klieo_embed_common/
fastembed.rs

1//! `FastEmbedEmbedder` — fastembed-rs ONNX-runtime CPU embedder.
2//!
3//! Default model: `AllMiniLML6V2` (384-dim, ~80 MB ONNX file
4//! auto-downloaded on first use to the user's cache directory).
5//!
6//! Heavy first-build (ONNX runtime crate is ~5-10 minutes from cold
7//! cache). Subsequent builds reuse the cached compile.
8
9use async_trait::async_trait;
10use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
11use klieo_core::error::MemoryError;
12use std::sync::{Arc, Mutex};
13
14use crate::Embedder;
15
16/// Default model: 384-dimensional `AllMiniLML6V2`.
17pub const DEFAULT_MODEL: EmbeddingModel = EmbeddingModel::AllMiniLML6V2;
18/// Output dimension of [`DEFAULT_MODEL`].
19pub const DEFAULT_DIM: usize = 384;
20
21/// fastembed-rs–backed `Embedder`.
22///
23/// fastembed 5.x requires `&mut TextEmbedding` for `embed`; the trait
24/// surface gives us `&self`. Interior `Mutex` serialises concurrent
25/// embed calls — acceptable because ONNX inference is already CPU-
26/// saturating.
27pub struct FastEmbedEmbedder {
28    inner: Arc<Mutex<TextEmbedding>>,
29    dim: usize,
30}
31
32impl FastEmbedEmbedder {
33    /// Build with [`DEFAULT_MODEL`]. Triggers model download on first
34    /// call if not already cached.
35    pub fn new() -> Result<Self, MemoryError> {
36        Self::with_model(DEFAULT_MODEL, DEFAULT_DIM)
37    }
38
39    /// Build with the supplied fastembed model. The caller asserts the
40    /// model's output dimension; mismatch is reported at first
41    /// `embed` call against a dimension-checked store.
42    pub fn with_model(model: EmbeddingModel, dim: usize) -> Result<Self, MemoryError> {
43        let inner = TextEmbedding::try_new(InitOptions::new(model))
44            .map_err(|e| MemoryError::Embedding(format!("fastembed init: {e}")))?;
45        Ok(Self {
46            inner: Arc::new(Mutex::new(inner)),
47            dim,
48        })
49    }
50}
51
52#[async_trait]
53impl Embedder for FastEmbedEmbedder {
54    fn dimension(&self) -> usize {
55        self.dim
56    }
57
58    async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, MemoryError> {
59        let texts_owned: Vec<String> = texts.to_vec();
60        let inner = self.inner.clone();
61        let dim = self.dim;
62        tokio::task::spawn_blocking(move || -> Result<Vec<Vec<f32>>, MemoryError> {
63            // `TextEmbedding`'s mutated state is the ort `Session`; a
64            // panic inside one inference does not leave any invariants
65            // broken for the next call. Recover from poison by taking
66            // the inner guard so a single panic does not permanently
67            // disable long-term memory writes.
68            //
69            // Held inside `spawn_blocking`, never across `.await` — use
70            // `std::sync::Mutex`, not `tokio::sync::Mutex`. Do NOT move
71            // this lock onto an async path without revisiting that
72            // choice.
73            let mut guard = inner.lock().unwrap_or_else(|p| p.into_inner());
74            let docs = guard
75                .embed(texts_owned, None)
76                .map_err(|e| MemoryError::Embedding(format!("fastembed embed: {e}")))?;
77            for v in &docs {
78                if v.len() != dim {
79                    return Err(MemoryError::Embedding(format!(
80                        "fastembed produced {}-dim vector, expected {dim}",
81                        v.len()
82                    )));
83                }
84            }
85            Ok(docs)
86        })
87        .await
88        .map_err(|e| MemoryError::Embedding(format!("blocking task join: {e}")))?
89    }
90}
klieo_embed_common/fastembed.rs

klieo_embed_common/
fastembed.rs
