mirror of
https://github.com/openai/codex.git
synced 2026-10-19 16:17:49 +00:00
Avoid full-string scans in token-budget truncation (#49712)
## Why Token-budget truncation uses byte counts to estimate tokens, but previously scanned the entire string to count characters, including the discarded middle. ## What changed Use UTF-8 boundary lookups to select the retained prefix and suffix for token-budget truncation. Skip character counting for a zero-token budget and defer it to paths that need character counts. ## Testing Add regression coverage for multibyte characters crossing both retained boundaries, checking the truncated output and original token estimate. GitOrigin-RevId: 37fee0c8d4ede8e74b491d3d641b0ea0e9003dd4
This commit is contained in:
@@ -40,9 +40,8 @@ fn truncate_with_byte_estimate(s: &str, max_bytes: usize, use_tokens: bool) -> S
|
||||
return String::new();
|
||||
}
|
||||
|
||||
let total_chars = s.chars().count();
|
||||
|
||||
if max_bytes == 0 {
|
||||
let total_chars = if use_tokens { 0 } else { s.chars().count() };
|
||||
return format_truncation_marker(
|
||||
use_tokens,
|
||||
removed_units(use_tokens, s.len(), total_chars),
|
||||
@@ -55,7 +54,14 @@ fn truncate_with_byte_estimate(s: &str, max_bytes: usize, use_tokens: bool) -> S
|
||||
|
||||
let total_bytes = s.len();
|
||||
let (left_budget, right_budget) = split_budget(max_bytes);
|
||||
let (removed_chars, left, right) = split_string(s, left_budget, right_budget);
|
||||
let (removed_chars, left, right) = if use_tokens {
|
||||
// Token markers depend only on byte counts; avoid scanning the discarded middle.
|
||||
let prefix_end = s.floor_char_boundary(left_budget);
|
||||
let suffix_start = s.ceil_char_boundary(total_bytes - right_budget);
|
||||
(0, &s[..prefix_end], &s[suffix_start..])
|
||||
} else {
|
||||
split_string(s, left_budget, right_budget)
|
||||
};
|
||||
let marker = format_truncation_marker(
|
||||
use_tokens,
|
||||
removed_units(
|
||||
|
||||
@@ -109,6 +109,20 @@ fn truncate_middle_tokens_handles_utf8_content() {
|
||||
assert_eq!(tokens, Some(16));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_middle_tokens_keeps_partial_characters_out_of_both_ends() {
|
||||
for (input, expected) in [
|
||||
("a😀b😀c", "a…2 tokens truncated…c"),
|
||||
("é中😀é", "é…2 tokens truncated…é"),
|
||||
("😀😀😀", "…2 tokens truncated…"),
|
||||
] {
|
||||
assert_eq!(
|
||||
truncate_middle_with_token_budget(input, /*max_tokens*/ 1),
|
||||
(expected.to_string(), Some(3)),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn truncate_middle_bytes_handles_utf8_content() {
|
||||
let s = "😀😀😀😀😀😀😀😀😀😀\nsecond line with text\n";
|
||||
|
||||
Reference in New Issue
Block a user