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    id: String,
75    dim: usize,
76}
77
78#[cfg(feature = "local-embed")]
79impl OnnxEmbedder {
80    pub fn new(cache_dir: &std::path::Path) -> Result<Self> {
81        Self::with_model(cache_dir, "bge-small-en-v1.5")
82    }
83
84    /// Open a specific local model by short name. Supported:
85    /// bge-small-en-v1.5 (384d, default), bge-base-en-v1.5 (768d),
86    /// nomic-embed-text-v1.5 (768d, long context).
87    pub fn with_model(cache_dir: &std::path::Path, name: &str) -> Result<Self> {
88        let (model, dim) = match name {
89            "bge-small-en-v1.5" => (fastembed::EmbeddingModel::BGESmallENV15, 384),
90            "bge-base-en-v1.5" => (fastembed::EmbeddingModel::BGEBaseENV15, 768),
91            "nomic-embed-text-v1.5" => (fastembed::EmbeddingModel::NomicEmbedTextV15, 768),
92            other => {
93                return Err(crate::SconeError::InvalidInput(format!(
94                    "unknown embed model {other:?}"
95                )));
96            }
97        };
98        let options = fastembed::InitOptions::new(model)
99            .with_cache_dir(cache_dir.to_path_buf())
100            .with_show_download_progress(false);
101        let model = fastembed::TextEmbedding::try_new(options)
102            .map_err(|e| crate::SconeError::Embed(e.to_string()))?;
103        Ok(Self {
104            model: std::cell::RefCell::new(model),
105            id: name.to_owned(),
106            dim,
107        })
108    }
109}
110
111#[cfg(feature = "local-embed")]
112impl EmbeddingProvider for OnnxEmbedder {
113    fn id(&self) -> &str {
114        &self.id
115    }
116
117    fn dim(&self) -> usize {
118        self.dim
119    }
120
121    fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
122        self.model
123            .borrow_mut()
124            .embed(texts, None)
125            .map_err(|e| crate::SconeError::Embed(e.to_string()))
126    }
127}