aimem_core/
embedder.rs

1//! Embedder — local text-to-vector using fastembed-rs.
2//!
3//! # Model
4//!
5//! Uses `all-MiniLM-L6-v2` (384 dimensions), the same model ChromaDB defaults to,
6//! so embeddings produced here are semantically comparable to an existing ChromaDB setup.
7//!
8//! The ONNX model (~23 MB) is downloaded from HuggingFace Hub on first call and
9//! cached at `~/.cache/huggingface/hub/`. Subsequent calls are instant.
10//!
11//! # Threading
12//!
13//! [`Embedder`] is `Clone + Send + Sync`. The inner ONNX session is wrapped in
14//! `Arc<Mutex<...>>` so multiple tasks can share one instance without contention —
15//! embed calls are batched internally by fastembed.
16//!
17//! # Example
18//!
19//! ```rust,no_run
20//! use aimem_core::Embedder;
21//!
22//! let embedder = Embedder::new()?;
23//!
24//! // Batch embedding (preferred — amortizes ONNX overhead)
25//! let vecs = embedder.embed(&["hello world", "turso vector search"])?;
26//! assert_eq!(vecs[0].len(), 384); // all-MiniLM-L6-v2 dim
27//!
28//! // Single embedding
29//! let v = embedder.embed_one("semantic memory")?;
30//! assert_eq!(v.len(), 384);
31//!
32//! // Serialize to Turso vector32 JSON for SQL params
33//! let json = Embedder::to_vector32_json(&v);
34//! assert!(json.starts_with('['));
35//! # anyhow::Ok(())
36//! ```
37//! ```
38
39use std::sync::{Arc, Mutex};
40
41use fastembed::{EmbeddingModel, TextEmbedding, TextInitOptions};
42use thiserror::Error;
43
44#[derive(Debug, Error)]
45pub enum EmbedError {
46    #[error("fastembed error: {0}")]
47    Fastembed(#[from] anyhow::Error),
48    #[error("mutex poisoned")]
49    Poison,
50}
51
52/// Dimension of the default embedding model (all-MiniLM-L6-v2).
53pub const EMBED_DIM: usize = 384;
54
55/// Thread-safe wrapper around a fastembed `TextEmbedding` instance.
56///
57/// Clone is cheap — inner state is Arc-wrapped.
58#[derive(Clone)]
59pub struct Embedder {
60    inner: Arc<Mutex<TextEmbedding>>,
61}
62
63impl std::fmt::Debug for Embedder {
64    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
65        f.debug_struct("Embedder")
66            .field("model", &"AllMiniLML6V2")
67            .finish()
68    }
69}
70
71impl Embedder {
72    /// Create a new embedder with the default model (all-MiniLM-L6-v2).
73    ///
74    /// On first call this will download the ONNX model (~23 MB) to
75    /// `~/.cache/huggingface/hub/` if not already cached.
76    pub fn new() -> Result<Self, EmbedError> {
77        let model = TextEmbedding::try_new(
78            TextInitOptions::new(EmbeddingModel::AllMiniLML6V2).with_show_download_progress(true),
79        )?;
80        Ok(Self {
81            inner: Arc::new(Mutex::new(model)),
82        })
83    }
84
85    /// Embed a slice of strings, returning one `Vec<f32>` per input.
86    pub fn embed(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, EmbedError> {
87        let mut guard = self.inner.lock().map_err(|_| EmbedError::Poison)?;
88        let results = guard.embed(texts, None)?;
89        Ok(results)
90    }
91
92    /// Convenience: embed a single string.
93    pub fn embed_one(&self, text: &str) -> Result<Vec<f32>, EmbedError> {
94        let mut results = self.embed(&[text])?;
95        results
96            .pop()
97            .ok_or_else(|| EmbedError::Fastembed(anyhow::anyhow!("empty embedding result")))
98    }
99
100    /// Serialize a `Vec<f32>` to a JSON array string suitable for `vector32(?)`
101    /// in Turso SQL.
102    pub fn to_vector32_json(v: &[f32]) -> String {
103        let nums: Vec<String> = v.iter().map(|f| f.to_string()).collect();
104        format!("[{}]", nums.join(","))
105    }
106}
aimem_core/embedder.rs

aimem_core/
embedder.rs
