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 {
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
33impl Embedder for FastEmbedder {
34    fn embed(&self, texts: Vec<&str>) -> Result<Vec<Vec<f32>>, GraphError> {
35        let docs: Vec<String> = texts.into_iter().map(|t| t.to_string()).collect();
36        let embeddings = self
37            .model
38            .embed(docs, None)
39            .map_err(|e| GraphError::Embed(e.to_string()))?;
40        Ok(embeddings)
41    }
42
43    fn embed_single(&self, text: &str) -> Result<Vec<f32>, GraphError> {
44        let embeddings = self
45            .model
46            .embed(vec![text.to_string()], None)
47            .map_err(|e| GraphError::Embed(e.to_string()))?;
48        embeddings
49            .into_iter()
50            .next()
51            .ok_or_else(|| GraphError::Embed("no embedding returned".into()))
52    }
53
54    fn dimensions(&self) -> usize {
55        384
56    }
57}
recall_echo/graph/embed.rs

recall_echo/graph/
embed.rs
