1#![deny(missing_docs)]
10#![warn(rust_2018_idioms)]
11
12mod backend;
13mod prompts;
14
15use std::sync::Arc;
16
17use async_trait::async_trait;
18use cel_memory::{
19 offdevice_hook::{OffdeviceCallDescriptor, OffdeviceCallHook, OffdeviceDecision},
20 summarizer::{Summarizer, SummarizerError, SummarizerResult, SummaryContext},
21 MemoryChunk,
22};
23
24pub use backend::{AnthropicBackend, CompletionBackend, MockBackend, OllamaBackend};
25
26pub const DEFAULT_ANTHROPIC_MODEL: &str = "claude-haiku-4-5";
28
29pub const DEFAULT_OLLAMA_MODEL: &str = "llama3.2:3b-instruct-q4_K_M";
31
32pub const PROVIDER_ENV: &str = "CEL_SUMMARIZER_PROVIDER";
35
36pub const ANTHROPIC_API_KEY_ENV: &str = "ANTHROPIC_API_KEY";
38
39pub const MODEL_ENV: &str = "CEL_SUMMARIZER_MODEL";
41
42pub struct AnthropicSummarizer {
44 name: String,
45 model: String,
46 backend: Arc<dyn CompletionBackend>,
47 offdevice_hook: Option<Arc<dyn OffdeviceCallHook>>,
48}
49
50impl AnthropicSummarizer {
51 pub fn new(backend: Arc<dyn CompletionBackend>, model: impl Into<String>) -> Self {
53 let model: String = model.into();
54 let name = format!("anthropic:{model}");
55 Self {
56 name,
57 model,
58 backend,
59 offdevice_hook: None,
60 }
61 }
62
63 pub fn with_offdevice_hook(mut self, hook: Arc<dyn OffdeviceCallHook>) -> Self {
65 self.offdevice_hook = Some(hook);
66 self
67 }
68
69 pub fn from_env() -> SummarizerResult<Self> {
72 let backend = Arc::new(AnthropicBackend::from_env()?);
73 let model =
74 std::env::var(MODEL_ENV).unwrap_or_else(|_| DEFAULT_ANTHROPIC_MODEL.to_string());
75 Ok(Self::new(backend, model))
76 }
77}
78
79#[async_trait]
80impl Summarizer for AnthropicSummarizer {
81 fn name(&self) -> &str {
82 &self.name
83 }
84
85 async fn summarize(
86 &self,
87 chunks: &[MemoryChunk],
88 ctx: &SummaryContext,
89 ) -> SummarizerResult<String> {
90 if chunks.is_empty() {
91 return Err(SummarizerError::NoInput);
92 }
93 if let Some(hook) = &self.offdevice_hook {
94 let mut desc = OffdeviceCallDescriptor::new(
95 "summarizer",
96 "anthropic",
97 &self.model,
98 "memory_summarizer",
99 );
100 if let Some(label) = &ctx.kind_label {
101 desc = desc.with_metadata("kind_label", label.clone());
102 }
103 desc = desc.with_metadata("chunk_count", chunks.len() as i64);
104 match hook.before_call(&desc).await {
105 OffdeviceDecision::Allow => {}
106 OffdeviceDecision::Veto { reason } => {
107 return Err(SummarizerError::Provider(format!(
108 "off-device call vetoed: {reason}"
109 )));
110 }
111 }
112 }
113 backend::summarize_with_backend(self.backend.as_ref(), &self.model, chunks, ctx).await
114 }
115}
116
117pub struct OllamaSummarizer {
119 name: String,
120 model: String,
121 backend: Arc<dyn CompletionBackend>,
122}
123
124impl OllamaSummarizer {
125 pub fn new(backend: Arc<dyn CompletionBackend>, model: impl Into<String>) -> Self {
127 let model: String = model.into();
128 let name = format!("ollama:{model}");
129 Self {
130 name,
131 model,
132 backend,
133 }
134 }
135
136 pub fn from_env() -> SummarizerResult<Self> {
138 let backend = Arc::new(OllamaBackend::from_env());
139 let model = std::env::var(MODEL_ENV).unwrap_or_else(|_| DEFAULT_OLLAMA_MODEL.to_string());
140 Ok(Self::new(backend, model))
141 }
142}
143
144#[async_trait]
145impl Summarizer for OllamaSummarizer {
146 fn name(&self) -> &str {
147 &self.name
148 }
149
150 async fn summarize(
151 &self,
152 chunks: &[MemoryChunk],
153 ctx: &SummaryContext,
154 ) -> SummarizerResult<String> {
155 backend::summarize_with_backend(self.backend.as_ref(), &self.model, chunks, ctx).await
156 }
157}
158
159pub fn build_default() -> SummarizerResult<Arc<dyn Summarizer>> {
165 let kind = std::env::var(PROVIDER_ENV)
166 .ok()
167 .map(|s| s.to_lowercase())
168 .unwrap_or_else(|| "anthropic".to_string());
169
170 match kind.as_str() {
171 "ollama" => Ok(Arc::new(OllamaSummarizer::from_env()?)),
172 _ => {
173 let key_present = std::env::var(ANTHROPIC_API_KEY_ENV)
174 .ok()
175 .filter(|s| !s.trim().is_empty())
176 .is_some();
177 if key_present {
178 match AnthropicSummarizer::from_env() {
179 Ok(s) => Ok(Arc::new(s)),
180 Err(_) => Ok(Arc::new(OllamaSummarizer::from_env()?)),
181 }
182 } else {
183 tracing::info!(
184 "{ANTHROPIC_API_KEY_ENV} not set; memory summarizer falling back to Ollama \
185 ({DEFAULT_OLLAMA_MODEL})"
186 );
187 Ok(Arc::new(OllamaSummarizer::from_env()?))
188 }
189 }
190 }
191}
192
193#[cfg(test)]
194mod tests {
195 use super::*;
196 use cel_memory::{ChunkKind, ChunkSource, ClosureOffdeviceHook, MemoryTier};
197 use chrono::Utc;
198 use serde_json::Value;
199
200 fn chunk(id: &str, content: &str) -> MemoryChunk {
201 MemoryChunk {
202 id: id.into(),
203 created_at: Utc::now(),
204 kind: ChunkKind::Chat,
205 tier: MemoryTier::Session,
206 source: ChunkSource::Embedded,
207 session_id: Some("s1".into()),
208 project_root: None,
209 caller_id: "embedded".into(),
210 content: content.into(),
211 metadata: Value::Null,
212 importance: 0.5,
213 pinned: false,
214 shareable: false,
215 superseded_by: None,
216 embedding_model: "mock".into(),
217 embedding_dim: 0,
218 }
219 }
220
221 #[tokio::test]
222 async fn anthropic_summarizer_returns_backend_text() {
223 let backend = Arc::new(MockBackend::with_text("the user closed the chat"));
224 let s = AnthropicSummarizer::new(backend, DEFAULT_ANTHROPIC_MODEL);
225 assert_eq!(s.name(), "anthropic:claude-haiku-4-5");
226 let out = s
227 .summarize(
228 &[chunk("a", "hi"), chunk("b", "bye")],
229 &SummaryContext {
230 kind_label: Some("session".into()),
231 ..Default::default()
232 },
233 )
234 .await
235 .unwrap();
236 assert_eq!(out, "the user closed the chat");
237 }
238
239 #[tokio::test]
240 async fn ollama_summarizer_returns_backend_text() {
241 let backend = Arc::new(MockBackend::with_text("local model summary"));
242 let s = OllamaSummarizer::new(backend, DEFAULT_OLLAMA_MODEL);
243 assert!(s.name().starts_with("ollama:"));
244 let out = s
245 .summarize(&[chunk("a", "x")], &SummaryContext::default())
246 .await
247 .unwrap();
248 assert_eq!(out, "local model summary");
249 }
250
251 #[tokio::test]
252 async fn empty_chunks_errors_no_input() {
253 let backend = Arc::new(MockBackend::with_text("unused"));
254 let s = AnthropicSummarizer::new(backend, DEFAULT_ANTHROPIC_MODEL);
255 let err = s
256 .summarize(&[], &SummaryContext::default())
257 .await
258 .unwrap_err();
259 assert!(matches!(err, SummarizerError::NoInput));
260 }
261
262 #[tokio::test]
263 async fn anthropic_offdevice_veto_short_circuits() {
264 let backend = Arc::new(MockBackend::with_text("should not be called"));
265 let hook: Arc<dyn OffdeviceCallHook> =
266 Arc::new(ClosureOffdeviceHook(|_d: &OffdeviceCallDescriptor| {
267 OffdeviceDecision::Veto {
268 reason: "work hours".into(),
269 }
270 }));
271 let s =
272 AnthropicSummarizer::new(backend, DEFAULT_ANTHROPIC_MODEL).with_offdevice_hook(hook);
273 let err = s
274 .summarize(&[chunk("a", "alpha")], &SummaryContext::default())
275 .await
276 .unwrap_err();
277 match err {
278 SummarizerError::Provider(msg) => {
279 assert!(msg.contains("off-device call vetoed"));
280 }
281 other => panic!("expected Provider, got {other:?}"),
282 }
283 }
284
285 #[tokio::test]
286 async fn build_default_falls_back_to_ollama_when_no_api_key() {
287 let prior_provider = std::env::var(PROVIDER_ENV).ok();
288 let prior_key = std::env::var(ANTHROPIC_API_KEY_ENV).ok();
289 std::env::remove_var(PROVIDER_ENV);
290 std::env::remove_var(ANTHROPIC_API_KEY_ENV);
291
292 let summ = build_default().expect("ollama fallback should construct");
293 assert!(summ.name().starts_with("ollama:"));
294
295 if let Some(v) = prior_provider {
296 std::env::set_var(PROVIDER_ENV, v);
297 }
298 if let Some(v) = prior_key {
299 std::env::set_var(ANTHROPIC_API_KEY_ENV, v);
300 }
301 }
302}