velesdb_memory/
embedder.rs

1//! Pluggable text → vector embedding.
2//!
3//! The Agent Memory SDK is *bring-your-own-vector*: it never generates
4//! embeddings. This crate mirrors the repo's established pattern (the Python
5//! SDK's `Embedder` protocol, the tauri-rag demo's `fastembed` backend): an
6//! [`Embedder`] trait with a default on-device model and a deterministic,
7//! network-free fallback for tests and air-gapped reproducibility.
8
9#[cfg(feature = "ollama")]
10use serde::Deserialize;
11
12/// Failure produced by an [`Embedder`] backend (e.g. a network-backed embedder
13/// that cannot reach its model). The in-memory [`HashEmbedder`] never fails.
14#[derive(Debug, thiserror::Error)]
15pub enum EmbedError {
16    /// The embedding backend (network, subprocess, …) returned an error.
17    #[error("embedding backend error: {0}")]
18    Backend(String),
19    /// The backend returned an empty embedding vector.
20    #[error("embedding backend returned an empty vector")]
21    Empty,
22}
23
24/// Turns text into a fixed-dimension embedding vector.
25pub trait Embedder {
26    /// Embedding dimension produced by [`Embedder::embed`].
27    fn dimension(&self) -> usize;
28
29    /// Embed `text` into a vector of length [`Embedder::dimension`].
30    ///
31    /// # Errors
32    /// Returns [`EmbedError`] if the backend cannot produce an embedding.
33    fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedError>;
34}
35
36/// Deterministic, network-free embedder (token-hashing into L2-normalized
37/// buckets). Not semantically strong — its purpose is reproducible tests and
38/// offline behavior, exactly like the `fake_embed` used in the repo's
39/// `agent_memory` examples. Swap in a real model (e.g. `fastembed`,
40/// all-MiniLM-L6-v2, 384-dim) for production recall quality.
41#[derive(Debug, Clone)]
42pub struct HashEmbedder {
43    dimension: usize,
44}
45
46impl HashEmbedder {
47    /// Create a [`HashEmbedder`] producing vectors of `dimension` length.
48    /// Use `384` to match the SDK's `DEFAULT_DIMENSION`.
49    #[must_use]
50    pub fn new(dimension: usize) -> Self {
51        Self { dimension }
52    }
53}
54
55impl Embedder for HashEmbedder {
56    fn dimension(&self) -> usize {
57        self.dimension
58    }
59
60    fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedError> {
61        let mut vector = vec![0.0_f32; self.dimension];
62        if self.dimension == 0 {
63            return Ok(vector);
64        }
65        let modulus = self.dimension as u64;
66        for token in text.split_whitespace() {
67            let bucket = usize::try_from(crate::id::stable_id(token) % modulus).unwrap_or(0);
68            vector[bucket] += 1.0;
69        }
70        velesdb_core::simd_native::normalize_inplace_native(&mut vector);
71        Ok(vector)
72    }
73}
74
75/// A boxed, object-safe embedder. Lets a non-generic `MemoryService<DynEmbedder>`
76/// be stored behind a concrete type — the MCP server and the language bindings
77/// both need this, since handler/pyclass types can't carry a generic `E`.
78pub type DynEmbedder = Box<dyn Embedder + Send + Sync>;
79
80/// Forward [`Embedder`] through a box, enabling a non-generic
81/// `MemoryService<DynEmbedder>` for the MCP server and bindings.
82impl<T: Embedder + ?Sized> Embedder for Box<T> {
83    fn dimension(&self) -> usize {
84        (**self).dimension()
85    }
86
87    fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedError> {
88        (**self).embed(text)
89    }
90}
91
92// --- Optional real-recall backend: a local Ollama embeddings endpoint --------
93//
94// Enabled with `--features ollama`. The default build omits this backend (and
95// its HTTP dependency) so the shipped binary stays tiny, zero-dependency, and
96// fully offline. This backend keeps the binary small too: it calls a model the
97// user already runs locally, so the memory still never leaves the machine.
98
99/// Default Ollama base URL.
100#[cfg(feature = "ollama")]
101pub const DEFAULT_OLLAMA_URL: &str = "http://localhost:11434";
102
103/// Default Ollama embedding model (384-dim; `ollama pull all-minilm`).
104#[cfg(feature = "ollama")]
105pub const DEFAULT_OLLAMA_MODEL: &str = "all-minilm";
106
107/// Embeds text through a local Ollama `/api/embeddings` endpoint — real
108/// semantic recall while the model stays on the user's own machine.
109#[cfg(feature = "ollama")]
110#[derive(Debug, Clone)]
111pub struct OllamaEmbedder {
112    base_url: String,
113    model: String,
114    dimension: usize,
115}
116
117#[cfg(feature = "ollama")]
118impl OllamaEmbedder {
119    /// Connect to Ollama at `base_url` using `model`, probing the embedding
120    /// dimension once so it adapts to whatever model is configured.
121    ///
122    /// # Errors
123    /// Returns [`EmbedError`] if Ollama is unreachable or the model does not
124    /// produce embeddings.
125    pub fn new(base_url: impl Into<String>, model: impl Into<String>) -> Result<Self, EmbedError> {
126        let base_url = base_url.into();
127        let model = model.into();
128        let dimension = request_embedding(&base_url, &model, "dimension probe")?.len();
129        if dimension == 0 {
130            return Err(EmbedError::Empty);
131        }
132        Ok(Self {
133            base_url,
134            model,
135            dimension,
136        })
137    }
138}
139
140#[cfg(feature = "ollama")]
141impl Embedder for OllamaEmbedder {
142    fn dimension(&self) -> usize {
143        self.dimension
144    }
145
146    fn embed(&self, text: &str) -> Result<Vec<f32>, EmbedError> {
147        request_embedding(&self.base_url, &self.model, text)
148    }
149}
150
151/// Build the JSON request body for the embeddings endpoint.
152#[cfg(feature = "ollama")]
153fn build_request_body(model: &str, text: &str) -> String {
154    serde_json::json!({ "model": model, "prompt": text }).to_string()
155}
156
157/// Ollama `/api/embeddings` response shape.
158#[cfg(feature = "ollama")]
159#[derive(Deserialize)]
160struct EmbeddingResponse {
161    embedding: Vec<f32>,
162}
163
164/// Parse an embeddings response body into a vector.
165#[cfg(feature = "ollama")]
166fn parse_embedding_response(body: &str) -> Result<Vec<f32>, EmbedError> {
167    let parsed: EmbeddingResponse = serde_json::from_str(body)
168        .map_err(|err| EmbedError::Backend(format!("invalid embeddings response: {err}")))?;
169    if parsed.embedding.is_empty() {
170        return Err(EmbedError::Empty);
171    }
172    Ok(parsed.embedding)
173}
174
175/// Perform one blocking embeddings request against a local Ollama.
176#[cfg(feature = "ollama")]
177fn request_embedding(base_url: &str, model: &str, text: &str) -> Result<Vec<f32>, EmbedError> {
178    let url = format!("{base_url}/api/embeddings");
179    let body = build_request_body(model, text);
180    let response = ureq::post(&url)
181        .set("Content-Type", "application/json")
182        .send_string(&body)
183        .map_err(|err| EmbedError::Backend(format!("ollama request failed: {err}")))?;
184    let payload = response
185        .into_string()
186        .map_err(|err| EmbedError::Backend(format!("reading ollama response failed: {err}")))?;
187    parse_embedding_response(&payload)
188}
189
190#[cfg(all(test, feature = "ollama"))]
191mod ollama_tests {
192    use super::*;
193
194    #[test]
195    fn request_body_carries_model_and_prompt() {
196        let body = build_request_body("all-minilm", "hello world");
197        let json: serde_json::Value = serde_json::from_str(&body).expect("valid json");
198        assert_eq!(json["model"], "all-minilm");
199        assert_eq!(json["prompt"], "hello world");
200    }
201
202    #[test]
203    fn parses_a_well_formed_embedding() {
204        let vector = parse_embedding_response(r#"{"embedding":[0.1,0.2,0.3]}"#).expect("parse");
205        assert_eq!(vector.len(), 3);
206        assert!((vector[0] - 0.1_f32).abs() < f32::EPSILON);
207    }
208
209    #[test]
210    fn rejects_an_empty_embedding() {
211        let parsed = parse_embedding_response(r#"{"embedding":[]}"#);
212        assert!(matches!(parsed, Err(EmbedError::Empty)));
213    }
214
215    #[test]
216    fn rejects_a_malformed_response() {
217        let parsed = parse_embedding_response(r#"{"oops":true}"#);
218        assert!(matches!(parsed, Err(EmbedError::Backend(_))));
219    }
220
221    #[test]
222    #[ignore = "requires a local Ollama with an embedding model (ollama pull all-minilm)"]
223    fn embeds_through_a_running_ollama() {
224        let embedder = OllamaEmbedder::new(DEFAULT_OLLAMA_URL, DEFAULT_OLLAMA_MODEL)
225            .expect("connect to ollama");
226        let vector = embedder
227            .embed("parking_lot avoids lock poisoning")
228            .expect("embed");
229        assert_eq!(vector.len(), embedder.dimension());
230        assert!(vector
231            .iter()
232            .any(|&component| component.abs() > f32::EPSILON));
233    }
234}
velesdb_memory/embedder.rs

velesdb_memory/
embedder.rs
