From 2a9aa97dc14347fcb43482949e2fda1d651ee7ff Mon Sep 17 00:00:00 2001 From: Salman Mohammed Date: Mon, 20 Jan 2025 15:22:09 -0500 Subject: [PATCH] feat: setup tokenizer on the provider model config (#656) --- .../goose/benches/tokenization_benchmark.rs | 12 +- crates/goose/build.rs | 12 +- crates/goose/src/agents/default.rs | 22 +- crates/goose/src/agents/reference.rs | 3 +- crates/goose/src/agents/truncate.rs | 20 +- crates/goose/src/providers/configs.rs | 24 ++ crates/goose/src/token_counter.rs | 287 ++++++++++-------- 7 files changed, 204 insertions(+), 176 deletions(-) diff --git a/crates/goose/benches/tokenization_benchmark.rs b/crates/goose/benches/tokenization_benchmark.rs index 11367c0ee7..708cc68bf5 100644 --- a/crates/goose/benches/tokenization_benchmark.rs +++ b/crates/goose/benches/tokenization_benchmark.rs @@ -2,15 +2,15 @@ use criterion::{black_box, criterion_group, criterion_main, Criterion}; use goose::token_counter::TokenCounter; fn benchmark_tokenization(c: &mut Criterion) { - let counter = TokenCounter::new(); let lengths = [1_000, 5_000, 10_000, 50_000, 100_000, 124_000, 200_000]; - let models = ["gpt-4o", "claude-3.5-sonnet"]; + let tokenizer_names = ["Xenova--gpt-4o", "Xenova--claude-tokenizer"]; - for &length in &lengths { - for model_name in models { + for tokenizer_name in tokenizer_names { + let counter = TokenCounter::new(tokenizer_name); + for &length in &lengths { let text = "hello ".repeat(length); - c.bench_function(&format!("{}_{}_tokens", model_name, length), |b| { - b.iter(|| counter.count_tokens(black_box(&text), Some(model_name))) + c.bench_function(&format!("{}_{}_tokens", tokenizer_name, length), |b| { + b.iter(|| counter.count_tokens(black_box(&text))) }); } } diff --git a/crates/goose/build.rs b/crates/goose/build.rs index ccfb369848..446a3d4896 100644 --- a/crates/goose/build.rs +++ b/crates/goose/build.rs @@ -3,13 +3,7 @@ use std::fs; use std::path::Path; const BASE_DIR: &str = "../../tokenizer_files"; -const MODELS: &[&str] = &[ - "Xenova/claude-tokenizer", - "Xenova/gemma-2-tokenizer", - "Xenova/gpt-4o", - "Qwen/Qwen2.5-Coder-32B-Instruct", - "Xenova/llama3-tokenizer", -]; +const TOKENIZERS: &[&str] = &["Xenova/gpt-4o", "Xenova/claude-tokenizer"]; #[tokio::main] async fn main() -> Result<(), Box> { @@ -18,8 +12,8 @@ async fn main() -> Result<(), Box> { println!("cargo:rerun-if-changed=build.rs"); println!("cargo:rerun-if-changed={}", BASE_DIR); - for model in MODELS { - download_tokenizer(model).await?; + for tokenizer_name in TOKENIZERS { + download_tokenizer(tokenizer_name).await?; } Ok(()) diff --git a/crates/goose/src/agents/default.rs b/crates/goose/src/agents/default.rs index a73f5d9d3b..897b89d46b 100644 --- a/crates/goose/src/agents/default.rs +++ b/crates/goose/src/agents/default.rs @@ -23,9 +23,10 @@ pub struct DefaultAgent { impl DefaultAgent { pub fn new(provider: Box) -> Self { + let token_counter = TokenCounter::new(provider.get_model_config().tokenizer_name()); Self { capabilities: Mutex::new(Capabilities::new(provider)), - token_counter: TokenCounter::new(), + token_counter: token_counter, } } @@ -37,7 +38,6 @@ impl DefaultAgent { messages: &[Message], pending: &[Message], target_limit: usize, - model_name: &str, resource_items: &mut [ResourceItem], ) -> SystemResult> { // Flatten all resource content into a vector of strings @@ -46,13 +46,9 @@ impl DefaultAgent { .map(|item| item.content.clone()) .collect(); - let approx_count = self.token_counter.count_everything( - system_prompt, - messages, - tools, - &resources, - Some(model_name), - ); + let approx_count = + self.token_counter + .count_everything(system_prompt, messages, tools, &resources); let mut status_content: Vec = Vec::new(); if approx_count > target_limit { @@ -60,10 +56,7 @@ impl DefaultAgent { for item in resource_items.iter_mut() { if item.token_count.is_none() { - let count = self - .token_counter - .count_tokens(&item.content, Some(model_name)) - as u32; + let count = self.token_counter.count_tokens(&item.content) as u32; item.token_count = Some(count); } } @@ -183,7 +176,6 @@ impl Agent for DefaultAgent { messages, &Vec::new(), estimated_limit, - &capabilities.provider().get_model_config().model_name, &mut capabilities.get_resources().await?, ) .await?; @@ -249,7 +241,7 @@ impl Agent for DefaultAgent { let pending = vec![response, message_tool_response]; - messages = self.prepare_inference(&system_prompt, &tools, &messages, &pending, estimated_limit, &capabilities.provider().get_model_config().model_name, &mut capabilities.get_resources().await?).await?; + messages = self.prepare_inference(&system_prompt, &tools, &messages, &pending, estimated_limit, &mut capabilities.get_resources().await?).await?; } })) } diff --git a/crates/goose/src/agents/reference.rs b/crates/goose/src/agents/reference.rs index b5e7c97a54..d11b399c4e 100644 --- a/crates/goose/src/agents/reference.rs +++ b/crates/goose/src/agents/reference.rs @@ -22,9 +22,10 @@ pub struct ReferenceAgent { impl ReferenceAgent { pub fn new(provider: Box) -> Self { + let token_counter = TokenCounter::new(provider.get_model_config().tokenizer_name()); Self { capabilities: Mutex::new(Capabilities::new(provider)), - _token_counter: TokenCounter::new(), + _token_counter: token_counter, } } } diff --git a/crates/goose/src/agents/truncate.rs b/crates/goose/src/agents/truncate.rs index 1d2d3199e3..f775727df8 100644 --- a/crates/goose/src/agents/truncate.rs +++ b/crates/goose/src/agents/truncate.rs @@ -23,9 +23,10 @@ pub struct TruncateAgent { impl TruncateAgent { pub fn new(provider: Box) -> Self { + let token_counter = TokenCounter::new(provider.get_model_config().tokenizer_name()); Self { capabilities: Mutex::new(Capabilities::new(provider)), - _token_counter: TokenCounter::new(), + _token_counter: token_counter, } } @@ -35,7 +36,6 @@ impl TruncateAgent { tools: &[Tool], messages: &[Message], target_limit: usize, - model_name: &str, resource_items: &mut [ResourceItem], ) -> SystemResult> { // Flatten all resource content into a vector of strings @@ -44,14 +44,13 @@ impl TruncateAgent { .map(|item| item.content.clone()) .collect(); - let model = Some(model_name); let approx_count = self._token_counter - .count_everything(system_prompt, messages, tools, &resources, model); + .count_everything(system_prompt, messages, tools, &resources); let mut new_messages = messages.to_vec(); if approx_count > target_limit { - let user_msg_size = self.text_content_size(new_messages.last(), model); + let user_msg_size = self.text_content_size(new_messages.last()); if user_msg_size > target_limit { debug!( "[WARNING] User message {} exceeds token budget {}.", @@ -61,7 +60,7 @@ impl TruncateAgent { return Err(SystemError::ContextLimit); } - new_messages = self.chop_front_messages(messages, approx_count, target_limit, model); + new_messages = self.chop_front_messages(messages, approx_count, target_limit); if new_messages.is_empty() { return Err(SystemError::ContextLimit); } @@ -70,13 +69,13 @@ impl TruncateAgent { Ok(new_messages) } - fn text_content_size(&self, message: Option<&Message>, model: Option<&str>) -> usize { + fn text_content_size(&self, message: Option<&Message>) -> usize { let text = message .and_then(|msg| msg.content.first()) .and_then(|c| c.as_text()); if let Some(txt) = text { - let count = self._token_counter.count_tokens(txt, model); + let count = self._token_counter.count_tokens(txt); return count; } @@ -90,7 +89,6 @@ impl TruncateAgent { messages: &[Message], approx_count: usize, target_limit: usize, - model: Option<&str>, ) -> Vec { debug!( "[WARNING] Conversation history has size: {} exceeding the token budget of {}. \ @@ -107,7 +105,7 @@ impl TruncateAgent { if current_tokens < target_limit || trimmed_items.is_empty() { break; } - let count = self.text_content_size(Some(msg), model); + let count = self.text_content_size(Some(msg)); let _ = trimmed_items.pop_front().unwrap(); // Subtract removed message’s token_count current_tokens = current_tokens.saturating_sub(count as usize); @@ -151,7 +149,6 @@ impl Agent for TruncateAgent { &tools, messages, _estimated_limit, - &capabilities.provider().get_model_config().model_name, &mut capabilities.get_resources().await?, ) .await?; @@ -209,7 +206,6 @@ impl Agent for TruncateAgent { &tools, &messages, _estimated_limit, - &capabilities.provider().get_model_config().model_name, &mut capabilities.get_resources().await? ).await?; diff --git a/crates/goose/src/providers/configs.rs b/crates/goose/src/providers/configs.rs index bbfb465023..9dca7c0bad 100644 --- a/crates/goose/src/providers/configs.rs +++ b/crates/goose/src/providers/configs.rs @@ -3,11 +3,19 @@ use serde::{Deserialize, Serialize}; const DEFAULT_CONTEXT_LIMIT: usize = 200_000; const DEFAULT_ESTIMATE_FACTOR: f32 = 0.8; +// Tokenizer names, used to infer from model name +pub const GPT_4O_TOKENIZER: &str = "Xenova--gpt-4o"; +pub const CLAUDE_TOKENIZER: &str = "Xenova--claude-tokenizer"; + /// Configuration for model-specific settings and limits #[derive(Debug, Clone, Serialize, Deserialize)] pub struct ModelConfig { /// The name of the model to use pub model_name: String, + // Optional tokenizer name (corresponds to the sanitized HuggingFace tokenizer name) + // "Xenova/gpt-4o" -> "Xenova/gpt-4o" + // If not provided, best attempt will be made to infer from model name or default + pub tokenizer_name: String, /// Optional explicit context limit that overrides any defaults pub context_limit: Option, /// Optional temperature setting (0.0 - 1.0) @@ -28,9 +36,11 @@ impl ModelConfig { /// 3. Global default (128_000) (in get_context_limit) pub fn new(model_name: String) -> Self { let context_limit = Self::get_model_specific_limit(&model_name); + let tokenizer_name = Self::infer_tokenizer_name(&model_name); Self { model_name, + tokenizer_name: tokenizer_name.to_string(), context_limit, temperature: None, max_tokens: None, @@ -38,6 +48,15 @@ impl ModelConfig { } } + fn infer_tokenizer_name(model_name: &str) -> &'static str { + if model_name.contains("claude") { + CLAUDE_TOKENIZER + } else { + // Default tokenizer + GPT_4O_TOKENIZER + } + } + /// Get model-specific context limit based on model name fn get_model_specific_limit(model_name: &str) -> Option { // Implement some sensible defaults @@ -85,6 +104,11 @@ impl ModelConfig { self } + // Get the tokenizer name + pub fn tokenizer_name(&self) -> &str { + &self.tokenizer_name + } + /// Get the context_limit for the current model /// If none are defined, use the DEFAULT_CONTEXT_LIMIT pub fn context_limit(&self) -> usize { diff --git a/crates/goose/src/token_counter.rs b/crates/goose/src/token_counter.rs index 28b49756ce..23e954c2e5 100644 --- a/crates/goose/src/token_counter.rs +++ b/crates/goose/src/token_counter.rs @@ -1,120 +1,118 @@ -use crate::message::Message; use include_dir::{include_dir, Dir}; -use indicatif::{ProgressBar, ProgressStyle}; -use mcp_core::tool::Tool; -use std::collections::HashMap; -use std::time::Duration; +use mcp_core::Tool; +use std::error::Error; +use std::fs; +use std::path::Path; use tokenizers::tokenizer::Tokenizer; -// Embed the tokenizer files directory +use crate::message::Message; + +// The embedded directory with all possible tokenizer files. +// If one of them doesn’t exist, we’ll download it at startup. static TOKENIZER_FILES: Dir = include_dir!("$CARGO_MANIFEST_DIR/../../tokenizer_files"); +/// The `TokenCounter` now stores exactly one `Tokenizer`. pub struct TokenCounter { - tokenizers: HashMap, -} - -const GPT_4O_TOKENIZER_KEY: &str = "Xenova--gpt-4o"; -const CLAUDE_TOKENIZER_KEY: &str = "Xenova--claude-tokenizer"; -const GOOGLE_TOKENIZER_KEY: &str = "Xenova--gemma-2-tokenizer"; -const QWEN_TOKENIZER_KEY: &str = "Qwen--Qwen2.5-Coder-32B-Instruct"; -const LLAMA_TOKENIZER_KEY: &str = "Xenova--llama3-tokenizer"; - -impl Default for TokenCounter { - fn default() -> Self { - Self::new() - } + tokenizer: Tokenizer, } impl TokenCounter { - fn load_tokenizer(&mut self, tokenizer_key: &str) { - // Create a spinner that will show during both loading and parsing - let pb = ProgressBar::new_spinner(); - pb.set_style( - ProgressStyle::default_spinner() - .template("{spinner:.green} {msg}") - .unwrap(), - ); - - // Start spinner for loading phase - pb.set_message(format!("Loading {} tokenizer from disk...", tokenizer_key)); - pb.enable_steady_tick(Duration::from_millis(120)); - - // Load from embedded tokenizer files. The tokenizer_key must match the directory name. - let tokenizer_path = format!("{}/tokenizer.json", tokenizer_key); - let file_content = TOKENIZER_FILES - .get_file(&tokenizer_path) - .map(|f| f.contents()) - .ok_or_else(|| format!("Embedded tokenizer file not found: {}", tokenizer_path)) - .unwrap(); - - // Update spinner for parsing phase - pb.set_message(format!("Initializing {} tokenizer...", tokenizer_key)); - - let tokenizer = Tokenizer::from_bytes(file_content); - - match tokenizer { - Ok(tokenizer) => { - self.tokenizers.insert(tokenizer_key.to_string(), tokenizer); - pb.finish_and_clear(); - } + /// Creates a new `TokenCounter` using the given HuggingFace tokenizer name. + /// + /// * `tokenizer_name` might look like "Xenova--gpt-4o" + /// or "Qwen--Qwen2.5-Coder-32B-Instruct", etc. + pub fn new(tokenizer_name: &str) -> Self { + match Self::load_from_embedded(tokenizer_name) { + Ok(tokenizer) => Self { tokenizer }, Err(e) => { - pb.finish_and_clear(); - eprintln!("Failed to load tokenizer {}: {}", tokenizer_key, e); + println!( + "Tokenizer '{}' not found in embedded dir: {}", + tokenizer_name, e + ); + println!("Attempting to download tokenizer and load..."); + // Fallback to download tokenizer and load from disk + match Self::download_and_load(tokenizer_name) { + Ok(counter) => counter, + Err(e) => panic!("Failed to initialize tokenizer: {}", e), + } } } } - pub fn new() -> Self { - let mut counter = TokenCounter { - tokenizers: HashMap::new(), - }; - // Add default tokenizers - for tokenizer_key in [ - GPT_4O_TOKENIZER_KEY, - CLAUDE_TOKENIZER_KEY, - GOOGLE_TOKENIZER_KEY, - QWEN_TOKENIZER_KEY, - LLAMA_TOKENIZER_KEY, - ] { - counter.load_tokenizer(tokenizer_key); + /// Load tokenizer bytes from the embedded directory (via `include_dir!`). + fn load_from_embedded(tokenizer_name: &str) -> Result> { + let tokenizer_file_path = format!("{}/tokenizer.json", tokenizer_name); + let file = TOKENIZER_FILES + .get_file(&tokenizer_file_path) + .ok_or_else(|| { + format!( + "Tokenizer file not found in embedded: {}", + tokenizer_file_path + ) + })?; + let contents = file.contents(); + let tokenizer = Tokenizer::from_bytes(contents) + .map_err(|e| format!("Failed to parse tokenizer bytes: {}", e))?; + Ok(tokenizer) + } + + /// Fallback: If not found in embedded, we look in `base_dir` on disk. + /// If not on disk, we download from Hugging Face, then load from disk. + fn download_and_load(tokenizer_name: &str) -> Result> { + let local_dir = std::env::temp_dir().join(tokenizer_name); + let local_json_path = local_dir.join("tokenizer.json"); + + // If the file doesn't already exist, we download from HF + if !Path::new(&local_json_path).exists() { + eprintln!("Tokenizer file not on disk, downloading…"); + let repo_id = tokenizer_name.replace("--", "/"); + // e.g. "Xenova--llama3-tokenizer" -> "Xenova/llama3-tokenizer" + Self::download_tokenizer(&repo_id, &local_dir)?; } - counter + + // Load from disk + let file_content = fs::read(&local_json_path)?; + let tokenizer = Tokenizer::from_bytes(&file_content) + .map_err(|e| format!("Failed to parse tokenizer after download: {}", e))?; + + Ok(Self { tokenizer }) } - pub fn add_tokenizer(&mut self, tokenizer_key: &str) { - self.load_tokenizer(tokenizer_key); + /// Download from Hugging Face into the local directory if not already present. + /// Synchronous version using a blocking runtime for simplicity. + fn download_tokenizer(repo_id: &str, download_dir: &Path) -> Result<(), Box> { + fs::create_dir_all(download_dir)?; + + let file_url = format!( + "https://huggingface.co/{}/resolve/main/tokenizer.json", + repo_id + ); + let file_path = download_dir.join("tokenizer.json"); + + // Blocking for example: just spawn a short-lived runtime + let content = tokio::runtime::Runtime::new()?.block_on(async { + let response = reqwest::get(&file_url).await?; + if !response.status().is_success() { + let error_msg = + format!("Failed to download tokenizer: status {}", response.status()); + return Err(Box::::from(error_msg)); + } + let bytes = response.bytes().await?; + Ok(bytes) + })?; + + fs::write(&file_path, content)?; + + Ok(()) } - fn model_to_tokenizer_key(model_name: Option<&str>) -> &str { - let model_name = model_name.unwrap_or("gpt-4o").to_lowercase(); - if model_name.contains("claude") { - CLAUDE_TOKENIZER_KEY - } else if model_name.contains("qwen") { - QWEN_TOKENIZER_KEY - } else if model_name.contains("gemini") { - GOOGLE_TOKENIZER_KEY - } else if model_name.contains("llama") { - LLAMA_TOKENIZER_KEY - } else { - // default - GPT_4O_TOKENIZER_KEY - } - } - - fn get_tokenizer(&self, model_name: Option<&str>) -> &Tokenizer { - let tokenizer_key = Self::model_to_tokenizer_key(model_name); - self.tokenizers - .get(tokenizer_key) - .expect("Tokenizer not found") - } - - pub fn count_tokens(&self, text: &str, model_name: Option<&str>) -> usize { - let tokenizer = self.get_tokenizer(model_name); - let encoding = tokenizer.encode(text, false).unwrap(); + /// Count tokens for a piece of text using our single tokenizer. + pub fn count_tokens(&self, text: &str) -> usize { + let encoding = self.tokenizer.encode(text, false).unwrap(); encoding.len() } - fn count_tokens_for_tools(&self, tools: &[Tool], model_name: Option<&str>) -> usize { + fn count_tokens_for_tools(&self, tools: &[Tool]) -> usize { // Token counts for different function components let func_init = 7; // Tokens for function initialization let prop_init = 3; // Tokens for properties initialization @@ -130,7 +128,7 @@ impl TokenCounter { let name = &tool.name; let description = &tool.description.trim_end_matches('.'); let line = format!("{}:{}", name, description); - func_token_count += self.count_tokens(&line, model_name); // Add tokens for name and description + func_token_count += self.count_tokens(&line); // Add tokens for name and description if let serde_json::Value::Object(properties) = &tool.input_schema["properties"] { if !properties.is_empty() { @@ -144,14 +142,14 @@ impl TokenCounter { .unwrap_or("") .trim_end_matches('.'); let line = format!("{}:{}:{}", p_name, p_type, p_desc); - func_token_count += self.count_tokens(&line, model_name); + func_token_count += self.count_tokens(&line); if let Some(enum_values) = value["enum"].as_array() { func_token_count = func_token_count.saturating_add_signed(enum_init); // Add tokens if property has enum list for item in enum_values { if let Some(item_str) = item.as_str() { func_token_count += enum_item; - func_token_count += self.count_tokens(item_str, model_name); + func_token_count += self.count_tokens(item_str); } } } @@ -170,7 +168,6 @@ impl TokenCounter { system_prompt: &str, messages: &[Message], tools: &[Tool], - model_name: Option<&str>, ) -> usize { // <|im_start|>ROLE<|im_sep|>MESSAGE<|im_end|> let tokens_per_message = 4; @@ -178,7 +175,7 @@ impl TokenCounter { // Count tokens in the system prompt let mut num_tokens = 0; if !system_prompt.is_empty() { - num_tokens += self.count_tokens(system_prompt, model_name) + tokens_per_message; + num_tokens += self.count_tokens(system_prompt) + tokens_per_message; } for message in messages { @@ -187,7 +184,7 @@ impl TokenCounter { for content in &message.content { // content can either be text response or tool request if let Some(content_text) = content.as_text() { - num_tokens += self.count_tokens(content_text, model_name); + num_tokens += self.count_tokens(content_text); } else if let Some(tool_request) = content.as_tool_request() { // TODO: count tokens for tool request let tool_call = tool_request.tool_call.as_ref().unwrap(); @@ -195,9 +192,9 @@ impl TokenCounter { "{}:{}:{}", tool_request.id, tool_call.name, tool_call.arguments ); - num_tokens += self.count_tokens(&text, model_name); + num_tokens += self.count_tokens(&text); } else if let Some(tool_response_text) = content.as_tool_response_text() { - num_tokens += self.count_tokens(&tool_response_text, model_name); + num_tokens += self.count_tokens(&tool_response_text); } else { // unsupported content type such as image - pass continue; @@ -207,7 +204,7 @@ impl TokenCounter { // Count tokens for tools if provided if !tools.is_empty() { - num_tokens += self.count_tokens_for_tools(tools, model_name); + num_tokens += self.count_tokens_for_tools(tools); } // Every reply is primed with <|start|>assistant<|message|> @@ -222,13 +219,12 @@ impl TokenCounter { messages: &[Message], tools: &[Tool], resources: &[String], - model_name: Option<&str>, ) -> usize { - let mut num_tokens = self.count_chat_tokens(system_prompt, messages, tools, model_name); + let mut num_tokens = self.count_chat_tokens(system_prompt, messages, tools); if !resources.is_empty() { for resource in resources { - num_tokens += self.count_tokens(resource, model_name); + num_tokens += self.count_tokens(resource); } } num_tokens @@ -238,40 +234,39 @@ impl TokenCounter { #[cfg(test)] mod tests { use super::*; - use crate::message::MessageContent; + use crate::message::{Message, MessageContent}; // or however your `Message` is imported + use crate::providers::configs::{CLAUDE_TOKENIZER, GPT_4O_TOKENIZER}; use mcp_core::role::Role; + use mcp_core::tool::Tool; use serde_json::json; #[test] - fn test_add_tokenizer_and_count_tokens() { - let mut counter = TokenCounter::new(); - counter.add_tokenizer(QWEN_TOKENIZER_KEY); - let text = "Hey there!"; - let count = counter.count_tokens(text, Some("qwen2.5-ollama")); - println!("Token count for '{}': {:?}", text, count); - assert_eq!(count, 3); - } + fn test_claude_tokenizer() { + let counter = TokenCounter::new(CLAUDE_TOKENIZER); - // Update the default tokenizer test similarly - #[test] - fn test_specific_claude_tokenizer() { - let counter = TokenCounter::new(); let text = "Hello, how are you?"; - let count = counter.count_tokens(text, Some("claude-3-5-sonnet-2")); + let count = counter.count_tokens(text); println!("Token count for '{}': {:?}", text, count); - assert_eq!(count, 6); + + // The old test expected 6 tokens + assert_eq!(count, 6, "Claude tokenizer token count mismatch"); } #[test] - fn test_default_gpt_4o_tokenizer() { - let counter = TokenCounter::new(); - let count = counter.count_tokens("Hey there!", None); - assert_eq!(count, 3); + fn test_gpt_4o_tokenizer() { + let counter = TokenCounter::new(GPT_4O_TOKENIZER); + + let text = "Hey there!"; + let count = counter.count_tokens(text); + println!("Token count for '{}': {:?}", text, count); + + // The old test expected 3 tokens + assert_eq!(count, 3, "GPT-4o tokenizer token count mismatch"); } #[test] fn test_count_chat_tokens() { - let token_counter = TokenCounter::new(); + let counter = TokenCounter::new(GPT_4O_TOKENIZER); let system_prompt = "You are a helpful assistant that can answer questions about the weather."; @@ -317,17 +312,43 @@ mod tests { }), }]; - let token_count_without_tools = - token_counter.count_chat_tokens(system_prompt, &messages, &[], Some("gpt-4o")); + let token_count_without_tools = counter.count_chat_tokens(system_prompt, &messages, &[]); println!("Total tokens without tools: {}", token_count_without_tools); - let token_count_with_tools = - token_counter.count_chat_tokens(system_prompt, &messages, &tools, Some("gpt-4o")); + let token_count_with_tools = counter.count_chat_tokens(system_prompt, &messages, &tools); println!("Total tokens with tools: {}", token_count_with_tools); - // The token count for messages without tools is calculated using the tokenizer - https://tiktokenizer.vercel.app/ - // The token count for messages with tools is taken from tiktoken github repo example (notebook) + // The old test used 56 / 124 for GPT-4o. Adjust if your actual tokenizer changes assert_eq!(token_count_without_tools, 56); assert_eq!(token_count_with_tools, 124); } + + #[test] + #[should_panic] + fn test_panic_if_provided_tokenizer_doesnt_exist() { + // This should panic because the tokenizer doesn't exist + // in the embedded directory and the download fails + + TokenCounter::new("nonexistent-tokenizer"); + } + + // Optional test to confirm that fallback download works if not found in embedded: + // Ignored cause this actually downloads a tokenizer from Hugging Face + #[test] + #[ignore] + fn test_download_tokenizer_successfully_if_not_embedded() { + let non_embedded_key = "openai-community/gpt2"; + let counter = TokenCounter::new(non_embedded_key); + + // If it downloads successfully, we can do a quick count to ensure it's valid + let text = "print('hello world')"; + let count = counter.count_tokens(text); + println!( + "Downloaded tokenizer, token count for '{}': {}", + text, count + ); + + // https://tiktokenizer.vercel.app/?model=gpt2 + assert!(count == 5, "Expected 5 tokens from downloaded tokenizer"); + } }