use anyhow::Result;
pub const DEFAULT_COMPRESSION_THRESHOLD: f64 = 0.5;
pub const MIN_MESSAGES_TO_KEEP: usize = 20;
pub const DEFAULT_TARGET_RATIO: f64 = 0.4;
pub const DEFAULT_COMPRESSOR_MODEL: &str = "claude-3-5-haiku-20241022";
pub const MAX_CONSECUTIVE_FAILURES: u32 = 3;
pub const AUTOCOMPACT_BUFFER_TOKENS: u32 = 13_000;
pub const WARNING_THRESHOLD_BUFFER_TOKENS: u32 = 20_000;
pub const ERROR_THRESHOLD_BUFFER_TOKENS: u32 = 20_000;
pub const MANUAL_COMPACT_BUFFER_TOKENS: u32 = 3_000;
pub const TIME_BASED_MC_GAP_THRESHOLD_MINUTES: u32 = 5;
pub const TIME_BASED_MC_CLEARED_MESSAGE: &str = "[Old tool result content cleared]";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum ThresholdLevel {
Normal,
Warning,
Error,
Blocking,
}
pub fn format_tokens(n: u32) -> String {
if n < 1_000 {
n.to_string()
} else if n < 10_000 {
format!("{:.1}K", n as f64 / 1_000.0)
} else {
format!("{:.0}K", n as f64 / 1_000.0)
}
}
#[derive(Debug, Clone, Default)]
pub struct CompressionBias {
pub preserve_tools: bool,
pub preserve_thinking: bool,
pub preserve_user_questions: bool,
pub compact_long_outputs: bool,
pub aggressive: bool,
pub preserve_keywords: Vec<String>,
}
impl CompressionBias {
pub fn balanced() -> Self {
Self {
preserve_tools: true,
preserve_thinking: false,
preserve_user_questions: true,
compact_long_outputs: false,
aggressive: false,
preserve_keywords: vec![
"决定".to_string(),
"decision".to_string(),
"重要".to_string(),
"important".to_string(),
"关键".to_string(),
"key".to_string(),
],
}
}
pub fn preserve_important() -> Self {
Self {
preserve_tools: true,
preserve_thinking: true,
preserve_user_questions: true,
compact_long_outputs: true,
aggressive: false,
preserve_keywords: vec![
"决定".to_string(),
"decision".to_string(),
"重要".to_string(),
"important".to_string(),
"关键".to_string(),
"key".to_string(),
"完成".to_string(),
"done".to_string(),
"成功".to_string(),
"success".to_string(),
],
}
}
pub fn aggressive() -> Self {
Self {
preserve_tools: false,
preserve_thinking: false,
preserve_user_questions: false,
compact_long_outputs: false,
aggressive: true,
preserve_keywords: vec![],
}
}
pub fn tool_focused() -> Self {
Self {
preserve_tools: true,
preserve_thinking: false,
preserve_user_questions: false,
compact_long_outputs: false,
aggressive: false,
preserve_keywords: vec![
"工具".to_string(),
"tool".to_string(),
"执行".to_string(),
"execute".to_string(),
"文件".to_string(),
"file".to_string(),
],
}
}
pub fn parse(spec: &str) -> Result<Self> {
let spec = spec.trim().to_lowercase();
if spec == "balanced" || spec == "default" || spec.is_empty() {
return Ok(Self::balanced());
}
if spec == "aggressive" {
return Ok(Self::aggressive());
}
if spec == "preserve_important" || spec == "important" {
return Ok(Self::preserve_important());
}
if spec == "tool_focused" || spec == "tools" {
return Ok(Self::tool_focused());
}
let mut bias = Self::default();
for part in spec.split_whitespace() {
if let Some(preserve_list) = part.strip_prefix("preserve:") {
for item in preserve_list.split(',') {
match item.trim() {
"tools" | "tool" => bias.preserve_tools = true,
"thinking" | "think" => bias.preserve_thinking = true,
"user" | "questions" => bias.preserve_user_questions = true,
"compact" | "long" => bias.compact_long_outputs = true,
_ => {}
}
}
} else if let Some(keyword_list) = part.strip_prefix("keywords:") {
bias.preserve_keywords = keyword_list
.split(',')
.map(|k| k.trim().to_string())
.filter(|k| !k.is_empty())
.collect();
} else if part == "aggressive" {
bias.aggressive = true;
}
}
Ok(bias)
}
pub fn format(&self) -> String {
let mut parts: Vec<String> = Vec::new();
if self.preserve_tools {
parts.push("tools".to_string());
}
if self.preserve_thinking {
parts.push("thinking".to_string());
}
if self.preserve_user_questions {
parts.push("user".to_string());
}
if self.compact_long_outputs {
parts.push("compact".to_string());
}
if self.aggressive {
parts.push("aggressive".to_string());
}
if !self.preserve_keywords.is_empty() {
parts.push(format!("keywords:{}", self.preserve_keywords.join(",")));
}
if parts.is_empty() {
"default".to_string()
} else {
parts.join(", ")
}
}
}
#[derive(Debug, Clone)]
pub struct CompressionConfig {
pub threshold: f64,
pub target_ratio: f64,
pub min_preserve_messages: usize,
pub use_summarization: bool,
pub compressor_model: Option<String>,
pub bias: CompressionBias,
}
impl Default for CompressionConfig {
fn default() -> Self {
Self {
threshold: DEFAULT_COMPRESSION_THRESHOLD,
target_ratio: DEFAULT_TARGET_RATIO,
min_preserve_messages: MIN_MESSAGES_TO_KEEP,
use_summarization: true,
compressor_model: None,
bias: CompressionBias::balanced(),
}
}
}
impl CompressionConfig {
pub fn compressor_model_name(&self) -> &str {
self.compressor_model
.as_deref()
.unwrap_or(DEFAULT_COMPRESSOR_MODEL)
}
pub fn calculate_threshold_level(
token_usage: u32,
context_window: u32,
) -> (ThresholdLevel, u32) {
let percent_left = if context_window > 0 {
((context_window - token_usage) as f64 / context_window as f64 * 100.0) as u32
} else {
0
};
let auto_threshold = context_window.saturating_sub(AUTOCOMPACT_BUFFER_TOKENS);
let warning_threshold = auto_threshold.saturating_sub(WARNING_THRESHOLD_BUFFER_TOKENS);
let error_threshold = auto_threshold.saturating_sub(ERROR_THRESHOLD_BUFFER_TOKENS);
let blocking_threshold = context_window.saturating_sub(MANUAL_COMPACT_BUFFER_TOKENS);
let level = if token_usage >= blocking_threshold {
ThresholdLevel::Blocking
} else if token_usage >= error_threshold {
ThresholdLevel::Error
} else if token_usage >= warning_threshold {
ThresholdLevel::Warning
} else {
ThresholdLevel::Normal
};
(level, percent_left.max(0))
}
}
#[derive(Debug, Clone, Default)]
pub struct CircuitBreakerState {
pub consecutive_failures: u32,
pub is_tripped: bool,
pub last_failure_time: Option<u64>,
}
impl CircuitBreakerState {
pub fn new() -> Self {
Self::default()
}
pub fn record_failure(&mut self) -> bool {
self.consecutive_failures += 1;
self.last_failure_time = Some(std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
.unwrap_or_default()
.as_secs());
if self.consecutive_failures >= MAX_CONSECUTIVE_FAILURES {
self.is_tripped = true;
return true;
}
false
}
pub fn record_success(&mut self) {
self.consecutive_failures = 0;
self.is_tripped = false;
self.last_failure_time = None;
}
pub fn should_skip(&self) -> bool {
self.is_tripped
}
pub fn reset(&mut self) {
self.consecutive_failures = 0;
self.is_tripped = false;
self.last_failure_time = None;
}
}