Skip to main content

cel_summarizer/
backend.rs

1//! HTTP completion backends for summarizer providers.
2
3use async_trait::async_trait;
4use cel_memory::SummarizerResult;
5use reqwest::Client;
6use serde::Deserialize;
7
8use crate::prompts;
9
10/// Abstraction over an LLM HTTP client — injectable for tests.
11#[async_trait]
12pub trait CompletionBackend: Send + Sync {
13    /// Run a single completion with system and user prompts.
14    async fn complete(
15        &self,
16        model: &str,
17        system: &str,
18        user: &str,
19        max_tokens: u32,
20    ) -> SummarizerResult<String>;
21}
22
23/// Anthropic Messages API backend.
24pub struct AnthropicBackend {
25    client: Client,
26    api_key: String,
27}
28
29impl AnthropicBackend {
30    /// Construct from [`crate::ANTHROPIC_API_KEY_ENV`].
31    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
114/// Ollama chat API backend.
115pub struct OllamaBackend {
116    client: Client,
117    base_url: String,
118}
119
120impl OllamaBackend {
121    /// Construct using `OLLAMA_BASE_URL` when set.
122    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
183/// Fixed-response backend for unit tests.
184pub struct MockBackend {
185    text: String,
186}
187
188impl MockBackend {
189    /// Return a fixed string for every completion call.
190    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
208/// Shared summarize path used by both provider types.
209pub(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}