use crate::segment::{ContextPriority, ContextSegment, ContextSegmentType};
use chrono::{DateTime, Utc};
use enact_core::kernel::ExecutionId;
use serde::{Deserialize, Serialize};
use thiserror::Error;
#[derive(Debug, Error)]
pub enum CompactionError {
#[error("Nothing to compact")]
NothingToCompact,
#[error("Target token count too low: {0}")]
TargetTooLow(usize),
#[error("Summarization failed: {0}")]
SummarizationFailed(String),
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum CompactionStrategyType {
Truncate,
Summarize,
ExtractKeyPoints,
SlidingWindow,
ImportanceWeighted,
Hybrid,
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CompactionStrategy {
#[serde(rename = "type")]
pub strategy_type: CompactionStrategyType,
pub target_tokens: usize,
pub min_preserve_percent: u8,
#[serde(skip_serializing_if = "Option::is_none")]
pub segments_to_compact: Option<Vec<ContextSegmentType>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub protected_segments: Option<Vec<ContextSegmentType>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub summary_max_tokens: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub window_size: Option<usize>,
#[serde(skip_serializing_if = "Option::is_none")]
pub min_importance_score: Option<f64>,
}
impl CompactionStrategy {
pub fn truncate(target_tokens: usize) -> Self {
Self {
strategy_type: CompactionStrategyType::Truncate,
target_tokens,
min_preserve_percent: 20,
segments_to_compact: None,
protected_segments: Some(vec![
ContextSegmentType::System,
ContextSegmentType::UserInput,
]),
summary_max_tokens: None,
window_size: None,
min_importance_score: None,
}
}
pub fn sliding_window(target_tokens: usize, window_size: usize) -> Self {
Self {
strategy_type: CompactionStrategyType::SlidingWindow,
target_tokens,
min_preserve_percent: 20,
segments_to_compact: Some(vec![ContextSegmentType::History]),
protected_segments: Some(vec![
ContextSegmentType::System,
ContextSegmentType::UserInput,
]),
summary_max_tokens: None,
window_size: Some(window_size),
min_importance_score: None,
}
}
pub fn summarize(target_tokens: usize, summary_max_tokens: usize) -> Self {
Self {
strategy_type: CompactionStrategyType::Summarize,
target_tokens,
min_preserve_percent: 30,
segments_to_compact: Some(vec![
ContextSegmentType::History,
ContextSegmentType::ToolResults,
]),
protected_segments: Some(vec![
ContextSegmentType::System,
ContextSegmentType::UserInput,
ContextSegmentType::Guidance,
]),
summary_max_tokens: Some(summary_max_tokens),
window_size: None,
min_importance_score: None,
}
}
pub fn is_protected(&self, segment_type: ContextSegmentType) -> bool {
self.protected_segments
.as_ref()
.map(|p| p.contains(&segment_type))
.unwrap_or(false)
}
pub fn should_compact(&self, segment_type: ContextSegmentType) -> bool {
if self.is_protected(segment_type) {
return false;
}
self.segments_to_compact
.as_ref()
.map(|s| s.contains(&segment_type))
.unwrap_or(true) }
}
#[derive(Debug, Clone, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct CompactionResult {
pub execution_id: ExecutionId,
pub strategy: CompactionStrategyType,
pub tokens_before: usize,
pub tokens_after: usize,
pub tokens_saved: usize,
pub compression_ratio: f64,
pub segments_compacted: usize,
pub duration_ms: u64,
pub success: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
pub compacted_at: DateTime<Utc>,
}
impl CompactionResult {
pub fn success(
execution_id: ExecutionId,
strategy: CompactionStrategyType,
tokens_before: usize,
tokens_after: usize,
segments_compacted: usize,
duration_ms: u64,
) -> Self {
let tokens_saved = tokens_before.saturating_sub(tokens_after);
let compression_ratio = if tokens_before > 0 {
tokens_after as f64 / tokens_before as f64
} else {
1.0
};
Self {
execution_id,
strategy,
tokens_before,
tokens_after,
tokens_saved,
compression_ratio,
segments_compacted,
duration_ms,
success: true,
error: None,
compacted_at: Utc::now(),
}
}
pub fn failure(
execution_id: ExecutionId,
strategy: CompactionStrategyType,
tokens_before: usize,
error: String,
duration_ms: u64,
) -> Self {
Self {
execution_id,
strategy,
tokens_before,
tokens_after: tokens_before,
tokens_saved: 0,
compression_ratio: 1.0,
segments_compacted: 0,
duration_ms,
success: false,
error: Some(error),
compacted_at: Utc::now(),
}
}
}
pub struct Compactor {
strategy: CompactionStrategy,
}
impl Compactor {
pub fn new(strategy: CompactionStrategy) -> Self {
Self { strategy }
}
pub fn truncate(target_tokens: usize) -> Self {
Self::new(CompactionStrategy::truncate(target_tokens))
}
pub fn sliding_window(target_tokens: usize, window_size: usize) -> Self {
Self::new(CompactionStrategy::sliding_window(
target_tokens,
window_size,
))
}
pub fn strategy(&self) -> &CompactionStrategy {
&self.strategy
}
pub fn compact_truncate(
&self,
segments: &mut Vec<ContextSegment>,
current_tokens: usize,
) -> Result<usize, CompactionError> {
if current_tokens <= self.strategy.target_tokens {
return Ok(0);
}
let tokens_to_remove = current_tokens - self.strategy.target_tokens;
let mut removed = 0;
segments.sort_by(|a, b| {
a.priority
.cmp(&b.priority)
.then(a.sequence.cmp(&b.sequence))
});
let mut i = 0;
while i < segments.len() && removed < tokens_to_remove {
let segment = &segments[i];
if !segment.compressible || self.strategy.is_protected(segment.segment_type) {
i += 1;
continue;
}
if segment.priority == ContextPriority::Critical {
i += 1;
continue;
}
removed += segment.token_count;
segments.remove(i);
}
Ok(removed)
}
pub fn compact_sliding_window(
&self,
segments: &mut Vec<ContextSegment>,
) -> Result<usize, CompactionError> {
let window_size = self.strategy.window_size.unwrap_or(10);
let history_indices: Vec<usize> = segments
.iter()
.enumerate()
.filter(|(_, s)| s.segment_type == ContextSegmentType::History)
.map(|(i, _)| i)
.collect();
if history_indices.len() <= window_size {
return Ok(0);
}
let to_remove = history_indices.len() - window_size;
let mut removed_tokens = 0;
for &idx in history_indices.iter().take(to_remove).rev() {
removed_tokens += segments[idx].token_count;
segments.remove(idx);
}
Ok(removed_tokens)
}
pub fn summarize(target_tokens: usize, summary_max_tokens: usize) -> Self {
Self::new(CompactionStrategy::summarize(
target_tokens,
summary_max_tokens,
))
}
pub fn compact_summarize(
&self,
segments: &mut Vec<ContextSegment>,
current_tokens: usize,
) -> Result<usize, CompactionError> {
if current_tokens <= self.strategy.target_tokens {
return Ok(0);
}
let summary_max = self.strategy.summary_max_tokens.unwrap_or(500);
let mut removed_tokens = 0;
let segments_to_compact = self
.strategy
.segments_to_compact
.clone()
.unwrap_or_else(|| vec![ContextSegmentType::History, ContextSegmentType::ToolResults]);
for segment_type in segments_to_compact {
let mut indices_to_summarize: Vec<usize> = Vec::new();
let mut combined_content = String::new();
let mut total_tokens_in_group = 0;
for (i, segment) in segments.iter().enumerate() {
if segment.segment_type == segment_type
&& segment.compressible
&& !self.strategy.is_protected(segment.segment_type)
&& segment.priority != ContextPriority::Critical
{
indices_to_summarize.push(i);
if !combined_content.is_empty() {
combined_content.push_str("\n---\n");
}
combined_content.push_str(&segment.content);
total_tokens_in_group += segment.token_count;
}
}
if indices_to_summarize.is_empty() || total_tokens_in_group <= summary_max {
continue;
}
let summary = self.extract_key_content(&combined_content, summary_max);
let summary_tokens = summary.len() / 4;
for &idx in indices_to_summarize.iter().rev() {
removed_tokens += segments[idx].token_count;
segments.remove(idx);
}
if !summary.is_empty() {
let summarized_segment = ContextSegment::new(
segment_type,
format!(
"[Summarized {}]\n{}",
segment_type_display(segment_type),
summary
),
summary_tokens,
0, )
.with_priority(ContextPriority::Low);
segments.push(summarized_segment);
removed_tokens = removed_tokens.saturating_sub(summary_tokens);
}
}
if removed_tokens == 0 {
return Err(CompactionError::NothingToCompact);
}
Ok(removed_tokens)
}
fn extract_key_content(&self, text: &str, max_tokens: usize) -> String {
let sentences: Vec<&str> = text
.split(&['.', '!', '?', '\n'][..])
.map(|s| s.trim())
.filter(|s| !s.is_empty() && s.len() > 10)
.collect();
if sentences.is_empty() {
return String::new();
}
let mut scored_sentences: Vec<(usize, &str, i32)> = sentences
.iter()
.enumerate()
.map(|(i, &s)| (i, s, self.score_sentence(s, i, sentences.len())))
.collect();
scored_sentences.sort_by(|a, b| b.2.cmp(&a.2));
let max_chars = max_tokens * 4; let mut summary_parts: Vec<(usize, &str)> = Vec::new();
let mut current_len = 0;
for (idx, sentence, _score) in scored_sentences {
if current_len + sentence.len() + 2 > max_chars {
break;
}
summary_parts.push((idx, sentence));
current_len += sentence.len() + 2;
}
summary_parts.sort_by_key(|(idx, _)| *idx);
summary_parts
.iter()
.map(|(_, s)| *s)
.collect::<Vec<_>>()
.join(". ")
+ "."
}
fn score_sentence(&self, sentence: &str, position: usize, total: usize) -> i32 {
let mut score = 0i32;
let lower = sentence.to_lowercase();
if position == 0 {
score += 10; }
if position == total - 1 {
score += 8; }
let important_keywords = [
("result", 5),
("output", 4),
("error", 6),
("success", 5),
("fail", 6),
("complete", 4),
("return", 3),
("created", 3),
("found", 3),
("important", 4),
("note", 3),
("warning", 5),
("summary", 4),
("conclusion", 5),
("decision", 4),
("because", 3),
("therefore", 3),
];
for (keyword, keyword_score) in important_keywords {
if lower.contains(keyword) {
score += keyword_score;
}
}
let len = sentence.len();
if len < 20 {
score -= 2;
} else if len > 200 {
score -= 1;
}
if sentence.contains('`') || sentence.contains("()") || sentence.contains("::") {
score += 2;
}
score
}
pub fn compact_extract_key_points(
&self,
_segments: &mut Vec<ContextSegment>,
_current_tokens: usize,
) -> Result<usize, CompactionError> {
Err(CompactionError::SummarizationFailed(
"ExtractKeyPoints strategy is not yet implemented. \
This strategy requires LLM integration for semantic key point extraction. \
Consider using 'Summarize' for extractive summarization or \
'SlidingWindow' for recency-based compaction."
.to_string(),
))
}
pub fn compact_importance_weighted(
&self,
_segments: &mut Vec<ContextSegment>,
_current_tokens: usize,
) -> Result<usize, CompactionError> {
Err(CompactionError::SummarizationFailed(
"ImportanceWeighted strategy is not yet implemented. \
This strategy requires embedding model integration for semantic importance scoring. \
Consider using 'Truncate' for priority-based removal or \
'SlidingWindow' for recency-based compaction."
.to_string(),
))
}
pub fn compact_hybrid(
&self,
_segments: &mut Vec<ContextSegment>,
_current_tokens: usize,
) -> Result<usize, CompactionError> {
Err(CompactionError::SummarizationFailed(
"Hybrid strategy is not yet implemented. \
This strategy combines multiple compaction approaches for optimal results. \
Consider using individual strategies: 'Summarize', 'Truncate', or 'SlidingWindow'."
.to_string(),
))
}
}
fn segment_type_display(segment_type: ContextSegmentType) -> &'static str {
match segment_type {
ContextSegmentType::System => "System",
ContextSegmentType::History => "Conversation History",
ContextSegmentType::WorkingMemory => "Working Memory",
ContextSegmentType::ToolResults => "Tool Results",
ContextSegmentType::RagContext => "Retrieved Context",
ContextSegmentType::UserInput => "User Input",
ContextSegmentType::AgentScratchpad => "Agent Notes",
ContextSegmentType::ChildSummary => "Child Execution",
ContextSegmentType::Guidance => "Guidance",
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_truncation_strategy() {
let strategy = CompactionStrategy::truncate(1000);
assert_eq!(strategy.strategy_type, CompactionStrategyType::Truncate);
assert!(strategy.is_protected(ContextSegmentType::System));
assert!(!strategy.is_protected(ContextSegmentType::History));
}
#[test]
fn test_compaction_result() {
let exec_id = ExecutionId::new();
let result = CompactionResult::success(
exec_id,
CompactionStrategyType::Truncate,
10000,
5000,
5,
100,
);
assert!(result.success);
assert_eq!(result.tokens_saved, 5000);
assert!((result.compression_ratio - 0.5).abs() < 0.01);
}
#[test]
fn test_summarize_strategy_creation() {
let strategy = CompactionStrategy::summarize(5000, 500);
assert_eq!(strategy.strategy_type, CompactionStrategyType::Summarize);
assert_eq!(strategy.target_tokens, 5000);
assert_eq!(strategy.summary_max_tokens, Some(500));
assert!(strategy.is_protected(ContextSegmentType::System));
assert!(strategy.is_protected(ContextSegmentType::UserInput));
assert!(!strategy.is_protected(ContextSegmentType::History));
}
#[test]
fn test_summarize_compactor_creation() {
let compactor = Compactor::summarize(5000, 500);
assert_eq!(
compactor.strategy().strategy_type,
CompactionStrategyType::Summarize
);
}
#[test]
fn test_compact_summarize_no_change_under_target() {
let compactor = Compactor::summarize(10000, 500);
let mut segments = vec![ContextSegment::history(
"Some history content here.",
100,
1,
)];
let result = compactor.compact_summarize(&mut segments, 100);
assert_eq!(result.unwrap(), 0);
assert_eq!(segments.len(), 1);
}
#[test]
fn test_compact_summarize_with_history() {
let compactor = Compactor::summarize(500, 100);
let mut segments = vec![
ContextSegment::system("You are a helpful assistant.", 10),
ContextSegment::history(
"The user asked about Rust programming. We discussed memory safety and ownership. \
The result was a successful explanation. The conclusion is that Rust is great.",
800,
1,
),
ContextSegment::history(
"Then we talked about error handling. The important point is that Result types are used. \
This is a note about the discussion. The output showed various patterns.",
700,
2,
),
];
let current_tokens = 10 + 800 + 700;
let result = compactor.compact_summarize(&mut segments, current_tokens);
assert!(result.is_ok());
let removed = result.unwrap();
assert!(removed > 0, "Should have removed some tokens");
assert!(segments
.iter()
.any(|s| s.segment_type == ContextSegmentType::System));
let summarized = segments
.iter()
.find(|s| s.segment_type == ContextSegmentType::History);
assert!(summarized.is_some());
assert!(summarized.unwrap().content.contains("[Summarized"));
}
#[test]
fn test_compact_summarize_preserves_protected() {
let compactor = Compactor::summarize(100, 50);
let mut segments = vec![
ContextSegment::system("System prompt", 20),
ContextSegment::user_input("User question", 15, 1),
ContextSegment::guidance("Important guidance", 25, 2),
ContextSegment::history("Some history that can be compressed.", 500, 3),
];
let current_tokens = 20 + 15 + 25 + 500;
let _ = compactor.compact_summarize(&mut segments, current_tokens);
assert!(segments
.iter()
.any(|s| s.segment_type == ContextSegmentType::System));
assert!(segments
.iter()
.any(|s| s.segment_type == ContextSegmentType::UserInput));
assert!(segments
.iter()
.any(|s| s.segment_type == ContextSegmentType::Guidance));
}
#[test]
fn test_extract_key_content_prioritizes_important_sentences() {
let compactor = Compactor::summarize(5000, 100);
let text = "This is the first sentence and sets context for the discussion. \
Some filler information here that is not particularly important. \
Another sentence with no real significance to the outcome. \
The result of the operation was successful and completed without errors. \
More random content follows that could be removed. \
Yet another sentence that adds little value to understanding. \
Some additional padding content here. \
In conclusion, this is the summary of our findings.";
let summary = compactor.extract_key_content(text, 20);
assert!(!summary.is_empty());
assert!(
summary.len() < text.len(),
"Summary ({} chars) should be shorter than original ({} chars)",
summary.len(),
text.len()
);
}
#[test]
fn test_extract_key_content_handles_empty() {
let compactor = Compactor::summarize(5000, 100);
let summary = compactor.extract_key_content("", 50);
assert!(summary.is_empty() || summary == ".");
}
#[test]
fn test_score_sentence_keywords() {
let compactor = Compactor::summarize(5000, 100);
let error_sentence = "There was an error in the process";
let normal_sentence = "The weather is nice today";
let error_score = compactor.score_sentence(error_sentence, 1, 5);
let normal_score = compactor.score_sentence(normal_sentence, 1, 5);
assert!(
error_score > normal_score,
"Error sentence should score higher"
);
}
#[test]
fn test_score_sentence_position() {
let compactor = Compactor::summarize(5000, 100);
let sentence = "This is a test sentence";
let first_score = compactor.score_sentence(sentence, 0, 5);
let middle_score = compactor.score_sentence(sentence, 2, 5);
let last_score = compactor.score_sentence(sentence, 4, 5);
assert!(
first_score > middle_score,
"First sentence should score higher"
);
assert!(
last_score > middle_score,
"Last sentence should score higher"
);
}
#[test]
fn test_extract_key_points_not_implemented() {
let strategy = CompactionStrategy {
strategy_type: CompactionStrategyType::ExtractKeyPoints,
target_tokens: 5000,
min_preserve_percent: 20,
segments_to_compact: None,
protected_segments: None,
summary_max_tokens: None,
window_size: None,
min_importance_score: None,
};
let compactor = Compactor::new(strategy);
let mut segments = vec![];
let result = compactor.compact_extract_key_points(&mut segments, 1000);
assert!(result.is_err());
let err = result.unwrap_err();
match err {
CompactionError::SummarizationFailed(msg) => {
assert!(msg.contains("ExtractKeyPoints"));
assert!(msg.contains("not yet implemented"));
}
_ => panic!("Expected SummarizationFailed error"),
}
}
#[test]
fn test_importance_weighted_not_implemented() {
let strategy = CompactionStrategy {
strategy_type: CompactionStrategyType::ImportanceWeighted,
target_tokens: 5000,
min_preserve_percent: 20,
segments_to_compact: None,
protected_segments: None,
summary_max_tokens: None,
window_size: None,
min_importance_score: Some(0.5),
};
let compactor = Compactor::new(strategy);
let mut segments = vec![];
let result = compactor.compact_importance_weighted(&mut segments, 1000);
assert!(result.is_err());
let err = result.unwrap_err();
match err {
CompactionError::SummarizationFailed(msg) => {
assert!(msg.contains("ImportanceWeighted"));
assert!(msg.contains("not yet implemented"));
assert!(msg.contains("embedding model"));
}
_ => panic!("Expected SummarizationFailed error"),
}
}
#[test]
fn test_hybrid_not_implemented() {
let strategy = CompactionStrategy {
strategy_type: CompactionStrategyType::Hybrid,
target_tokens: 5000,
min_preserve_percent: 20,
segments_to_compact: None,
protected_segments: None,
summary_max_tokens: None,
window_size: None,
min_importance_score: None,
};
let compactor = Compactor::new(strategy);
let mut segments = vec![];
let result = compactor.compact_hybrid(&mut segments, 1000);
assert!(result.is_err());
let err = result.unwrap_err();
match err {
CompactionError::SummarizationFailed(msg) => {
assert!(msg.contains("Hybrid"));
assert!(msg.contains("not yet implemented"));
}
_ => panic!("Expected SummarizationFailed error"),
}
}
#[test]
fn test_segment_type_display() {
assert_eq!(
segment_type_display(ContextSegmentType::History),
"Conversation History"
);
assert_eq!(
segment_type_display(ContextSegmentType::ToolResults),
"Tool Results"
);
assert_eq!(segment_type_display(ContextSegmentType::System), "System");
}
}