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

lean_ctx/proxy/
tool_output.rs
