llm_token/
count.rs

1//! Pure token estimation for LLM requests
2//!
3//! Estimation layer built on top of [`TokenEstimator`]: counts messages,
4//! tools, and multimodal content. Transport-level counting (provider
5//! count-tokens API first, estimation as fallback) lives in
6//! `client::LlmClientImpl::count_tokens_inner`; this module never touches
7//! HTTP clients or codecs.
8
9use super::estimation::{estimate_tokens, TokenEstimator, MESSAGE_OVERHEAD_TOKENS};
10use llm_types::llm::LlmRequest;
11use llm_types::message::{Message, MessageContent, MessageContentValue};
12
13/// Estimate tokens for a single message: content (text / rich blocks),
14/// thinking, tool calls, plus a fixed metadata overhead per message.
15pub fn estimate_message_tokens(msg: &Message) -> u32 {
16    let mut total = 0u32;
17
18    match &msg.content {
19        MessageContentValue::Text(text) => total += estimate_tokens(text) as u32,
20        MessageContentValue::Rich(blocks) => {
21            for block in blocks {
22                total += match block {
23                    MessageContent::Text { text } => estimate_tokens(text) as u32,
24                    MessageContent::Thinking { thinking, .. } => estimate_tokens(thinking) as u32,
25                    MessageContent::ToolUse { tool_use } => {
26                        let json = serde_json::to_string(&tool_use.input).unwrap_or_default();
27                        estimate_tokens(&json) as u32
28                    }
29                    MessageContent::ToolResult { tool_result } => {
30                        estimate_tokens(&tool_result.content) as u32
31                    }
32                    MessageContent::ImageUrl { image_url } => estimate_image_tokens(&image_url.url),
33                };
34            }
35        }
36    }
37
38    if let Some(thinking) = &msg.thinking {
39        total += estimate_tokens(thinking) as u32;
40    }
41
42    if let Some(tool_calls) = &msg.tool_calls {
43        if let Ok(json) = serde_json::to_string(tool_calls) {
44            total += estimate_tokens(&json) as u32;
45        }
46    }
47
48    total + MESSAGE_OVERHEAD_TOKENS
49}
50
51/// Estimate tokens for tool declarations alone (the per-request dynamic
52/// overhead that never accumulates into a message-array ledger).
53pub fn estimate_tool_declarations(tools: Option<&[llm_types::tool::Tool]>) -> u32 {
54    let mut total = 0u32;
55    for tool in tools.unwrap_or_default() {
56        total += estimate_tokens(&tool.name) as u32;
57        total += estimate_tokens(&tool.description) as u32;
58        if let Some(params) = &tool.parameters {
59            if let Ok(json) = serde_json::to_string(params) {
60                total += estimate_tokens(&json) as u32;
61            }
62        }
63    }
64    total
65}
66
67/// Estimate tokens for a full LLM request: messages + tool declarations.
68pub fn estimate_request_tokens(request: &LlmRequest) -> u32 {
69    let mut total = 0u32;
70
71    for msg in &request.messages {
72        total += estimate_message_tokens(msg);
73    }
74
75    total += estimate_tool_declarations(request.tools.as_deref());
76
77    total
78}
79
80/// Estimate tokens for an image reference.
81///
82/// Heuristic:
83/// - data URI: derive pixel count from base64 length, `ceil(pixels/750) + 200`
84/// - remote URL: conservative fixed estimate of 170 tokens
85pub fn estimate_image_tokens(url_or_data_uri: &str) -> u32 {
86    if let Some(rest) = url_or_data_uri.strip_prefix("data:") {
87        let b64 = if let Some(semicolon) = rest.find(';') {
88            rest[semicolon + 1..]
89                .strip_prefix("base64,")
90                .unwrap_or(&rest[semicolon + 1..])
91        } else {
92            rest
93        };
94        // Raw bytes ≈ 3/4 of base64 length; ~3 bytes per pixel (RGB)
95        let bytes = b64.len() as f64 * 3.0 / 4.0;
96        let pixels = bytes / 3.0;
97        (pixels / 750.0).ceil() as u32 + 200
98    } else {
99        170
100    }
101}
102
103/// Estimate the token count of the current conversation messages.
104pub fn estimate_messages(messages: &[Message]) -> u32 {
105    messages.iter().map(estimate_message_tokens).sum()
106}
107
108/// Estimator used by counting helpers.
109pub fn default_estimator() -> TokenEstimator {
110    TokenEstimator::default()
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116    use llm_types::message::{
117        ImageUrlContent, Message, MessageContent, MessageContentValue, MessageRole,
118        ToolResultContent,
119    };
120
121    fn text_message(role: MessageRole, text: &str) -> Message {
122        Message {
123            id: "id".to_string(),
124            role,
125            content: MessageContentValue::Text(text.to_string()),
126            timestamp: 0,
127            tool_call_id: None,
128            tool_name: None,
129            tool_calls: None,
130            thinking: None,
131            metadata: None,
132        }
133    }
134
135    #[test]
136    fn test_estimate_simple_message() {
137        let msg = text_message(MessageRole::User, "Hello");
138        // 5 * 0.25 = 1.25 -> 1 + 4 overhead = 5
139        assert_eq!(estimate_message_tokens(&msg), 5);
140    }
141
142    #[test]
143    fn test_estimate_cjk_message() {
144        let msg = text_message(MessageRole::User, "你好世界");
145        // 4 CJK = 4 + 4 overhead = 8
146        assert_eq!(estimate_message_tokens(&msg), 8);
147    }
148
149    #[test]
150    fn test_estimate_message_with_thinking() {
151        let mut msg = text_message(MessageRole::Assistant, "result");
152        msg.thinking = Some("思考过程".to_string());
153        // content "result" 6 * 0.25 = 1.5 -> 2 + thinking 4 CJK = 4 + 4 overhead = 10
154        assert_eq!(estimate_message_tokens(&msg), 10);
155    }
156
157    #[test]
158    fn test_estimate_message_with_tool_calls() {
159        let mut msg = text_message(MessageRole::Assistant, "");
160        msg.tool_calls = Some(vec![llm_types::message::LlmToolCall {
161            id: "call_1".to_string(),
162            r#type: "function".to_string(),
163            function: llm_types::message::LlmFunctionCall {
164                name: "get_weather".to_string(),
165                arguments: r#"{"city":"Beijing"}"#.to_string(),
166            },
167        }]);
168        let total = estimate_message_tokens(&msg);
169        // JSON estimate + 4 overhead; must be at least overhead
170        assert!(total >= 4, "tokens: {}", total);
171    }
172
173    #[test]
174    fn test_estimate_rich_content() {
175        let msg = Message {
176            id: "id".to_string(),
177            role: MessageRole::User,
178            content: MessageContentValue::Rich(vec![
179                MessageContent::Text {
180                    text: "look at this".to_string(),
181                },
182                MessageContent::ImageUrl {
183                    image_url: ImageUrlContent {
184                        url: "https://example.com/img.png".to_string(),
185                        detail: None,
186                    },
187                },
188            ]),
189            timestamp: 0,
190            tool_call_id: None,
191            tool_name: None,
192            tool_calls: None,
193            thinking: None,
194            metadata: None,
195        };
196        let total = estimate_message_tokens(&msg);
197        // text ~3 + image 170 + 4 overhead
198        assert!(total >= 170, "tokens: {}", total);
199    }
200
201    #[test]
202    fn test_estimate_rich_tool_result() {
203        let msg = Message {
204            id: "id".to_string(),
205            role: MessageRole::Tool,
206            content: MessageContentValue::Rich(vec![MessageContent::ToolResult {
207                tool_result: ToolResultContent {
208                    tool_use_id: "call_1".to_string(),
209                    content: "the weather is sunny".to_string(),
210                    is_error: None,
211                },
212            }]),
213            timestamp: 0,
214            tool_call_id: None,
215            tool_name: None,
216            tool_calls: None,
217            thinking: None,
218            metadata: None,
219        };
220        let total = estimate_message_tokens(&msg);
221        assert!(total >= 4, "tokens: {}", total);
222    }
223
224    #[test]
225    fn test_estimate_image_tokens_url() {
226        assert_eq!(estimate_image_tokens("https://example.com/a.png"), 170);
227    }
228
229    #[test]
230    fn test_estimate_image_tokens_data_uri() {
231        // 900 pixels -> ceil(900/750) + 200 = 202
232        // b64 length for 900*3 bytes = 2700 bytes -> ceil(2700/3)=900 px
233        // base64 of 2700 bytes is 3600 chars
234        let uri = format!("data:image/png;base64,{}", "A".repeat(3600));
235        assert_eq!(estimate_image_tokens(&uri), 202);
236    }
237
238    #[test]
239    fn test_estimate_request_tokens() {
240        let req = LlmRequest {
241            profile_id: "default".to_string(),
242            messages: vec![
243                text_message(MessageRole::System, "You are a helper"),
244                text_message(MessageRole::User, "你好"),
245            ],
246            parameters: None,
247            generation: None,
248            tools: None,
249            tool_call_protocol: None,
250            locked_tool_call_protocol: None,
251            violation_policy: None,
252            execution_id: None,
253            stream: None,
254            dead_loop_detection: None,
255            protocol_auto_converted: None,
256            timeout_ms: None,
257        };
258        // system: 13 letters * 0.25 + 3 spaces * 0.5 = 4.75 -> 5 + 4 = 9;
259        // user: 2 + 4 = 6; total 15
260        assert_eq!(estimate_request_tokens(&req), 15);
261    }
262
263    #[test]
264    fn test_estimate_messages_empty() {
265        assert_eq!(estimate_messages(&[]), 0);
266    }
267}
llm_token/count.rs

llm_token/
count.rs
