use alloc::{
collections::BTreeMap,
format,
string::{String, ToString},
vec::Vec,
};
use super::policy::TruncationPolicy;
pub const DEFAULT_BYTES_PER_TOKEN: usize = 4;
pub const MIN_ADAPTIVE_TOOL_OUTPUT_BYTES: usize = 4 * 1024;
pub const MAX_ADAPTIVE_TOOL_OUTPUT_BYTES: usize = 64 * 1024;
pub const DEFAULT_TOOL_OUTPUT_CONTEXT_NUMERATOR: usize = 3;
pub const DEFAULT_TOOL_OUTPUT_CONTEXT_DENOMINATOR: usize = 10;
#[must_use]
pub const fn compute_token_adaptive_budget(
remaining_context_tokens: usize,
min_bytes: usize,
max_bytes: usize,
) -> usize {
let ceiling = if max_bytes < min_bytes {
min_bytes
} else {
max_bytes
};
let allocated_tokens = remaining_context_tokens
.saturating_mul(DEFAULT_TOOL_OUTPUT_CONTEXT_NUMERATOR)
/ DEFAULT_TOOL_OUTPUT_CONTEXT_DENOMINATOR;
let raw_bytes = allocated_tokens.saturating_mul(DEFAULT_BYTES_PER_TOKEN);
if raw_bytes < min_bytes {
min_bytes
} else if raw_bytes > ceiling {
ceiling
} else {
raw_bytes
}
}
impl TruncationPolicy {
#[must_use]
pub const fn for_context_window(remaining_context_tokens: usize) -> Self {
Self::new(compute_token_adaptive_budget(
remaining_context_tokens,
MIN_ADAPTIVE_TOOL_OUTPUT_BYTES,
MAX_ADAPTIVE_TOOL_OUTPUT_BYTES,
))
}
#[must_use]
pub const fn for_context_window_clamped(
remaining_context_tokens: usize,
min_bytes: usize,
max_bytes: usize,
) -> Self {
Self::new(compute_token_adaptive_budget(
remaining_context_tokens,
min_bytes,
max_bytes,
))
}
}
#[must_use]
pub fn truncate_head_chars(text: &str, max_chars: usize, marker: &str) -> String {
if text.chars().count() <= max_chars {
return text.to_string();
}
let keep = max_chars.saturating_sub(marker.chars().count());
let end = text
.char_indices()
.nth(keep)
.map_or(text.len(), |(idx, _)| idx);
format!("{}{marker}", &text[..end])
}
#[must_use]
pub fn truncate_tail_chars(text: &str, max_chars: usize, prefix: &str) -> String {
let total = text.chars().count();
if total <= max_chars {
return text.to_string();
}
let keep = max_chars.saturating_sub(prefix.chars().count());
let skip = total - keep;
let start = text
.char_indices()
.nth(skip)
.map_or(text.len(), |(idx, _)| idx);
format!("{prefix}{}", &text[start..])
}
#[must_use]
pub fn truncate_middle_chars(
text: &str,
head_chars: usize,
tail_chars: usize,
marker: &str,
) -> String {
let total = text.chars().count();
if total <= head_chars.saturating_add(tail_chars) {
return text.to_string();
}
let head_end = text
.char_indices()
.nth(head_chars)
.map_or(text.len(), |(idx, _)| idx);
let tail_start = text
.char_indices()
.nth(total - tail_chars)
.map_or(text.len(), |(idx, _)| idx);
format!("{}{marker}{}", &text[..head_end], &text[tail_start..])
}
#[must_use]
pub fn canonical_json(value: &serde_json::Value) -> serde_json::Value {
match value {
serde_json::Value::Object(map) => {
let mut btree = BTreeMap::new();
for (k, v) in map {
btree.insert(k.clone(), canonical_json(v));
}
let mut sorted_map = serde_json::Map::with_capacity(btree.len());
for (k, v) in btree {
sorted_map.insert(k, v);
}
serde_json::Value::Object(sorted_map)
}
serde_json::Value::Array(arr) => {
let canonical_arr: Vec<serde_json::Value> = arr.iter().map(canonical_json).collect();
serde_json::Value::Array(canonical_arr)
}
_ => value.clone(),
}
}