Skip to main content

apollo/agent/
compaction.rs

1//! Pluggable context compaction for long conversations.
2//!
3//! Compactor trait: decides when to compact + performs the compaction.
4//! DefaultCompactor: summarization-based using a fast model.
5//!
6//! Inspired by hermes-agent context_engine.py ABC pattern.
7//!
8//! pontytail: single compactor, no plugin discovery. Trait-based so plugins
9//! can register via the PluginRegistry when multi-engine support is needed.
10
11use std::sync::Arc;
12
13use async_trait::async_trait;
14
15use crate::providers::ChatMessage;
16use crate::text::truncate_chars_counted;
17
18// ── Config ────────────────────────────────────────────────────────────────
19
20#[derive(Debug, Clone)]
21pub struct ContextInfo {
22    pub message_count: usize,
23    pub total_chars: usize,
24    pub max_chars: usize,
25    pub compactions_done: usize,
26}
27
28#[derive(Debug, Clone)]
29pub struct CompressResult {
30    pub did_compact: bool,
31    pub messages: Vec<ChatMessage>,
32}
33
34// ── Compactor trait ───────────────────────────────────────────────────────
35
36#[async_trait]
37pub trait Compactor: Send + Sync {
38    fn name(&self) -> &str;
39    fn should_compress(&self, info: &ContextInfo) -> bool;
40    async fn compress(&self, messages: &[ChatMessage], task: Option<&str>) -> CompressResult;
41}
42
43// ── Default compactor ─────────────────────────────────────────────────────
44
45pub struct DefaultCompactor;
46
47impl DefaultCompactor {
48    pub fn new() -> Self {
49        Self
50    }
51}
52
53impl Default for DefaultCompactor {
54    fn default() -> Self {
55        Self::new()
56    }
57}
58
59#[async_trait]
60impl Compactor for DefaultCompactor {
61    fn name(&self) -> &str {
62        "default_compactor"
63    }
64
65    fn should_compress(&self, info: &ContextInfo) -> bool {
66        // Compress when we exceed 75% of max context
67        let threshold = (0.75 * info.max_chars as f64) as usize;
68        info.total_chars > threshold
69    }
70
71    async fn compress(&self, messages: &[ChatMessage], task: Option<&str>) -> CompressResult {
72        let keep_recent = 6;
73
74        if messages.len() <= keep_recent + 2 {
75            return CompressResult {
76                did_compact: false,
77                messages: messages.to_vec(),
78            };
79        }
80
81        let system_msgs: Vec<&ChatMessage> =
82            messages.iter().filter(|m| m.role == "system").collect();
83        let non_system: Vec<&ChatMessage> =
84            messages.iter().filter(|m| m.role != "system").collect();
85
86        if non_system.len() <= keep_recent {
87            return CompressResult {
88                did_compact: false,
89                messages: messages.to_vec(),
90            };
91        }
92
93        let (old_msgs, recent_msgs) = non_system.split_at(non_system.len() - keep_recent);
94
95        let mut summary_input = String::new();
96        for m in old_msgs {
97            let role_label = match m.role.as_str() {
98                "user" => "User",
99                "assistant" | "assistant_tool_use" => "Assistant",
100                "tool_result" => "Tool Result",
101                _ => &m.role,
102            };
103            let content = match truncate_chars_counted(&m.content, 500) {
104                Some((head, _)) => format!("{head}..."),
105                None => m.content.clone(),
106            };
107            summary_input.push_str(&format!("[{}]: {}\n", role_label, content));
108        }
109
110        let original_task = task.unwrap_or("unknown");
111
112        let compact_prompt = format!(
113            "Summarize this conversation concisely. Original task: \"{}\"\n\n\
114             Focus on: what was accomplished, key results, what's still pending.\n\n\
115             Conversation:\n{}",
116            original_task, summary_input
117        );
118
119        let _compact_msgs = [ChatMessage {
120            role: "user".into(),
121            content: compact_prompt,
122            tool_use_id: None,
123        }];
124
125        // TODO: wire in provider for LLM-based summarization.
126        // For now, fall back to simple truncation summary.
127        let summary = format!(
128            "[Compacted — {} earlier messages. Prior context summarized.]",
129            old_msgs.len()
130        );
131
132        let mut compacted = Vec::new();
133        for sm in &system_msgs {
134            compacted.push((*sm).clone());
135        }
136        compacted.push(ChatMessage {
137            role: "user".into(),
138            content: format!(
139                "[Compacted — {} earlier messages summarized]\n\n{}",
140                old_msgs.len(),
141                summary
142            ),
143            tool_use_id: None,
144        });
145        compacted.push(ChatMessage {
146            role: "assistant".into(),
147            content: "Understood, continuing from summary.".into(),
148            tool_use_id: None,
149        });
150        for rm in recent_msgs {
151            compacted.push((*rm).clone());
152        }
153
154        CompressResult {
155            did_compact: old_msgs.len() > 2,
156            messages: compacted,
157        }
158    }
159}
160
161// ── Convenience factory ───────────────────────────────────────────────────
162
163pub fn default_compactor() -> Arc<dyn Compactor> {
164    Arc::new(DefaultCompactor::new())
165}