cascade_agent/knowledge/
embeddings.rs

1//! Fastembed-based text embedding with thread-safe Mutex wrapper.
2
3use std::sync::Mutex;
4
5use fastembed::{EmbeddingModel, TextEmbedding, TextInitOptions};
6
7use crate::error::{AgentError, Result};
8
9/// Maps a human-readable model name to the fastembed enum.
10fn model_from_name(name: &str) -> Result<EmbeddingModel> {
11    match name {
12        "all-MiniLM-L6-v2" => Ok(EmbeddingModel::AllMiniLML6V2),
13        "all-MiniLM-L6-v2-q" => Ok(EmbeddingModel::AllMiniLML6V2Q),
14        "all-MiniLM-L12-v2" => Ok(EmbeddingModel::AllMiniLML12V2),
15        "all-mpnet-base-v2" => Ok(EmbeddingModel::AllMpnetBaseV2),
16        "bge-base-en-v1.5" => Ok(EmbeddingModel::BGEBaseENV15),
17        "bge-base-en-v1.5-q" => Ok(EmbeddingModel::BGEBaseENV15Q),
18        "bge-large-en-v1.5" => Ok(EmbeddingModel::BGELargeENV15),
19        "bge-large-en-v1.5-q" => Ok(EmbeddingModel::BGELargeENV15Q),
20        "bge-small-en-v1.5" => Ok(EmbeddingModel::BGESmallENV15),
21        "bge-small-en-v1.5-q" => Ok(EmbeddingModel::BGESmallENV15Q),
22        "nomic-embed-text-v1" => Ok(EmbeddingModel::NomicEmbedTextV1),
23        "nomic-embed-text-v1.5" => Ok(EmbeddingModel::NomicEmbedTextV15),
24        "nomic-embed-text-v1.5-q" => Ok(EmbeddingModel::NomicEmbedTextV15Q),
25        "paraphrase-multilingual-MiniLM-L12-v2" => Ok(EmbeddingModel::ParaphraseMLMiniLML12V2),
26        "paraphrase-multilingual-mpnet-base-v2" => Ok(EmbeddingModel::ParaphraseMLMpnetBaseV2),
27        "bgem3" | "BAAI/bgem3" => Ok(EmbeddingModel::BGEM3),
28        "multilingual-e5-small" | "intfloat/multilingual-e5-small" => {
29            Ok(EmbeddingModel::MultilingualE5Small)
30        }
31        "multilingual-e5-base" | "intfloat/multilingual-e5-base" => {
32            Ok(EmbeddingModel::MultilingualE5Base)
33        }
34        "multilingual-e5-large" | "intfloat/multilingual-e5-large" => {
35            Ok(EmbeddingModel::MultilingualE5Large)
36        }
37        "mxbai-embed-large-v1" => Ok(EmbeddingModel::MxbaiEmbedLargeV1),
38        "mxbai-embed-large-v1-q" => Ok(EmbeddingModel::MxbaiEmbedLargeV1Q),
39        "gte-base-en-v1.5" => Ok(EmbeddingModel::GTEBaseENV15),
40        _ => Err(AgentError::EmbeddingError(format!(
41            "Unknown embedding model: '{}'. \
42             See docs for supported models.",
43            name
44        ))),
45    }
46}
47
48/// Thread-safe embedding wrapper using fastembed.
49///
50/// The inner `TextEmbedding` is wrapped in a `Mutex` because `embed(&mut self)`
51/// requires mutable access.
52pub struct Embedder {
53    model: Mutex<TextEmbedding>,
54    dimension: usize,
55}
56
57impl std::fmt::Debug for Embedder {
58    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
59        f.debug_struct("Embedder")
60            .field("dimension", &self.dimension)
61            .finish_non_exhaustive()
62    }
63}
64
65impl Embedder {
66    /// Create a new embedder, downloading the model if necessary.
67    pub fn new(model_name: &str) -> Result<Self> {
68        let model_enum = model_from_name(model_name)?;
69        let opts = TextInitOptions::new(model_enum).with_show_download_progress(true);
70
71        let mut text_embedding = TextEmbedding::try_new(opts)
72            .map_err(|e| AgentError::EmbeddingError(format!("Failed to load model: {}", e)))?;
73
74        // Determine dimension by embedding a dummy text.
75        let dimension = {
76            let mut te = text_embedding;
77            let result = te
78                .embed(["dim_probe"], None)
79                .map_err(|e| AgentError::EmbeddingError(format!("Probe embed failed: {}", e)))?;
80            let dim = result.into_iter().next().unwrap().len();
81            text_embedding = te;
82            dim
83        };
84
85        Ok(Self {
86            model: Mutex::new(text_embedding),
87            dimension,
88        })
89    }
90
91    /// Embed a single query string (prepends "query: " prefix for E5 models).
92    pub fn embed_query(&self, text: &str) -> Result<Vec<f32>> {
93        let prefixed = format!("query: {}", text);
94        let mut guard = self
95            .model
96            .lock()
97            .map_err(|e| AgentError::EmbeddingError(format!("Embedder lock poisoned: {}", e)))?;
98        let results = guard
99            .embed([&prefixed], None)
100            .map_err(|e| AgentError::EmbeddingError(format!("Query embed failed: {}", e)))?;
101        Ok(results.into_iter().next().unwrap())
102    }
103
104    /// Embed a single passage string (prepends "passage: " prefix for E5 models).
105    pub fn embed_passage(&self, text: &str) -> Result<Vec<f32>> {
106        let prefixed = format!("passage: {}", text);
107        let mut guard = self
108            .model
109            .lock()
110            .map_err(|e| AgentError::EmbeddingError(format!("Embedder lock poisoned: {}", e)))?;
111        let results = guard
112            .embed([&prefixed], None)
113            .map_err(|e| AgentError::EmbeddingError(format!("Passage embed failed: {}", e)))?;
114        Ok(results.into_iter().next().unwrap())
115    }
116
117    /// Embed a batch of passage strings (prepends "passage: " prefix).
118    pub fn embed_batch_passages(&self, texts: &[String]) -> Result<Vec<Vec<f32>>> {
119        if texts.is_empty() {
120            return Ok(Vec::new());
121        }
122        let prefixed: Vec<String> = texts.iter().map(|t| format!("passage: {}", t)).collect();
123        let mut guard = self
124            .model
125            .lock()
126            .map_err(|e| AgentError::EmbeddingError(format!("Embedder lock poisoned: {}", e)))?;
127        guard
128            .embed(&prefixed, None)
129            .map_err(|e| AgentError::EmbeddingError(format!("Batch embed failed: {}", e)))
130    }
131
132    /// Returns the embedding dimensionality.
133    pub fn dimension(&self) -> usize {
134        self.dimension
135    }
136}
cascade_agent/knowledge/embeddings.rs

cascade_agent/knowledge/
embeddings.rs
