use serde_json::Value;
use crate::core::tokens::TokenizerFamily;
use super::compress::{compress_tool_result, compress_tool_result_for};
use super::tool_kind::{ToolResultKind, should_protect};
enum JsonRewrite {
NotJson,
Unchanged,
Changed(String),
}
pub(super) fn compress_text(
text: &mut String,
tool_name: Option<&str>,
kind: ToolResultKind,
) -> bool {
match rewrite_json_payload_text(text, kind, |inner| {
if should_protect(kind, inner) {
return None;
}
let compressed = compress_tool_result(inner, tool_name);
(compressed.len() < inner.len()).then_some(compressed)
}) {
JsonRewrite::Changed(compressed) => {
*text = compressed;
return true;
}
JsonRewrite::Unchanged => return false,
JsonRewrite::NotJson => {}
}
if should_protect(kind, text) {
return false;
}
let compressed = compress_tool_result(text, tool_name);
if compressed.len() < text.len() {
*text = compressed;
return true;
}
false
}
#[allow(dead_code)]
pub(super) fn compress_text_for(
text: &mut String,
tool_name: Option<&str>,
kind: ToolResultKind,
family: TokenizerFamily,
) -> bool {
match rewrite_json_payload_text(text, kind, |inner| {
if should_protect(kind, inner) {
return None;
}
let compressed = compress_tool_result_for(inner, tool_name, family);
(compressed.len() < inner.len()).then_some(compressed)
}) {
JsonRewrite::Changed(compressed) => {
*text = compressed;
return true;
}
JsonRewrite::Unchanged => return false,
JsonRewrite::NotJson => {}
}
if should_protect(kind, text) {
return false;
}
let compressed = compress_tool_result_for(text, tool_name, family);
if compressed.len() < text.len() {
*text = compressed;
return true;
}
false
}
pub(super) fn compress_value(
value: &mut Value,
tool_name: Option<&str>,
kind: ToolResultKind,
) -> bool {
match value {
Value::String(text) => compress_text(text, tool_name, kind),
Value::Array(parts) => {
let mut changed = false;
for part in parts.iter_mut() {
if let Some(Value::String(text)) = part.get_mut("text") {
changed |= compress_text(text, tool_name, kind);
}
}
changed
}
_ => false,
}
}
#[allow(dead_code)]
pub(super) fn compress_value_for(
value: &mut Value,
tool_name: Option<&str>,
kind: ToolResultKind,
family: TokenizerFamily,
) -> bool {
match value {
Value::String(text) => compress_text_for(text, tool_name, kind, family),
Value::Array(parts) => {
let mut changed = false;
for part in parts.iter_mut() {
if let Some(Value::String(text)) = part.get_mut("text") {
changed |= compress_text_for(text, tool_name, kind, family);
}
}
changed
}
_ => false,
}
}
pub(super) fn prune_text(text: &mut String, kind: ToolResultKind) -> bool {
match rewrite_json_payload_text(text, kind, |inner| {
super::history_prune::prune_output_text(inner, kind)
}) {
JsonRewrite::Changed(pruned) => {
*text = pruned;
return true;
}
JsonRewrite::Unchanged => return false,
JsonRewrite::NotJson => {}
}
if let Some(pruned) = super::history_prune::prune_output_text(text, kind) {
*text = pruned;
return true;
}
false
}
pub(super) fn prune_value(value: &mut Value, kind: ToolResultKind) -> bool {
match value {
Value::String(text) => prune_text(text, kind),
Value::Array(parts) => {
let mut changed = false;
for part in parts.iter_mut() {
if let Some(Value::String(text)) = part.get_mut("text") {
changed |= prune_text(text, kind);
}
}
changed
}
_ => false,
}
}
fn rewrite_json_payload_text(
text: &str,
kind: ToolResultKind,
mut rewrite: impl FnMut(&str) -> Option<String>,
) -> JsonRewrite {
let trimmed = text.trim();
if !(trimmed.starts_with('{') || trimmed.starts_with('[')) {
return JsonRewrite::NotJson;
}
let Ok(mut value) = serde_json::from_str::<Value>(trimmed) else {
return JsonRewrite::NotJson;
};
let mut touched = false;
let mut changed = false;
rewrite_json_text_values(&mut value, kind, &mut rewrite, &mut touched, &mut changed);
if !touched || !changed {
return JsonRewrite::Unchanged;
}
match serde_json::to_string(&value) {
Ok(serialized) if serialized.len() < text.len() => JsonRewrite::Changed(serialized),
_ => JsonRewrite::Unchanged,
}
}
fn rewrite_json_text_values(
value: &mut Value,
kind: ToolResultKind,
rewrite: &mut impl FnMut(&str) -> Option<String>,
touched: &mut bool,
changed: &mut bool,
) {
match value {
Value::Object(map) => {
let is_text_part = map
.get("type")
.and_then(Value::as_str)
.is_some_and(|t| matches!(t, "text" | "input_text" | "output_text"));
let rewrite_all_strings =
matches!(kind, ToolResultKind::Shell | ToolResultKind::Search);
for (key, child) in map.iter_mut() {
if let Value::String(s) = child
&& (rewrite_all_strings || (is_text_part && key == "text"))
{
*touched = true;
if let Some(next) = rewrite(s) {
*s = next;
*changed = true;
}
continue;
}
rewrite_json_text_values(child, kind, rewrite, touched, changed);
}
}
Value::Array(items) => {
for item in items {
rewrite_json_text_values(item, kind, rewrite, touched, changed);
}
}
_ => {}
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::core::tokens::count_tokens_for;
#[test]
fn compress_text_for_uses_selected_tokenizer_family() {
let paragraph = "Grüezi 世界 means hello in Swiss German, and this repeated sentence carries multilingual text for tokenizer accounting. 🤝";
let input = format!("{}\n", [paragraph; 20].join("\n\n"));
let mut cl100k = input.clone();
let mut o200k = input;
assert!(compress_text_for(
&mut cl100k,
Some("shell"),
ToolResultKind::Shell,
TokenizerFamily::Cl100k,
));
assert!(compress_text_for(
&mut o200k,
Some("shell"),
ToolResultKind::Shell,
TokenizerFamily::O200kBase,
));
assert_ne!(
count_tokens_for(&cl100k, TokenizerFamily::Cl100k),
count_tokens_for(&o200k, TokenizerFamily::O200kBase),
"the selected tokenizer family must produce model-specific accounting"
);
}
}