llm_kernel/embedding/

fastembed.rs

//! Local ONNX embedding provider via fastembed-rs.
//!
//! Wraps [`fastembed::TextEmbedding`] behind the [`EmbeddingProvider`] trait.
//! Models are downloaded from HuggingFace on first use and cached locally.
//!
//! ```ignore
//! use llm_kernel::embedding::{EmbeddingModel, FastembedProvider};
//! use llm_kernel::embedding::EmbeddingProvider;
//!
//! let provider = FastembedProvider::new(EmbeddingModel::BGESmallENV15, None)?;
//! let result = provider.embed("hello world")?;
//! assert_eq!(result.vector.len(), 384);
//! ```

use std::path::PathBuf;
use std::sync::Mutex;

use crate::embedding::catalog::EmbeddingModel;
use crate::embedding::types::{EmbeddingProvider, EmbeddingResult};
use crate::error::{KernelError, Result};

/// Local ONNX embedding provider backed by fastembed-rs.
///
/// `TextEmbedding::embed()` requires `&mut self`, so the inner instance is
/// protected by a `Mutex`. Thread-safety is guaranteed by the `Send + Sync`
/// bounds on `EmbeddingProvider`.
pub struct FastembedProvider {
    inner: Mutex<fastembed::TextEmbedding>,
    model: EmbeddingModel,
}

impl FastembedProvider {
    /// Create a new provider.
    ///
    /// `cache_dir` overrides the HuggingFace model cache directory.
    /// Pass `None` to use the default cache location.
    pub fn new(model: EmbeddingModel, cache_dir: Option<PathBuf>) -> Result<Self> {
        let mut options = fastembed::TextInitOptions::new(model.as_fastembed())
            .with_show_download_progress(false);
        if let Some(dir) = cache_dir {
            options = options.with_cache_dir(dir);
        }
        let te = fastembed::TextEmbedding::try_new(options).map_err(KernelError::embedding)?;
        Ok(Self {
            inner: Mutex::new(te),
            model,
        })
    }

    /// Create with DirectML GPU execution on Windows.
    ///
    /// Requires the `embedding-fastembed-directml` feature and Windows OS.
    /// The DirectML runtime DLL must be present on the target system.
    ///
    /// **Initialization cost:** the first call initialises the D3D12 device and
    /// loads the DirectML DLL, which can take hundreds of milliseconds to
    /// several seconds. Create the provider once and reuse it across requests.
    ///
    /// `cache_dir` overrides the HuggingFace model cache directory.
    #[cfg(all(feature = "embedding-fastembed-directml", target_os = "windows"))]
    pub fn new_with_directml(model: EmbeddingModel, cache_dir: Option<PathBuf>) -> Result<Self> {
        use ort::execution_providers::DirectMLExecutionProvider;
        let mut options = fastembed::TextInitOptions::new(model.as_fastembed())
            .with_show_download_progress(false)
            .with_execution_providers(vec![DirectMLExecutionProvider::default().build()]);
        if let Some(dir) = cache_dir {
            options = options.with_cache_dir(dir);
        }
        let te = fastembed::TextEmbedding::try_new(options).map_err(KernelError::embedding)?;
        Ok(Self {
            inner: Mutex::new(te),
            model,
        })
    }

    /// Create with CoreML GPU/ANE execution on macOS (feature `embedding-fastembed-coreml`).
    ///
    /// Accelerates ONNX inference via the CoreML execution provider — Neural Engine
    /// / GPU on Apple Silicon. The CoreML runtime is bundled with macOS; no extra
    /// dylib needed (unlike DirectML on Windows).
    ///
    /// Compiled CoreML models are cached under `<cache_dir>/coreml` and reused
    /// across sessions, so repeated provider creation does not re-compile for
    /// the Neural Engine.
    #[cfg(all(feature = "embedding-fastembed-coreml", target_os = "macos"))]
    pub fn new_with_coreml(model: EmbeddingModel, cache_dir: Option<PathBuf>) -> Result<Self> {
        use ort::execution_providers::CoreMLExecutionProvider;
        // Default to fastembed's own cache location so the CoreML model cache
        // lives next to the downloaded ONNX models.
        let hf_cache = cache_dir
            .clone()
            .unwrap_or_else(|| PathBuf::from(fastembed::get_cache_dir()));
        let coreml_cache = hf_cache.join("coreml");
        let _ = std::fs::create_dir_all(&coreml_cache);
        let coreml_ep = CoreMLExecutionProvider::default()
            // Reuse the compiled CoreML model across sessions. Without this,
            // every provider creation re-compiles for ANE — slow and
            // memory-heavy for per-request providers (research-agent incident,
            // RSS 22.69GB / 35 threads).
            .with_model_cache_dir(coreml_cache.to_string_lossy())
            .build();
        let mut options = fastembed::TextInitOptions::new(model.as_fastembed())
            .with_show_download_progress(false)
            // CoreML EP does the heavy compute; cap CPU intra threads instead
            // of fastembed's available_parallelism default.
            // ponytail: fixed 4 — raise if CPU-only fallback throughput matters.
            .with_intra_threads(4)
            .with_execution_providers(vec![coreml_ep]);
        options = options.with_cache_dir(hf_cache);
        let te = fastembed::TextEmbedding::try_new(options).map_err(KernelError::embedding)?;
        Ok(Self {
            inner: Mutex::new(te),
            model,
        })
    }

    /// Create with a custom maximum sequence length.
    pub fn with_max_length(
        model: EmbeddingModel,
        cache_dir: Option<PathBuf>,
        max_length: usize,
    ) -> Result<Self> {
        let mut options = fastembed::TextInitOptions::new(model.as_fastembed())
            .with_show_download_progress(false)
            .with_max_length(max_length);
        if let Some(dir) = cache_dir {
            options = options.with_cache_dir(dir);
        }
        let te = fastembed::TextEmbedding::try_new(options).map_err(KernelError::embedding)?;
        Ok(Self {
            inner: Mutex::new(te),
            model,
        })
    }
}

use super::types::text_preview;

impl EmbeddingProvider for FastembedProvider {
    fn dim(&self) -> usize {
        self.model.dimension()
    }

    fn name(&self) -> &str {
        self.model.as_str()
    }

    fn embed(&self, text: &str) -> Result<EmbeddingResult> {
        self.embed_with_prefix(text, self.model.query_prefix())
    }

    fn embed_document(&self, text: &str) -> Result<EmbeddingResult> {
        // Asymmetric models (E5/BGE) train with a passage prefix for corpus text.
        // query_prefix() is correct for queries; doc_prefix() is correct here.
        self.embed_with_prefix(text, self.model.doc_prefix())
    }

    fn embed_batch(&self, texts: &[&str]) -> Result<Vec<EmbeddingResult>> {
        self.embed_batch_with_prefix(texts, self.model.query_prefix())
    }

    fn embed_documents(&self, texts: &[&str]) -> Result<Vec<EmbeddingResult>> {
        self.embed_batch_with_prefix(texts, self.model.doc_prefix())
    }
}

impl FastembedProvider {
    /// Single-text embed with a caller-chosen prefix (query vs document).
    fn embed_with_prefix(&self, text: &str, prefix: Option<&str>) -> Result<EmbeddingResult> {
        let owned = match prefix {
            Some(p) => format!("{p}{text}"),
            None => text.to_string(),
        };
        let mut te = self
            .inner
            .lock()
            .map_err(|e| KernelError::Embedding(format!("lock: {e}")))?;
        let embeddings = te
            .embed(vec![owned], None)
            .map_err(KernelError::embedding)?;
        let vector = embeddings
            .into_iter()
            .next()
            .ok_or_else(|| KernelError::Embedding("empty embedding output".into()))?;

        Ok(EmbeddingResult {
            vector,
            text_preview: text_preview(text),
        })
    }

    /// Batch embed with a caller-chosen prefix (query vs document).
    fn embed_batch_with_prefix(
        &self,
        texts: &[&str],
        prefix: Option<&str>,
    ) -> Result<Vec<EmbeddingResult>> {
        if texts.is_empty() {
            return Ok(vec![]);
        }
        let prepared: Vec<String> = texts
            .iter()
            .map(|t| match prefix {
                Some(p) => format!("{p}{t}"),
                None => t.to_string(),
            })
            .collect();

        let mut te = self
            .inner
            .lock()
            .map_err(|e| KernelError::Embedding(format!("lock: {e}")))?;
        let embeddings = te.embed(prepared, None).map_err(KernelError::embedding)?;

        Ok(embeddings
            .into_iter()
            .zip(texts.iter())
            .map(|(vector, &text)| EmbeddingResult {
                vector,
                text_preview: text_preview(text),
            })
            .collect())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn provider_name_matches_model() {
        // Doesn't need a model download — just checks the constructor doesn't
        // change the name mapping.
        for &m in EmbeddingModel::ALL {
            // Verify as_str() round-trips through as_fastembed()
            let fe = m.as_fastembed();
            assert_eq!(format!("{fe:?}"), m.as_str());
        }
    }

    #[test]
    #[ignore = "requires model download"]
    fn embed_single_text() {
        let dir = tempfile::tempdir().unwrap();
        let provider = FastembedProvider::new(
            EmbeddingModel::BGESmallENV15,
            Some(dir.path().to_path_buf()),
        )
        .unwrap();
        let result = provider.embed("hello world").unwrap();
        assert_eq!(result.vector.len(), 384);
        assert!(!result.vector.is_empty());
    }

    #[test]
    #[ignore = "requires model download"]
    fn embed_batch_texts() {
        let dir = tempfile::tempdir().unwrap();
        let provider = FastembedProvider::new(
            EmbeddingModel::BGESmallENV15,
            Some(dir.path().to_path_buf()),
        )
        .unwrap();
        let results = provider
            .embed_batch(&["hello", "world", "foo bar"])
            .unwrap();
        assert_eq!(results.len(), 3);
        for r in &results {
            assert_eq!(r.vector.len(), 384);
        }
    }

    // Research-agent incident regression: providers are created per request /
    // per MCP call, so repeated create→embed→drop cycles must stay stable and
    // reuse the CoreML model cache (second cycle compiles no fresh ANE model).
    #[cfg(all(feature = "embedding-fastembed-coreml", target_os = "macos"))]
    #[test]
    #[ignore = "requires model download"]
    fn coreml_repeated_session_cycles() {
        let dir = tempfile::tempdir().unwrap();
        let cache_dir = Some(dir.path().to_path_buf());
        for _ in 0..10 {
            let provider = FastembedProvider::new_with_coreml(
                EmbeddingModel::BGESmallENV15,
                cache_dir.clone(),
            )
            .unwrap();
            drop(provider);
        }
        // Cached compile artifacts must exist after the cycles.
        let cache = dir.path().join("coreml");
        std::fs::read_dir(&cache).expect("coreml cache dir populated");
        let provider =
            FastembedProvider::new_with_coreml(EmbeddingModel::BGESmallENV15, cache_dir).unwrap();
        let result = provider.embed("hello world").unwrap();
        assert_eq!(result.vector.len(), 384);
    }
}