diff options
Diffstat (limited to 'src')
| -rw-r--r-- | src/config/conversation.rs | 3 | ||||
| -rw-r--r-- | src/config/message.rs | 13 | ||||
| -rw-r--r-- | src/config/mod.rs | 8 |
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) } |
