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;
16
17// ── Config ────────────────────────────────────────────────────────────────
18
19#[derive(Debug, Clone)]
20pub struct ContextInfo {
21    pub message_count: usize,
22    pub total_chars: usize,
23    pub max_chars: usize,
24    pub compactions_done: usize,
25}
26
27#[derive(Debug, Clone)]
28pub struct CompressResult {
29    pub did_compact: bool,
30    pub messages: Vec<ChatMessage>,
31}
32
33// ── Compactor trait ───────────────────────────────────────────────────────
34
35#[async_trait]
36pub trait Compactor: Send + Sync {
37    fn name(&self) -> &str;
38    fn should_compress(&self, info: &ContextInfo) -> bool;
39    async fn compress(&self, messages: &[ChatMessage], task: Option<&str>) -> CompressResult;
40}
41
42// ── Default compactor ─────────────────────────────────────────────────────
43
44pub struct DefaultCompactor;
45
46impl DefaultCompactor {
47    pub fn new() -> Self {
48        Self
49    }
50}
51
52impl Default for DefaultCompactor {
53    fn default() -> Self {
54        Self::new()
55    }
56}
57
58#[async_trait]
59impl Compactor for DefaultCompactor {
60    fn name(&self) -> &str {
61        "default_compactor"
62    }
63
64    fn should_compress(&self, info: &ContextInfo) -> bool {
65        // Compress when we exceed 75% of max context
66        let threshold = (0.75 * info.max_chars as f64) as usize;
67        info.total_chars > threshold
68    }
69
70    async fn compress(&self, messages: &[ChatMessage], task: Option<&str>) -> CompressResult {
71        let keep_recent = 6;
72
73        if messages.len() <= keep_recent + 2 {
74            return CompressResult {
75                did_compact: false,
76                messages: messages.to_vec(),
77            };
78        }
79
80        let system_msgs: Vec<&ChatMessage> =
81            messages.iter().filter(|m| m.role == "system").collect();
82        let non_system: Vec<&ChatMessage> =
83            messages.iter().filter(|m| m.role != "system").collect();
84
85        if non_system.len() <= keep_recent {
86            return CompressResult {
87                did_compact: false,
88                messages: messages.to_vec(),
89            };
90        }
91
92        let (old_msgs, recent_msgs) = non_system.split_at(non_system.len() - keep_recent);
93
94        let mut summary_input = String::new();
95        for m in old_msgs {
96            let role_label = match m.role.as_str() {
97                "user" => "User",
98                "assistant" | "assistant_tool_use" => "Assistant",
99                "tool_result" => "Tool Result",
100                _ => &m.role,
101            };
102            let content = if m.content.len() > 500 {
103                format!("{}...", &m.content[..500])
104            } else {
105                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}