recall_echo/graph/
embed.rs

1//! Text embedding via fastembed (BGE-Small-EN-v1.5, 384 dimensions).
2
3use std::path::Path;
4
5use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
6
7use super::error::GraphError;
8
9/// Trait for embedding text into vectors.
10pub trait Embedder: Send + Sync {
11    fn embed(&self, texts: Vec<&str>) -> Result<Vec<Vec<f32>>, GraphError>;
12    fn embed_single(&self, text: &str) -> Result<Vec<f32>, GraphError>;
13    fn dimensions(&self) -> usize;
14}
15
16/// Local embedding using fastembed (BGE-Small-EN-v1.5, 384 dimensions).
17pub struct FastEmbedder {
18    model: TextEmbedding,
19}
20
21impl FastEmbedder {
22    pub fn new(cache_dir: &Path) -> Result<Self, GraphError> {
23        let options = InitOptions::new(EmbeddingModel::BGESmallENV15)
24            .with_cache_dir(cache_dir.to_path_buf())
25            .with_show_download_progress(true);
26
27        let model =
28            TextEmbedding::try_new(options).map_err(|e| GraphError::Embed(e.to_string()))?;
29        Ok(Self { model })
30    }
31}
32
33/// Lazily-initialized [`FastEmbedder`].
34///
35/// Construction is free — the ONNX model is loaded (and downloaded on first
36/// ever use) only when an operation actually needs an embedding. Operations
37/// that never embed (schema init, CRUD reads, GC, status) never touch the
38/// network or pay the model-load cost. This also keeps unit tests that open
39/// a graph store fully offline.
40pub struct LazyEmbedder {
41    cache_dir: std::path::PathBuf,
42    cell: std::sync::OnceLock<FastEmbedder>,
43    init_lock: std::sync::Mutex<()>,
44}
45
46impl LazyEmbedder {
47    pub fn new(cache_dir: &Path) -> Self {
48        Self {
49            cache_dir: cache_dir.to_path_buf(),
50            cell: std::sync::OnceLock::new(),
51            init_lock: std::sync::Mutex::new(()),
52        }
53    }
54
55    /// Get the embedder, initializing it on first use.
56    pub fn get(&self) -> Result<&FastEmbedder, GraphError> {
57        if let Some(e) = self.cell.get() {
58            return Ok(e);
59        }
60        // Serialize initialization; losers of the race find the cell filled.
61        let _guard = self
62            .init_lock
63            .lock()
64            .map_err(|_| GraphError::Embed("embedder init lock poisoned".into()))?;
65        if self.cell.get().is_none() {
66            let embedder = FastEmbedder::new(&self.cache_dir)?;
67            let _ = self.cell.set(embedder);
68        }
69        self.cell
70            .get()
71            .ok_or_else(|| GraphError::Embed("embedder cell empty after init".into()))
72    }
73}
74
75impl Embedder for FastEmbedder {
76    fn embed(&self, texts: Vec<&str>) -> Result<Vec<Vec<f32>>, GraphError> {
77        let docs: Vec<String> = texts.into_iter().map(|t| t.to_string()).collect();
78        let embeddings = self
79            .model
80            .embed(docs, None)
81            .map_err(|e| GraphError::Embed(e.to_string()))?;
82        Ok(embeddings)
83    }
84
85    fn embed_single(&self, text: &str) -> Result<Vec<f32>, GraphError> {
86        let embeddings = self
87            .model
88            .embed(vec![text.to_string()], None)
89            .map_err(|e| GraphError::Embed(e.to_string()))?;
90        embeddings
91            .into_iter()
92            .next()
93            .ok_or_else(|| GraphError::Embed("no embedding returned".into()))
94    }
95
96    fn dimensions(&self) -> usize {
97        384
98    }
99}
recall_echo/graph/embed.rs

recall_echo/graph/
embed.rs
