From aab3170dd0a62a41746138dc6ce06ccee37ef586 Mon Sep 17 00:00:00 2001 From: bhavnick Date: Mon, 20 Apr 2026 18:25:18 -0700 Subject: [PATCH] feat: use tokie for accurate token counting in Cloudflare and Gemini providers Replace the rough len/5 heuristic with real tokenizer-based counting via tokie. Tokenizers are lazily loaded from HuggingFace Hub and cached. For models without a known tokenizer, falls back to cl100k (GPT-4). Co-Authored-By: Claude Opus 4.6 --- Cargo.toml | 1 + src/catalog.rs | 5 ++- src/lib.rs | 3 +- src/models.rs | 9 +++++ src/providers/cloudflare.rs | 20 ++++++---- src/providers/gemini.rs | 20 ++++++---- src/tokenizer.rs | 80 +++++++++++++++++++++++++++++++++++++ 7 files changed, 122 insertions(+), 16 deletions(-) create mode 100644 src/tokenizer.rs diff --git a/Cargo.toml b/Cargo.toml index 7e84abe..cb9c88b 100644 --- a/Cargo.toml +++ b/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" diff --git a/src/catalog.rs b/src/catalog.rs index c9776c3..467a0c9 100644 --- a/src/catalog.rs +++ b/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, } impl From for ModelInfo { @@ -36,6 +38,7 @@ impl From 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, } } } diff --git a/src/lib.rs b/src/lib.rs index 8463d18..8bb3529 100644 --- a/src/lib.rs +++ b/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; diff --git a/src/models.rs b/src/models.rs index 358d02f..01e5409 100644 --- a/src/models.rs +++ b/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, + /// Tokenizer configuration. + pub tokenizer: Option, } diff --git a/src/providers/cloudflare.rs b/src/providers/cloudflare.rs index 8024eb8..4b2ffd4 100644 --- a/src/providers/cloudflare.rs +++ b/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); diff --git a/src/providers/gemini.rs b/src/providers/gemini.rs index fabd3af..4504941 100644 --- a/src/providers/gemini.rs +++ b/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); diff --git a/src/tokenizer.rs b/src/tokenizer.rs new file mode 100644 index 0000000..8bbe60c --- /dev/null +++ b/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>>> = + 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::() as u64 + } + None => fallback_count(texts), + } +} + +fn get_or_load_tokenizer(name: &str) -> Option> { + 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::() as u64 + } + None => texts.iter().map(|s| (s.len() as u64 / 5).max(1)).sum(), + } +}