Skip to main content

scone_core/
embed.rs

1//! Embedding providers (spec §9).
2//!
3//! The engine never talks to a model directly — only through this trait, so
4//! local ONNX, remote endpoints, and the deterministic test embedder are
5//! interchangeable, and the engine works offline by construction.
6
7use crate::error::Result;
8
9pub trait EmbeddingProvider: Send {
10    /// Stable identity, pinned into the index metadata; changing providers
11    /// requires `doctor --rebuild` (spec §9).
12    fn id(&self) -> &str;
13    fn dim(&self) -> usize;
14    fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>>;
15}
16
17/// Deterministic, model-free embedder: hashed bag-of-words, L2-normalized.
18///
19/// Exists so tests and degraded mode need no model download; not a semantic
20/// embedder.
21pub struct HashEmbedder {
22    dim: usize,
23}
24
25impl HashEmbedder {
26    pub fn new(dim: usize) -> Self {
27        Self { dim: dim.max(1) }
28    }
29}
30
31impl EmbeddingProvider for HashEmbedder {
32    fn id(&self) -> &str {
33        "hash-v1"
34    }
35
36    fn dim(&self) -> usize {
37        self.dim
38    }
39
40    fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
41        Ok(texts
42            .iter()
43            .map(|text| {
44                let mut v = vec![0.0f32; self.dim];
45                for token in text.to_lowercase().split_whitespace() {
46                    // FNV-1a over the token bytes.
47                    let mut h: u64 = 0xcbf2_9ce4_8422_2325;
48                    for b in token.as_bytes() {
49                        h ^= u64::from(*b);
50                        h = h.wrapping_mul(0x0000_0100_0000_01b3);
51                    }
52                    let bucket = (h % self.dim as u64) as usize;
53                    // Second hash bit decides sign, reducing bucket bias.
54                    let sign = if h & (1 << 63) == 0 { 1.0 } else { -1.0 };
55                    v[bucket] += sign;
56                }
57                let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
58                if norm > 0.0 {
59                    for x in &mut v {
60                        *x /= norm;
61                    }
62                }
63                v
64            })
65            .collect())
66    }
67}
68
69/// Local ONNX embedder (spec §9 default): bge-small-en-v1.5, 384 dims,
70/// downloaded once into the data dir and cached — offline thereafter.
71#[cfg(feature = "local-embed")]
72pub struct OnnxEmbedder {
73    model: std::cell::RefCell<fastembed::TextEmbedding>,
74}
75
76#[cfg(feature = "local-embed")]
77impl OnnxEmbedder {
78    pub fn new(cache_dir: &std::path::Path) -> Result<Self> {
79        let options = fastembed::InitOptions::new(fastembed::EmbeddingModel::BGESmallENV15)
80            .with_cache_dir(cache_dir.to_path_buf())
81            .with_show_download_progress(false);
82        let model = fastembed::TextEmbedding::try_new(options)
83            .map_err(|e| crate::SconeError::Embed(e.to_string()))?;
84        Ok(Self {
85            model: std::cell::RefCell::new(model),
86        })
87    }
88}
89
90#[cfg(feature = "local-embed")]
91impl EmbeddingProvider for OnnxEmbedder {
92    fn id(&self) -> &str {
93        "bge-small-en-v1.5"
94    }
95
96    fn dim(&self) -> usize {
97        384
98    }
99
100    fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
101        self.model
102            .borrow_mut()
103            .embed(texts, None)
104            .map_err(|e| crate::SconeError::Embed(e.to_string()))
105    }
106}