Skip to main content

klieo_embed_common/
lib.rs

1#![deny(missing_docs)]
2#![deny(rust_2018_idioms)]
3#![deny(rustdoc::broken_intra_doc_links)]
4
5//! Shared [`Embedder`] trait + dummy/fake implementations for klieo
6//! memory backends.
7//!
8//! Before W3.A17 the trait was duplicated byte-for-byte across
9//! `klieo-memory-sqlite::embedder` and `klieo-memory-qdrant::embedder`
10//! — any downstream embedder (Ollama, OpenAI, fastembed) had to
11//! `impl Embedder` twice. This crate is the single home; both memory
12//! crates re-export from here to keep their public APIs source-stable.
13//!
14//! # Features
15//!
16//! - **Default** — `Embedder` trait + `DummyEmbedder` (zero vectors).
17//! - **`test-utils`** — adds `FakeEmbedder` (deterministic per-text
18//!   hashing) for downstream test harnesses.
19//! - **`fastembed`** — adds `FastEmbedEmbedder` (fastembed-rs ONNX CPU
20//!   embeddings). Heavy; pulls the ONNX runtime.
21
22use async_trait::async_trait;
23use klieo_core::error::MemoryError;
24
25#[cfg(feature = "fastembed")]
26mod fastembed;
27#[cfg(feature = "fastembed")]
28pub use fastembed::{FastEmbedEmbedder, DEFAULT_DIM, DEFAULT_MODEL};
29
30/// Compute embeddings for one or more texts.
31///
32/// Output vectors must each be of length [`Embedder::dimension`].
33/// Implementations must be deterministic at the type level — the
34/// dimensionality cannot vary across calls on the same instance.
35#[async_trait]
36pub trait Embedder: Send + Sync {
37    /// Embedding dimensionality. Must be constant for a given
38    /// `Embedder` instance — long-term memory backends reject vectors
39    /// of the wrong length at runtime.
40    fn dimension(&self) -> usize;
41
42    /// Compute one embedding per input text.
43    async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, MemoryError>;
44}
45
46/// Default embedder that returns zero vectors.
47///
48/// Lets long-term memory backends store and retrieve facts but cosine
49/// similarity is always 1.0, so recall is FIFO order, not semantic.
50pub struct DummyEmbedder;
51
52#[async_trait]
53impl Embedder for DummyEmbedder {
54    fn dimension(&self) -> usize {
55        384
56    }
57
58    async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, MemoryError> {
59        Ok(texts.iter().map(|_| vec![0.0f32; 384]).collect())
60    }
61}
62
63/// Test-only embedder that hashes each input text into a deterministic
64/// vector. Identical texts produce identical embeddings, so cosine
65/// recall behaves predictably under test.
66#[cfg(any(test, feature = "test-utils"))]
67pub struct FakeEmbedder {
68    dim: usize,
69}
70
71#[cfg(any(test, feature = "test-utils"))]
72impl FakeEmbedder {
73    /// Build a deterministic embedder of the given dimensionality.
74    pub fn new(dim: usize) -> Self {
75        Self { dim }
76    }
77}
78
79#[cfg(any(test, feature = "test-utils"))]
80impl Default for FakeEmbedder {
81    fn default() -> Self {
82        Self::new(8)
83    }
84}
85
86#[cfg(any(test, feature = "test-utils"))]
87#[async_trait]
88impl Embedder for FakeEmbedder {
89    fn dimension(&self) -> usize {
90        self.dim
91    }
92
93    async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, MemoryError> {
94        let dim = self.dim;
95        Ok(texts
96            .iter()
97            .map(|text| {
98                // Deterministic per-text vector via FNV-1a per slot —
99                // toolchain-stable across rustc/std hasher upgrades.
100                let mut v = vec![0.0f32; dim];
101                let bytes = text.as_bytes();
102                for (i, slot) in v.iter_mut().enumerate() {
103                    const FNV_OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
104                    const FNV_PRIME: u64 = 0x0000_0001_0000_01b3;
105                    let mut h: u64 = FNV_OFFSET;
106                    h ^= i as u64;
107                    h = h.wrapping_mul(FNV_PRIME);
108                    for &b in bytes {
109                        h ^= b as u64;
110                        h = h.wrapping_mul(FNV_PRIME);
111                    }
112                    *slot = (h as f32 / u64::MAX as f32) - 0.5;
113                }
114                let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
115                if norm > 0.0 {
116                    for x in &mut v {
117                        *x /= norm;
118                    }
119                }
120                v
121            })
122            .collect())
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129
130    #[tokio::test]
131    async fn dummy_returns_zero_vectors_of_dim_384() {
132        let e = DummyEmbedder;
133        assert_eq!(e.dimension(), 384);
134        let out = e.embed(&["a".into(), "b".into()]).await.unwrap();
135        assert_eq!(out.len(), 2);
136        assert_eq!(out[0].len(), 384);
137        assert!(out[0].iter().all(|x| *x == 0.0));
138    }
139
140    #[tokio::test]
141    async fn fake_embedder_is_deterministic() {
142        let e = FakeEmbedder::new(16);
143        let a = e.embed(&["hello".into()]).await.unwrap();
144        let b = e.embed(&["hello".into()]).await.unwrap();
145        assert_eq!(a, b);
146    }
147
148    #[tokio::test]
149    async fn fake_embedder_distinguishes_inputs() {
150        let e = FakeEmbedder::new(16);
151        let a = e.embed(&["alpha".into()]).await.unwrap();
152        let b = e.embed(&["beta".into()]).await.unwrap();
153        assert_ne!(a, b);
154    }
155
156    #[tokio::test]
157    async fn fake_embedder_outputs_unit_vectors() {
158        let e = FakeEmbedder::new(8);
159        let v = e.embed(&["hello".into()]).await.unwrap();
160        let norm: f32 = v[0].iter().map(|x| x * x).sum::<f32>().sqrt();
161        assert!(
162            (norm - 1.0).abs() < 1e-5,
163            "fake embedder must produce unit vectors, got norm={norm}"
164        );
165    }
166}