sara-tasks 0.2.2

Sara — folder-aware task manager (successor to tk)
use anyhow::{Context, Result};
use serde_json::json;
use std::time::Duration;

use super::{
    EnrichmentRequest, EnrichmentResponse, ItemEnrichmentRequest, LlmProvider, system_prompt,
    user_prompt,
};
use crate::config::LlmConfig;

pub struct OllamaProvider {
    base_url: String,
    model: String,
    timeout: Duration,
}

impl OllamaProvider {
    pub fn new(cfg: &LlmConfig) -> Self {
        OllamaProvider {
            base_url: cfg
                .base_url
                .clone()
                .unwrap_or_else(|| "http://localhost:11434".to_string()),
            model: cfg.model.clone(),
            timeout: Duration::from_secs(cfg.timeout_secs),
        }
    }
}

impl LlmProvider for OllamaProvider {
    fn enrich(&self, req: &EnrichmentRequest) -> Result<EnrichmentResponse> {
        let schema = schemars::schema_for!(EnrichmentResponse);
        let schema_val = serde_json::to_value(&schema)?;

        let body = json!({
            "model": self.model,
            "messages": [
                {"role": "system", "content": system_prompt(req)},
                {"role": "user",   "content": user_prompt(req)}
            ],
            "format": schema_val,
            "options": { "temperature": 0 },
            "stream": false
        });

        let client = reqwest::blocking::Client::builder()
            .timeout(self.timeout)
            .build()?;

        let url = format!("{}/api/chat", self.base_url);
        let resp = client
            .post(&url)
            .json(&body)
            .send()
            .context("Ollama request failed — is Ollama running?")?;

        if resp.status() == 404 {
            anyhow::bail!(
                "Ollama model '{}' not found. Run: ollama pull {}",
                self.model,
                self.model
            );
        }

        let resp_json: serde_json::Value = resp
            .error_for_status()
            .context("Ollama returned an error")?
            .json()
            .context("Ollama response was not valid JSON")?;

        let content = resp_json["message"]["content"]
            .as_str()
            .ok_or_else(|| anyhow::anyhow!("Unexpected Ollama response shape"))?;

        let enrichment: EnrichmentResponse =
            serde_json::from_str(content).context("Could not parse Ollama enrichment JSON")?;
        Ok(enrichment)
    }

    fn enrich_item(&self, req: &ItemEnrichmentRequest) -> Result<super::ItemEnrichmentResponse> {
        use super::{ItemEnrichmentResponse, item_system_prompt, item_user_prompt};
        let schema = schemars::schema_for!(ItemEnrichmentResponse);
        let schema_val = serde_json::to_value(&schema)?;

        let body = json!({
            "model": self.model,
            "messages": [
                {"role": "system", "content": item_system_prompt(req)},
                {"role": "user",   "content": item_user_prompt(req)}
            ],
            "format": schema_val,
            "options": { "temperature": 0 },
            "stream": false
        });

        let client = reqwest::blocking::Client::builder()
            .timeout(self.timeout)
            .build()?;

        let url = format!("{}/api/chat", self.base_url);
        let resp = client
            .post(&url)
            .json(&body)
            .send()
            .context("Ollama request failed — is Ollama running?")?;

        let resp_json: serde_json::Value = resp.error_for_status()?.json()?;
        let content = resp_json["message"]["content"]
            .as_str()
            .ok_or_else(|| anyhow::anyhow!("Unexpected Ollama response shape"))?;

        serde_json::from_str(content).context("Could not parse Ollama item enrichment JSON")
    }

    fn chat(&self, system: &str, user: &str) -> Result<String> {
        let body = json!({
            "model": self.model,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user",   "content": user}
            ],
            "options": { "temperature": 0.3 },
            "stream": false
        });

        let client = reqwest::blocking::Client::builder()
            .timeout(self.timeout)
            .build()?;

        let url = format!("{}/api/chat", self.base_url);
        let resp = client
            .post(&url)
            .json(&body)
            .send()
            .context("Ollama request failed — is Ollama running?")?;

        let resp_json: serde_json::Value = resp.error_for_status()?.json()?;
        resp_json["message"]["content"]
            .as_str()
            .map(|s| s.to_string())
            .ok_or_else(|| anyhow::anyhow!("Unexpected Ollama response shape"))
    }
}