From 19ce62ab769b74a8e417cd1213b7711a5ed55638 Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Mon, 20 Mar 2023 22:51:51 +0800
Subject: feat: check token usage in dry_run mode (#82)

---
 src/config/mod.rs | 17 +++++++++++++++--
 1 file changed, 15 insertions(+), 2 deletions(-)

(limited to 'src/config/mod.rs')

diff --git a/src/config/mod.rs b/src/config/mod.rs
index 147a817..e3d81cb 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -2,10 +2,11 @@ mod conversation;
 mod message;
 mod role;
 
+use self::conversation::Conversation;
 use self::message::Message;
 use self::role::Role;
-use self::{conversation::Conversation, message::within_max_tokens_limit};
 
+use crate::config::message::num_tokens_from_messages;
 use crate::utils::now;
 
 use anyhow::{anyhow, bail, Context, Result};
@@ -292,7 +293,10 @@ impl Config {
             let message = Message::new(content);
             vec![message]
         };
-        within_max_tokens_limit(&messages, self.model.1)?;
+        let tokens = num_tokens_from_messages(&messages);
+        if tokens >= self.model.1 {
+            bail!("Exceed max tokens limit")
+        }
 
         Ok(messages)
     }
@@ -434,6 +438,15 @@ impl Config {
         (self.highlight, self.light_theme)
     }
 
+    pub fn maybe_print_send_tokens(&self, input: &str) {
+        if self.dry_run {
+            if let Ok(messages) = self.build_messages(input) {
+                let tokens = num_tokens_from_messages(&messages);
+                println!(">>> The following message consumes {tokens} tokens.")
+            }
+        }
+    }
+
     fn open_message_file(&self) -> Result<File> {
         let path = Config::messages_file()?;
         ensure_parent_exists(&path)?;
-- 
cgit v1.2.3