engram/
embedding_ollama.rs

1//! Ollama-backed `EmbeddingProvider` implementation.
2//!
3//! Uses the Ollama `/api/embed` endpoint to embed text via a locally-running
4//! Ollama instance. Defaults to `nomic-embed-text` (768 dimensions).
5
6use crate::embedding::EmbeddingProvider;
7use crate::store::MemoryError;
8use async_trait::async_trait;
9use serde::{Deserialize, Serialize};
10
11// ---------------------------------------------------------------------------
12// Request / response types
13// ---------------------------------------------------------------------------
14
15#[derive(Serialize)]
16struct EmbedRequest<'a> {
17    model: &'a str,
18    input: Vec<&'a str>,
19}
20
21#[derive(Deserialize)]
22struct EmbedResponse {
23    embeddings: Vec<Vec<f32>>,
24}
25
26// ---------------------------------------------------------------------------
27// OllamaEmbeddingProvider
28// ---------------------------------------------------------------------------
29
30/// `EmbeddingProvider` backed by a locally-running Ollama instance.
31///
32/// # Construction
33///
34/// ```no_run
35/// use engram::OllamaEmbeddingProvider;
36///
37/// // Use defaults: http://localhost:11434, nomic-embed-text, 768 dims
38/// let provider = OllamaEmbeddingProvider::new();
39///
40/// // Or supply custom config
41/// let provider = OllamaEmbeddingProvider::with_config(
42///     "http://localhost:11434",
43///     "nomic-embed-text",
44///     768,
45/// );
46/// ```
47pub struct OllamaEmbeddingProvider {
48    base_url: String,
49    model: String,
50    dims: usize,
51    client: reqwest::Client,
52}
53
54impl OllamaEmbeddingProvider {
55    /// Create a provider with default settings:
56    /// - base URL: `http://localhost:11434`
57    /// - model: `nomic-embed-text`
58    /// - dims: `768`
59    pub fn new() -> Self {
60        Self::with_config("http://localhost:11434", "nomic-embed-text", 768)
61    }
62
63    /// Create a provider with explicit configuration.
64    pub fn with_config(base_url: &str, model: &str, dims: usize) -> Self {
65        Self {
66            base_url: base_url.trim_end_matches('/').to_string(),
67            model: model.to_string(),
68            dims,
69            client: reqwest::Client::new(),
70        }
71    }
72}
73
74impl Default for OllamaEmbeddingProvider {
75    fn default() -> Self {
76        Self::new()
77    }
78}
79
80#[async_trait]
81impl EmbeddingProvider for OllamaEmbeddingProvider {
82    async fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, MemoryError> {
83        let url = format!("{}/api/embed", self.base_url);
84        let body = EmbedRequest {
85            model: &self.model,
86            input: texts.to_vec(),
87        };
88
89        let response = self
90            .client
91            .post(&url)
92            .json(&body)
93            .send()
94            .await
95            .map_err(|e| MemoryError::Embedding(format!("Ollama request failed: {e}")))?;
96
97        if !response.status().is_success() {
98            let status = response.status();
99            let text = response
100                .text()
101                .await
102                .unwrap_or_else(|_| "<no body>".to_string());
103            return Err(MemoryError::Embedding(format!(
104                "Ollama returned {status}: {text}"
105            )));
106        }
107
108        let embed_response: EmbedResponse = response
109            .json()
110            .await
111            .map_err(|e| MemoryError::Embedding(format!("Failed to parse Ollama response: {e}")))?;
112
113        Ok(embed_response.embeddings)
114    }
115
116    fn dimensions(&self) -> usize {
117        self.dims
118    }
119}
engram/embedding_ollama.rs

engram/
embedding_ollama.rs
