cel_summarizer/
backend.rs1use async_trait::async_trait;
4use cel_memory::SummarizerResult;
5use reqwest::Client;
6use serde::Deserialize;
7
8use crate::prompts;
9
10#[async_trait]
12pub trait CompletionBackend: Send + Sync {
13 async fn complete(
15 &self,
16 model: &str,
17 system: &str,
18 user: &str,
19 max_tokens: u32,
20 ) -> SummarizerResult<String>;
21}
22
23pub struct AnthropicBackend {
25 client: Client,
26 api_key: String,
27}
28
29impl AnthropicBackend {
30 pub fn from_env() -> SummarizerResult<Self> {
32 let api_key = std::env::var(crate::ANTHROPIC_API_KEY_ENV).map_err(|_| {
33 cel_memory::SummarizerError::InvalidConfig(format!(
34 "{} is not set",
35 crate::ANTHROPIC_API_KEY_ENV
36 ))
37 })?;
38 if api_key.trim().is_empty() {
39 return Err(cel_memory::SummarizerError::InvalidConfig(format!(
40 "{} is empty",
41 crate::ANTHROPIC_API_KEY_ENV
42 )));
43 }
44 Ok(Self {
45 client: Client::new(),
46 api_key,
47 })
48 }
49}
50
51#[derive(Deserialize)]
52struct AnthropicResponse {
53 content: Vec<AnthropicBlock>,
54}
55
56#[derive(Deserialize)]
57struct AnthropicBlock {
58 #[serde(rename = "type")]
59 kind: String,
60 text: Option<String>,
61}
62
63#[async_trait]
64impl CompletionBackend for AnthropicBackend {
65 async fn complete(
66 &self,
67 model: &str,
68 system: &str,
69 user: &str,
70 max_tokens: u32,
71 ) -> SummarizerResult<String> {
72 let body = serde_json::json!({
73 "model": model,
74 "max_tokens": max_tokens,
75 "system": system,
76 "messages": [{"role": "user", "content": user}]
77 });
78 let resp = self
79 .client
80 .post("https://api.anthropic.com/v1/messages")
81 .header("x-api-key", &self.api_key)
82 .header("anthropic-version", "2023-06-01")
83 .header("content-type", "application/json")
84 .json(&body)
85 .send()
86 .await
87 .map_err(|e| cel_memory::SummarizerError::Provider(e.to_string()))?;
88 if !resp.status().is_success() {
89 let status = resp.status();
90 let text = resp.text().await.unwrap_or_default();
91 return Err(cel_memory::SummarizerError::Provider(format!(
92 "anthropic HTTP {status}: {text}"
93 )));
94 }
95 let parsed: AnthropicResponse = resp
96 .json()
97 .await
98 .map_err(|e| cel_memory::SummarizerError::Provider(e.to_string()))?;
99 let mut out = String::new();
100 for block in parsed.content {
101 if block.kind == "text" {
102 if let Some(text) = block.text {
103 if !out.is_empty() {
104 out.push('\n');
105 }
106 out.push_str(text.trim());
107 }
108 }
109 }
110 Ok(out)
111 }
112}
113
114pub struct OllamaBackend {
116 client: Client,
117 base_url: String,
118}
119
120impl OllamaBackend {
121 pub fn from_env() -> Self {
123 let base_url = std::env::var("OLLAMA_BASE_URL")
124 .unwrap_or_else(|_| "http://127.0.0.1:11434".to_string());
125 Self {
126 client: Client::new(),
127 base_url: base_url.trim_end_matches('/').to_string(),
128 }
129 }
130}
131
132#[derive(Deserialize)]
133struct OllamaResponse {
134 message: OllamaMessage,
135}
136
137#[derive(Deserialize)]
138struct OllamaMessage {
139 content: String,
140}
141
142#[async_trait]
143impl CompletionBackend for OllamaBackend {
144 async fn complete(
145 &self,
146 model: &str,
147 system: &str,
148 user: &str,
149 max_tokens: u32,
150 ) -> SummarizerResult<String> {
151 let url = format!("{}/api/chat", self.base_url);
152 let body = serde_json::json!({
153 "model": model,
154 "stream": false,
155 "options": { "num_predict": max_tokens },
156 "messages": [
157 {"role": "system", "content": system},
158 {"role": "user", "content": user}
159 ]
160 });
161 let resp = self
162 .client
163 .post(url)
164 .json(&body)
165 .send()
166 .await
167 .map_err(|e| cel_memory::SummarizerError::Provider(e.to_string()))?;
168 if !resp.status().is_success() {
169 let status = resp.status();
170 let text = resp.text().await.unwrap_or_default();
171 return Err(cel_memory::SummarizerError::Provider(format!(
172 "ollama HTTP {status}: {text}"
173 )));
174 }
175 let parsed: OllamaResponse = resp
176 .json()
177 .await
178 .map_err(|e| cel_memory::SummarizerError::Provider(e.to_string()))?;
179 Ok(parsed.message.content.trim().to_string())
180 }
181}
182
183pub struct MockBackend {
185 text: String,
186}
187
188impl MockBackend {
189 pub fn with_text(text: impl Into<String>) -> Self {
191 Self { text: text.into() }
192 }
193}
194
195#[async_trait]
196impl CompletionBackend for MockBackend {
197 async fn complete(
198 &self,
199 _model: &str,
200 _system: &str,
201 _user: &str,
202 _max_tokens: u32,
203 ) -> SummarizerResult<String> {
204 Ok(self.text.clone())
205 }
206}
207
208pub(crate) async fn summarize_with_backend(
210 backend: &dyn CompletionBackend,
211 model: &str,
212 chunks: &[cel_memory::MemoryChunk],
213 ctx: &cel_memory::SummaryContext,
214) -> SummarizerResult<String> {
215 if chunks.is_empty() {
216 return Err(cel_memory::SummarizerError::NoInput);
217 }
218 let system = prompts::build_system_prompt(ctx);
219 let user = prompts::build_user_prompt(chunks, ctx);
220 let max_tokens = prompts::max_tokens_for_context(ctx);
221 backend.complete(model, &system, &user, max_tokens).await
222}