feat: use tokie for accurate token counting

Cargo.toml

@@ -38,6 +38,7 @@ rand = "0.9"
 tracing = "0.1"
 once_cell = "1.21"
 async-trait = "0.1"
+tokie = { version = "0.0.8", features = ["hf"] }
 
 [dev-dependencies]
 tokio-test = "0.4"

src/catalog.rs

@@ -7,7 +7,7 @@ use std::collections::HashMap;
 use once_cell::sync::Lazy;
 use serde::Deserialize;
 
-use crate::models::ModelInfo;
+use crate::models::{ModelInfo, TokenizerInfo};
 
 /// Raw JSON model data embedded at compile time.
 const MODELS_JSON: &str = include_str!("data/models.json");
@@ -24,6 +24,8 @@ struct CatalogEntry {
     supports_dimensions: bool,
     #[serde(default)]
     supports_input_type: bool,
+    #[serde(default)]
+    tokenizer: Option<TokenizerInfo>,
 }
 
 impl From<CatalogEntry> for ModelInfo {
@@ -36,6 +38,7 @@ impl From<CatalogEntry> for ModelInfo {
             supports_dimensions: entry.supports_dimensions,
             supports_input_type: entry.supports_input_type,
             cost_per_million_tokens: entry.cost_per_million_tokens,
+            tokenizer: entry.tokenizer,
         }
     }
 }

src/lib.rs

@@ -54,11 +54,12 @@ pub mod errors;
 pub mod http;
 pub mod models;
 pub mod providers;
+pub(crate) mod tokenizer;
 
 // Re-exports
 pub use catalog::{find_model_by_name, get_model, list_catalog_providers, list_models};
 pub use client::Client;
 pub use errors::ClientError;
 pub use http::{HttpClient, HttpConfig};
-pub use models::{EmbedRequest, EmbedResponse, InputType, ModelInfo, Usage};
+pub use models::{EmbedRequest, EmbedResponse, InputType, ModelInfo, TokenizerInfo, Usage};
 pub use providers::EmbeddingProvider;

src/models.rs

@@ -58,6 +58,13 @@ pub enum InputType {
     Document,
 }
 
+/// Tokenizer configuration from the model catalog.
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct TokenizerInfo {
+    pub engine: String,
+    pub name: String,
+}
+
 /// Model information from the catalog.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct ModelInfo {
@@ -75,4 +82,6 @@ pub struct ModelInfo {
     pub supports_input_type: bool,
     /// Cost per 1M tokens in USD.
     pub cost_per_million_tokens: Option<f64>,
+    /// Tokenizer configuration.
+    pub tokenizer: Option<TokenizerInfo>,
 }

src/providers/cloudflare.rs

@@ -162,13 +162,19 @@ impl EmbeddingProvider for CloudflareProvider {
             .copied()
             .unwrap_or_else(|| embedding_vectors.first().map(|e| e.len()).unwrap_or(0));
 
-        // Cloudflare doesn't return token count, estimate based on input
-        // Rough estimate: 1 token per 5 characters
-        let estimated_tokens: u64 = request
-            .inputs
-            .iter()
-            .map(|s| (s.len() as u64 / 5).max(1))
-            .sum();
+        let inputs_for_counting = request.inputs.clone();
+        let model_for_counting = request.model.clone();
+        let estimated_tokens = tokio::task::spawn_blocking(move || {
+            crate::tokenizer::count_tokens_for_model(
+                "cloudflare",
+                &model_for_counting,
+                &inputs_for_counting,
+            )
+        })
+        .await
+        .unwrap_or_else(|_| {
+            crate::tokenizer::fallback_count(&request.inputs)
+        });
 
         let cost = calculate_cost(&request.model, estimated_tokens);
 

src/providers/gemini.rs

@@ -163,13 +163,19 @@ impl EmbeddingProvider for GeminiProvider {
         let latency_ms = start.elapsed().as_secs_f64() * 1000.0;
         let dimensions = embedding_vectors.first().map(|e| e.len()).unwrap_or(0);
 
-        // Gemini doesn't return token count, estimate based on input
-        // Rough estimate: 1 token per 5 characters
-        let estimated_tokens: u64 = request
-            .inputs
-            .iter()
-            .map(|s| (s.len() as u64 / 5).max(1))
-            .sum();
+        let inputs_for_counting = request.inputs.clone();
+        let model_for_counting = request.model.clone();
+        let estimated_tokens = tokio::task::spawn_blocking(move || {
+            crate::tokenizer::count_tokens_for_model(
+                "gemini",
+                &model_for_counting,
+                &inputs_for_counting,
+            )
+        })
+        .await
+        .unwrap_or_else(|_| {
+            crate::tokenizer::fallback_count(&request.inputs)
+        });
 
         let cost = calculate_cost(&request.model, estimated_tokens);
 

src/tokenizer.rs

@@ -0,0 +1,80 @@
+use std::collections::HashMap;
+use std::sync::{Arc, Mutex};
+
+use once_cell::sync::Lazy;
+use tokie::Tokenizer;
+use tracing::warn;
+
+use crate::catalog::get_model;
+
+static TOKENIZER_CACHE: Lazy<Mutex<HashMap<String, Arc<Tokenizer>>>> =
+    Lazy::new(|| Mutex::new(HashMap::new()));
+
+pub fn count_tokens_for_model(provider: &str, model: &str, texts: &[String]) -> u64 {
+    let tokenizer_info = match get_model(provider, model) {
+        Some(info) => match info.tokenizer {
+            Some(t) => t,
+            None => return fallback_count(texts),
+        },
+        None => return fallback_count(texts),
+    };
+
+    let repo_name = match tokenizer_info.engine.as_str() {
+        "huggingface" => tokenizer_info.name.clone(),
+        "tiktoken" => match tiktoken_to_repo(&tokenizer_info.name) {
+            Some(repo) => repo.to_string(),
+            None => return fallback_count(texts),
+        },
+        _ => return fallback_count(texts),
+    };
+
+    match get_or_load_tokenizer(&repo_name) {
+        Some(tokenizer) => {
+            let refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
+            tokenizer.count_tokens_batch(&refs).iter().sum::<usize>() as u64
+        }
+        None => fallback_count(texts),
+    }
+}
+
+fn get_or_load_tokenizer(name: &str) -> Option<Arc<Tokenizer>> {
+    let mut cache = TOKENIZER_CACHE.lock().ok()?;
+
+    if let Some(tok) = cache.get(name) {
+        return Some(Arc::clone(tok));
+    }
+
+    match Tokenizer::from_pretrained(name) {
+        Ok(tok) => {
+            let tok = Arc::new(tok);
+            cache.insert(name.to_string(), Arc::clone(&tok));
+            Some(tok)
+        }
+        Err(e) => {
+            warn!(
+                tokenizer = name,
+                error = %e,
+                "Failed to load tokenizer, falling back to estimate"
+            );
+            None
+        }
+    }
+}
+
+fn tiktoken_to_repo(name: &str) -> Option<&'static str> {
+    match name {
+        "cl100k_base" => Some("xenova/gpt-4"),
+        "o200k_base" => Some("xenova/gpt-4o"),
+        _ => None,
+    }
+}
+
+pub fn fallback_count(texts: &[String]) -> u64 {
+    match get_or_load_tokenizer("xenova/gpt-4") {
+        Some(tokenizer) => {
+            let refs: Vec<&str> = texts.iter().map(|s| s.as_str()).collect();
+            tokenizer.count_tokens_batch(&refs).iter().sum::<usize>() as u64
+        }
+        None => texts.iter().map(|s| (s.len() as u64 / 5).max(1)).sum(),
+    }
+}