use std::cmp::Ordering;
use std::collections::BTreeSet;
use std::ops::Range;
use super::token_engine::ContextTokenEngine;
use super::units::unit_boundaries;
use crate::lexical::{overlap_count, terms};
use crate::types::message::{Content, ContentPart, Message};
pub struct UtilitySelectionContext<'a> {
pub goal: &'a str,
pub criteria: &'a [String],
pub preserved_refs: &'a [String],
pub active_directives: &'a [String],
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct UtilityUnitScore {
pub range: Range<usize>,
pub tokens: u32,
pub mandatory: bool,
pub goal_overlap: u32,
pub has_unresolved: bool,
pub referenced_later: bool,
pub is_error_or_decision: bool,
pub recency: u32,
pub token_cost: u32,
pub prefix_invalidation_cost: u32,
pub utility: i64,
}
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct UtilityArchivePlan {
pub archived_ranges: Vec<Range<usize>>,
pub retained_ranges: Vec<Range<usize>>,
pub archived_tokens: u32,
pub retained_tokens: u32,
pub scores: Vec<UtilityUnitScore>,
}
pub fn plan_utility_archive(
messages: &[Message],
total_tokens: u32,
target_tokens: u32,
preserve_recent_units: usize,
engine: &ContextTokenEngine,
context: &UtilitySelectionContext<'_>,
) -> UtilityArchivePlan {
let ranges = unit_boundaries(messages);
if ranges.is_empty() {
return UtilityArchivePlan::default();
}
let unit_texts = ranges
.iter()
.map(|range| unit_text(&messages[range.clone()]))
.collect::<Vec<_>>();
let goal_terms = terms(
std::iter::once(context.goal)
.chain(context.criteria.iter().map(String::as_str))
.collect::<Vec<_>>()
.join(" ")
.as_str(),
);
let recent_start = ranges.len().saturating_sub(preserve_recent_units);
let denominator = total_tokens.max(1);
let unit_count = ranges.len().max(1) as u32;
let mut scores = Vec::with_capacity(ranges.len());
for (index, range) in ranges.iter().enumerate() {
let slice = &messages[range.clone()];
let text = &unit_texts[index];
let tokens = slice
.iter()
.map(|message| {
message
.token_count
.unwrap_or_else(|| engine.count_message(message))
})
.sum::<u32>();
let goal_overlap = overlap_count(&terms(text), &goal_terms);
let has_unresolved = has_unresolved(slice, text);
let referenced_later = unit_referenced_later(slice, text, &unit_texts[index + 1..]);
let is_error_or_decision = is_error_or_decision(slice, text);
let dependency = context
.preserved_refs
.iter()
.any(|reference| contains_folded(text, reference))
|| context
.active_directives
.iter()
.any(|directive| directive_dependency(text, directive));
let mandatory = index >= recent_start || has_unresolved || dependency;
let recency = ((index as u64 + 1) * 1_000 / u64::from(unit_count)) as u32;
let token_cost = (u64::from(tokens) * 1_000 / u64::from(denominator)) as u32;
let prefix_invalidation_cost =
((ranges.len() - index) as u64 * 1_000 / u64::from(unit_count)) as u32;
let utility = i64::from(goal_overlap) * 4_000
+ if has_unresolved { 20_000 } else { 0 }
+ if referenced_later { 5_000 } else { 0 }
+ if is_error_or_decision { 6_000 } else { 0 }
+ i64::from(recency) * 2
- i64::from(token_cost) * 2
- i64::from(prefix_invalidation_cost);
scores.push(UtilityUnitScore {
range: range.clone(),
tokens,
mandatory,
goal_overlap,
has_unresolved,
referenced_later,
is_error_or_decision,
recency,
token_cost,
prefix_invalidation_cost,
utility,
});
}
if total_tokens <= target_tokens {
return UtilityArchivePlan {
archived_ranges: Vec::new(),
retained_ranges: ranges,
archived_tokens: 0,
retained_tokens: scores.iter().map(|score| score.tokens).sum(),
scores,
};
}
let mut retained = scores
.iter()
.enumerate()
.filter_map(|(index, score)| score.mandatory.then_some(index))
.collect::<BTreeSet<_>>();
let mut retained_tokens = retained
.iter()
.map(|index| scores[*index].tokens)
.sum::<u32>();
let mut optional = scores
.iter()
.enumerate()
.filter_map(|(index, score)| (!score.mandatory).then_some(index))
.collect::<Vec<_>>();
optional.sort_by(|left, right| compare_density(&scores[*right], &scores[*left]));
for index in optional {
let tokens = scores[index].tokens;
if retained_tokens.saturating_add(tokens) <= target_tokens {
retained.insert(index);
retained_tokens = retained_tokens.saturating_add(tokens);
}
}
let retained_ranges = ranges
.iter()
.enumerate()
.filter_map(|(index, range)| retained.contains(&index).then_some(range.clone()))
.collect::<Vec<_>>();
let archived_ranges = ranges
.iter()
.enumerate()
.filter_map(|(index, range)| (!retained.contains(&index)).then_some(range.clone()))
.collect::<Vec<_>>();
let archived_tokens = scores
.iter()
.enumerate()
.filter_map(|(index, score)| (!retained.contains(&index)).then_some(score.tokens))
.sum();
UtilityArchivePlan {
archived_ranges,
retained_ranges,
archived_tokens,
retained_tokens,
scores,
}
}
fn compare_density(left: &UtilityUnitScore, right: &UtilityUnitScore) -> Ordering {
let left_density = i128::from(left.utility) * i128::from(right.tokens.max(1));
let right_density = i128::from(right.utility) * i128::from(left.tokens.max(1));
left_density
.cmp(&right_density)
.then_with(|| left.utility.cmp(&right.utility))
.then_with(|| left.range.start.cmp(&right.range.start))
}
fn unit_text(messages: &[Message]) -> String {
let mut parts = Vec::new();
for message in messages {
match &message.content {
Content::Text(text) => parts.push(text.clone()),
Content::Parts(content_parts) => {
for part in content_parts {
match part {
ContentPart::Text { text } => parts.push(text.clone()),
ContentPart::ToolResult {
call_id, output, ..
} => parts.push(format!("{call_id} {output}")),
ContentPart::Image { url, .. } => {
parts.push(url.clone().unwrap_or_default())
}
ContentPart::Audio { .. } => parts.push("audio".into()),
}
}
}
}
for call in &message.tool_calls {
parts.push(format!("{} {} {}", call.id, call.name, call.arguments));
}
}
parts.join("\n")
}
fn contains_folded(text: &str, pattern: &str) -> bool {
!pattern.trim().is_empty() && text.to_lowercase().contains(&pattern.to_lowercase())
}
fn directive_dependency(text: &str, directive: &str) -> bool {
if contains_folded(text, directive) {
return true;
}
let directive_terms = terms(directive);
if directive_terms.is_empty() {
return false;
}
let threshold = directive_terms.len().min(2);
terms(text).intersection(&directive_terms).count() >= threshold
}
fn has_unresolved(messages: &[Message], text: &str) -> bool {
let mut opened = BTreeSet::new();
let mut resolved = BTreeSet::new();
for message in messages {
for call in &message.tool_calls {
opened.insert(call.id.to_string());
}
if let Content::Parts(parts) = &message.content {
for part in parts {
if let ContentPart::ToolResult {
call_id, is_error, ..
} = part
{
if *is_error {
return true;
}
resolved.insert(call_id.to_string());
}
}
}
}
opened.iter().any(|call_id| !resolved.contains(call_id))
|| marker(
text,
&[
"unresolved",
"open question",
"retry",
"blocked",
"待确认",
"未解决",
"重试",
"阻塞",
],
)
}
fn is_error_or_decision(messages: &[Message], text: &str) -> bool {
messages.iter().any(|message| {
matches!(&message.content, Content::Parts(parts) if parts.iter().any(|part| matches!(part, ContentPart::ToolResult { is_error: true, .. })))
}) || marker(
text,
&[
"error", "failed", "failure", "exception", "decision", "decided", "must", "should",
"错误", "失败", "异常", "决定", "选择", "必须", "应当",
],
)
}
fn marker(text: &str, markers: &[&str]) -> bool {
markers.iter().any(|marker| contains_folded(text, marker))
}
fn unit_referenced_later(messages: &[Message], text: &str, later: &[String]) -> bool {
let mut references = messages
.iter()
.flat_map(|message| message.tool_calls.iter().map(|call| call.id.to_string()))
.collect::<BTreeSet<_>>();
references.extend(
text.split_whitespace()
.map(|token| token.trim_matches(|character: char| character.is_ascii_punctuation()))
.filter(|token| token.contains('/') || token.contains("://"))
.filter(|token| token.len() > 3)
.map(str::to_string),
);
references.iter().any(|reference| {
later
.iter()
.any(|later_text| contains_folded(later_text, reference))
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::message::{ContentPart, ToolCall};
#[test]
fn unresolved_tool_unit_is_mandatory() {
let mut call = Message::assistant("working");
call.tool_calls.push(ToolCall {
id: "call-1".into(),
name: "read".into(),
arguments: serde_json::json!({"path": "/work/a"}),
});
call.token_count = Some(20);
let mut recent = Message::user("recent");
recent.token_count = Some(20);
let messages = vec![call, recent];
let plan = plan_utility_archive(
&messages,
40,
20,
1,
&ContextTokenEngine::char_approx(),
&UtilitySelectionContext {
goal: "",
criteria: &[],
preserved_refs: &[],
active_directives: &[],
},
);
assert!(plan.scores[0].mandatory);
assert!(plan.scores[0].has_unresolved);
assert_eq!(plan.retained_tokens, 40);
}
#[test]
fn chinese_directive_dependency_requires_bigram_overlap_not_shared_characters() {
let mut unrelated = Message::assistant("我们在文中回顾了天气");
unrelated.token_count = Some(30);
let mut on_topic = Message::user("已按要求保持中文回答");
on_topic.token_count = Some(30);
let mut recent = Message::user("recent");
recent.token_count = Some(10);
let messages = vec![unrelated, on_topic, recent];
let plan = plan_utility_archive(
&messages,
70,
10,
1,
&ContextTokenEngine::char_approx(),
&UtilitySelectionContext {
goal: "",
criteria: &[],
preserved_refs: &[],
active_directives: &["必须用中文回答".into()],
},
);
assert!(
!plan.scores[0].mandatory,
"unrelated Chinese text must not bind to the directive"
);
assert!(
plan.scores[1].mandatory,
"text restating the directive must stay mandatory"
);
}
#[test]
fn preserved_ref_keeps_complete_tool_unit() {
let mut call = Message::assistant("read artifact");
call.tool_calls.push(ToolCall {
id: "call-keep".into(),
name: "read".into(),
arguments: serde_json::json!({}),
});
call.token_count = Some(20);
let mut result = Message::tool(vec![ContentPart::ToolResult {
call_id: "call-keep".into(),
output: "artifact".into(),
is_error: false,
}]);
result.token_count = Some(20);
let messages = vec![call, result];
let plan = plan_utility_archive(
&messages,
40,
0,
0,
&ContextTokenEngine::char_approx(),
&UtilitySelectionContext {
goal: "",
criteria: &[],
preserved_refs: &["call-keep".into()],
active_directives: &[],
},
);
assert!(plan.scores[0].mandatory);
assert_eq!(plan.archived_ranges, Vec::<Range<usize>>::new());
}
}