trusty_embedder/
lib.rs

1//! Shared text-embedding abstraction for trusty-* projects.
2//!
3//! Why: trusty-memory and trusty-search both shipped near-identical
4//! `Embedder` traits and `FastEmbedder` implementations, with subtle
5//! drift (cache vs no-cache, sync vs async warmup, `dim()` vs `dimension()`).
6//! Centralising fixes one bug in one place and lets future consumers pick up
7//! the embedder for free.
8//!
9//! What: an async `Embedder` trait with `embed_batch` as the single primitive
10//! (single-text embed is a free helper), plus a production `FastEmbedder`
11//! (fastembed-rs, all-MiniLM-L6-v2, 384-d) with LRU caching and ORT warmup,
12//! and a `MockEmbedder` test double behind the `test-support` feature.
13//!
14//! Test: `cargo test -p trusty-embedder` covers shape, cache hits, and the
15//! mock embedder. ONNX-backed tests are `#[ignore]` to keep CI under one
16//! cargo-feature umbrella.
17
18use std::num::NonZeroUsize;
19use std::sync::Arc;
20
21use anyhow::{Context, Result};
22use async_trait::async_trait;
23use fastembed::{EmbeddingModel, TextEmbedding, TextInitOptions};
24use lru::LruCache;
25use parking_lot::Mutex;
26
27/// Output dimension of the all-MiniLM-L6-v2 model.
28///
29/// Note: we now load the INT8-quantised variant (`AllMiniLML6V2Q`) which
30/// produces identical 384-dim vectors but runs ~3-4× faster on CPU ONNX
31/// and ships as a ~22MB file (vs 86MB for the f32 model).
32pub const EMBED_DIM: usize = 384;
33
34/// Default LRU cache capacity. Picked to be large enough to keep the
35/// hot working set of repeat queries in memory but small enough that the
36/// cache itself fits well inside L2/L3 on a typical developer machine.
37pub const DEFAULT_CACHE_CAPACITY: usize = 256;
38
39/// Abstraction over embedding backends.
40///
41/// Why: Decouple consumers from any one model so we can swap in remote APIs,
42/// quantised models, or deterministic mocks without changing call sites.
43/// What: a single primitive — `embed_batch` — plus a dimension accessor.
44/// Single-text callers should use the [`embed_one`] convenience helper.
45/// Test: covered by `FastEmbedder` and `MockEmbedder` tests below.
46#[async_trait]
47pub trait Embedder: Send + Sync {
48    /// Embed a batch of texts. Returns one `Vec<f32>` per input, each of
49    /// length `self.dimension()`. An empty input batch returns an empty Vec.
50    async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>>;
51
52    /// Output dimension of the produced embeddings.
53    fn dimension(&self) -> usize;
54}
55
56/// Convenience helper: embed a single text via `embed_batch` and return the
57/// lone vector.
58///
59/// Why: Most call sites only need one embedding at a time and writing
60/// `.embed_batch(&[text]).await?.into_iter().next()` everywhere is noise.
61/// What: builds a 1-element batch, calls `embed_batch`, returns the first
62/// vector (or errors if the embedder produced nothing).
63/// Test: covered indirectly by `mock_embedder_round_trip`.
64pub async fn embed_one(embedder: &dyn Embedder, text: &str) -> Result<Vec<f32>> {
65    let mut v = embedder.embed_batch(&[text.to_string()]).await?;
66    v.pop()
67        .context("embedder returned no embedding for non-empty input")
68}
69
70/// Local CPU embedder backed by fastembed-rs (ONNX runtime, all-MiniLM-L6-v2).
71///
72/// Why: Default to local-only embeddings so consumers have zero external
73/// network dependency and predictable latency. The LRU cache keeps the hot
74/// path free of redundant ONNX work for repeat strings (queries, common
75/// chunks).
76/// What: wraps a single `TextEmbedding` behind a `parking_lot::Mutex` (the
77/// underlying `embed` requires `&mut self`) and an `LruCache<String, Vec<f32>>`.
78/// Initialisation warms the ORT graph with a small batch so the first user
79/// query doesn't pay the one-shot compile cost.
80/// Test: `embed_batch_returns_correct_dim` and `cache_hit_is_idempotent`
81/// (marked `#[ignore]` — they download a real model).
82pub struct FastEmbedder {
83    model: Arc<Mutex<TextEmbedding>>,
84    cache: Arc<Mutex<LruCache<String, Vec<f32>>>>,
85    dim: usize,
86}
87
88impl FastEmbedder {
89    /// Construct a new `FastEmbedder` with the default cache size.
90    pub async fn new() -> Result<Self> {
91        Self::with_cache_size(DEFAULT_CACHE_CAPACITY).await
92    }
93
94    /// Construct with an explicit LRU capacity.
95    pub async fn with_cache_size(capacity: usize) -> Result<Self> {
96        let capacity =
97            NonZeroUsize::new(capacity.max(1)).expect("capacity.max(1) is always non-zero");
98
99        // fastembed's `try_new` downloads + builds an ONNX session — blocking
100        // work that must run off the async reactor.
101        let model = tokio::task::spawn_blocking(|| -> Result<TextEmbedding> {
102            let mut m =
103                TextEmbedding::try_new(TextInitOptions::new(EmbeddingModel::AllMiniLML6V2Q))
104                    .or_else(|q_err| {
105                        tracing::warn!(
106                            "AllMiniLML6V2Q init failed ({q_err:#}), falling back to AllMiniLML6V2"
107                        );
108                        TextEmbedding::try_new(TextInitOptions::new(EmbeddingModel::AllMiniLML6V2))
109                    })
110                    .context(
111                        "failed to initialise fastembed (tried AllMiniLML6V2Q and AllMiniLML6V2)",
112                    )?;
113
114            // Warm the graph so the first real user query is hot.
115            let warmup: Vec<&str> = vec![
116                "hello world",
117                "the quick brown fox",
118                "memory palace warmup",
119                "embedding model ready",
120                "trusty common warmup",
121            ];
122            let _ = m
123                .embed(warmup, None)
124                .context("fastembed warmup batch failed")?;
125            Ok(m)
126        })
127        .await
128        .context("spawn_blocking joined with error during embedder init")??;
129
130        Ok(Self {
131            model: Arc::new(Mutex::new(model)),
132            cache: Arc::new(Mutex::new(LruCache::new(capacity))),
133            dim: EMBED_DIM,
134        })
135    }
136}
137
138#[async_trait]
139impl Embedder for FastEmbedder {
140    async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
141        if texts.is_empty() {
142            return Ok(Vec::new());
143        }
144
145        // Split into cached hits vs misses.
146        let mut results: Vec<Option<Vec<f32>>> = vec![None; texts.len()];
147        let mut to_compute: Vec<(usize, String)> = Vec::new();
148        {
149            let mut cache = self.cache.lock();
150            for (i, t) in texts.iter().enumerate() {
151                if let Some(v) = cache.get(t) {
152                    results[i] = Some(v.clone());
153                } else {
154                    to_compute.push((i, t.clone()));
155                }
156            }
157        }
158
159        if !to_compute.is_empty() {
160            let model = Arc::clone(&self.model);
161            let owned: Vec<String> = to_compute.iter().map(|(_, s)| s.clone()).collect();
162            let computed = tokio::task::spawn_blocking(move || -> Result<Vec<Vec<f32>>> {
163                let mut guard = model.lock();
164                guard
165                    .embed(owned, None)
166                    .context("fastembed embed call failed")
167            })
168            .await
169            .context("spawn_blocking joined with error during embed")??;
170
171            if computed.len() != to_compute.len() {
172                anyhow::bail!(
173                    "fastembed returned {} embeddings, expected {}",
174                    computed.len(),
175                    to_compute.len()
176                );
177            }
178
179            let mut cache = self.cache.lock();
180            for ((idx, key), vector) in to_compute.into_iter().zip(computed.into_iter()) {
181                cache.put(key, vector.clone());
182                results[idx] = Some(vector);
183            }
184        }
185
186        results
187            .into_iter()
188            .map(|opt| opt.context("missing embedding slot after batch"))
189            .collect()
190    }
191
192    fn dimension(&self) -> usize {
193        self.dim
194    }
195}
196
197/// Deterministic test double — hashes input bytes into a fixed-dim vector.
198///
199/// Why: ONNX model downloads dominate test runtime and can race on cold
200/// caches when multiple tests construct embedders in parallel. The mock
201/// gives integration tests a "rank by similarity" surface without any I/O.
202/// What: a tiny per-byte hash spread across `dim` slots, with the first byte
203/// always contributing so short/empty strings still differ.
204/// Test: `mock_embedder_round_trip` confirms shape + determinism.
205#[cfg(any(test, feature = "test-support"))]
206pub struct MockEmbedder {
207    dim: usize,
208}
209
210#[cfg(any(test, feature = "test-support"))]
211impl MockEmbedder {
212    pub fn new(dim: usize) -> Self {
213        Self { dim }
214    }
215
216    fn hash_to_vec(&self, text: &str) -> Vec<f32> {
217        let mut v = vec![0.0_f32; self.dim];
218        for (i, b) in text.bytes().enumerate() {
219            let slot = (i + b as usize) % self.dim;
220            v[slot] += (b as f32) / 255.0;
221        }
222        if let Some(first) = text.bytes().next() {
223            v[0] += first as f32 / 255.0;
224        }
225        v
226    }
227}
228
229#[cfg(any(test, feature = "test-support"))]
230#[async_trait]
231impl Embedder for MockEmbedder {
232    async fn embed_batch(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
233        Ok(texts.iter().map(|t| self.hash_to_vec(t)).collect())
234    }
235
236    fn dimension(&self) -> usize {
237        self.dim
238    }
239}
240
241#[cfg(test)]
242mod tests {
243    use super::*;
244
245    #[tokio::test]
246    async fn mock_embedder_round_trip() {
247        let e = MockEmbedder::new(EMBED_DIM);
248        assert_eq!(e.dimension(), EMBED_DIM);
249        let v = embed_one(&e, "hello").await.unwrap();
250        assert_eq!(v.len(), EMBED_DIM);
251        let batch = e
252            .embed_batch(&["a".to_string(), "b".to_string()])
253            .await
254            .unwrap();
255        assert_eq!(batch.len(), 2);
256        assert_ne!(batch[0], batch[1]);
257    }
258
259    #[tokio::test]
260    async fn mock_embedder_empty_input_returns_empty() {
261        let e = MockEmbedder::new(EMBED_DIM);
262        let v = e.embed_batch(&[]).await.unwrap();
263        assert!(v.is_empty());
264    }
265
266    // ONNX-backed test: downloads ~23MB on first run. Marked ignored so default
267    // `cargo test` stays offline; run with `cargo test -- --ignored` when needed.
268    #[tokio::test]
269    #[ignore]
270    async fn fastembed_returns_correct_dim() {
271        let e = FastEmbedder::new().await.unwrap();
272        assert_eq!(e.dimension(), 384);
273        let v = embed_one(&e, "fn authenticate(user: &str) -> bool")
274            .await
275            .unwrap();
276        assert_eq!(v.len(), 384);
277        assert!(v.iter().any(|x| *x != 0.0));
278    }
279
280    #[tokio::test]
281    #[ignore]
282    async fn fastembed_cache_hit_is_idempotent() {
283        let e = FastEmbedder::new().await.unwrap();
284        let v1 = embed_one(&e, "cached").await.unwrap();
285        let v2 = embed_one(&e, "cached").await.unwrap();
286        assert_eq!(v1, v2);
287    }
288}
trusty_embedder/lib.rs

trusty_embedder/
lib.rs
