cel_memory_sqlite/
embedder.rs

1//! Embedder trait + reference implementations.
2//!
3//! The [`Embedder`] trait abstracts over local (fastembed/ONNX) and cloud
4//! (OpenAI/Voyage) embedding backends. The [`SqliteMemoryProvider`] takes
5//! a `Box<dyn Embedder>` at construction and uses it for write-time
6//! embedding and retrieval-time query embedding.
7//!
8//! v1 ships two implementations:
9//!
10//! - [`MockEmbedder`] — deterministic small-dim vectors for tests. No
11//!   external dependencies; always available.
12//! - `FastEmbedEmbedder` — real `bge-small-en-v1.5` model (384 dim) via
13//!   the [`fastembed`] crate. Gated behind the `fastembed` feature
14//!   because the model file downloads on first instantiation (~130 MB)
15//!   and onnxruntime is a heavy dep.
16//!
17//! [`SqliteMemoryProvider`]: crate::SqliteMemoryProvider
18//! [`fastembed`]: https://crates.io/crates/fastembed
19
20use async_trait::async_trait;
21
22use crate::error::SqliteMemoryError;
23
24/// Embedding result alias.
25pub type EmbedResult<T> = Result<T, SqliteMemoryError>;
26
27/// An embedder turns text into a fixed-dimension vector.
28///
29/// Implementations must produce vectors of [`Embedder::dim`] length on
30/// every call; the storage layer validates and rejects mismatches.
31#[async_trait]
32pub trait Embedder: Send + Sync {
33    /// Vector dimensionality.
34    fn dim(&self) -> usize;
35
36    /// Stable model identifier (e.g., `"bge-small-en-v1.5"`).
37    fn model_name(&self) -> &str;
38
39    /// Embed one piece of text.
40    async fn embed(&self, text: &str) -> EmbedResult<Vec<f32>>;
41
42    /// Embed a batch of texts. Default implementation calls [`embed`]
43    /// sequentially; production embedders should override for batching.
44    ///
45    /// [`embed`]: Embedder::embed
46    async fn embed_batch(&self, texts: &[String]) -> EmbedResult<Vec<Vec<f32>>> {
47        let mut out = Vec::with_capacity(texts.len());
48        for t in texts {
49            out.push(self.embed(t).await?);
50        }
51        Ok(out)
52    }
53}
54
55/// Deterministic test embedder. Hashes the input text to a small vector
56/// of pseudo-random floats. **Never use in production** — produces
57/// meaningless vectors.
58///
59/// Useful for unit tests of the SQLite layer where we just need *some*
60/// vector to round-trip through `memory_vec`.
61#[derive(Debug, Clone)]
62pub struct MockEmbedder {
63    dim: usize,
64    model: String,
65}
66
67impl MockEmbedder {
68    /// Mock embedder with the default `384` dimension, matching the
69    /// `memory_vec` schema produced by the initial migration.
70    pub fn new() -> Self {
71        Self {
72            dim: 384,
73            model: "mock-384".into(),
74        }
75    }
76
77    /// Mock embedder with an arbitrary dim. Use only for tests that
78    /// override the migration schema.
79    pub fn with_dim(dim: usize) -> Self {
80        Self {
81            dim,
82            model: format!("mock-{dim}"),
83        }
84    }
85}
86
87impl Default for MockEmbedder {
88    fn default() -> Self {
89        Self::new()
90    }
91}
92
93#[async_trait]
94impl Embedder for MockEmbedder {
95    fn dim(&self) -> usize {
96        self.dim
97    }
98
99    fn model_name(&self) -> &str {
100        &self.model
101    }
102
103    async fn embed(&self, text: &str) -> EmbedResult<Vec<f32>> {
104        // Tiny deterministic hash → seed → fill. Not cryptographic; just
105        // makes equal inputs produce equal outputs and slightly differs
106        // across short inputs.
107        let mut seed: u64 = 0xcbf29ce484222325;
108        for b in text.bytes() {
109            seed ^= b as u64;
110            seed = seed.wrapping_mul(0x100000001b3);
111        }
112        let mut out = Vec::with_capacity(self.dim);
113        for i in 0..self.dim {
114            seed = seed
115                .wrapping_mul(6364136223846793005)
116                .wrapping_add(1442695040888963407);
117            // Map to [-1, 1].
118            let f = ((seed >> (i % 32)) as i32) as f32 / i32::MAX as f32;
119            out.push(f.clamp(-1.0, 1.0));
120        }
121        Ok(out)
122    }
123}
124
125#[cfg(test)]
126mod tests {
127    use super::*;
128
129    #[tokio::test]
130    async fn mock_embedder_deterministic() {
131        let e = MockEmbedder::new();
132        let a = e.embed("hello").await.unwrap();
133        let b = e.embed("hello").await.unwrap();
134        assert_eq!(a, b);
135        assert_eq!(a.len(), 384);
136    }
137
138    #[tokio::test]
139    async fn mock_embedder_different_for_different_input() {
140        let e = MockEmbedder::new();
141        let a = e.embed("hello").await.unwrap();
142        let b = e.embed("world").await.unwrap();
143        assert_ne!(a, b);
144    }
145
146    #[tokio::test]
147    async fn mock_embedder_with_dim_honors_dim() {
148        let e = MockEmbedder::with_dim(8);
149        let v = e.embed("hi").await.unwrap();
150        assert_eq!(v.len(), 8);
151        assert_eq!(e.dim(), 8);
152    }
153
154    #[tokio::test]
155    async fn batch_default_works() {
156        let e = MockEmbedder::new();
157        let v = e
158            .embed_batch(&["a".to_string(), "b".to_string()])
159            .await
160            .unwrap();
161        assert_eq!(v.len(), 2);
162        assert_eq!(v[0].len(), 384);
163    }
164}
cel_memory_sqlite/embedder.rs

cel_memory_sqlite/
embedder.rs
