mw_memory/
embed.rs

1//! Pluggable text embedding — so similarity can be *semantic*, not just lexical.
2//!
3//! Local-first: the default real implementation talks to **Ollama** on
4//! `localhost` (e.g. `nomic-embed-text`), which is part of the user's machine —
5//! no external API. When no embedder is configured (or it fails), the scorer
6//! falls back to lexical term overlap, so retrieval always works offline.
7
8/// Anything that can turn text into a vector.
9pub trait Embedder: Send + Sync {
10    fn embed(&self, text: &str) -> anyhow::Result<Vec<f32>>;
11    fn name(&self) -> &str;
12}
13
14/// Cosine similarity in [0,1] (negatives clamped to 0 — we only care about
15/// "how related," not "how opposite").
16pub fn cosine(a: &[f32], b: &[f32]) -> f32 {
17    if a.is_empty() || b.is_empty() || a.len() != b.len() {
18        return 0.0;
19    }
20    let (mut dot, mut na, mut nb) = (0.0f32, 0.0f32, 0.0f32);
21    for i in 0..a.len() {
22        dot += a[i] * b[i];
23        na += a[i] * a[i];
24        nb += b[i] * b[i];
25    }
26    if na == 0.0 || nb == 0.0 {
27        0.0
28    } else {
29        (dot / (na.sqrt() * nb.sqrt())).clamp(0.0, 1.0)
30    }
31}
32
33/// Encode an embedding as little-endian f32 bytes (for caching as a SQLite BLOB).
34pub fn vec_to_bytes(v: &[f32]) -> Vec<u8> {
35    let mut out = Vec::with_capacity(v.len() * 4);
36    for f in v {
37        out.extend_from_slice(&f.to_le_bytes());
38    }
39    out
40}
41
42/// Decode little-endian f32 bytes back into an embedding. Returns an empty vec
43/// if the byte length isn't a multiple of 4 (corrupt/foreign blob).
44pub fn bytes_to_vec(b: &[u8]) -> Vec<f32> {
45    if b.len() % 4 != 0 {
46        return Vec::new();
47    }
48    b.chunks_exact(4)
49        .map(|c| f32::from_le_bytes([c[0], c[1], c[2], c[3]]))
50        .collect()
51}
52
53/// Local-first embedder backed by Ollama's `/api/embeddings`.
54/// Run e.g. `ollama pull nomic-embed-text` first.
55pub struct OllamaEmbedder {
56    endpoint: String,
57    model: String,
58}
59
60impl OllamaEmbedder {
61    pub fn new(model: impl Into<String>) -> Self {
62        Self {
63            endpoint: "http://localhost:11434".into(),
64            model: model.into(),
65        }
66    }
67
68    pub fn with_endpoint(mut self, endpoint: impl Into<String>) -> Self {
69        self.endpoint = endpoint.into();
70        self
71    }
72}
73
74impl Default for OllamaEmbedder {
75    fn default() -> Self {
76        Self::new("nomic-embed-text")
77    }
78}
79
80impl Embedder for OllamaEmbedder {
81    fn name(&self) -> &str {
82        "ollama"
83    }
84
85    fn embed(&self, text: &str) -> anyhow::Result<Vec<f32>> {
86        let url = format!("{}/api/embeddings", self.endpoint);
87        let resp = ureq::post(&url)
88            .send_json(ureq::json!({ "model": self.model, "prompt": text }))
89            .map_err(|e| anyhow::anyhow!("ollama request failed: {e}"))?;
90        let body: serde_json::Value = resp
91            .into_json()
92            .map_err(|e| anyhow::anyhow!("ollama response not JSON: {e}"))?;
93        let arr = body
94            .get("embedding")
95            .and_then(|v| v.as_array())
96            .ok_or_else(|| anyhow::anyhow!("ollama response missing 'embedding'"))?;
97        let vec: Vec<f32> = arr
98            .iter()
99            .filter_map(|x| x.as_f64().map(|f| f as f32))
100            .collect();
101        if vec.is_empty() {
102            anyhow::bail!("ollama returned an empty embedding");
103        }
104        Ok(vec)
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
113    fn cosine_basics() {
114        assert!((cosine(&[1.0, 0.0], &[1.0, 0.0]) - 1.0).abs() < 1e-6);
115        assert!(cosine(&[1.0, 0.0], &[0.0, 1.0]).abs() < 1e-6);
116        assert_eq!(cosine(&[], &[1.0]), 0.0);
117        assert_eq!(cosine(&[1.0, 2.0], &[1.0]), 0.0); // dim mismatch
118    }
119
120    #[test]
121    fn vec_bytes_roundtrip() {
122        let v = vec![0.0f32, 1.5, -2.25, 3.125e9, f32::MIN];
123        assert_eq!(bytes_to_vec(&vec_to_bytes(&v)), v);
124        assert_eq!(vec_to_bytes(&v).len(), v.len() * 4);
125        // corrupt length -> empty, no panic
126        assert!(bytes_to_vec(&[1, 2, 3]).is_empty());
127        assert!(bytes_to_vec(&[]).is_empty());
128    }
129}
mw_memory/embed.rs

mw_memory/
embed.rs
