recall_echo/graph/
embed.rs

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

recall_echo/graph/
embed.rs
