lean_ctx/proxy/
tool_output.rs

1use serde_json::Value;
2
3use crate::core::tokens::TokenizerFamily;
4
5use super::compress::{compress_tool_result, compress_tool_result_for};
6use super::tool_kind::{ToolResultKind, should_protect};
7
8enum JsonRewrite {
9    NotJson,
10    Unchanged,
11    Changed(String),
12}
13
14pub(super) fn compress_text(
15    text: &mut String,
16    tool_name: Option<&str>,
17    kind: ToolResultKind,
18) -> bool {
19    match rewrite_json_payload_text(text, kind, |inner| {
20        if should_protect(kind, inner) {
21            return None;
22        }
23        let compressed = compress_tool_result(inner, tool_name);
24        (compressed.len() < inner.len()).then_some(compressed)
25    }) {
26        JsonRewrite::Changed(compressed) => {
27            *text = compressed;
28            return true;
29        }
30        JsonRewrite::Unchanged => return false,
31        JsonRewrite::NotJson => {}
32    }
33
34    if should_protect(kind, text) {
35        return false;
36    }
37    let compressed = compress_tool_result(text, tool_name);
38    if compressed.len() < text.len() {
39        *text = compressed;
40        return true;
41    }
42    false
43}
44
45pub(super) fn compress_text_for(
46    text: &mut String,
47    tool_name: Option<&str>,
48    kind: ToolResultKind,
49    family: TokenizerFamily,
50) -> bool {
51    match rewrite_json_payload_text(text, kind, |inner| {
52        if should_protect(kind, inner) {
53            return None;
54        }
55        let compressed = compress_tool_result_for(inner, tool_name, family);
56        (compressed.len() < inner.len()).then_some(compressed)
57    }) {
58        JsonRewrite::Changed(compressed) => {
59            *text = compressed;
60            return true;
61        }
62        JsonRewrite::Unchanged => return false,
63        JsonRewrite::NotJson => {}
64    }
65
66    if should_protect(kind, text) {
67        return false;
68    }
69    let compressed = compress_tool_result_for(text, tool_name, family);
70    if compressed.len() < text.len() {
71        *text = compressed;
72        return true;
73    }
74    false
75}
76
77pub(super) fn compress_value(
78    value: &mut Value,
79    tool_name: Option<&str>,
80    kind: ToolResultKind,
81) -> bool {
82    match value {
83        Value::String(text) => compress_text(text, tool_name, kind),
84        Value::Array(parts) => {
85            let mut changed = false;
86            for part in parts.iter_mut() {
87                if let Some(Value::String(text)) = part.get_mut("text") {
88                    changed |= compress_text(text, tool_name, kind);
89                }
90            }
91            changed
92        }
93        _ => false,
94    }
95}
96
97// TODO(#1354): remove dead code or implement
98pub(super) fn compress_value_for(
99    value: &mut Value,
100    tool_name: Option<&str>,
101    kind: ToolResultKind,
102    family: TokenizerFamily,
103) -> bool {
104    match value {
105        Value::String(text) => compress_text_for(text, tool_name, kind, family),
106        Value::Array(parts) => {
107            let mut changed = false;
108            for part in parts.iter_mut() {
109                if let Some(Value::String(text)) = part.get_mut("text") {
110                    changed |= compress_text_for(text, tool_name, kind, family);
111                }
112            }
113            changed
114        }
115        _ => false,
116    }
117}
118
119pub(super) fn prune_text(text: &mut String, kind: ToolResultKind) -> bool {
120    match rewrite_json_payload_text(text, kind, |inner| {
121        super::history_prune::prune_output_text(inner, kind)
122    }) {
123        JsonRewrite::Changed(pruned) => {
124            *text = pruned;
125            return true;
126        }
127        JsonRewrite::Unchanged => return false,
128        JsonRewrite::NotJson => {}
129    }
130
131    if let Some(pruned) = super::history_prune::prune_output_text(text, kind) {
132        *text = pruned;
133        return true;
134    }
135    false
136}
137
138pub(super) fn prune_value(value: &mut Value, kind: ToolResultKind) -> bool {
139    match value {
140        Value::String(text) => prune_text(text, kind),
141        Value::Array(parts) => {
142            let mut changed = false;
143            for part in parts.iter_mut() {
144                if let Some(Value::String(text)) = part.get_mut("text") {
145                    changed |= prune_text(text, kind);
146                }
147            }
148            changed
149        }
150        _ => false,
151    }
152}
153
154fn rewrite_json_payload_text(
155    text: &str,
156    kind: ToolResultKind,
157    mut rewrite: impl FnMut(&str) -> Option<String>,
158) -> JsonRewrite {
159    let trimmed = text.trim();
160    if !(trimmed.starts_with('{') || trimmed.starts_with('[')) {
161        return JsonRewrite::NotJson;
162    }
163    let Ok(mut value) = serde_json::from_str::<Value>(trimmed) else {
164        return JsonRewrite::NotJson;
165    };
166    let mut touched = false;
167    let mut changed = false;
168    rewrite_json_text_values(&mut value, kind, &mut rewrite, &mut touched, &mut changed);
169    if !touched || !changed {
170        return JsonRewrite::Unchanged;
171    }
172    match serde_json::to_string(&value) {
173        Ok(serialized) if serialized.len() < text.len() => JsonRewrite::Changed(serialized),
174        _ => JsonRewrite::Unchanged,
175    }
176}
177
178fn rewrite_json_text_values(
179    value: &mut Value,
180    kind: ToolResultKind,
181    rewrite: &mut impl FnMut(&str) -> Option<String>,
182    touched: &mut bool,
183    changed: &mut bool,
184) {
185    match value {
186        Value::Object(map) => {
187            let is_text_part = map
188                .get("type")
189                .and_then(Value::as_str)
190                .is_some_and(|t| matches!(t, "text" | "input_text" | "output_text"));
191            let rewrite_all_strings =
192                matches!(kind, ToolResultKind::Shell | ToolResultKind::Search);
193            for (key, child) in map.iter_mut() {
194                if let Value::String(s) = child
195                    && (rewrite_all_strings || (is_text_part && key == "text"))
196                {
197                    *touched = true;
198                    if let Some(next) = rewrite(s) {
199                        *s = next;
200                        *changed = true;
201                    }
202                    continue;
203                }
204                rewrite_json_text_values(child, kind, rewrite, touched, changed);
205            }
206        }
207        Value::Array(items) => {
208            for item in items {
209                rewrite_json_text_values(item, kind, rewrite, touched, changed);
210            }
211        }
212        _ => {}
213    }
214}
215
216#[cfg(test)]
217mod tests {
218    use super::*;
219    use crate::core::tokens::count_tokens_for;
220
221    #[test]
222    fn compress_text_for_uses_selected_tokenizer_family() {
223        let paragraph = "Grüezi 世界 means hello in Swiss German, and this repeated sentence carries multilingual text for tokenizer accounting. 🤝";
224        let input = format!("{}\n", [paragraph; 20].join("\n\n"));
225        let mut cl100k = input.clone();
226        let mut o200k = input;
227
228        assert!(compress_text_for(
229            &mut cl100k,
230            Some("shell"),
231            ToolResultKind::Shell,
232            TokenizerFamily::Cl100k,
233        ));
234        assert!(compress_text_for(
235            &mut o200k,
236            Some("shell"),
237            ToolResultKind::Shell,
238            TokenizerFamily::O200kBase,
239        ));
240        assert_ne!(
241            count_tokens_for(&cl100k, TokenizerFamily::Cl100k),
242            count_tokens_for(&o200k, TokenizerFamily::O200kBase),
243            "the selected tokenizer family must produce model-specific accounting"
244        );
245    }
246}
lean_ctx/proxy/tool_output.rs

lean_ctx/proxy/
tool_output.rs
