fix: dry run on role or session (#181)

1 files changed, 11 insertions, 11 deletions

diff --git a/src/utils/mod.rs b/src/utils/mod.rs
index 699f8a0..db53fd6 100644
--- a/src/utils/mod.rs
+++ b/src/utils/mod.rs
@@ -1,10 +1,8 @@
 mod tiktoken;
 
-use self::tiktoken::cl100k_base;
-pub use self::tiktoken::{cl100k_base_singleton, count_tokens, text_to_tokens, tokens_to_text};
+pub use self::tiktoken::cl100k_base_singleton;
 
 use chrono::prelude::*;
-use crossterm::style::{Color, Stylize};
 use std::io::{stdout, Write};
 
 #[macro_export]
@@ -32,16 +30,18 @@ pub fn get_env_name(key: &str) -> String {
     )
 }
 
-#[allow(unused)]
-pub fn emphasis(text: &str) -> String {
-    text.stylize().with(Color::White).to_string()
+/// Split text to tokens
+pub fn tokenize(text: &str) -> Vec<String> {
+    let tokens = cl100k_base_singleton().lock().tokenize(text);
+    tokens.into_iter().map(|(_, text)| text).collect()
 }
 
-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
+/// Count how many tokens a piece of text needs to consume
+pub fn count_tokens(text: &str) -> usize {
+    cl100k_base_singleton()
+        .lock()
+        .encode_with_special_tokens(text)
+        .len()
 }
 
 pub fn light_theme_from_colorfgbg(colorfgbg: &str) -> Option<bool> {