convergio_knowledge/
embedder.rs

1//! Pure Rust embedding via fastembed (ONNX Runtime, no Python).
2//!
3//! Downloads model on first use (~25MB for AllMiniLML6V2).
4//! Produces 384-dim embeddings in-process with zero external dependencies.
5
6use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
7use std::sync::OnceLock;
8use tokio::sync::Mutex;
9use tracing::{info, warn};
10
11/// Shared model instance — loaded once, reused across calls.
12static MODEL: OnceLock<Mutex<TextEmbedding>> = OnceLock::new();
13
14fn get_model() -> &'static Mutex<TextEmbedding> {
15    MODEL.get_or_init(|| {
16        info!("loading embedding model AllMiniLML6V2...");
17        let model = TextEmbedding::try_new(
18            InitOptions::new(EmbeddingModel::AllMiniLML6V2).with_show_download_progress(true),
19        )
20        .expect("failed to load embedding model");
21        info!("embedding model loaded (384 dims)");
22        Mutex::new(model)
23    })
24}
25
26/// Embed a single text. Returns 384-dim f32 vector.
27pub async fn embed(text: &str) -> Vec<f32> {
28    let model = get_model();
29    let mut guard = model.lock().await;
30    match guard.embed(vec![text.to_string()], None) {
31        Ok(mut results) if !results.is_empty() => results.remove(0),
32        Ok(_) => {
33            warn!("embedding returned empty results");
34            vec![0.0; 384]
35        }
36        Err(e) => {
37            warn!(error = %e, "embedding failed");
38            vec![0.0; 384]
39        }
40    }
41}
42
43/// Embed multiple texts in batch. More efficient than individual calls.
44pub async fn embed_batch(texts: &[String]) -> Vec<Vec<f32>> {
45    if texts.is_empty() {
46        return vec![];
47    }
48    let model = get_model();
49    let mut guard = model.lock().await;
50    match guard.embed(texts, None) {
51        Ok(results) => results,
52        Err(e) => {
53            warn!(error = %e, "batch embedding failed");
54            texts.iter().map(|_| vec![0.0; 384]).collect()
55        }
56    }
57}
58
59/// Embedding dimension (AllMiniLML6V2 = 384).
60pub const fn embedding_dim() -> usize {
61    384
62}
63
64#[cfg(test)]
65mod tests {
66    use super::*;
67
68    #[tokio::test]
69    async fn embed_produces_384_dims() {
70        let v = embed("test convergio knowledge store").await;
71        assert_eq!(v.len(), 384);
72        // Should not be all zeros (model loaded successfully)
73        let nonzero = v.iter().any(|x| *x != 0.0);
74        assert!(nonzero, "embedding should have non-zero values");
75    }
76
77    #[tokio::test]
78    async fn similar_texts_have_high_similarity() {
79        let a = embed("rate limiter fix for authenticated agents").await;
80        let b = embed("fixing the rate limit for auth users").await;
81        let c = embed("chocolate cake recipe with vanilla frosting").await;
82
83        let sim_ab = cosine(&a, &b);
84        let sim_ac = cosine(&a, &c);
85        assert!(
86            sim_ab > sim_ac,
87            "similar texts should score higher: ab={sim_ab:.3} ac={sim_ac:.3}"
88        );
89    }
90
91    fn cosine(a: &[f32], b: &[f32]) -> f32 {
92        let dot: f32 = a.iter().zip(b).map(|(x, y)| x * y).sum();
93        let na: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
94        let nb: f32 = b.iter().map(|x| x * x).sum::<f32>().sqrt();
95        if na == 0.0 || nb == 0.0 {
96            0.0
97        } else {
98            dot / (na * nb)
99        }
100    }
101}
convergio_knowledge/embedder.rs

convergio_knowledge/
embedder.rs
