utils/output-truncation/src/lib.rs
157 lines
//! Helpers for truncating tool and exec output using [`TruncationPolicy`](codex_protocol::protocol::TruncationPolicy).use codex_protocol::models::FunctionCallOutputContentItem;pub use codex_utils_string::approx_bytes_for_tokens;pub use codex_utils_string::approx_token_count;pub use codex_utils_string::approx_tokens_from_byte_count;use codex_utils_string::truncate_middle_chars;use codex_utils_string::truncate_middle_with_token_budget;pub use codex_protocol::protocol::TruncationPolicy;pub fn formatted_truncate_text(content: &str, policy: TruncationPolicy) -> String { if content.len() <= policy.byte_budget() { return content.to_string(); } let original_token_count = approx_token_count(content); let total_lines = content.lines().count(); let result = truncate_text(content, policy); format!( "Warning: truncated output (original token count: {original_token_count})\nTotal output lines: {total_lines}\n\n{result}" )}pub fn truncate_text(content: &str, policy: TruncationPolicy) -> String { match policy { TruncationPolicy::Bytes(bytes) => truncate_middle_chars(content, bytes), TruncationPolicy::Tokens(tokens) => truncate_middle_with_token_budget(content, tokens).0, }}pub fn formatted_truncate_text_content_items_with_policy( items: &[FunctionCallOutputContentItem], policy: TruncationPolicy,) -> (Vec<FunctionCallOutputContentItem>, Option<usize>) { let text_segments = items .iter() .filter_map(|item| match item { FunctionCallOutputContentItem::InputText { text } => Some(text.as_str()), FunctionCallOutputContentItem::InputImage { .. } | FunctionCallOutputContentItem::EncryptedContent { .. } => None, }) .collect::<Vec<_>>(); if text_segments.is_empty() { return (items.to_vec(), None); } let mut combined = String::new(); for text in &text_segments { if !combined.is_empty() { combined.push('\n'); } combined.push_str(text); } if combined.len() <= policy.byte_budget() { return (items.to_vec(), None); } let original_token_count = approx_token_count(&combined); let mut out = vec![FunctionCallOutputContentItem::InputText { text: formatted_truncate_text(&combined, policy), }]; out.extend(items.iter().filter_map(|item| match item { FunctionCallOutputContentItem::InputImage { image_url, detail } => { Some(FunctionCallOutputContentItem::InputImage { image_url: image_url.clone(), detail: *detail, }) } FunctionCallOutputContentItem::EncryptedContent { encrypted_content } => { Some(FunctionCallOutputContentItem::EncryptedContent { encrypted_content: encrypted_content.clone(), }) } FunctionCallOutputContentItem::InputText { .. } => None, })); (out, Some(original_token_count))}pub fn truncate_function_output_items_with_policy( items: &[FunctionCallOutputContentItem], policy: TruncationPolicy,) -> Vec<FunctionCallOutputContentItem> { let mut out: Vec<FunctionCallOutputContentItem> = Vec::with_capacity(items.len()); let mut remaining_budget = match policy { TruncationPolicy::Bytes(_) => policy.byte_budget(), TruncationPolicy::Tokens(_) => policy.token_budget(), }; let mut omitted_text_items = 0usize; for item in items { match item { FunctionCallOutputContentItem::InputText { text } => { if remaining_budget == 0 { omitted_text_items += 1; continue; } let cost = match policy { TruncationPolicy::Bytes(_) => text.len(), TruncationPolicy::Tokens(_) => approx_token_count(text), }; if cost <= remaining_budget { out.push(FunctionCallOutputContentItem::InputText { text: text.clone() }); remaining_budget = remaining_budget.saturating_sub(cost); } else { let snippet_policy = match policy { TruncationPolicy::Bytes(_) => TruncationPolicy::Bytes(remaining_budget), TruncationPolicy::Tokens(_) => TruncationPolicy::Tokens(remaining_budget), }; let snippet = truncate_text(text, snippet_policy); if snippet.is_empty() { omitted_text_items += 1; } else { out.push(FunctionCallOutputContentItem::InputText { text: snippet }); } remaining_budget = 0; } } FunctionCallOutputContentItem::InputImage { image_url, detail } => { out.push(FunctionCallOutputContentItem::InputImage { image_url: image_url.clone(), detail: *detail, }); } FunctionCallOutputContentItem::EncryptedContent { encrypted_content } => { out.push(FunctionCallOutputContentItem::EncryptedContent { encrypted_content: encrypted_content.clone(), }); } } } if omitted_text_items > 0 { out.push(FunctionCallOutputContentItem::InputText { text: format!("[omitted {omitted_text_items} text items ...]"), }); } out}pub fn approx_tokens_from_byte_count_i64(bytes: i64) -> i64 { if bytes <= 0 { return 0; } let bytes = usize::try_from(bytes).unwrap_or(usize::MAX); i64::try_from(approx_tokens_from_byte_count(bytes)).unwrap_or(i64::MAX)}#[cfg(test)]mod truncate_tests;