bamboo_memory/memory_store/embedding.rs
1//! Optional semantic-retrieval seam for durable-memory recall ("L1" of the
2//! memory redesign): the [`MemoryEmbedder`] interface a backend implements, plus
3//! the hybrid-score building blocks ([`cosine`], [`HYBRID_COSINE_WEIGHT`],
4//! [`SEMANTIC_FLOOR`]) that `lexical_bm25` combines as `bm25 + β·cosine`.
5//!
6//! **Embedding is OPTIONAL.** With no embedder configured, recall stays pure
7//! BM25 — this seam is entirely inert (query vectors are `None`, so the cosine
8//! term drops out and scoring is byte-identical to L0). Only the interface and
9//! the hybrid-scoring path land here; the concrete backend (a local multilingual
10//! ONNX embedder, provider-independent — or an optional hosted embedding
11//! endpoint) is a deferred follow-up that just implements this trait, populates
12//! `LexicalIndexItem.embedding` at index-build time, and embeds the query at
13//! recall time. BM25 stays the PRIMARY term so the #61 cache-stable ordering and
14//! the CJK/lexical guarantees survive; the vector term only re-ranks within that
15//! and surfaces paraphrase matches lexical scoring would miss.
16
17/// A text embedder. Backends SHOULD return a unit-normalized vector; an empty
18/// vec means "no embedding available" (recall falls back to pure BM25 for that
19/// item/query). Must be cheap enough to call once per query at recall time
20/// (document vectors are precomputed at write/index time).
21pub trait MemoryEmbedder: Send + Sync {
22 /// Embed `text` into a vector. Empty return = no embedding.
23 fn embed(&self, text: &str) -> Vec<f32>;
24}
25
26/// Weight (β) on the cosine term in the hybrid score `bm25 + β·cosine`.
27pub const HYBRID_COSINE_WEIGHT: f64 = 1.0;
28
29/// Minimum cosine for a NON-lexical (pure-semantic) match to be recalled — stops
30/// the vector term from surfacing every embedded doc as a weak match. A doc still
31/// recalls on lexical BM25 alone below this; the floor only gates cosine-only hits.
32pub const SEMANTIC_FLOOR: f64 = 0.25;
33
34/// Cosine similarity in [-1, 1]. Returns 0 for empty, mismatched-length,
35/// degenerate (zero-norm), or non-finite (NaN/inf) vectors, so a missing/absent
36/// or malformed embedding is a no-op that can never leak NaN/inf into a score.
37pub fn cosine(a: &[f32], b: &[f32]) -> f64 {
38 if a.is_empty() || a.len() != b.len() {
39 return 0.0;
40 }
41 let mut dot = 0.0f64;
42 let mut na = 0.0f64;
43 let mut nb = 0.0f64;
44 for (x, y) in a.iter().zip(b) {
45 let (x, y) = (*x as f64, *y as f64);
46 dot += x * y;
47 na += x * x;
48 nb += y * y;
49 }
50 let denom = na.sqrt() * nb.sqrt();
51 if denom <= f64::EPSILON {
52 return 0.0;
53 }
54 let sim = dot / denom;
55 // Backends only SHOULD unit-normalize; a stray NaN/inf component would make
56 // `sim` non-finite. Clamp it out rather than trust the contract.
57 if sim.is_finite() {
58 sim
59 } else {
60 0.0
61 }
62}
63
64#[cfg(test)]
65mod tests {
66 use super::*;
67
68 #[test]
69 fn cosine_identical_is_one_orthogonal_is_zero() {
70 assert!((cosine(&[1.0, 0.0], &[2.0, 0.0]) - 1.0).abs() < 1e-9);
71 assert!(cosine(&[1.0, 0.0], &[0.0, 1.0]).abs() < 1e-9);
72 }
73
74 #[test]
75 fn cosine_degenerate_and_mismatch_are_zero() {
76 assert_eq!(cosine(&[], &[1.0]), 0.0);
77 assert_eq!(cosine(&[1.0, 2.0], &[1.0]), 0.0);
78 assert_eq!(cosine(&[0.0, 0.0], &[1.0, 1.0]), 0.0);
79 }
80
81 #[test]
82 fn cosine_non_finite_component_is_zero() {
83 assert_eq!(cosine(&[f32::NAN, 1.0], &[1.0, 1.0]), 0.0);
84 assert_eq!(cosine(&[f32::INFINITY, 0.0], &[1.0, 1.0]), 0.0);
85 }
86}