aether_core/context/
compaction.rs

1use std::sync::Arc;
2
3use tokio_stream::StreamExt;
4
5use llm::{ChatMessage, Context, LlmResponse, StreamingModelProvider, TokenUsage};
6
7const SUMMARIZATION_PROMPT: &str = include_str!("prompts/summarization.md");
8
9/// Result of a compaction operation
10#[derive(Debug, Clone)]
11pub struct CompactionResult {
12    /// The summary text that replaces the compacted messages
13    pub summary: String,
14    /// Number of messages that were removed/compacted
15    pub messages_removed: usize,
16    /// Token usage reported by the summarization LLM call, if any
17    pub usage: Option<TokenUsage>,
18}
19
20/// Errors that can occur during compaction
21#[derive(Debug, Clone, thiserror::Error)]
22pub enum CompactionError {
23    /// The LLM failed to generate a summary
24    #[error("summarization failed: {0}")]
25    SummarizationFailed(String),
26    /// No messages to compact
27    #[error("nothing to compact")]
28    NothingToCompact,
29}
30
31/// Configuration for context compaction
32#[derive(Debug, Clone)]
33pub struct CompactionConfig {
34    /// Threshold (0.0-1.0) at which to trigger compaction
35    pub threshold: f64,
36}
37
38impl Default for CompactionConfig {
39    fn default() -> Self {
40        Self { threshold: super::DEFAULT_COMPACTION_THRESHOLD }
41    }
42}
43
44impl CompactionConfig {
45    /// Create a new compaction config with the given threshold
46    pub fn with_threshold(threshold: f64) -> Self {
47        Self { threshold }
48    }
49}
50
51/// Compacts context by generating an LLM summary
52pub struct Compactor {
53    llm: Arc<dyn StreamingModelProvider>,
54}
55
56impl Compactor {
57    pub fn new(llm: Arc<dyn StreamingModelProvider>) -> Self {
58        Self { llm }
59    }
60
61    /// Generate a structured summary of the conversation.
62    ///
63    /// Takes the context snapshot by value since the caller applies the summary
64    /// to its live context (which may have changed while the summarization
65    /// request was in flight) via [`Context::with_compacted_summary`].
66    pub async fn compact(&self, mut context: Context) -> Result<CompactionResult, CompactionError> {
67        let messages_to_summarize = context.messages_for_summary();
68        if messages_to_summarize.is_empty() {
69            return Err(CompactionError::NothingToCompact);
70        }
71
72        let messages_removed = messages_to_summarize.len();
73
74        context.add_message(ChatMessage::user(format!(
75            "{SUMMARIZATION_PROMPT}\n\nPlease perform a structured handoff of the conversation above."
76        )));
77
78        let mut stream = self.llm.stream_response(&context);
79        let mut summary = String::new();
80        let mut usage = None;
81
82        while let Some(result) = stream.next().await {
83            match result {
84                Ok(LlmResponse::Text { chunk }) => {
85                    summary.push_str(&chunk);
86                }
87                Ok(LlmResponse::Usage { tokens }) => usage = Some(tokens),
88                Ok(LlmResponse::Done { .. }) => break,
89                Ok(LlmResponse::Error { message }) => {
90                    return Err(CompactionError::SummarizationFailed(message));
91                }
92                Err(e) => {
93                    return Err(CompactionError::SummarizationFailed(e.to_string()));
94                }
95                _ => {}
96            }
97        }
98
99        if summary.is_empty() {
100            return Err(CompactionError::SummarizationFailed("LLM returned empty summary".to_string()));
101        }
102
103        Ok(CompactionResult { summary, messages_removed, usage })
104    }
105}
106
107#[cfg(test)]
108mod tests {
109    use super::*;
110
111    #[test]
112    fn test_compaction_config_default() {
113        let config = CompactionConfig::default();
114        assert!((config.threshold - 0.85).abs() < 0.001);
115    }
116
117    #[test]
118    fn test_compaction_config_with_threshold() {
119        let config = CompactionConfig::with_threshold(0.9);
120        assert!((config.threshold - 0.9).abs() < 0.001);
121    }
122
123    #[tokio::test]
124    async fn test_compactor_generates_summary() {
125        use llm::testing::FakeLlmProvider;
126
127        let summary_response = vec![
128            LlmResponse::Start,
129            LlmResponse::text(
130                "## Primary Goal\nTest the compaction feature\n\n## Completed Work\n- Wrote initial tests\n\n## File Changes\n- `src/main.rs` — added entry point\n\n## Key Decisions\n- Use structured handoff — preserves context better\n\n## Current State\nRunning compaction tests\n\n## Next Steps\n1. Verify all tests pass\n\n## Open Questions\n(none)\n\n## Constraints\n(none)",
131            ),
132            LlmResponse::done(),
133        ];
134
135        let fake_llm = Arc::new(FakeLlmProvider::with_single_response(summary_response));
136        let compactor = Compactor::new(fake_llm);
137
138        let context = Context::new(vec![ChatMessage::system("System"), ChatMessage::user("Test message")], vec![]);
139
140        let result = compactor.compact(context).await;
141        assert!(result.is_ok());
142
143        let result = result.unwrap();
144        assert!(result.summary.contains("Primary Goal"));
145        assert!(result.summary.contains("File Changes"));
146        assert!(result.summary.contains("Next Steps"));
147        assert_eq!(result.messages_removed, 1);
148    }
149
150    #[tokio::test]
151    async fn test_compactor_handles_error() {
152        use llm::testing::FakeLlmProvider;
153
154        let error_response = vec![LlmResponse::Error { message: "API error".to_string() }];
155
156        let fake_llm = Arc::new(FakeLlmProvider::with_single_response(error_response));
157        let compactor = Compactor::new(fake_llm);
158
159        let context = Context::new(vec![ChatMessage::system("System"), ChatMessage::user("Test")], vec![]);
160
161        let result = compactor.compact(context).await;
162        assert!(matches!(result, Err(CompactionError::SummarizationFailed(_))));
163    }
164
165    #[tokio::test]
166    async fn test_compactor_empty_context() {
167        use llm::testing::FakeLlmProvider;
168
169        let fake_llm = Arc::new(FakeLlmProvider::with_single_response(vec![]));
170        let compactor = Compactor::new(fake_llm);
171
172        let context = Context::new(vec![ChatMessage::system("System")], vec![]);
173
174        let result = compactor.compact(context).await;
175        assert!(matches!(result, Err(CompactionError::NothingToCompact)));
176    }
177}
aether_core/context/compaction.rs

aether_core/context/
compaction.rs
