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] = &[
8    "text",
9    "content",
10    "input",
11    "instructions",
12    "system",
13    "reasoning",
14    "reasoning_content",
15    "arguments",
16    "partial_json",
17];
18/// Keys whose non-string values (tool defs, structured system) are counted
19/// by serializing the whole subtree.
20const SERIALIZE_KEYS: &[&str] = &[
21    "tools",
22    "tool_choice",
23    "system",
24    "response_format",
25    "json_schema",
26    "schema",
27    "generation_config",
28];
29/// Keys whose array length approximates the message count.
30const MESSAGE_KEYS: &[&str] = &["messages", "contents", "input"];
31
32/// Harvest human-text from any provider-native request JSON: walks the value,
33/// collecting strings under text-ish keys (`text`, `content`, `instructions`,
34/// string-form `system`, gemini parts text) plus tool definitions serialized.
35/// Returns `(texts, message_count)` where `message_count` is the length of
36/// the largest `messages` / `contents` / `input` array found (0 if none).
37pub fn harvest(body: &[u8]) -> (Vec<String>, u64) {
38    try_harvest(body).unwrap_or_default()
39}
40
41/// Fallible harvesting for callers that must distinguish malformed JSON from
42/// an intentionally empty request.
43pub fn try_harvest(body: &[u8]) -> Result<(Vec<String>, u64), serde_json::Error> {
44    let root = serde_json::from_slice::<Value>(body)?;
45    let mut texts = Vec::new();
46    let mut messages = 0u64;
47    walk(&root, &mut texts, &mut messages);
48    Ok((texts, messages))
49}
50
51fn walk(value: &Value, texts: &mut Vec<String>, messages: &mut u64) {
52    match value {
53        Value::Object(map) => {
54            for (key, val) in map {
55                match val {
56                    Value::String(s) if TEXT_KEYS.contains(&key.as_str()) => {
57                        texts.push(s.clone());
58                    }
59                    _ if SERIALIZE_KEYS.contains(&key.as_str()) && !val.is_null() => {
60                        texts.push(val.to_string());
61                    }
62                    Value::Array(arr) => {
63                        if MESSAGE_KEYS.contains(&key.as_str()) {
64                            *messages = (*messages).max(arr.len() as u64);
65                        }
66                        for item in arr {
67                            walk(item, texts, messages);
68                        }
69                    }
70                    Value::Object(_) => walk(val, texts, messages),
71                    _ => {}
72                }
73            }
74        }
75        Value::Array(arr) => {
76            for item in arr {
77                walk(item, texts, messages);
78            }
79        }
80        _ => {}
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::{harvest, try_harvest};
87
88    #[test]
89    fn harvest_claude_body() {
90        let body = serde_json::json!({
91            "model": "claude-sonnet-4",
92            "system": "be terse",
93            "messages": [
94                { "role": "user", "content": "hello there" },
95                { "role": "assistant", "content": [
96                    { "type": "text", "text": "hi!" }
97                ]}
98            ],
99            "tools": [{ "name": "get_weather", "description": "weather" }]
100        })
101        .to_string();
102        let (texts, messages) = harvest(body.as_bytes());
103        assert_eq!(messages, 2);
104        assert!(texts.iter().any(|t| t == "hello there"));
105        assert!(texts.iter().any(|t| t == "hi!"));
106        assert!(texts.iter().any(|t| t == "be terse"));
107        assert!(texts.iter().any(|t| t.contains("get_weather")));
108    }
109
110    #[test]
111    fn harvest_openai_responses_string_input() {
112        let body = serde_json::json!({
113            "model": "deepseek-v4-flash",
114            "input": "Count these words."
115        })
116        .to_string();
117        let (texts, messages) = harvest(body.as_bytes());
118        assert_eq!(texts, ["Count these words."]);
119        assert_eq!(messages, 0);
120    }
121
122    #[test]
123    fn invalid_json_is_explicit_on_fallible_path() {
124        assert!(try_harvest(br#"{"messages":["#).is_err());
125        assert_eq!(harvest(br#"{"messages":["#), (Vec::new(), 0));
126    }
127
128    #[test]
129    fn harvests_reasoning_arguments_and_schema() {
130        let body = serde_json::json!({
131            "messages": [{
132                "role": "assistant",
133                "reasoning_content": "think",
134                "tool_calls": [{"function": {"arguments": "{\"city\":\"Paris\"}"}}]
135            }],
136            "response_format": {"type": "json_schema", "json_schema": {"name": "answer"}}
137        });
138        let (texts, _) = try_harvest(body.to_string().as_bytes()).unwrap();
139        assert!(texts.iter().any(|text| text == "think"));
140        assert!(texts.iter().any(|text| text.contains("Paris")));
141        assert!(texts.iter().any(|text| text.contains("json_schema")));
142    }
143}
gproxy_tokenize/tokenize/extract.rs

gproxy_tokenize/tokenize/
extract.rs
