second_brain_core/
embedding.rs

1use std::sync::Mutex;
2
3use anyhow::{Context, Result};
4use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
5
6pub struct Embedder {
7    model: Mutex<TextEmbedding>,
8}
9
10impl Embedder {
11    pub fn new() -> Result<Self> {
12        let model = TextEmbedding::try_new(
13            InitOptions::new(EmbeddingModel::BGESmallENV15).with_show_download_progress(true),
14        )
15        .context("initializing embedding model")?;
16
17        Ok(Self {
18            model: Mutex::new(model),
19        })
20    }
21
22    pub fn embed(&self, text: &str) -> Result<Vec<f32>> {
23        let mut model = self.model.lock().unwrap();
24        let results = model
25            .embed(vec![text], None)
26            .context("generating embedding")?;
27
28        results
29            .into_iter()
30            .next()
31            .ok_or_else(|| anyhow::anyhow!("no embedding returned"))
32    }
33
34    pub fn embed_batch(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>> {
35        let mut model = self.model.lock().unwrap();
36        let owned: Vec<String> = texts.iter().map(|t| t.to_string()).collect();
37        let refs: Vec<&str> = owned.iter().map(|s| s.as_str()).collect();
38        model
39            .embed(refs, None)
40            .context("generating batch embeddings")
41    }
42
43    pub fn dimension(&self) -> usize {
44        384
45    }
46}
second_brain_core/embedding.rs

second_brain_core/
embedding.rs
