infiniloom_engine/tokenizer/
counts.rs

1//! Token count aggregation for multiple models
2//!
3//! This module provides the TokenCounts struct for storing token counts
4//! across multiple LLM model families.
5
6use super::models::TokenModel;
7
8/// Token counts for multiple models
9///
10/// Counts are grouped by tokenizer encoding family:
11/// - `o200k`: OpenAI modern models (GPT-5.x, GPT-4o, O1/O3/O4) - EXACT
12/// - `cl100k`: OpenAI legacy models (GPT-4, GPT-3.5-turbo) - EXACT
13/// - Other fields: Estimation-based counts for non-OpenAI vendors
14#[derive(Debug, Clone, Copy, Default, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
15pub struct TokenCounts {
16    /// OpenAI o200k_base encoding (GPT-5.2, GPT-5.1, GPT-5, GPT-4o, O1, O3, O4)
17    pub o200k: u32,
18    /// OpenAI cl100k_base encoding (GPT-4, GPT-3.5-turbo) - legacy
19    pub cl100k: u32,
20    /// Anthropic Claude (all versions)
21    pub claude: u32,
22    /// Google Gemini (all versions)
23    pub gemini: u32,
24    /// Meta Llama (3, 4, CodeLlama)
25    pub llama: u32,
26    /// Mistral AI (Large, Medium, Small, Codestral)
27    pub mistral: u32,
28    /// DeepSeek (V3, R1, Coder)
29    pub deepseek: u32,
30    /// Alibaba Qwen (Qwen3, Qwen2.5)
31    pub qwen: u32,
32    /// Cohere (Command R+, Command R)
33    pub cohere: u32,
34    /// xAI Grok (Grok 2, Grok 3)
35    pub grok: u32,
36}
37
38impl TokenCounts {
39    /// Create zero counts
40    pub fn zero() -> Self {
41        Self::default()
42    }
43
44    /// Create counts with all fields set to the same value
45    /// Useful for testing
46    #[cfg(test)]
47    pub fn default_with_value(value: u32) -> Self {
48        Self {
49            o200k: value,
50            cl100k: value,
51            claude: value,
52            gemini: value,
53            llama: value,
54            mistral: value,
55            deepseek: value,
56            qwen: value,
57            cohere: value,
58            grok: value,
59        }
60    }
61
62    /// Get count for a specific model
63    pub fn get(&self, model: TokenModel) -> u32 {
64        match model {
65            // OpenAI o200k_base models (all share same encoding)
66            TokenModel::Gpt52
67            | TokenModel::Gpt52Pro
68            | TokenModel::Gpt51
69            | TokenModel::Gpt51Mini
70            | TokenModel::Gpt51Codex
71            | TokenModel::Gpt5
72            | TokenModel::Gpt5Mini
73            | TokenModel::Gpt5Nano
74            | TokenModel::O4Mini
75            | TokenModel::O3
76            | TokenModel::O3Mini
77            | TokenModel::O1
78            | TokenModel::O1Mini
79            | TokenModel::O1Preview
80            | TokenModel::Gpt4o
81            | TokenModel::Gpt4oMini => self.o200k,
82            // OpenAI cl100k_base models (legacy, same encoding)
83            TokenModel::Gpt4 | TokenModel::Gpt35Turbo => self.cl100k,
84            // Other vendors
85            TokenModel::Claude => self.claude,
86            TokenModel::Gemini => self.gemini,
87            TokenModel::Llama | TokenModel::CodeLlama => self.llama,
88            TokenModel::Mistral => self.mistral,
89            TokenModel::DeepSeek => self.deepseek,
90            TokenModel::Qwen => self.qwen,
91            TokenModel::Cohere => self.cohere,
92            TokenModel::Grok => self.grok,
93        }
94    }
95
96    /// Set count for a specific model
97    pub fn set(&mut self, model: TokenModel, count: u32) {
98        match model {
99            // OpenAI o200k_base models
100            TokenModel::Gpt52
101            | TokenModel::Gpt52Pro
102            | TokenModel::Gpt51
103            | TokenModel::Gpt51Mini
104            | TokenModel::Gpt51Codex
105            | TokenModel::Gpt5
106            | TokenModel::Gpt5Mini
107            | TokenModel::Gpt5Nano
108            | TokenModel::O4Mini
109            | TokenModel::O3
110            | TokenModel::O3Mini
111            | TokenModel::O1
112            | TokenModel::O1Mini
113            | TokenModel::O1Preview
114            | TokenModel::Gpt4o
115            | TokenModel::Gpt4oMini => self.o200k = count,
116            // OpenAI cl100k_base models (legacy)
117            TokenModel::Gpt4 | TokenModel::Gpt35Turbo => self.cl100k = count,
118            // Other vendors
119            TokenModel::Claude => self.claude = count,
120            TokenModel::Gemini => self.gemini = count,
121            TokenModel::Llama | TokenModel::CodeLlama => self.llama = count,
122            TokenModel::Mistral => self.mistral = count,
123            TokenModel::DeepSeek => self.deepseek = count,
124            TokenModel::Qwen => self.qwen = count,
125            TokenModel::Cohere => self.cohere = count,
126            TokenModel::Grok => self.grok = count,
127        }
128    }
129
130    /// Sum all counts (useful for aggregate statistics)
131    pub fn total(&self) -> u64 {
132        self.o200k as u64
133            + self.cl100k as u64
134            + self.claude as u64
135            + self.gemini as u64
136            + self.llama as u64
137            + self.mistral as u64
138            + self.deepseek as u64
139            + self.qwen as u64
140            + self.cohere as u64
141            + self.grok as u64
142    }
143
144    /// Add counts from another TokenCounts
145    pub fn add(&mut self, other: &TokenCounts) {
146        self.o200k += other.o200k;
147        self.cl100k += other.cl100k;
148        self.claude += other.claude;
149        self.gemini += other.gemini;
150        self.llama += other.llama;
151        self.mistral += other.mistral;
152        self.deepseek += other.deepseek;
153        self.qwen += other.qwen;
154        self.cohere += other.cohere;
155        self.grok += other.grok;
156    }
157
158    /// Get the minimum token count across all models
159    pub fn min(&self) -> u32 {
160        [
161            self.o200k,
162            self.cl100k,
163            self.claude,
164            self.gemini,
165            self.llama,
166            self.mistral,
167            self.deepseek,
168            self.qwen,
169            self.cohere,
170            self.grok,
171        ]
172        .into_iter()
173        .min()
174        .unwrap_or(0)
175    }
176
177    /// Get the maximum token count across all models
178    pub fn max(&self) -> u32 {
179        [
180            self.o200k,
181            self.cl100k,
182            self.claude,
183            self.gemini,
184            self.llama,
185            self.mistral,
186            self.deepseek,
187            self.qwen,
188            self.cohere,
189            self.grok,
190        ]
191        .into_iter()
192        .max()
193        .unwrap_or(0)
194    }
195}
196
197impl std::ops::Add for TokenCounts {
198    type Output = Self;
199
200    fn add(self, rhs: Self) -> Self::Output {
201        Self {
202            o200k: self.o200k + rhs.o200k,
203            cl100k: self.cl100k + rhs.cl100k,
204            claude: self.claude + rhs.claude,
205            gemini: self.gemini + rhs.gemini,
206            llama: self.llama + rhs.llama,
207            mistral: self.mistral + rhs.mistral,
208            deepseek: self.deepseek + rhs.deepseek,
209            qwen: self.qwen + rhs.qwen,
210            cohere: self.cohere + rhs.cohere,
211            grok: self.grok + rhs.grok,
212        }
213    }
214}
215
216impl std::iter::Sum for TokenCounts {
217    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
218        iter.fold(Self::zero(), |acc, x| acc + x)
219    }
220}
infiniloom_engine/tokenizer/counts.rs

infiniloom_engine/tokenizer/
counts.rs
