1use serde::Serialize;
2use serde_json::Value;
3use std::sync::atomic::{AtomicU64, Ordering};
4
5#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
6#[serde(rename_all = "snake_case")]
7pub enum Provider {
8 Anthropic,
9 OpenAi,
10 Gemini,
11}
12
13#[derive(Debug, Clone, Serialize)]
14pub struct RequestBreakdown {
15 pub provider: Provider,
16 pub model: String,
17 pub system_prompt_tokens: usize,
18 pub user_message_tokens: usize,
19 pub assistant_message_tokens: usize,
20 pub tool_definition_tokens: usize,
21 pub tool_definition_count: usize,
22 pub tool_result_tokens: usize,
23 pub image_count: usize,
24 pub total_input_tokens: usize,
25 pub message_count: usize,
26}
27
28pub fn analyze_request(body: &Value, provider: Provider) -> RequestBreakdown {
29 match provider {
30 Provider::Anthropic => analyze_anthropic(body),
31 Provider::OpenAi => analyze_openai(body),
32 Provider::Gemini => analyze_gemini(body),
33 }
34}
35
36fn analyze_anthropic(body: &Value) -> RequestBreakdown {
37 let model = body
38 .get("model")
39 .and_then(|m| m.as_str())
40 .unwrap_or("unknown")
41 .to_string();
42
43 let system_prompt_tokens = match body.get("system") {
44 Some(Value::String(s)) => chars_to_tokens(s.len()),
45 Some(Value::Array(arr)) => {
46 arr.iter()
47 .map(|block| {
48 block
49 .get("text")
50 .and_then(|t| t.as_str())
51 .map_or(0, str::len)
52 })
53 .sum::<usize>()
54 / 4
55 }
56 _ => 0,
57 };
58
59 let tool_definition_tokens = body
60 .get("tools")
61 .and_then(|t| t.as_array())
62 .map_or(0, |arr| json_chars(arr) / 4);
63
64 let tool_definition_count = body
65 .get("tools")
66 .and_then(|t| t.as_array())
67 .map_or(0, Vec::len);
68
69 let mut user_message_tokens = 0;
70 let mut assistant_message_tokens = 0;
71 let mut tool_result_tokens = 0;
72 let mut image_count = 0;
73 let mut message_count = 0;
74
75 if let Some(messages) = body.get("messages").and_then(|m| m.as_array()) {
76 message_count = messages.len();
77 for msg in messages {
78 let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
79 let content_tokens = estimate_content_tokens(msg.get("content"));
80 let has_images = count_images(msg.get("content"));
81 image_count += has_images;
82
83 match role {
84 "user" => {
85 if has_tool_results(msg.get("content")) {
86 tool_result_tokens += content_tokens;
87 } else {
88 user_message_tokens += content_tokens;
89 }
90 }
91 "assistant" => assistant_message_tokens += content_tokens,
92 _ => user_message_tokens += content_tokens,
93 }
94 }
95 }
96
97 let total_input_tokens = system_prompt_tokens
98 + user_message_tokens
99 + assistant_message_tokens
100 + tool_definition_tokens
101 + tool_result_tokens;
102
103 RequestBreakdown {
104 provider: Provider::Anthropic,
105 model,
106 system_prompt_tokens,
107 user_message_tokens,
108 assistant_message_tokens,
109 tool_definition_tokens,
110 tool_definition_count,
111 tool_result_tokens,
112 image_count,
113 total_input_tokens,
114 message_count,
115 }
116}
117
118fn analyze_openai(body: &Value) -> RequestBreakdown {
119 let model = body
120 .get("model")
121 .and_then(|m| m.as_str())
122 .unwrap_or("unknown")
123 .to_string();
124
125 let mut system_prompt_tokens = 0;
126 let mut user_message_tokens = 0;
127 let mut assistant_message_tokens = 0;
128 let mut tool_result_tokens = 0;
129 let mut image_count = 0;
130 let mut message_count = 0;
131
132 if let Some(messages) = body.get("messages").and_then(|m| m.as_array()) {
133 message_count = messages.len();
134 for msg in messages {
135 let role = msg.get("role").and_then(|r| r.as_str()).unwrap_or("");
136 let content_tokens = estimate_content_tokens(msg.get("content"));
137 image_count += count_images(msg.get("content"));
138
139 match role {
140 "system" | "developer" => system_prompt_tokens += content_tokens,
141 "assistant" => assistant_message_tokens += content_tokens,
142 "tool" => tool_result_tokens += content_tokens,
143 _ => user_message_tokens += content_tokens,
144 }
145 }
146 }
147
148 let tool_definition_tokens = body
149 .get("tools")
150 .and_then(|t| t.as_array())
151 .map_or(0, |arr| json_chars(arr) / 4);
152
153 let tool_definition_count = body
154 .get("tools")
155 .and_then(|t| t.as_array())
156 .map_or(0, Vec::len);
157
158 let total_input_tokens = system_prompt_tokens
159 + user_message_tokens
160 + assistant_message_tokens
161 + tool_definition_tokens
162 + tool_result_tokens;
163
164 RequestBreakdown {
165 provider: Provider::OpenAi,
166 model,
167 system_prompt_tokens,
168 user_message_tokens,
169 assistant_message_tokens,
170 tool_definition_tokens,
171 tool_definition_count,
172 tool_result_tokens,
173 image_count,
174 total_input_tokens,
175 message_count,
176 }
177}
178
179fn analyze_gemini(body: &Value) -> RequestBreakdown {
180 let model = "gemini".to_string();
181
182 let system_prompt_tokens = body
183 .get("systemInstruction")
184 .and_then(|si| si.get("parts"))
185 .and_then(|p| p.as_array())
186 .map_or(0, |parts| {
187 parts
188 .iter()
189 .map(|p| p.get("text").and_then(|t| t.as_str()).map_or(0, str::len))
190 .sum::<usize>()
191 / 4
192 });
193
194 let mut user_message_tokens = 0;
195 let mut assistant_message_tokens = 0;
196 let mut tool_result_tokens = 0;
197 let mut message_count = 0;
198
199 if let Some(contents) = body.get("contents").and_then(|c| c.as_array()) {
200 message_count = contents.len();
201 for content in contents {
202 let role = content
203 .get("role")
204 .and_then(|r| r.as_str())
205 .unwrap_or("user");
206 let parts_tokens = content
207 .get("parts")
208 .and_then(|p| p.as_array())
209 .map_or(0, |parts| {
210 parts
211 .iter()
212 .map(|p| {
213 if p.get("functionResponse").is_some() {
214 json_chars(std::slice::from_ref(p)) / 4
215 } else {
216 p.get("text")
217 .and_then(|t| t.as_str())
218 .map_or(0, |s| chars_to_tokens(s.len()))
219 }
220 })
221 .sum::<usize>()
222 });
223
224 let has_fn_response = content
225 .get("parts")
226 .and_then(|p| p.as_array())
227 .is_some_and(|parts| parts.iter().any(|p| p.get("functionResponse").is_some()));
228
229 if has_fn_response {
230 tool_result_tokens += parts_tokens;
231 } else {
232 match role {
233 "model" => assistant_message_tokens += parts_tokens,
234 _ => user_message_tokens += parts_tokens,
235 }
236 }
237 }
238 }
239
240 let tool_definition_tokens = body
241 .get("tools")
242 .and_then(|t| t.as_array())
243 .map_or(0, |arr| json_chars(arr) / 4);
244
245 let tool_definition_count = body
246 .get("tools")
247 .and_then(|t| t.as_array())
248 .map_or(0, |arr| {
249 arr.iter()
250 .filter_map(|t| t.get("functionDeclarations").and_then(|f| f.as_array()))
251 .map(Vec::len)
252 .sum()
253 });
254
255 let total_input_tokens = system_prompt_tokens
256 + user_message_tokens
257 + assistant_message_tokens
258 + tool_definition_tokens
259 + tool_result_tokens;
260
261 RequestBreakdown {
262 provider: Provider::Gemini,
263 model,
264 system_prompt_tokens,
265 user_message_tokens,
266 assistant_message_tokens,
267 tool_definition_tokens,
268 tool_definition_count,
269 tool_result_tokens,
270 image_count: 0,
271 total_input_tokens,
272 message_count,
273 }
274}
275
276fn chars_to_tokens(chars: usize) -> usize {
277 chars / 4
278}
279
280fn json_chars(arr: &[Value]) -> usize {
281 arr.iter().map(|v| v.to_string().len()).sum()
282}
283
284fn estimate_content_tokens(content: Option<&Value>) -> usize {
285 match content {
286 Some(Value::String(s)) => chars_to_tokens(s.len()),
287 Some(Value::Array(arr)) => arr
288 .iter()
289 .map(|block| {
290 if let Some(text) = block.get("text").and_then(|t| t.as_str()) {
291 chars_to_tokens(text.len())
292 } else {
293 block.to_string().len() / 4
294 }
295 })
296 .sum(),
297 Some(v) => v.to_string().len() / 4,
298 None => 0,
299 }
300}
301
302fn count_images(content: Option<&Value>) -> usize {
303 match content {
304 Some(Value::Array(arr)) => arr
305 .iter()
306 .filter(|block| {
307 block.get("type").and_then(|t| t.as_str()) == Some("image")
308 || block.get("type").and_then(|t| t.as_str()) == Some("image_url")
309 })
310 .count(),
311 _ => 0,
312 }
313}
314
315fn has_tool_results(content: Option<&Value>) -> bool {
316 match content {
317 Some(Value::Array(arr)) => arr
318 .iter()
319 .any(|block| block.get("type").and_then(|t| t.as_str()) == Some("tool_result")),
320 _ => false,
321 }
322}
323
324pub struct IntrospectState {
326 pub last_breakdown: std::sync::Mutex<Option<RequestBreakdown>>,
327 pub total_system_prompt_tokens: AtomicU64,
328 pub total_requests: AtomicU64,
329}
330
331impl Default for IntrospectState {
332 fn default() -> Self {
333 Self {
334 last_breakdown: std::sync::Mutex::new(None),
335 total_system_prompt_tokens: AtomicU64::new(0),
336 total_requests: AtomicU64::new(0),
337 }
338 }
339}
340
341impl IntrospectState {
342 pub fn record(&self, breakdown: RequestBreakdown) {
343 self.total_system_prompt_tokens
344 .fetch_add(breakdown.system_prompt_tokens as u64, Ordering::Relaxed);
345 self.total_requests.fetch_add(1, Ordering::Relaxed);
346 if let Ok(mut last) = self.last_breakdown.lock() {
347 *last = Some(breakdown);
348 }
349 }
350}
351
352#[cfg(test)]
353mod tests {
354 use super::*;
355
356 #[test]
357 fn anthropic_basic() {
358 let body = serde_json::json!({
359 "model": "claude-sonnet-4-20250514",
360 "system": "You are a helpful assistant.",
361 "messages": [
362 {"role": "user", "content": "Hello"},
363 {"role": "assistant", "content": "Hi there!"}
364 ],
365 "tools": [{"name": "read", "description": "Read a file", "input_schema": {}}]
366 });
367 let b = analyze_request(&body, Provider::Anthropic);
368 assert_eq!(b.provider, Provider::Anthropic);
369 assert!(b.system_prompt_tokens > 0);
370 assert_eq!(b.message_count, 2);
371 assert!(b.user_message_tokens > 0);
372 assert!(b.assistant_message_tokens > 0);
373 assert_eq!(b.tool_definition_count, 1);
374 assert!(b.tool_definition_tokens > 0);
375 }
376
377 #[test]
378 fn openai_system_message() {
379 let body = serde_json::json!({
380 "model": "gpt-4o",
381 "messages": [
382 {"role": "system", "content": "System prompt here"},
383 {"role": "user", "content": "Hello"},
384 {"role": "tool", "content": "tool result data", "tool_call_id": "x"}
385 ]
386 });
387 let b = analyze_request(&body, Provider::OpenAi);
388 assert!(b.system_prompt_tokens > 0);
389 assert!(b.user_message_tokens > 0);
390 assert!(b.tool_result_tokens > 0);
391 assert_eq!(b.message_count, 3);
392 }
393
394 #[test]
395 fn gemini_system_instruction() {
396 let body = serde_json::json!({
397 "systemInstruction": {
398 "parts": [{"text": "Be concise and helpful to the user at all times."}]
399 },
400 "contents": [
401 {"role": "user", "parts": [{"text": "What is the meaning of life and everything?"}]},
402 {"role": "model", "parts": [{"text": "The answer is 42 according to Douglas Adams."}]}
403 ]
404 });
405 let b = analyze_request(&body, Provider::Gemini);
406 assert!(b.system_prompt_tokens > 0);
407 assert!(b.user_message_tokens > 0);
408 assert!(b.assistant_message_tokens > 0);
409 assert_eq!(b.message_count, 2);
410 }
411}