serializer/llm/
tokens.rs

1//! Token counting for LLM serializer output
2//!
3//! Provides token counting and measurement for various LLM models:
4//! - **GPT-4o / o1**: Uses `o200k_base` tokenizer via tiktoken-rs
5//! - **Gemini 3**: Uses SentencePiece tokenizer via tokenizers crate
6//! - **Claude Opus 4.5**: Uses approximate BPE tokenizer
7//!
8//! ## Usage
9//!
10//! ```rust
11//! use serializer::llm::tokens::{TokenCounter, ModelType, TokenInfo};
12//!
13//! let counter = TokenCounter::new();
14//! let text = "Hello, world!";
15//!
16//! // Count tokens for GPT-4o
17//! let info = counter.count(text, ModelType::Gpt4o);
18//! assert!(info.count > 0);
19//! println!("Token count: {}", info.count);
20//! ```
21
22use std::collections::HashMap;
23
24/// Supported LLM model types for token counting
25#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
26pub enum ModelType {
27    /// OpenAI GPT-4o (uses o200k_base tokenizer)
28    Gpt4o,
29    /// OpenAI o1 model (uses o200k_base tokenizer)
30    O1,
31    /// OpenAI GPT-4 (uses cl100k_base tokenizer)
32    Gpt4,
33    /// Google Gemini 3 (uses SentencePiece tokenizer)
34    Gemini3,
35    /// Anthropic Claude Opus 4.5 (uses BPE tokenizer)
36    ClaudeOpus45,
37    /// Anthropic Claude Sonnet 4 (uses BPE tokenizer)
38    ClaudeSonnet4,
39    /// Generic "Other" model (uses average tokenization)
40    Other,
41}
42
43impl std::fmt::Display for ModelType {
44    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
45        match self {
46            ModelType::Gpt4o => write!(f, "GPT-4o"),
47            ModelType::O1 => write!(f, "o1"),
48            ModelType::Gpt4 => write!(f, "GPT-4"),
49            ModelType::Gemini3 => write!(f, "Gemini 3"),
50            ModelType::ClaudeOpus45 => write!(f, "Claude Opus 4.5"),
51            ModelType::ClaudeSonnet4 => write!(f, "Claude Sonnet 4"),
52            ModelType::Other => write!(f, "Other"),
53        }
54    }
55}
56
57/// Token information returned from counting
58#[derive(Debug, Clone)]
59pub struct TokenInfo {
60    /// Total number of tokens
61    pub count: usize,
62    /// Token IDs (if available)
63    pub ids: Vec<u32>,
64    /// Token strings (decoded tokens)
65    pub tokens: Vec<String>,
66    /// Model used for counting
67    pub model: ModelType,
68}
69
70impl TokenInfo {
71    /// Create a new TokenInfo
72    pub fn new(count: usize, ids: Vec<u32>, tokens: Vec<String>, model: ModelType) -> Self {
73        Self {
74            count,
75            ids,
76            tokens,
77            model,
78        }
79    }
80
81    /// Create TokenInfo with just count (for models without ID access)
82    pub fn count_only(count: usize, model: ModelType) -> Self {
83        Self {
84            count,
85            ids: Vec::new(),
86            tokens: Vec::new(),
87            model,
88        }
89    }
90}
91
92/// Token counter for multiple LLM models
93///
94/// Provides unified interface for counting tokens across different models.
95/// Uses model-specific tokenizers internally.
96pub struct TokenCounter {
97    // Note: Caching was considered but removed as token counting is fast enough
98    // that the overhead of cache management outweighs the benefits for typical use cases.
99    // If profiling shows token counting as a bottleneck, caching can be re-added.
100}
101
102impl TokenCounter {
103    /// Create a new token counter
104    pub fn new() -> Self {
105        Self {}
106    }
107
108    /// Count tokens for the given text and model
109    ///
110    /// # Arguments
111    /// * `text` - The text to tokenize
112    /// * `model` - The model type to use for tokenization
113    ///
114    /// # Returns
115    /// TokenInfo containing count, IDs, and decoded tokens
116    pub fn count(&self, text: &str, model: ModelType) -> TokenInfo {
117        match model {
118            ModelType::Gpt4o | ModelType::O1 => self.count_openai_o200k(text, model),
119            ModelType::Gpt4 => self.count_openai_cl100k(text, model),
120            ModelType::Gemini3 => self.count_gemini(text, model),
121            ModelType::ClaudeOpus45 | ModelType::ClaudeSonnet4 => self.count_claude(text, model),
122            ModelType::Other => self.count_other(text, model),
123        }
124    }
125
126    /// Count tokens using OpenAI o200k_base tokenizer (GPT-4o, o1)
127    fn count_openai_o200k(&self, text: &str, model: ModelType) -> TokenInfo {
128        // Use tiktoken-rs o200k_base tokenizer
129        #[cfg(feature = "tiktoken")]
130        {
131            use tiktoken_rs::o200k_base;
132            if let Ok(bpe) = o200k_base() {
133                let tokens = bpe.encode_with_special_tokens(text);
134                let decoded: Vec<String> = tokens
135                    .iter()
136                    .filter_map(|&id| bpe.decode(vec![id]).ok())
137                    .collect();
138                return TokenInfo::new(tokens.len(), tokens, decoded, model);
139            }
140        }
141
142        // Fallback: approximate token count (4 chars per token average)
143        self.approximate_token_count(text, model, 4.0)
144    }
145
146    /// Count tokens using OpenAI cl100k_base tokenizer (GPT-4)
147    fn count_openai_cl100k(&self, text: &str, model: ModelType) -> TokenInfo {
148        #[cfg(feature = "tiktoken")]
149        {
150            use tiktoken_rs::cl100k_base;
151            if let Ok(bpe) = cl100k_base() {
152                let tokens = bpe.encode_with_special_tokens(text);
153                let decoded: Vec<String> = tokens
154                    .iter()
155                    .filter_map(|&id| bpe.decode(vec![id]).ok())
156                    .collect();
157                return TokenInfo::new(tokens.len(), tokens, decoded, model);
158            }
159        }
160
161        // Fallback: approximate token count
162        self.approximate_token_count(text, model, 4.0)
163    }
164
165    /// Count tokens using Gemini/Gemma tokenizer
166    fn count_gemini(&self, text: &str, model: ModelType) -> TokenInfo {
167        #[cfg(feature = "tokenizers")]
168        {
169            use tokenizers::Tokenizer;
170            // Try to load Gemma tokenizer from known paths
171            let paths = [
172                "tokenizers/gemma-tokenizer.json",
173                "~/.cache/huggingface/tokenizers/google/gemma-3/tokenizer.json",
174            ];
175
176            for path in &paths {
177                if let Ok(tokenizer) = Tokenizer::from_file(path) {
178                    if let Ok(encoding) = tokenizer.encode(text, false) {
179                        let ids: Vec<u32> = encoding.get_ids().to_vec();
180                        let tokens: Vec<String> = encoding
181                            .get_tokens()
182                            .iter()
183                            .map(|s| s.to_string())
184                            .collect();
185                        return TokenInfo::new(ids.len(), ids, tokens, model);
186                    }
187                }
188            }
189        }
190
191        // Fallback: Gemini uses ~3.5 chars per token
192        self.approximate_token_count(text, model, 3.5)
193    }
194
195    /// Count tokens using Claude tokenizer
196    fn count_claude(&self, text: &str, model: ModelType) -> TokenInfo {
197        // Claude tokenizer is approximated since no official local tokenizer exists
198        // Use HuggingFace tokenizers crate with community Claude tokenizer if available
199        #[cfg(feature = "tokenizers-hf")]
200        {
201            use tokenizers::Tokenizer;
202            // Try to load Claude tokenizer from known paths
203            let paths = [
204                "tokenizers/claude-tokenizer.json",
205                "~/.cache/huggingface/tokenizers/anthropic/claude/tokenizer.json",
206            ];
207
208            for path in &paths {
209                if let Ok(tokenizer) = Tokenizer::from_file(path) {
210                    if let Ok(encoding) = tokenizer.encode(text, false) {
211                        let ids: Vec<u32> = encoding.get_ids().to_vec();
212                        let tokens: Vec<String> = encoding
213                            .get_tokens()
214                            .iter()
215                            .map(|s| s.to_string())
216                            .collect();
217                        return TokenInfo::new(ids.len(), ids, tokens, model);
218                    }
219                }
220            }
221        }
222
223        // Fallback: Claude uses ~3.8 chars per token
224        self.approximate_token_count(text, model, 3.8)
225    }
226
227    /// Count tokens using generic "Other" model (average tokenization)
228    fn count_other(&self, text: &str, model: ModelType) -> TokenInfo {
229        // Use an average of ~3.7 chars per token (between Claude and OpenAI)
230        self.approximate_token_count(text, model, 3.7)
231    }
232
233    /// Approximate token count when tokenizer is not available
234    fn approximate_token_count(
235        &self,
236        text: &str,
237        model: ModelType,
238        chars_per_token: f64,
239    ) -> TokenInfo {
240        let count = (text.len() as f64 / chars_per_token).ceil() as usize;
241        TokenInfo::count_only(count.max(1), model)
242    }
243
244    /// Count tokens for all supported models
245    pub fn count_all(&self, text: &str) -> HashMap<ModelType, TokenInfo> {
246        let models = [
247            ModelType::Gpt4o,
248            ModelType::O1,
249            ModelType::Gpt4,
250            ModelType::Gemini3,
251            ModelType::ClaudeOpus45,
252            ModelType::ClaudeSonnet4,
253            ModelType::Other,
254        ];
255
256        models
257            .iter()
258            .map(|&model| (model, self.count(text, model)))
259            .collect()
260    }
261
262    /// Count tokens for the 4 primary models (OpenAI, Claude, Gemini, Other)
263    /// as required by the token efficiency display feature.
264    ///
265    /// Returns counts for:
266    /// - OpenAI (GPT-4o)
267    /// - Claude (Sonnet 4)
268    /// - Gemini (Gemini 3)
269    /// - Other (generic model)
270    pub fn count_primary_models(&self, text: &str) -> HashMap<ModelType, TokenInfo> {
271        let models = [
272            ModelType::Gpt4o,         // OpenAI representative
273            ModelType::ClaudeSonnet4, // Claude representative
274            ModelType::Gemini3,       // Gemini representative
275            ModelType::Other,         // Generic model
276        ];
277
278        models
279            .iter()
280            .map(|&model| (model, self.count(text, model)))
281            .collect()
282    }
283
284    /// Get a summary of token counts for all models
285    pub fn summary(&self, text: &str) -> String {
286        let counts = self.count_all(text);
287        let mut lines = vec![format!("Token counts for {} chars:", text.len())];
288
289        for model in [
290            ModelType::Gpt4o,
291            ModelType::Gemini3,
292            ModelType::ClaudeOpus45,
293        ] {
294            if let Some(info) = counts.get(&model) {
295                lines.push(format!("  {}: {} tokens", model, info.count));
296            }
297        }
298
299        lines.join("\n")
300    }
301}
302
303impl Default for TokenCounter {
304    fn default() -> Self {
305        Self::new()
306    }
307}
308
309/// Measure token efficiency of dx format vs other formats
310pub struct TokenEfficiencyMeasurement {
311    /// Original text (e.g., JSON)
312    pub original: TokenInfo,
313    /// DX format text
314    pub dx_format: TokenInfo,
315    /// Savings percentage
316    pub savings_percent: f64,
317}
318
319impl TokenEfficiencyMeasurement {
320    /// Calculate token savings
321    pub fn calculate(original: TokenInfo, dx_format: TokenInfo) -> Self {
322        let savings = if original.count > 0 {
323            ((original.count as f64 - dx_format.count as f64) / original.count as f64) * 100.0
324        } else {
325            0.0
326        };
327
328        Self {
329            original,
330            dx_format,
331            savings_percent: savings,
332        }
333    }
334}
335
336/// Extension trait for adding token counting to dx format serializer
337pub trait TokenCountExt {
338    /// Get token count for the serialized output
339    fn token_count(&self, model: ModelType) -> TokenInfo;
340
341    /// Get token counts for all models
342    fn token_counts(&self) -> HashMap<ModelType, TokenInfo>;
343}
344
345impl TokenCountExt for String {
346    fn token_count(&self, model: ModelType) -> TokenInfo {
347        let counter = TokenCounter::new();
348        counter.count(self, model)
349    }
350
351    fn token_counts(&self) -> HashMap<ModelType, TokenInfo> {
352        let counter = TokenCounter::new();
353        counter.count_all(self)
354    }
355}
356
357impl TokenCountExt for str {
358    fn token_count(&self, model: ModelType) -> TokenInfo {
359        let counter = TokenCounter::new();
360        counter.count(self, model)
361    }
362
363    fn token_counts(&self) -> HashMap<ModelType, TokenInfo> {
364        let counter = TokenCounter::new();
365        counter.count_all(self)
366    }
367}
368
369#[cfg(test)]
370mod tests {
371    use super::*;
372
373    #[test]
374    fn test_token_counter_creation() {
375        let counter = TokenCounter::new();
376        let info = counter.count("Hello, world!", ModelType::Gpt4o);
377        assert!(info.count > 0);
378    }
379
380    #[test]
381    fn test_count_all_models() {
382        let counter = TokenCounter::new();
383        let counts = counter.count_all("Hello, world!");
384
385        assert!(counts.contains_key(&ModelType::Gpt4o));
386        assert!(counts.contains_key(&ModelType::Gemini3));
387        assert!(counts.contains_key(&ModelType::ClaudeOpus45));
388        assert!(counts.contains_key(&ModelType::Other));
389    }
390
391    #[test]
392    fn test_count_primary_models() {
393        let counter = TokenCounter::new();
394        let counts = counter.count_primary_models("Hello, world!");
395
396        // Should have exactly 4 models
397        assert_eq!(counts.len(), 4);
398        assert!(counts.contains_key(&ModelType::Gpt4o));
399        assert!(counts.contains_key(&ModelType::ClaudeSonnet4));
400        assert!(counts.contains_key(&ModelType::Gemini3));
401        assert!(counts.contains_key(&ModelType::Other));
402
403        // All counts should be non-zero
404        for info in counts.values() {
405            assert!(info.count > 0, "Token count should be non-zero");
406        }
407    }
408
409    #[test]
410    fn test_other_model() {
411        let counter = TokenCounter::new();
412        let info = counter.count("Hello, world!", ModelType::Other);
413        assert!(info.count > 0);
414        assert_eq!(info.model, ModelType::Other);
415    }
416
417    #[test]
418    fn test_token_efficiency_measurement() {
419        let original = TokenInfo::count_only(100, ModelType::Gpt4o);
420        let dx = TokenInfo::count_only(73, ModelType::Gpt4o);
421        let measurement = TokenEfficiencyMeasurement::calculate(original, dx);
422
423        assert!((measurement.savings_percent - 27.0).abs() < 0.1);
424    }
425
426    #[test]
427    fn test_empty_string() {
428        let counter = TokenCounter::new();
429        let info = counter.count("", ModelType::Gpt4o);
430        assert_eq!(info.count, 1); // Minimum 1 token
431    }
432
433    #[test]
434    fn test_model_display() {
435        assert_eq!(format!("{}", ModelType::Gpt4o), "GPT-4o");
436        assert_eq!(format!("{}", ModelType::Gemini3), "Gemini 3");
437        assert_eq!(format!("{}", ModelType::ClaudeOpus45), "Claude Opus 4.5");
438    }
439
440    #[test]
441    fn test_summary() {
442        let counter = TokenCounter::new();
443        let summary = counter.summary("Hello, world!");
444        assert!(summary.contains("Token counts"));
445        assert!(summary.contains("GPT-4o"));
446    }
447
448    #[test]
449    fn test_token_count_extension() {
450        let text = "Hello, world!";
451        let info = text.token_count(ModelType::Gpt4o);
452        assert!(info.count > 0);
453    }
454
455    #[test]
456    fn test_dx_format_token_efficiency() {
457        // Compare JSON vs DX format for same data
458        // Use a larger example to show more significant savings
459        let json = r#"{"name":"dx-serializer","version":"0.1.0","description":"Binary-first serialization format for LLMs","workspace":["frontend/www","frontend/mobile","backend/api","backend/workers"],"dependencies":{"serde":"1.0","bincode":"2.0","tokio":"1.0"},"enabled":true,"count":42}"#;
460        let dx = "nm=dx-serializer ver=0.1.0 ds=\"Binary-first serialization format for LLMs\" ws:frontend/www,frontend/mobile,backend/api,backend/workers deps.serde=1.0 deps.bincode=2.0 deps.tokio=1.0 en=true ct=42";
461
462        let counter = TokenCounter::new();
463        let json_tokens = counter.count(json, ModelType::Gpt4o);
464        let dx_tokens = counter.count(dx, ModelType::Gpt4o);
465
466        // Store counts before moving
467        let json_count = json_tokens.count;
468        let dx_count = dx_tokens.count;
469
470        // DX format should use fewer tokens
471        let measurement = TokenEfficiencyMeasurement::calculate(json_tokens, dx_tokens);
472        println!(
473            "JSON: {} tokens, DX: {} tokens, Savings: {:.1}%",
474            measurement.original.count, measurement.dx_format.count, measurement.savings_percent
475        );
476
477        // DX should be more efficient (or at least not worse)
478        // Small examples may not show significant savings due to tokenizer overhead
479        assert!(
480            measurement.savings_percent >= 0.0 || dx_count <= json_count,
481            "DX format should not be worse than JSON"
482        );
483    }
484}
serializer/llm/tokens.rs

serializer/llm/
tokens.rs
