cel-summarizer 0.2.0

//! HTTP completion backends for summarizer providers.

use async_trait::async_trait;
use cel_memory::SummarizerResult;
use reqwest::Client;
use serde::Deserialize;

use crate::prompts;

/// Abstraction over an LLM HTTP client — injectable for tests.
#[async_trait]
pub trait CompletionBackend: Send + Sync {
    /// Run a single completion with system and user prompts.
    async fn complete(
        &self,
        model: &str,
        system: &str,
        user: &str,
        max_tokens: u32,
    ) -> SummarizerResult<String>;
}

/// Anthropic Messages API backend.
pub struct AnthropicBackend {
    client: Client,
    api_key: String,
}

impl AnthropicBackend {
    /// Construct from [`crate::ANTHROPIC_API_KEY_ENV`].
    pub fn from_env() -> SummarizerResult<Self> {
        let api_key = std::env::var(crate::ANTHROPIC_API_KEY_ENV).map_err(|_| {
            cel_memory::SummarizerError::InvalidConfig(format!(
                "{} is not set",
                crate::ANTHROPIC_API_KEY_ENV
            ))
        })?;
        if api_key.trim().is_empty() {
            return Err(cel_memory::SummarizerError::InvalidConfig(format!(
                "{} is empty",
                crate::ANTHROPIC_API_KEY_ENV
            )));
        }
        Ok(Self {
            client: Client::new(),
            api_key,
        })
    }
}

#[derive(Deserialize)]
struct AnthropicResponse {
    content: Vec<AnthropicBlock>,
}

#[derive(Deserialize)]
struct AnthropicBlock {
    #[serde(rename = "type")]
    kind: String,
    text: Option<String>,
}

#[async_trait]
impl CompletionBackend for AnthropicBackend {
    async fn complete(
        &self,
        model: &str,
        system: &str,
        user: &str,
        max_tokens: u32,
    ) -> SummarizerResult<String> {
        let body = serde_json::json!({
            "model": model,
            "max_tokens": max_tokens,
            "system": system,
            "messages": [{"role": "user", "content": user}]
        });
        let resp = self
            .client
            .post("https://api.anthropic.com/v1/messages")
            .header("x-api-key", &self.api_key)
            .header("anthropic-version", "2023-06-01")
            .header("content-type", "application/json")
            .json(&body)
            .send()
            .await
            .map_err(|e| cel_memory::SummarizerError::Provider(e.to_string()))?;
        if !resp.status().is_success() {
            let status = resp.status();
            let text = resp.text().await.unwrap_or_default();
            return Err(cel_memory::SummarizerError::Provider(format!(
                "anthropic HTTP {status}: {text}"
            )));
        }
        let parsed: AnthropicResponse = resp
            .json()
            .await
            .map_err(|e| cel_memory::SummarizerError::Provider(e.to_string()))?;
        let mut out = String::new();
        for block in parsed.content {
            if block.kind == "text" {
                if let Some(text) = block.text {
                    if !out.is_empty() {
                        out.push('\n');
                    }
                    out.push_str(text.trim());
                }
            }
        }
        Ok(out)
    }
}

/// Ollama chat API backend.
pub struct OllamaBackend {
    client: Client,
    base_url: String,
}

impl OllamaBackend {
    /// Construct using `OLLAMA_BASE_URL` when set.
    pub fn from_env() -> Self {
        let base_url = std::env::var("OLLAMA_BASE_URL")
            .unwrap_or_else(|_| "http://127.0.0.1:11434".to_string());
        Self {
            client: Client::new(),
            base_url: base_url.trim_end_matches('/').to_string(),
        }
    }
}

#[derive(Deserialize)]
struct OllamaResponse {
    message: OllamaMessage,
}

#[derive(Deserialize)]
struct OllamaMessage {
    content: String,
}

#[async_trait]
impl CompletionBackend for OllamaBackend {
    async fn complete(
        &self,
        model: &str,
        system: &str,
        user: &str,
        max_tokens: u32,
    ) -> SummarizerResult<String> {
        let url = format!("{}/api/chat", self.base_url);
        let body = serde_json::json!({
            "model": model,
            "stream": false,
            "options": { "num_predict": max_tokens },
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user}
            ]
        });
        let resp = self
            .client
            .post(url)
            .json(&body)
            .send()
            .await
            .map_err(|e| cel_memory::SummarizerError::Provider(e.to_string()))?;
        if !resp.status().is_success() {
            let status = resp.status();
            let text = resp.text().await.unwrap_or_default();
            return Err(cel_memory::SummarizerError::Provider(format!(
                "ollama HTTP {status}: {text}"
            )));
        }
        let parsed: OllamaResponse = resp
            .json()
            .await
            .map_err(|e| cel_memory::SummarizerError::Provider(e.to_string()))?;
        Ok(parsed.message.content.trim().to_string())
    }
}

/// Fixed-response backend for unit tests.
pub struct MockBackend {
    text: String,
}

impl MockBackend {
    /// Return a fixed string for every completion call.
    pub fn with_text(text: impl Into<String>) -> Self {
        Self { text: text.into() }
    }
}

#[async_trait]
impl CompletionBackend for MockBackend {
    async fn complete(
        &self,
        _model: &str,
        _system: &str,
        _user: &str,
        _max_tokens: u32,
    ) -> SummarizerResult<String> {
        Ok(self.text.clone())
    }
}

/// Shared summarize path used by both provider types.
pub(crate) async fn summarize_with_backend(
    backend: &dyn CompletionBackend,
    model: &str,
    chunks: &[cel_memory::MemoryChunk],
    ctx: &cel_memory::SummaryContext,
) -> SummarizerResult<String> {
    if chunks.is_empty() {
        return Err(cel_memory::SummarizerError::NoInput);
    }
    let system = prompts::build_system_prompt(ctx);
    let user = prompts::build_user_prompt(chunks, ctx);
    let max_tokens = prompts::max_tokens_for_context(ctx);
    backend.complete(model, &system, &user, max_tokens).await
}