gproxy_tokenize/tokenize/
extract.rs

1//! Protocol-agnostic text harvesting from provider-native request JSON.
2//! Compiled on every target (the edge estimate path needs it too).
3
4use serde_json::Value;
5
6/// Keys whose string values are human text worth counting.
7const TEXT_KEYS: &[&str] = &["text", "content", "input", "instructions", "system"];
8/// Keys whose non-string values (tool defs, structured system) are counted
9/// by serializing the whole subtree.
10const SERIALIZE_KEYS: &[&str] = &["tools", "tool_choice", "system"];
11/// Keys whose array length approximates the message count.
12const MESSAGE_KEYS: &[&str] = &["messages", "contents", "input"];
13
14/// Harvest human-text from any provider-native request JSON: walks the value,
15/// collecting strings under text-ish keys (`text`, `content`, `instructions`,
16/// string-form `system`, gemini parts text) plus tool definitions serialized.
17/// Returns `(texts, message_count)` where `message_count` is the length of
18/// the largest `messages` / `contents` / `input` array found (0 if none).
19pub fn harvest(body: &[u8]) -> (Vec<String>, u64) {
20    let Ok(root) = serde_json::from_slice::<Value>(body) else {
21        return (Vec::new(), 0);
22    };
23    let mut texts = Vec::new();
24    let mut messages = 0u64;
25    walk(&root, &mut texts, &mut messages);
26    (texts, messages)
27}
28
29fn walk(value: &Value, texts: &mut Vec<String>, messages: &mut u64) {
30    match value {
31        Value::Object(map) => {
32            for (key, val) in map {
33                match val {
34                    Value::String(s) if TEXT_KEYS.contains(&key.as_str()) => {
35                        texts.push(s.clone());
36                    }
37                    _ if SERIALIZE_KEYS.contains(&key.as_str()) && !val.is_null() => {
38                        texts.push(val.to_string());
39                    }
40                    Value::Array(arr) => {
41                        if MESSAGE_KEYS.contains(&key.as_str()) {
42                            *messages = (*messages).max(arr.len() as u64);
43                        }
44                        for item in arr {
45                            walk(item, texts, messages);
46                        }
47                    }
48                    Value::Object(_) => walk(val, texts, messages),
49                    _ => {}
50                }
51            }
52        }
53        Value::Array(arr) => {
54            for item in arr {
55                walk(item, texts, messages);
56            }
57        }
58        _ => {}
59    }
60}
61
62#[cfg(test)]
63mod tests {
64    use super::harvest;
65
66    #[test]
67    fn harvest_claude_body() {
68        let body = serde_json::json!({
69            "model": "claude-sonnet-4",
70            "system": "be terse",
71            "messages": [
72                { "role": "user", "content": "hello there" },
73                { "role": "assistant", "content": [
74                    { "type": "text", "text": "hi!" }
75                ]}
76            ],
77            "tools": [{ "name": "get_weather", "description": "weather" }]
78        })
79        .to_string();
80        let (texts, messages) = harvest(body.as_bytes());
81        assert_eq!(messages, 2);
82        assert!(texts.iter().any(|t| t == "hello there"));
83        assert!(texts.iter().any(|t| t == "hi!"));
84        assert!(texts.iter().any(|t| t == "be terse"));
85        assert!(texts.iter().any(|t| t.contains("get_weather")));
86    }
87
88    #[test]
89    fn harvest_openai_responses_string_input() {
90        let body = serde_json::json!({
91            "model": "deepseek-v4-flash",
92            "input": "Count these words."
93        })
94        .to_string();
95        let (texts, messages) = harvest(body.as_bytes());
96        assert_eq!(texts, ["Count these words."]);
97        assert_eq!(messages, 0);
98    }
99}
gproxy_tokenize/tokenize/extract.rs

gproxy_tokenize/tokenize/
extract.rs
