1use crate::segment::{ContextPriority, ContextSegment, ContextSegmentType};
8use chrono::{DateTime, Utc};
9use enact_core::kernel::ExecutionId;
10use serde::{Deserialize, Serialize};
11use thiserror::Error;
12
13#[derive(Debug, Error)]
15pub enum CompactionError {
16 #[error("Nothing to compact")]
17 NothingToCompact,
18
19 #[error("Target token count too low: {0}")]
20 TargetTooLow(usize),
21
22 #[error("Summarization failed: {0}")]
23 SummarizationFailed(String),
24}
25
26#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
30#[serde(rename_all = "snake_case")]
31pub enum CompactionStrategyType {
32 Truncate,
34 Summarize,
36 ExtractKeyPoints,
38 SlidingWindow,
40 ImportanceWeighted,
42 Hybrid,
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize)]
50#[serde(rename_all = "camelCase")]
51pub struct CompactionStrategy {
52 #[serde(rename = "type")]
54 pub strategy_type: CompactionStrategyType,
55
56 pub target_tokens: usize,
58
59 pub min_preserve_percent: u8,
61
62 #[serde(skip_serializing_if = "Option::is_none")]
64 pub segments_to_compact: Option<Vec<ContextSegmentType>>,
65
66 #[serde(skip_serializing_if = "Option::is_none")]
68 pub protected_segments: Option<Vec<ContextSegmentType>>,
69
70 #[serde(skip_serializing_if = "Option::is_none")]
72 pub summary_max_tokens: Option<usize>,
73
74 #[serde(skip_serializing_if = "Option::is_none")]
76 pub window_size: Option<usize>,
77
78 #[serde(skip_serializing_if = "Option::is_none")]
80 pub min_importance_score: Option<f64>,
81}
82
83impl CompactionStrategy {
84 pub fn truncate(target_tokens: usize) -> Self {
86 Self {
87 strategy_type: CompactionStrategyType::Truncate,
88 target_tokens,
89 min_preserve_percent: 20,
90 segments_to_compact: None,
91 protected_segments: Some(vec![
92 ContextSegmentType::System,
93 ContextSegmentType::UserInput,
94 ]),
95 summary_max_tokens: None,
96 window_size: None,
97 min_importance_score: None,
98 }
99 }
100
101 pub fn sliding_window(target_tokens: usize, window_size: usize) -> Self {
103 Self {
104 strategy_type: CompactionStrategyType::SlidingWindow,
105 target_tokens,
106 min_preserve_percent: 20,
107 segments_to_compact: Some(vec![ContextSegmentType::History]),
108 protected_segments: Some(vec![
109 ContextSegmentType::System,
110 ContextSegmentType::UserInput,
111 ]),
112 summary_max_tokens: None,
113 window_size: Some(window_size),
114 min_importance_score: None,
115 }
116 }
117
118 pub fn summarize(target_tokens: usize, summary_max_tokens: usize) -> Self {
120 Self {
121 strategy_type: CompactionStrategyType::Summarize,
122 target_tokens,
123 min_preserve_percent: 30,
124 segments_to_compact: Some(vec![
125 ContextSegmentType::History,
126 ContextSegmentType::ToolResults,
127 ]),
128 protected_segments: Some(vec![
129 ContextSegmentType::System,
130 ContextSegmentType::UserInput,
131 ContextSegmentType::Guidance,
132 ]),
133 summary_max_tokens: Some(summary_max_tokens),
134 window_size: None,
135 min_importance_score: None,
136 }
137 }
138
139 pub fn is_protected(&self, segment_type: ContextSegmentType) -> bool {
141 self.protected_segments
142 .as_ref()
143 .map(|p| p.contains(&segment_type))
144 .unwrap_or(false)
145 }
146
147 pub fn should_compact(&self, segment_type: ContextSegmentType) -> bool {
149 if self.is_protected(segment_type) {
150 return false;
151 }
152
153 self.segments_to_compact
154 .as_ref()
155 .map(|s| s.contains(&segment_type))
156 .unwrap_or(true) }
158}
159
160#[derive(Debug, Clone, Serialize, Deserialize)]
164#[serde(rename_all = "camelCase")]
165pub struct CompactionResult {
166 pub execution_id: ExecutionId,
168
169 pub strategy: CompactionStrategyType,
171
172 pub tokens_before: usize,
174
175 pub tokens_after: usize,
177
178 pub tokens_saved: usize,
180
181 pub compression_ratio: f64,
183
184 pub segments_compacted: usize,
186
187 pub duration_ms: u64,
189
190 pub success: bool,
192
193 #[serde(skip_serializing_if = "Option::is_none")]
195 pub error: Option<String>,
196
197 pub compacted_at: DateTime<Utc>,
199}
200
201impl CompactionResult {
202 pub fn success(
204 execution_id: ExecutionId,
205 strategy: CompactionStrategyType,
206 tokens_before: usize,
207 tokens_after: usize,
208 segments_compacted: usize,
209 duration_ms: u64,
210 ) -> Self {
211 let tokens_saved = tokens_before.saturating_sub(tokens_after);
212 let compression_ratio = if tokens_before > 0 {
213 tokens_after as f64 / tokens_before as f64
214 } else {
215 1.0
216 };
217
218 Self {
219 execution_id,
220 strategy,
221 tokens_before,
222 tokens_after,
223 tokens_saved,
224 compression_ratio,
225 segments_compacted,
226 duration_ms,
227 success: true,
228 error: None,
229 compacted_at: Utc::now(),
230 }
231 }
232
233 pub fn failure(
235 execution_id: ExecutionId,
236 strategy: CompactionStrategyType,
237 tokens_before: usize,
238 error: String,
239 duration_ms: u64,
240 ) -> Self {
241 Self {
242 execution_id,
243 strategy,
244 tokens_before,
245 tokens_after: tokens_before,
246 tokens_saved: 0,
247 compression_ratio: 1.0,
248 segments_compacted: 0,
249 duration_ms,
250 success: false,
251 error: Some(error),
252 compacted_at: Utc::now(),
253 }
254 }
255}
256
257pub struct Compactor {
259 strategy: CompactionStrategy,
260}
261
262impl Compactor {
263 pub fn new(strategy: CompactionStrategy) -> Self {
265 Self { strategy }
266 }
267
268 pub fn truncate(target_tokens: usize) -> Self {
270 Self::new(CompactionStrategy::truncate(target_tokens))
271 }
272
273 pub fn sliding_window(target_tokens: usize, window_size: usize) -> Self {
275 Self::new(CompactionStrategy::sliding_window(
276 target_tokens,
277 window_size,
278 ))
279 }
280
281 pub fn strategy(&self) -> &CompactionStrategy {
283 &self.strategy
284 }
285
286 pub fn compact_truncate(
290 &self,
291 segments: &mut Vec<ContextSegment>,
292 current_tokens: usize,
293 ) -> Result<usize, CompactionError> {
294 if current_tokens <= self.strategy.target_tokens {
295 return Ok(0);
296 }
297
298 let tokens_to_remove = current_tokens - self.strategy.target_tokens;
299 let mut removed = 0;
300
301 segments.sort_by(|a, b| {
303 a.priority
304 .cmp(&b.priority)
305 .then(a.sequence.cmp(&b.sequence))
306 });
307
308 let mut i = 0;
310 while i < segments.len() && removed < tokens_to_remove {
311 let segment = &segments[i];
312
313 if !segment.compressible || self.strategy.is_protected(segment.segment_type) {
315 i += 1;
316 continue;
317 }
318
319 if segment.priority == ContextPriority::Critical {
321 i += 1;
322 continue;
323 }
324
325 removed += segment.token_count;
326 segments.remove(i);
327 }
328
329 Ok(removed)
330 }
331
332 pub fn compact_sliding_window(
336 &self,
337 segments: &mut Vec<ContextSegment>,
338 ) -> Result<usize, CompactionError> {
339 let window_size = self.strategy.window_size.unwrap_or(10);
340
341 let history_indices: Vec<usize> = segments
343 .iter()
344 .enumerate()
345 .filter(|(_, s)| s.segment_type == ContextSegmentType::History)
346 .map(|(i, _)| i)
347 .collect();
348
349 if history_indices.len() <= window_size {
350 return Ok(0);
351 }
352
353 let to_remove = history_indices.len() - window_size;
355 let mut removed_tokens = 0;
356
357 for &idx in history_indices.iter().take(to_remove).rev() {
359 removed_tokens += segments[idx].token_count;
360 segments.remove(idx);
361 }
362
363 Ok(removed_tokens)
364 }
365}
366
367#[cfg(test)]
368mod tests {
369 use super::*;
370
371 #[test]
372 fn test_truncation_strategy() {
373 let strategy = CompactionStrategy::truncate(1000);
374 assert_eq!(strategy.strategy_type, CompactionStrategyType::Truncate);
375 assert!(strategy.is_protected(ContextSegmentType::System));
376 assert!(!strategy.is_protected(ContextSegmentType::History));
377 }
378
379 #[test]
380 fn test_compaction_result() {
381 let exec_id = ExecutionId::new();
382 let result = CompactionResult::success(
383 exec_id,
384 CompactionStrategyType::Truncate,
385 10000,
386 5000,
387 5,
388 100,
389 );
390
391 assert!(result.success);
392 assert_eq!(result.tokens_saved, 5000);
393 assert!((result.compression_ratio - 0.5).abs() < 0.01);
394 }
395}