chore: opitmize code, move MAX_TOKENS related to config/message.rs

3 files changed, 15 insertions, 9 deletions

diff --git a/src/config/conversation.rs b/src/config/conversation.rs
index 93382a3..c166195 100644
--- a/src/config/conversation.rs
+++ b/src/config/conversation.rs
@@ -1,6 +1,5 @@
-use super::message::{num_tokens_from_messages, Message, MessageRole};
+use super::message::{num_tokens_from_messages, Message, MessageRole, MAX_TOKENS};
 use super::role::Role;
-use super::MAX_TOKENS;
 
 use anyhow::{bail, Result};
 use serde::{Deserialize, Serialize};
diff --git a/src/config/message.rs b/src/config/message.rs
index 514a2fa..3fd5935 100644
--- a/src/config/message.rs
+++ b/src/config/message.rs
@@ -1,6 +1,9 @@
+use crate::utils::count_tokens;
+
+use anyhow::{bail, Result};
 use serde::{Deserialize, Serialize};
 
-use crate::utils::count_tokens;
+pub const MAX_TOKENS: usize = 4096;
 
 #[derive(Debug, Clone, Deserialize, Serialize)]
 pub struct Message {
@@ -25,6 +28,14 @@ pub enum MessageRole {
     User,
 }
 
+pub fn within_max_tokens_limit(messages: &[Message]) -> Result<()> {
+    let tokens = num_tokens_from_messages(messages);
+    if tokens >= MAX_TOKENS {
+        bail!("Exceed max tokens limit")
+    }
+    Ok(())
+}
+
 pub fn num_tokens_from_messages(messages: &[Message]) -> usize {
     let mut num_tokens = 0;
     for message in messages.iter() {
diff --git a/src/config/mod.rs b/src/config/mod.rs
index ff7acae..147d944 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -4,7 +4,7 @@ mod role;
 
 use self::message::Message;
 use self::role::Role;
-use self::{conversation::Conversation, message::num_tokens_from_messages};
+use self::{conversation::Conversation, message::within_max_tokens_limit};
 
 use crate::utils::now;
 
@@ -21,7 +21,6 @@ use std::{
     sync::Arc,
 };
 
-const MAX_TOKENS: usize = 4096;
 const CONFIG_FILE_NAME: &str = "config.yaml";
 const ROLES_FILE_NAME: &str = "roles.yaml";
 const HISTORY_FILE_NAME: &str = "history.txt";
@@ -229,10 +228,7 @@ impl Config {
             let message = Message::new(content);
             vec![message]
         };
-        let tokens = num_tokens_from_messages(&messages);
-        if tokens >= MAX_TOKENS {
-            bail!("Exceed max tokens limit")
-        }
+        within_max_tokens_limit(&messages)?;
 
         Ok(messages)
     }