refactor: streaming input text in dry-run mode (#154)

3 files changed, 15 insertions, 2 deletions

diff --git a/src/client/mod.rs b/src/client/mod.rs
index bf83c60..1f9e66c 100644
--- a/src/client/mod.rs
+++ b/src/client/mod.rs
@@ -17,6 +17,7 @@ use crate::{
     client::localai::LocalAIClient,
     config::{Config, SharedConfig},
     repl::{ReplyStreamHandler, SharedAbortSignal},
+    utils::split_text,
 };
 
 #[allow(clippy::struct_excessive_bools)]
@@ -92,7 +93,11 @@ pub trait Client {
             tokio::select! {
                 ret = async {
                     if self.get_config().read().dry_run {
-                        handler.text(&self.get_config().read().echo_messages(content))?;
+                        let words = split_text(content)?;
+                        for word in words {
+                            tokio::time::sleep(Duration::from_millis(25)).await;
+                            handler.text(&self.get_config().read().echo_messages(&word))?;
+                        }
                         return Ok(());
                     }
                     self.send_message_streaming_inner(content, handler).await
diff --git a/src/config/mod.rs b/src/config/mod.rs
index 4a0a126..4360dcf 100644
--- a/src/config/mod.rs
+++ b/src/config/mod.rs
@@ -412,7 +412,7 @@ impl Config {
         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.");
+                println!(">>> This message consumes {tokens} tokens. <<<");
             }
         }
     }
diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index ae75969..3ae1004 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -1,5 +1,6 @@
 mod tiktoken;
 
+use self::tiktoken::cl100k_base;
 pub use self::tiktoken::{cl100k_base_singleton, count_tokens, text_to_tokens, tokens_to_text};
 
 use arboard::Clipboard;
@@ -41,3 +42,10 @@ pub fn copy(src: &str) -> Result<(), arboard::Error> {
     let mut clipboard = Clipboard::new()?;
     clipboard.set_text(src)
 }
+
+pub fn split_text(text: &str) -> Result<Vec<String>, anyhow::Error> {
+    let bpe = cl100k_base()?;
+    let tokens = bpe.encode_with_special_tokens(text);
+    let data: Result<Vec<String>, _> = tokens.into_iter().map(|v| bpe.decode(&[v])).collect();
+    data
+}