refactor: optimize counting tokens (#53)

1 files changed, 7 insertions, 14 deletions

diff --git a/src/config/conversation.rs b/src/config/conversation.rs
index c3e9bed..4472781 100644
--- a/src/config/conversation.rs
+++ b/src/config/conversation.rs
@@ -1,9 +1,7 @@
-use super::message::{Message, MessageRole, MESSAGE_EXTRA_TOKENS};
+use super::message::{num_tokens_from_messages, Message, MessageRole};
 use super::role::Role;
 use super::MAX_TOKENS;
 
-use crate::utils::count_tokens;
-
 use anyhow::Result;
 use serde::{Deserialize, Serialize};
 
@@ -16,21 +14,17 @@ pub struct Conversation {
 
 impl Conversation {
     pub fn new(role: Option<Role>) -> Self {
-        let tokens = if let Some(role) = role.as_ref() {
-            role.consume_tokens()
-        } else {
-            0
-        };
-        Self {
-            tokens,
+        let mut value = Self {
+            tokens: 0,
             role,
             messages: vec![],
-        }
+        };
+        value.tokens = num_tokens_from_messages(&value.build_emssages(""));
+        value
     }
 
     pub fn add_message(&mut self, input: &str, output: &str) -> Result<()> {
         let mut need_add_msg = true;
-        let mut input_tokens = count_tokens(input);
         if self.messages.is_empty() {
             if let Some(role) = self.role.as_ref() {
                 self.messages.extend(role.build_emssages(input));
@@ -42,13 +36,12 @@ impl Conversation {
                 role: MessageRole::User,
                 content: input.to_string(),
             });
-            input_tokens += MESSAGE_EXTRA_TOKENS;
         }
         self.messages.push(Message {
             role: MessageRole::Assistant,
             content: output.to_string(),
         });
-        self.tokens += input_tokens + count_tokens(output) + MESSAGE_EXTRA_TOKENS;
+        self.tokens = num_tokens_from_messages(&self.messages);
         Ok(())
     }