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    /// Get count for a specific model
45    pub fn get(&self, model: TokenModel) -> u32 {
46        match model {
47            // OpenAI o200k_base models (all share same encoding)
48            TokenModel::Gpt52
49            | TokenModel::Gpt52Pro
50            | TokenModel::Gpt51
51            | TokenModel::Gpt51Mini
52            | TokenModel::Gpt51Codex
53            | TokenModel::Gpt5
54            | TokenModel::Gpt5Mini
55            | TokenModel::Gpt5Nano
56            | TokenModel::O4Mini
57            | TokenModel::O3
58            | TokenModel::O3Mini
59            | TokenModel::O1
60            | TokenModel::O1Mini
61            | TokenModel::O1Preview
62            | TokenModel::Gpt4o
63            | TokenModel::Gpt4oMini => self.o200k,
64            // OpenAI cl100k_base models (legacy, same encoding)
65            TokenModel::Gpt4 | TokenModel::Gpt35Turbo => self.cl100k,
66            // Other vendors
67            TokenModel::Claude => self.claude,
68            TokenModel::Gemini => self.gemini,
69            TokenModel::Llama | TokenModel::CodeLlama => self.llama,
70            TokenModel::Mistral => self.mistral,
71            TokenModel::DeepSeek => self.deepseek,
72            TokenModel::Qwen => self.qwen,
73            TokenModel::Cohere => self.cohere,
74            TokenModel::Grok => self.grok,
75        }
76    }
77
78    /// Set count for a specific model
79    pub fn set(&mut self, model: TokenModel, count: u32) {
80        match model {
81            // OpenAI o200k_base models
82            TokenModel::Gpt52
83            | TokenModel::Gpt52Pro
84            | TokenModel::Gpt51
85            | TokenModel::Gpt51Mini
86            | TokenModel::Gpt51Codex
87            | TokenModel::Gpt5
88            | TokenModel::Gpt5Mini
89            | TokenModel::Gpt5Nano
90            | TokenModel::O4Mini
91            | TokenModel::O3
92            | TokenModel::O3Mini
93            | TokenModel::O1
94            | TokenModel::O1Mini
95            | TokenModel::O1Preview
96            | TokenModel::Gpt4o
97            | TokenModel::Gpt4oMini => self.o200k = count,
98            // OpenAI cl100k_base models (legacy)
99            TokenModel::Gpt4 | TokenModel::Gpt35Turbo => self.cl100k = count,
100            // Other vendors
101            TokenModel::Claude => self.claude = count,
102            TokenModel::Gemini => self.gemini = count,
103            TokenModel::Llama | TokenModel::CodeLlama => self.llama = count,
104            TokenModel::Mistral => self.mistral = count,
105            TokenModel::DeepSeek => self.deepseek = count,
106            TokenModel::Qwen => self.qwen = count,
107            TokenModel::Cohere => self.cohere = count,
108            TokenModel::Grok => self.grok = count,
109        }
110    }
111
112    /// Sum all counts (useful for aggregate statistics)
113    pub fn total(&self) -> u64 {
114        self.o200k as u64
115            + self.cl100k as u64
116            + self.claude as u64
117            + self.gemini as u64
118            + self.llama as u64
119            + self.mistral as u64
120            + self.deepseek as u64
121            + self.qwen as u64
122            + self.cohere as u64
123            + self.grok as u64
124    }
125
126    /// Add counts from another TokenCounts
127    pub fn add(&mut self, other: &TokenCounts) {
128        self.o200k += other.o200k;
129        self.cl100k += other.cl100k;
130        self.claude += other.claude;
131        self.gemini += other.gemini;
132        self.llama += other.llama;
133        self.mistral += other.mistral;
134        self.deepseek += other.deepseek;
135        self.qwen += other.qwen;
136        self.cohere += other.cohere;
137        self.grok += other.grok;
138    }
139
140    /// Get the minimum token count across all models
141    pub fn min(&self) -> u32 {
142        [
143            self.o200k,
144            self.cl100k,
145            self.claude,
146            self.gemini,
147            self.llama,
148            self.mistral,
149            self.deepseek,
150            self.qwen,
151            self.cohere,
152            self.grok,
153        ]
154        .into_iter()
155        .min()
156        .unwrap_or(0)
157    }
158
159    /// Get the maximum token count across all models
160    pub fn max(&self) -> u32 {
161        [
162            self.o200k,
163            self.cl100k,
164            self.claude,
165            self.gemini,
166            self.llama,
167            self.mistral,
168            self.deepseek,
169            self.qwen,
170            self.cohere,
171            self.grok,
172        ]
173        .into_iter()
174        .max()
175        .unwrap_or(0)
176    }
177}
178
179impl std::ops::Add for TokenCounts {
180    type Output = Self;
181
182    fn add(self, rhs: Self) -> Self::Output {
183        Self {
184            o200k: self.o200k + rhs.o200k,
185            cl100k: self.cl100k + rhs.cl100k,
186            claude: self.claude + rhs.claude,
187            gemini: self.gemini + rhs.gemini,
188            llama: self.llama + rhs.llama,
189            mistral: self.mistral + rhs.mistral,
190            deepseek: self.deepseek + rhs.deepseek,
191            qwen: self.qwen + rhs.qwen,
192            cohere: self.cohere + rhs.cohere,
193            grok: self.grok + rhs.grok,
194        }
195    }
196}
197
198impl std::iter::Sum for TokenCounts {
199    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
200        iter.fold(Self::zero(), |acc, x| acc + x)
201    }
202}
infiniloom_engine/tokenizer/counts.rs

infiniloom_engine/tokenizer/
counts.rs
