semantic_memory/
embedder.rs

1//! Embedding trait and implementations.
2//!
3//! Provides the [`Embedder`] trait for text-to-vector conversion,
4//! with [`OllamaEmbedder`] (production) and [`MockEmbedder`] (testing).
5
6use crate::config::EmbeddingConfig;
7use crate::error::MemoryError;
8use std::future::Future;
9use std::hash::{Hash, Hasher};
10use std::pin::Pin;
11
12/// Boxed future type alias for single embedding results.
13pub type EmbedFuture<'a> = Pin<Box<dyn Future<Output = Result<Vec<f32>, MemoryError>> + Send + 'a>>;
14
15/// Boxed future type alias for batch embedding results.
16pub type EmbedBatchFuture<'a> =
17    Pin<Box<dyn Future<Output = Result<Vec<Vec<f32>>, MemoryError>> + Send + 'a>>;
18
19/// Trait for embedding text into vectors.
20///
21/// Implement this to swap embedding providers.
22pub trait Embedder: Send + Sync {
23    /// Embed a single text. Returns a vector of f32.
24    fn embed<'a>(&'a self, text: &'a str) -> EmbedFuture<'a>;
25
26    /// Embed multiple texts in a batch.
27    ///
28    /// Takes owned strings to avoid lifetime issues across async boundaries.
29    fn embed_batch<'a>(&'a self, texts: Vec<String>) -> EmbedBatchFuture<'a>;
30
31    /// The model name this embedder uses.
32    fn model_name(&self) -> &str;
33
34    /// Expected embedding dimensions.
35    fn dimensions(&self) -> usize;
36}
37
38// ─── OllamaEmbedder ─────────────────────────────────────────────
39
40/// Embedding provider that calls Ollama's `/api/embed` endpoint.
41pub struct OllamaEmbedder {
42    client: reqwest::Client,
43    base_url: String,
44    model: String,
45    dimensions: usize,
46    batch_size: usize,
47}
48
49impl OllamaEmbedder {
50    /// Create a new OllamaEmbedder from config.
51    ///
52    /// Returns an error if the HTTP client cannot be constructed (e.g. TLS backend
53    /// is unavailable).
54    pub fn try_new(config: &EmbeddingConfig) -> Result<Self, MemoryError> {
55        let client = reqwest::Client::builder()
56            .timeout(std::time::Duration::from_secs(config.timeout_secs))
57            .build()
58            .map_err(|e| {
59                MemoryError::EmbedderUnavailable(format!("failed to build HTTP client: {e}"))
60            })?;
61
62        Ok(Self {
63            client,
64            base_url: config.ollama_url.trim_end_matches('/').to_string(),
65            model: config.model.clone(),
66            dimensions: config.dimensions,
67            batch_size: config.batch_size,
68        })
69    }
70
71    // GOV-005: Deprecated `new()` method removed — all consumers should use `try_new()`.
72}
73
74impl Embedder for OllamaEmbedder {
75    fn embed<'a>(&'a self, text: &'a str) -> EmbedFuture<'a> {
76        Box::pin(async move {
77            let mut results = self.embed_batch(vec![text.to_string()]).await?;
78            results.pop().ok_or_else(|| {
79                MemoryError::Other("Ollama returned empty embeddings for single text".to_string())
80            })
81        })
82    }
83
84    fn embed_batch<'a>(&'a self, texts: Vec<String>) -> EmbedBatchFuture<'a> {
85        Box::pin(async move {
86            let mut all_embeddings = Vec::with_capacity(texts.len());
87
88            for batch in texts.chunks(self.batch_size) {
89                let input: Vec<&str> = batch.iter().map(|s| s.as_str()).collect();
90                let body = serde_json::json!({
91                    "model": self.model,
92                    "input": input
93                });
94
95                let url = format!("{}/api/embed", self.base_url);
96                let response = self
97                    .client
98                    .post(&url)
99                    .json(&body)
100                    .send()
101                    .await
102                    .map_err(|e| {
103                        if e.is_connect() {
104                            MemoryError::EmbedderUnavailable(format!(
105                                "Ollama not running at {}",
106                                self.base_url
107                            ))
108                        } else if e.is_timeout() {
109                            MemoryError::EmbedderUnavailable(format!(
110                                "Ollama embedding timed out: {}",
111                                e
112                            ))
113                        } else {
114                            MemoryError::EmbeddingRequest(e)
115                        }
116                    })?;
117
118                if response.status() == reqwest::StatusCode::NOT_FOUND {
119                    return Err(MemoryError::EmbedderUnavailable(format!(
120                        "Model '{}' not available in Ollama. Run: ollama pull {}",
121                        self.model, self.model
122                    )));
123                }
124
125                if !response.status().is_success() {
126                    let status = response.status();
127                    let body = response.text().await.unwrap_or_default();
128                    return Err(MemoryError::Other(format!(
129                        "Ollama returned HTTP {}: {}",
130                        status,
131                        &body[..body.len().min(500)]
132                    )));
133                }
134
135                let resp_body: serde_json::Value = response.json().await?;
136                let batch_embeddings = parse_embedding_response(&resp_body, self.dimensions)?;
137                all_embeddings.extend(batch_embeddings);
138            }
139
140            Ok(all_embeddings)
141        })
142    }
143
144    fn model_name(&self) -> &str {
145        &self.model
146    }
147
148    fn dimensions(&self) -> usize {
149        self.dimensions
150    }
151}
152
153/// Parse an Ollama embedding response body into vectors.
154///
155/// Validates that all values are numeric and dimensions match.
156#[doc(hidden)]
157pub fn parse_embedding_response(
158    body: &serde_json::Value,
159    expected_dims: usize,
160) -> Result<Vec<Vec<f32>>, MemoryError> {
161    let embeddings = body["embeddings"].as_array().ok_or_else(|| {
162        MemoryError::Other("Ollama response missing 'embeddings' field".to_string())
163    })?;
164
165    let mut result = Vec::with_capacity(embeddings.len());
166    for embedding_val in embeddings {
167        let raw_array = embedding_val
168            .as_array()
169            .ok_or_else(|| MemoryError::Other("Embedding is not an array".to_string()))?;
170
171        let mut embedding = Vec::with_capacity(raw_array.len());
172        for (i, v) in raw_array.iter().enumerate() {
173            let val = v.as_f64().ok_or_else(|| {
174                MemoryError::Other(format!(
175                    "Embedding dimension {} contains non-numeric value: {}",
176                    i, v
177                ))
178            })?;
179            embedding.push(val as f32);
180        }
181
182        if embedding.len() != expected_dims {
183            return Err(MemoryError::DimensionMismatch {
184                expected: expected_dims,
185                actual: embedding.len(),
186            });
187        }
188
189        result.push(embedding);
190    }
191
192    Ok(result)
193}
194
195// ─── MockEmbedder ────────────────────────────────────────────────
196
197/// Deterministic embedder for unit tests.
198///
199/// Generates consistent embeddings based on a hash of the input text.
200/// Same text always produces the same embedding. Output is normalized.
201pub struct MockEmbedder {
202    dimensions: usize,
203}
204
205impl MockEmbedder {
206    /// Create a new MockEmbedder with the given dimensions.
207    pub fn new(dimensions: usize) -> Self {
208        Self { dimensions }
209    }
210}
211
212impl Embedder for MockEmbedder {
213    fn embed<'a>(&'a self, text: &'a str) -> EmbedFuture<'a> {
214        let embedding = deterministic_embedding(text, self.dimensions);
215        Box::pin(async move { Ok(embedding) })
216    }
217
218    fn embed_batch<'a>(&'a self, texts: Vec<String>) -> EmbedBatchFuture<'a> {
219        let embeddings: Vec<Vec<f32>> = texts
220            .iter()
221            .map(|t| deterministic_embedding(t, self.dimensions))
222            .collect();
223        Box::pin(async move { Ok(embeddings) })
224    }
225
226    fn model_name(&self) -> &str {
227        "mock-embedder"
228    }
229
230    fn dimensions(&self) -> usize {
231        self.dimensions
232    }
233}
234
235/// Generate a deterministic embedding from text using a hash-seeded xorshift RNG.
236fn deterministic_embedding(text: &str, dimensions: usize) -> Vec<f32> {
237    let mut hasher = std::hash::DefaultHasher::new();
238    text.hash(&mut hasher);
239    let mut state = hasher.finish();
240    if state == 0 {
241        state = 1;
242    }
243
244    let mut values = Vec::with_capacity(dimensions);
245    for _ in 0..dimensions {
246        // xorshift64
247        state ^= state << 13;
248        state ^= state >> 7;
249        state ^= state << 17;
250        let val = ((state as f64) / (u64::MAX as f64)) * 2.0 - 1.0;
251        values.push(val as f32);
252    }
253
254    // Normalize to unit length
255    let magnitude: f32 = values.iter().map(|v| v * v).sum::<f32>().sqrt();
256    if magnitude > 0.0 {
257        for v in &mut values {
258            *v /= magnitude;
259        }
260    }
261
262    values
263}
semantic_memory/embedder.rs

semantic_memory/
embedder.rs
