klieo-memory-sqlite 0.4.0

//! `FastEmbedEmbedder` — fastembed-rs ONNX-runtime CPU embedder.
//!
//! Default model: `AllMiniLML6V2` (384-dim, ~80 MB ONNX file
//! auto-downloaded on first use to the user's cache directory).
//!
//! Heavy first-build (ONNX runtime crate is ~5-10 minutes from cold
//! cache). Subsequent builds reuse the cached compile.

use async_trait::async_trait;
use fastembed::{EmbeddingModel, InitOptions, TextEmbedding};
use klieo_core::error::MemoryError;
use std::sync::{Arc, Mutex};

use crate::embedder::Embedder;

/// Default model: 384-dimensional `AllMiniLML6V2`.
pub const DEFAULT_MODEL: EmbeddingModel = EmbeddingModel::AllMiniLML6V2;
/// Output dimension of [`DEFAULT_MODEL`].
pub const DEFAULT_DIM: usize = 384;

/// fastembed-rs–backed `Embedder`.
///
/// fastembed 5.x requires `&mut TextEmbedding` for `embed`; the trait
/// surface gives us `&self`. Interior `Mutex` serialises concurrent
/// embed calls — acceptable because ONNX inference is already CPU-
/// saturating.
pub struct FastEmbedEmbedder {
    inner: Arc<Mutex<TextEmbedding>>,
    dim: usize,
}

impl FastEmbedEmbedder {
    /// Build with [`DEFAULT_MODEL`]. Triggers model download on first
    /// call if not already cached.
    pub fn new() -> Result<Self, MemoryError> {
        Self::with_model(DEFAULT_MODEL, DEFAULT_DIM)
    }

    /// Build with the supplied fastembed model. The caller asserts the
    /// model's output dimension; mismatch is reported at first
    /// `embed` call when sqlite-vec is on.
    pub fn with_model(model: EmbeddingModel, dim: usize) -> Result<Self, MemoryError> {
        let inner = TextEmbedding::try_new(InitOptions::new(model))
            .map_err(|e| MemoryError::Embedding(format!("fastembed init: {e}")))?;
        Ok(Self {
            inner: Arc::new(Mutex::new(inner)),
            dim,
        })
    }
}

#[async_trait]
impl Embedder for FastEmbedEmbedder {
    fn dimension(&self) -> usize {
        self.dim
    }

    async fn embed(&self, texts: &[String]) -> Result<Vec<Vec<f32>>, MemoryError> {
        let texts_owned: Vec<String> = texts.to_vec();
        let inner = self.inner.clone();
        let dim = self.dim;
        tokio::task::spawn_blocking(move || -> Result<Vec<Vec<f32>>, MemoryError> {
            // `TextEmbedding`'s mutated state is the ort `Session`; a
            // panic inside one inference does not leave any invariants
            // broken for the next call. Recover from poison by taking
            // the inner guard so a single panic does not permanently
            // disable long-term memory writes.
            //
            // Held inside `spawn_blocking`, never across `.await` — use
            // `std::sync::Mutex`, not `tokio::sync::Mutex`. Do NOT move
            // this lock onto an async path without revisiting that
            // choice.
            let mut guard = inner.lock().unwrap_or_else(|p| p.into_inner());
            let docs = guard
                .embed(texts_owned, None)
                .map_err(|e| MemoryError::Embedding(format!("fastembed embed: {e}")))?;
            for v in &docs {
                if v.len() != dim {
                    return Err(MemoryError::Embedding(format!(
                        "fastembed produced {}-dim vector, expected {dim}",
                        v.len()
                    )));
                }
            }
            Ok(docs)
        })
        .await
        .map_err(|e| MemoryError::Embedding(format!("blocking task join: {e}")))?
    }
}