Skip to main content

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}