agent_io/agent/
compaction.rs

1//! Context compaction for long-running conversations
2
3use derive_builder::Builder;
4use serde::{Deserialize, Serialize};
5
6use crate::llm::Message;
7
8/// Default summary prompt for compaction
9pub const DEFAULT_SUMMARY_PROMPT: &str = r#"Summarize the conversation so far, preserving:
101. Key decisions and their reasons
112. Important facts learned
123. Current task state and next steps
134. Any user preferences or constraints
14
15Format the summary as a structured markdown document."#;
16
17/// Compaction configuration
18#[derive(Debug, Clone, Builder)]
19#[builder(pattern = "owned")]
20pub struct CompactionConfig {
21    /// Enable compaction
22    #[builder(default = "true")]
23    pub enabled: bool,
24
25    /// Threshold ratio of context window to trigger compaction
26    #[builder(default = "0.80")]
27    pub threshold_ratio: f32,
28
29    /// Model to use for generating summaries
30    #[builder(setter(into, strip_option), default = "None")]
31    pub model: Option<String>,
32
33    /// Custom summary prompt
34    #[builder(default = "DEFAULT_SUMMARY_PROMPT.to_string()")]
35    pub summary_prompt: String,
36}
37
38impl Default for CompactionConfig {
39    fn default() -> Self {
40        Self {
41            enabled: true,
42            threshold_ratio: 0.80,
43            model: None,
44            summary_prompt: DEFAULT_SUMMARY_PROMPT.to_string(),
45        }
46    }
47}
48
49/// Token usage tracking
50#[derive(Debug, Clone, Default, Serialize, Deserialize)]
51pub struct TokenUsage {
52    pub input_tokens: u64,
53    pub output_tokens: u64,
54    pub cache_creation_tokens: u64,
55    pub cache_read_tokens: u64,
56}
57
58impl TokenUsage {
59    pub fn new() -> Self {
60        Self::default()
61    }
62
63    pub fn total(&self) -> u64 {
64        self.input_tokens + self.output_tokens
65    }
66
67    pub fn add_input(&mut self, tokens: u64, cached: bool) {
68        if cached {
69            self.cache_read_tokens += tokens;
70        } else {
71            self.input_tokens += tokens;
72        }
73    }
74
75    pub fn add_output(&mut self, tokens: u64) {
76        self.output_tokens += tokens;
77    }
78
79    pub fn add_cache_creation(&mut self, tokens: u64) {
80        self.cache_creation_tokens += tokens;
81    }
82}
83
84/// Compaction service
85pub struct CompactionService {
86    config: CompactionConfig,
87    usage: TokenUsage,
88}
89
90impl CompactionService {
91    pub fn new(config: CompactionConfig) -> Self {
92        Self {
93            config,
94            usage: TokenUsage::new(),
95        }
96    }
97
98    /// Check if compaction is needed based on current usage
99    pub fn should_compact(&self, current_tokens: u64, context_window: u64) -> bool {
100        if !self.config.enabled || context_window == 0 {
101            return false;
102        }
103
104        let threshold = (context_window as f32 * self.config.threshold_ratio) as u64;
105        current_tokens >= threshold
106    }
107
108    /// Get the summary prompt
109    pub fn summary_prompt(&self) -> &str {
110        &self.config.summary_prompt
111    }
112
113    /// Update token usage
114    pub fn update_usage(&mut self, usage: &TokenUsage) {
115        self.usage.input_tokens += usage.input_tokens;
116        self.usage.output_tokens += usage.output_tokens;
117        self.usage.cache_creation_tokens += usage.cache_creation_tokens;
118        self.usage.cache_read_tokens += usage.cache_read_tokens;
119    }
120
121    /// Get current usage
122    pub fn get_usage(&self) -> &TokenUsage {
123        &self.usage
124    }
125}
126
127/// Result of compaction
128#[derive(Debug, Clone)]
129pub struct CompactionResult {
130    /// The summary message
131    pub summary: Message,
132    /// Number of messages removed
133    pub messages_removed: usize,
134    /// Tokens saved (approximate)
135    pub tokens_saved: u64,
136}
137
138impl CompactionResult {
139    pub fn new(summary: impl Into<String>, messages_removed: usize, tokens_saved: u64) -> Self {
140        Self {
141            summary: Message::system(summary.into()),
142            messages_removed,
143            tokens_saved,
144        }
145    }
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn test_should_compact() {
154        let config = CompactionConfig::default();
155        let service = CompactionService::new(config);
156
157        // Below threshold
158        assert!(!service.should_compact(50, 100));
159
160        // At threshold
161        assert!(service.should_compact(80, 100));
162
163        // Above threshold
164        assert!(service.should_compact(90, 100));
165    }
166
167    #[test]
168    fn test_disabled_compaction() {
169        let config = CompactionConfig {
170            enabled: false,
171            ..Default::default()
172        };
173        let service = CompactionService::new(config);
174
175        assert!(!service.should_compact(99, 100));
176    }
177
178    #[test]
179    fn test_token_usage() {
180        let mut usage = TokenUsage::new();
181        usage.add_input(100, false);
182        usage.add_input(50, true);
183        usage.add_output(75);
184
185        assert_eq!(usage.input_tokens, 100);
186        assert_eq!(usage.cache_read_tokens, 50);
187        assert_eq!(usage.output_tokens, 75);
188        assert_eq!(usage.total(), 175);
189    }
190}
agent_io/agent/compaction.rs

agent_io/agent/
compaction.rs
