mirror of
https://github.com/openai/codex.git
synced 2026-10-28 18:58:46 +00:00
Preserve output boundaries through token truncation (#57584)
## What changed Add `OutputBoundary` and token-budget formatting APIs that map item indices and UTF-8 text offsets to their positions after truncation. Normalize input positions, remap boundaries within omitted text, and keep end boundaries before generated omission notices. Unify code mode output formatting: merge text-only output for truncation with a warning, and budget mixed text and audio in input order while preserving images and encrypted content. Accept owned output items to reuse allocations and move retained media without cloning. ## Testing Add coverage for boundary normalization and remapping, UTF-8 offsets, zero and nonzero budgets, owned and borrowed output equivalence, and allocation reuse. GitOrigin-RevId: 598407caa88b5a1f84130c4bc2273b85ab02f7f2
This commit is contained in:
Sean Huang
@@ -46,8 +46,7 @@ use codex_protocol::openai_models::ToolMode;
|
||||
use codex_tools::ToolName;
|
||||
use codex_utils_audio::estimate_audio_token_count;
|
||||
use codex_utils_output_truncation::TruncationPolicy;
|
||||
use codex_utils_output_truncation::formatted_truncate_text_content_items_with_policy;
|
||||
use codex_utils_output_truncation::truncate_function_output_items_with_policy;
|
||||
use codex_utils_output_truncation::formatted_truncate_function_output_items_with_policy;
|
||||
|
||||
use delegate::CodeModeCellDelegate;
|
||||
use delegate::CodeModeDispatchBroker;
|
||||
@@ -321,16 +320,7 @@ fn truncate_code_mode_result(
|
||||
) -> Vec<FunctionCallOutputContentItem> {
|
||||
let max_output_tokens = resolve_max_tokens(max_output_tokens);
|
||||
let policy = TruncationPolicy::Tokens(max_output_tokens);
|
||||
if items
|
||||
.iter()
|
||||
.all(|item| matches!(item, FunctionCallOutputContentItem::InputText { .. }))
|
||||
{
|
||||
let (truncated_items, _) =
|
||||
formatted_truncate_text_content_items_with_policy(&items, policy);
|
||||
return truncated_items;
|
||||
}
|
||||
|
||||
truncate_function_output_items_with_policy(&items, policy, estimate_audio_token_count)
|
||||
formatted_truncate_function_output_items_with_policy(items, policy, estimate_audio_token_count)
|
||||
}
|
||||
|
||||
// Submit synchronously so the recorder sees the call before the cell's dispatch gate closes.
|
||||
|
||||
@@ -6,11 +6,12 @@ use codex_protocol::mcp::CallToolResult;
|
||||
use codex_protocol::models::FunctionCallOutputBody;
|
||||
use codex_protocol::models::FunctionCallOutputContentItem;
|
||||
use codex_protocol::models::FunctionCallOutputPayload;
|
||||
use codex_utils_string::TextTruncation;
|
||||
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::plan_token_truncation;
|
||||
use codex_utils_string::truncate_middle_chars;
|
||||
use codex_utils_string::truncate_middle_with_token_budget;
|
||||
|
||||
pub use codex_protocol::protocol::TruncationPolicy;
|
||||
|
||||
@@ -25,18 +26,43 @@ pub fn formatted_truncate_text(content: &str, policy: TruncationPolicy) -> Strin
|
||||
return content.to_string();
|
||||
}
|
||||
|
||||
format_truncated_text(content, policy, &mut [])
|
||||
}
|
||||
|
||||
fn format_truncated_text(
|
||||
content: &str,
|
||||
policy: TruncationPolicy,
|
||||
boundaries: &mut [OutputBoundary],
|
||||
) -> 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}"
|
||||
)
|
||||
let header = format!(
|
||||
"Warning: truncated output (original token count: {original_token_count})\nTotal output lines: {total_lines}\n\n"
|
||||
);
|
||||
let (snippet, edit) = truncate_text_with_plan(content, policy);
|
||||
if let Some(edit) = edit {
|
||||
for boundary in boundaries {
|
||||
boundary.text_offset = rendered_offset(&edit, boundary.text_offset, header.len());
|
||||
}
|
||||
}
|
||||
header + &snippet
|
||||
}
|
||||
|
||||
pub fn truncate_text(content: &str, policy: TruncationPolicy) -> String {
|
||||
truncate_text_with_plan(content, policy).0
|
||||
}
|
||||
|
||||
// Only token-budget callers supply boundaries that need remapping.
|
||||
fn truncate_text_with_plan(
|
||||
content: &str,
|
||||
policy: TruncationPolicy,
|
||||
) -> (String, Option<TextTruncation>) {
|
||||
match policy {
|
||||
TruncationPolicy::Bytes(bytes) => truncate_middle_chars(content, bytes),
|
||||
TruncationPolicy::Tokens(tokens) => truncate_middle_with_token_budget(content, tokens).0,
|
||||
TruncationPolicy::Bytes(bytes) => (truncate_middle_chars(content, bytes), None),
|
||||
TruncationPolicy::Tokens(tokens) => match plan_token_truncation(content, tokens) {
|
||||
Some(edit) => (edit.apply(content), Some(edit)),
|
||||
None => (content.to_owned(), None),
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
@@ -104,67 +130,169 @@ pub fn truncate_function_output_payload(
|
||||
}
|
||||
}
|
||||
|
||||
pub fn formatted_truncate_text_content_items_with_policy(
|
||||
items: &[FunctionCallOutputContentItem],
|
||||
/// An uncharged boundary in the original logical output. Token-boundary
|
||||
/// formatters update the item index and UTF-8 text offset to the rendered
|
||||
/// position, including when content at that position was omitted.
|
||||
///
|
||||
/// At entry, item indices are clamped to the end of the input, text offsets
|
||||
/// round down to an in-bounds UTF-8 boundary, and offsets on non-text items
|
||||
/// are zeroed. Positions preceding an earlier boundary move to that position,
|
||||
/// preserving caller order. On return, boundaries refer to the rendered items;
|
||||
/// an end boundary precedes any generated omission notices.
|
||||
pub struct OutputBoundary {
|
||||
pub item_index: usize,
|
||||
pub text_offset: usize,
|
||||
}
|
||||
|
||||
/// Format text-only output as a whole observation, including the truncation
|
||||
/// warning. Charge mixed text and audio against the budget in order while
|
||||
/// preserving images.
|
||||
pub fn formatted_truncate_function_output_items_with_policy<'a>(
|
||||
items: impl Into<Cow<'a, [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::InputAudio { .. }
|
||||
| FunctionCallOutputContentItem::EncryptedContent { .. } => None,
|
||||
})
|
||||
.collect::<Vec<_>>();
|
||||
estimate_audio_token_count: impl Fn(&str) -> usize,
|
||||
) -> Vec<FunctionCallOutputContentItem> {
|
||||
formatted_output(items.into(), policy, estimate_audio_token_count, &mut [])
|
||||
}
|
||||
|
||||
if text_segments.is_empty() {
|
||||
return (items.to_vec(), None);
|
||||
/// As above, with zero-width boundaries for callers that have already dropped
|
||||
/// output. Input positions are normalized as documented on [`OutputBoundary`].
|
||||
pub fn formatted_output_at_token_boundaries<'a>(
|
||||
items: impl Into<Cow<'a, [FunctionCallOutputContentItem]>>,
|
||||
max_tokens: usize,
|
||||
estimate_audio_token_count: impl Fn(&str) -> usize,
|
||||
boundaries: &mut [OutputBoundary],
|
||||
) -> Vec<FunctionCallOutputContentItem> {
|
||||
let items = items.into();
|
||||
normalize_boundaries(&items, boundaries);
|
||||
formatted_output(
|
||||
items,
|
||||
TruncationPolicy::Tokens(max_tokens),
|
||||
estimate_audio_token_count,
|
||||
boundaries,
|
||||
)
|
||||
}
|
||||
|
||||
fn formatted_output(
|
||||
items: Cow<'_, [FunctionCallOutputContentItem]>,
|
||||
policy: TruncationPolicy,
|
||||
estimate_audio_token_count: impl Fn(&str) -> usize,
|
||||
boundaries: &mut [OutputBoundary],
|
||||
) -> Vec<FunctionCallOutputContentItem> {
|
||||
let mut combined_len = 0;
|
||||
for item in items.iter() {
|
||||
let FunctionCallOutputContentItem::InputText { text } = item else {
|
||||
return truncate_output_items(items, policy, estimate_audio_token_count, boundaries);
|
||||
};
|
||||
combined_len += usize::from(combined_len != 0) + text.len();
|
||||
}
|
||||
if combined_len <= policy.byte_budget() {
|
||||
return items.into_owned();
|
||||
}
|
||||
|
||||
let mut combined = String::new();
|
||||
for text in &text_segments {
|
||||
if !combined.is_empty() {
|
||||
combined.push('\n');
|
||||
let mut combined = String::with_capacity(combined_len);
|
||||
let mut next_boundary = 0;
|
||||
for (index, item) in items.iter().enumerate() {
|
||||
if let FunctionCallOutputContentItem::InputText { text } = item {
|
||||
if !combined.is_empty() {
|
||||
combined.push('\n');
|
||||
}
|
||||
while let Some(boundary) = boundaries.get_mut(next_boundary)
|
||||
&& boundary.item_index == index
|
||||
{
|
||||
boundary.item_index = 0;
|
||||
boundary.text_offset += combined.len();
|
||||
next_boundary += 1;
|
||||
}
|
||||
combined.push_str(text);
|
||||
}
|
||||
combined.push_str(text);
|
||||
}
|
||||
|
||||
if combined.len() <= policy.byte_budget() {
|
||||
return (items.to_vec(), None);
|
||||
// A boundary after the last item belongs at the end of the merged text.
|
||||
for boundary in &mut boundaries[next_boundary..] {
|
||||
boundary.item_index = 0;
|
||||
boundary.text_offset = combined.len();
|
||||
}
|
||||
|
||||
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, detail } => {
|
||||
Some(FunctionCallOutputContentItem::InputImage {
|
||||
image: image.clone(),
|
||||
detail: *detail,
|
||||
})
|
||||
}
|
||||
FunctionCallOutputContentItem::InputAudio { audio_url } => {
|
||||
Some(FunctionCallOutputContentItem::InputAudio {
|
||||
audio_url: audio_url.clone(),
|
||||
})
|
||||
}
|
||||
FunctionCallOutputContentItem::EncryptedContent { encrypted_content } => {
|
||||
Some(FunctionCallOutputContentItem::EncryptedContent {
|
||||
encrypted_content: encrypted_content.clone(),
|
||||
})
|
||||
}
|
||||
FunctionCallOutputContentItem::InputText { .. } => None,
|
||||
}));
|
||||
|
||||
(out, Some(original_token_count))
|
||||
vec![FunctionCallOutputContentItem::InputText {
|
||||
text: format_truncated_text(&combined, policy, boundaries),
|
||||
}]
|
||||
}
|
||||
|
||||
pub fn truncate_function_output_items_with_policy(
|
||||
items: &[FunctionCallOutputContentItem],
|
||||
policy: TruncationPolicy,
|
||||
estimate_audio_token_count: impl Fn(&str) -> usize,
|
||||
) -> Vec<FunctionCallOutputContentItem> {
|
||||
truncate_output_items(
|
||||
Cow::Borrowed(items),
|
||||
policy,
|
||||
estimate_audio_token_count,
|
||||
&mut [],
|
||||
)
|
||||
}
|
||||
|
||||
/// Budget in input order even if a caller previously dropped every media item.
|
||||
/// Input positions are normalized as documented on [`OutputBoundary`].
|
||||
pub fn mixed_output_at_token_boundaries<'a>(
|
||||
items: impl Into<Cow<'a, [FunctionCallOutputContentItem]>>,
|
||||
max_tokens: usize,
|
||||
estimate_audio_token_count: impl Fn(&str) -> usize,
|
||||
boundaries: &mut [OutputBoundary],
|
||||
) -> Vec<FunctionCallOutputContentItem> {
|
||||
let items = items.into();
|
||||
normalize_boundaries(&items, boundaries);
|
||||
truncate_output_items(
|
||||
items,
|
||||
TruncationPolicy::Tokens(max_tokens),
|
||||
estimate_audio_token_count,
|
||||
boundaries,
|
||||
)
|
||||
}
|
||||
|
||||
fn normalize_boundaries(
|
||||
items: &[FunctionCallOutputContentItem],
|
||||
boundaries: &mut [OutputBoundary],
|
||||
) {
|
||||
let mut previous = (0, 0);
|
||||
for boundary in boundaries {
|
||||
let item_index = boundary.item_index.min(items.len());
|
||||
let text_offset = match items.get(item_index) {
|
||||
Some(FunctionCallOutputContentItem::InputText { text }) => {
|
||||
text.floor_char_boundary(boundary.text_offset.min(text.len()))
|
||||
}
|
||||
_ => 0,
|
||||
};
|
||||
previous = previous.max((item_index, text_offset));
|
||||
(boundary.item_index, boundary.text_offset) = previous;
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate_output_items(
|
||||
items: Cow<'_, [FunctionCallOutputContentItem]>,
|
||||
policy: TruncationPolicy,
|
||||
estimate_audio_token_count: impl Fn(&str) -> usize,
|
||||
boundaries: &mut [OutputBoundary],
|
||||
) -> Vec<FunctionCallOutputContentItem> {
|
||||
match items {
|
||||
Cow::Borrowed(items) => truncate_output_iter(
|
||||
items.iter().map(Cow::Borrowed),
|
||||
policy,
|
||||
estimate_audio_token_count,
|
||||
boundaries,
|
||||
),
|
||||
Cow::Owned(items) => truncate_output_iter(
|
||||
items.into_iter().map(Cow::Owned),
|
||||
policy,
|
||||
estimate_audio_token_count,
|
||||
boundaries,
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
fn truncate_output_iter<'a>(
|
||||
items: impl ExactSizeIterator<Item = Cow<'a, FunctionCallOutputContentItem>>,
|
||||
policy: TruncationPolicy,
|
||||
estimate_audio_token_count: impl Fn(&str) -> usize,
|
||||
boundaries: &mut [OutputBoundary],
|
||||
) -> Vec<FunctionCallOutputContentItem> {
|
||||
let mut out: Vec<FunctionCallOutputContentItem> = Vec::with_capacity(items.len());
|
||||
let mut remaining_budget = match policy {
|
||||
@@ -173,47 +301,44 @@ pub fn truncate_function_output_items_with_policy(
|
||||
};
|
||||
let mut omitted_text_items = 0usize;
|
||||
let mut omitted_audio_items = 0usize;
|
||||
let mut next_boundary = 0;
|
||||
|
||||
for item in items {
|
||||
match item {
|
||||
for (index, item) in items.enumerate() {
|
||||
let item_out = out.len();
|
||||
let item_policy = match policy {
|
||||
TruncationPolicy::Bytes(_) => TruncationPolicy::Bytes(remaining_budget),
|
||||
TruncationPolicy::Tokens(_) => TruncationPolicy::Tokens(remaining_budget),
|
||||
};
|
||||
let mut text_edit = None;
|
||||
let keep = match item.as_ref() {
|
||||
// Empty text contributes no model content but still consumes an API array slot.
|
||||
FunctionCallOutputContentItem::InputText { text } if text.is_empty() => false,
|
||||
FunctionCallOutputContentItem::InputText { .. } if remaining_budget == 0 => {
|
||||
omitted_text_items += 1;
|
||||
false
|
||||
}
|
||||
FunctionCallOutputContentItem::InputText { text } => {
|
||||
// Empty text contributes no model content but still consumes an API array slot.
|
||||
if text.is_empty() {
|
||||
continue;
|
||||
}
|
||||
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);
|
||||
true
|
||||
} 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);
|
||||
let (snippet, edit) = truncate_text_with_plan(text, item_policy);
|
||||
text_edit = edit;
|
||||
if snippet.is_empty() {
|
||||
omitted_text_items += 1;
|
||||
} else {
|
||||
out.push(FunctionCallOutputContentItem::InputText { text: snippet });
|
||||
}
|
||||
remaining_budget = 0;
|
||||
false
|
||||
}
|
||||
}
|
||||
FunctionCallOutputContentItem::InputImage { image, detail } => {
|
||||
out.push(FunctionCallOutputContentItem::InputImage {
|
||||
image: image.clone(),
|
||||
detail: *detail,
|
||||
});
|
||||
}
|
||||
FunctionCallOutputContentItem::InputImage { .. }
|
||||
| FunctionCallOutputContentItem::EncryptedContent { .. } => true,
|
||||
FunctionCallOutputContentItem::InputAudio { audio_url } => {
|
||||
let token_cost = estimate_audio_token_count(audio_url);
|
||||
let cost = match policy {
|
||||
@@ -221,20 +346,35 @@ pub fn truncate_function_output_items_with_policy(
|
||||
TruncationPolicy::Tokens(_) => token_cost,
|
||||
};
|
||||
if cost <= remaining_budget {
|
||||
out.push(FunctionCallOutputContentItem::InputAudio {
|
||||
audio_url: audio_url.clone(),
|
||||
});
|
||||
remaining_budget = remaining_budget.saturating_sub(cost);
|
||||
true
|
||||
} else {
|
||||
omitted_audio_items += 1;
|
||||
false
|
||||
}
|
||||
}
|
||||
FunctionCallOutputContentItem::EncryptedContent { encrypted_content } => {
|
||||
out.push(FunctionCallOutputContentItem::EncryptedContent {
|
||||
encrypted_content: encrypted_content.clone(),
|
||||
});
|
||||
};
|
||||
while let Some(boundary) = boundaries.get_mut(next_boundary)
|
||||
&& boundary.item_index == index
|
||||
{
|
||||
boundary.item_index = item_out;
|
||||
if let Some(edit) = &text_edit {
|
||||
boundary.text_offset =
|
||||
rendered_offset(edit, boundary.text_offset, /*header_bytes*/ 0);
|
||||
} else if !keep {
|
||||
boundary.text_offset = 0;
|
||||
}
|
||||
next_boundary += 1;
|
||||
}
|
||||
if keep {
|
||||
out.push(item.into_owned());
|
||||
}
|
||||
}
|
||||
// End boundaries precede generated omission counts: counts are not original
|
||||
// content and must never move a collector notice away from its source.
|
||||
for boundary in &mut boundaries[next_boundary..] {
|
||||
boundary.item_index = out.len();
|
||||
boundary.text_offset = 0;
|
||||
}
|
||||
|
||||
if omitted_text_items > 0 {
|
||||
@@ -247,10 +387,22 @@ pub fn truncate_function_output_items_with_policy(
|
||||
text: format!("[omitted {omitted_audio_items} audio items ...]"),
|
||||
});
|
||||
}
|
||||
|
||||
out
|
||||
}
|
||||
|
||||
fn rendered_offset(edit: &TextTruncation, boundary: usize, header_bytes: usize) -> usize {
|
||||
if boundary == 0 {
|
||||
0
|
||||
} else if boundary <= edit.removed.start {
|
||||
header_bytes + boundary
|
||||
} else {
|
||||
header_bytes
|
||||
+ edit.removed.start
|
||||
+ edit.marker.len()
|
||||
+ boundary.saturating_sub(edit.removed.end)
|
||||
}
|
||||
}
|
||||
|
||||
pub fn approx_tokens_from_byte_count_i64(bytes: i64) -> i64 {
|
||||
if bytes <= 0 {
|
||||
return 0;
|
||||
|
||||
@@ -1,8 +1,11 @@
|
||||
use crate::OutputBoundary;
|
||||
use crate::TruncationPolicy;
|
||||
use crate::approx_token_count;
|
||||
use crate::approx_tokens_from_byte_count_i64;
|
||||
use crate::formatted_output_at_token_boundaries;
|
||||
use crate::formatted_truncate_function_output_items_with_policy;
|
||||
use crate::formatted_truncate_text;
|
||||
use crate::formatted_truncate_text_content_items_with_policy;
|
||||
use crate::mixed_output_at_token_boundaries;
|
||||
use crate::truncate_function_output_items_with_policy;
|
||||
use crate::truncate_mcp_tool_result;
|
||||
use crate::truncate_text;
|
||||
@@ -278,39 +281,28 @@ fn truncate_function_output_items_with_policy_discards_empty_text() {
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formatted_truncate_text_content_items_with_policy_returns_original_under_limit() {
|
||||
let items = vec![
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "alpha".to_string(),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: String::new(),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "beta".to_string(),
|
||||
},
|
||||
];
|
||||
fn formatted_text_output_reuses_owned_items_under_limit() {
|
||||
let items = vec![text("alpha"), text(""), text("beta")];
|
||||
let policy = TruncationPolicy::Bytes(32);
|
||||
|
||||
let (output, original_token_count) =
|
||||
formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(32));
|
||||
let borrowed = formatted_truncate_function_output_items_with_policy(&items, policy, |_| 0);
|
||||
assert_eq!(borrowed, items);
|
||||
|
||||
assert_eq!(output, items);
|
||||
assert_eq!(original_token_count, None);
|
||||
let allocation = items.as_ptr();
|
||||
let owned = formatted_truncate_function_output_items_with_policy(items, policy, |_| 0);
|
||||
assert_eq!(owned, borrowed);
|
||||
assert_eq!(owned.as_ptr(), allocation);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formatted_truncate_text_content_items_with_policy_preserves_empty_leading_text_behavior() {
|
||||
let items = vec![
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: String::new(),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "abc".to_string(),
|
||||
},
|
||||
];
|
||||
fn formatted_text_output_preserves_empty_leading_text_behavior() {
|
||||
let items = vec![text(""), text("abc")];
|
||||
|
||||
let (output, original_token_count) =
|
||||
formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(0));
|
||||
let output = formatted_truncate_function_output_items_with_policy(
|
||||
items,
|
||||
TruncationPolicy::Bytes(0),
|
||||
|_| 0,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
output,
|
||||
@@ -318,101 +310,63 @@ fn formatted_truncate_text_content_items_with_policy_preserves_empty_leading_tex
|
||||
text: "Warning: truncated output (original token count: 1)\nTotal output lines: 1\n\n…3 chars truncated…".to_string(),
|
||||
}]
|
||||
);
|
||||
assert_eq!(original_token_count, Some(1));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formatted_truncate_text_content_items_with_policy_merges_text_and_appends_media() {
|
||||
fn formatted_mixed_output_budgets_text_and_audio_in_order() {
|
||||
let audio = FunctionCallOutputContentItem::InputAudio {
|
||||
audio_url: "audio:one".to_owned(),
|
||||
};
|
||||
let file_image = FunctionCallOutputContentItem::InputImage {
|
||||
image: ImageReference::File {
|
||||
file_id: "file_123".to_owned(),
|
||||
},
|
||||
detail: Some(DEFAULT_IMAGE_DETAIL),
|
||||
};
|
||||
let items = vec![
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "abcd".to_string(),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputImage {
|
||||
image: ImageReference::Inline {
|
||||
image_url: "img:one".to_string(),
|
||||
},
|
||||
detail: Some(DEFAULT_IMAGE_DETAIL),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "efgh".to_string(),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputAudio {
|
||||
audio_url: "audio:one".to_string(),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputImage {
|
||||
image: ImageReference::File {
|
||||
file_id: "file_123".to_string(),
|
||||
},
|
||||
detail: Some(DEFAULT_IMAGE_DETAIL),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "ijkl".to_string(),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputImage {
|
||||
image: ImageReference::Inline {
|
||||
image_url: "img:two".to_string(),
|
||||
},
|
||||
detail: Some(DEFAULT_IMAGE_DETAIL),
|
||||
},
|
||||
text("abcd"),
|
||||
inline_image("img:one".to_owned()),
|
||||
text("efgh"),
|
||||
audio.clone(),
|
||||
file_image.clone(),
|
||||
text("ijkl"),
|
||||
inline_image("img:two".to_owned()),
|
||||
];
|
||||
|
||||
let (output, original_token_count) =
|
||||
formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(8));
|
||||
let output = formatted_truncate_function_output_items_with_policy(
|
||||
items,
|
||||
TruncationPolicy::Bytes(12),
|
||||
|_| 1,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
output,
|
||||
vec![
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "Warning: truncated output (original token count: 4)\nTotal output lines: 3\n\nabcd…6 chars truncated…ijkl".to_string(),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputImage {
|
||||
image: ImageReference::Inline { image_url: "img:one".to_string() },
|
||||
detail: Some(DEFAULT_IMAGE_DETAIL),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputAudio {
|
||||
audio_url: "audio:one".to_string(),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputImage {
|
||||
image: ImageReference::File {
|
||||
file_id: "file_123".to_string(),
|
||||
},
|
||||
detail: Some(DEFAULT_IMAGE_DETAIL),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputImage {
|
||||
image: ImageReference::Inline { image_url: "img:two".to_string() },
|
||||
detail: Some(DEFAULT_IMAGE_DETAIL),
|
||||
},
|
||||
text("abcd"),
|
||||
inline_image("img:one".to_owned()),
|
||||
text("efgh"),
|
||||
audio,
|
||||
file_image,
|
||||
inline_image("img:two".to_owned()),
|
||||
text("[omitted 1 text items ...]"),
|
||||
]
|
||||
);
|
||||
assert_eq!(original_token_count, Some(4));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formatted_truncate_text_content_items_with_policy_preserves_encrypted_content() {
|
||||
let items = vec![
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "abcdefgh".to_string(),
|
||||
},
|
||||
FunctionCallOutputContentItem::EncryptedContent {
|
||||
encrypted_content: "enc_opaque".to_string(),
|
||||
},
|
||||
];
|
||||
fn formatted_mixed_output_preserves_encrypted_content() {
|
||||
let encrypted = FunctionCallOutputContentItem::EncryptedContent {
|
||||
encrypted_content: "enc_opaque".to_owned(),
|
||||
};
|
||||
let items = vec![text("abcdefgh"), encrypted.clone()];
|
||||
|
||||
let (output, original_token_count) =
|
||||
formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Bytes(2));
|
||||
|
||||
assert_eq!(
|
||||
output,
|
||||
vec![
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "Warning: truncated output (original token count: 2)\nTotal output lines: 1\n\na…6 chars truncated…h".to_string(),
|
||||
},
|
||||
FunctionCallOutputContentItem::EncryptedContent {
|
||||
encrypted_content: "enc_opaque".to_string(),
|
||||
},
|
||||
]
|
||||
let output = formatted_truncate_function_output_items_with_policy(
|
||||
items,
|
||||
TruncationPolicy::Bytes(2),
|
||||
|_| 0,
|
||||
);
|
||||
assert_eq!(original_token_count, Some(2));
|
||||
|
||||
assert_eq!(output, vec![text("a…6 chars truncated…h"), encrypted]);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -475,18 +429,14 @@ fn truncate_function_output_items_with_policy_charges_audio_against_byte_budget(
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn formatted_truncate_text_content_items_with_policy_merges_all_text_for_token_budget() {
|
||||
let items = vec![
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "abcdefgh".to_string(),
|
||||
},
|
||||
FunctionCallOutputContentItem::InputText {
|
||||
text: "ijklmnop".to_string(),
|
||||
},
|
||||
];
|
||||
fn formatted_text_output_merges_all_text_for_token_budget() {
|
||||
let items = vec![text("abcdefgh"), text("ijklmnop")];
|
||||
|
||||
let (output, original_token_count) =
|
||||
formatted_truncate_text_content_items_with_policy(&items, TruncationPolicy::Tokens(2));
|
||||
let output = formatted_truncate_function_output_items_with_policy(
|
||||
items,
|
||||
TruncationPolicy::Tokens(2),
|
||||
|_| 0,
|
||||
);
|
||||
|
||||
assert_eq!(
|
||||
output,
|
||||
@@ -494,7 +444,6 @@ fn formatted_truncate_text_content_items_with_policy_merges_all_text_for_token_b
|
||||
text: "Warning: truncated output (original token count: 5)\nTotal output lines: 2\n\nabcd…3 tokens truncated…mnop".to_string(),
|
||||
}]
|
||||
);
|
||||
assert_eq!(original_token_count, Some(5));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -503,3 +452,113 @@ fn byte_count_conversion_clamps_non_positive_values() {
|
||||
assert_eq!(approx_tokens_from_byte_count_i64(/*bytes*/ 0), 0);
|
||||
assert_eq!(approx_tokens_from_byte_count_i64(/*bytes*/ 5), 2);
|
||||
}
|
||||
|
||||
fn boundaries<const N: usize>(positions: [(usize, usize); N]) -> [OutputBoundary; N] {
|
||||
positions.map(|(item_index, text_offset)| OutputBoundary {
|
||||
item_index,
|
||||
text_offset,
|
||||
})
|
||||
}
|
||||
|
||||
fn positions<const N: usize>(boundaries: [OutputBoundary; N]) -> [(usize, usize); N] {
|
||||
boundaries.map(|b| (b.item_index, b.text_offset))
|
||||
}
|
||||
|
||||
fn text(s: &str) -> FunctionCallOutputContentItem {
|
||||
FunctionCallOutputContentItem::InputText { text: s.to_owned() }
|
||||
}
|
||||
|
||||
fn inline_image(image_url: String) -> FunctionCallOutputContentItem {
|
||||
FunctionCallOutputContentItem::InputImage {
|
||||
image: ImageReference::Inline { image_url },
|
||||
detail: Some(DEFAULT_IMAGE_DETAIL),
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn boundary_formatters_normalize_input_and_remap_truncated_text() {
|
||||
let audio_cost: fn(&str) -> usize = |_| 1;
|
||||
let items = [text("left-head"), text("right-tail")];
|
||||
let max_tokens = 2;
|
||||
let mut boundary = boundaries([(0, 0), (0, 2), (1, 0), (1, 8), (2, 0)]);
|
||||
let output =
|
||||
formatted_output_at_token_boundaries(&items, max_tokens, audio_cost, &mut boundary);
|
||||
let [FunctionCallOutputContentItem::InputText { text: rendered }] = output.as_slice() else {
|
||||
panic!("expected one merged text item")
|
||||
};
|
||||
let head = rendered.find("left").expect("expected retained prefix");
|
||||
let tail = rendered
|
||||
.strip_suffix("tail")
|
||||
.expect("expected retained suffix")
|
||||
.len();
|
||||
assert_eq!(
|
||||
positions(boundary),
|
||||
[
|
||||
(0, 0),
|
||||
(0, head + 2),
|
||||
(0, tail),
|
||||
(0, tail + 2),
|
||||
(0, rendered.len())
|
||||
]
|
||||
);
|
||||
|
||||
for (middle, middle_end) in [(text("after"), 5), (inline_image("img:one".to_owned()), 0)] {
|
||||
let items = [text("🙂é boundary"), middle, text("last")];
|
||||
for formatter in [
|
||||
formatted_output_at_token_boundaries,
|
||||
mixed_output_at_token_boundaries,
|
||||
] {
|
||||
for tokens in [0, 1, 2, 256] {
|
||||
let mut b = boundaries([(0, 1), (1, usize::MAX), (0, 8), (usize::MAX, usize::MAX)]);
|
||||
let output = formatter(&items, tokens, audio_cost, &mut b);
|
||||
assert_eq!(output, formatter(&items, tokens, audio_cost, &mut []));
|
||||
let result = positions(b);
|
||||
assert!(result.is_sorted());
|
||||
assert!(
|
||||
result
|
||||
.iter()
|
||||
.all(|&(index, offset)| match output.get(index) {
|
||||
Some(FunctionCallOutputContentItem::InputText { text }) =>
|
||||
text.is_char_boundary(offset),
|
||||
_ => index <= output.len() && offset == 0,
|
||||
})
|
||||
);
|
||||
if tokens == 256 {
|
||||
assert_eq!(result, [(0, 0), (1, middle_end), (1, middle_end), (3, 0)]);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn owned_mixed_output_moves_a_preserved_image_across_text_truncation() {
|
||||
let image_url = "data:image/png;base64,".to_owned() + &"a".repeat(4096);
|
||||
let url_ptr = image_url.as_ptr();
|
||||
let items = vec![
|
||||
text("long prefix"),
|
||||
inline_image(image_url),
|
||||
text("omitted"),
|
||||
];
|
||||
let mut borrowed = boundaries([(0, 5), (2, 3)]);
|
||||
let mut owned = boundaries([(0, 5), (2, 3)]);
|
||||
let max_tokens = 1;
|
||||
let expected = formatted_output_at_token_boundaries(&items, max_tokens, |_| 0, &mut borrowed);
|
||||
let actual = formatted_output_at_token_boundaries(items, max_tokens, |_| 0, &mut owned);
|
||||
assert_eq!(actual, expected);
|
||||
let projected = positions(owned);
|
||||
assert_eq!(projected, positions(borrowed));
|
||||
let [(index, offset), omitted] = projected;
|
||||
let FunctionCallOutputContentItem::InputText { text: snippet } = &actual[index] else {
|
||||
panic!("expected boundary inside truncated text");
|
||||
};
|
||||
assert_eq!((index, &snippet[offset..], omitted), (0, "ix", (2, 0)));
|
||||
let FunctionCallOutputContentItem::InputImage {
|
||||
image: ImageReference::Inline { image_url },
|
||||
..
|
||||
} = &actual[1]
|
||||
else {
|
||||
panic!("expected preserved image")
|
||||
};
|
||||
assert_eq!(image_url.as_ptr(), url_ptr);
|
||||
}
|
||||
|
||||
@@ -12,6 +12,17 @@ pub struct TextTruncation {
|
||||
pub marker: String,
|
||||
}
|
||||
|
||||
impl TextTruncation {
|
||||
/// Applies this edit to the string it was planned from.
|
||||
pub fn apply(&self, original: &str) -> String {
|
||||
assemble_truncated_output(
|
||||
&original[..self.removed.start],
|
||||
&original[self.removed.end..],
|
||||
&self.marker,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
/// Plans the existing head/tail token truncation. The marker is outside the
|
||||
/// retained-text budget, just as it is in [`truncate_middle_with_token_budget`].
|
||||
pub fn plan_token_truncation(s: &str, max_tokens: usize) -> Option<TextTruncation> {
|
||||
@@ -70,11 +81,7 @@ pub fn truncate_middle_with_token_budget(s: &str, max_tokens: usize) -> (String,
|
||||
let Some(truncation) = plan_token_truncation(s, max_tokens) else {
|
||||
return (s.to_string(), None);
|
||||
};
|
||||
let truncated = assemble_truncated_output(
|
||||
&s[..truncation.removed.start],
|
||||
&s[truncation.removed.end..],
|
||||
&truncation.marker,
|
||||
);
|
||||
let truncated = truncation.apply(s);
|
||||
let total_tokens = u64::try_from(approx_token_count(s)).unwrap_or(u64::MAX);
|
||||
|
||||
if truncated == s {
|
||||
|
||||
Reference in New Issue
Block a user