cognee_embedding/
engine.rs

1use crate::error::EmbeddingResult;
2use async_trait::async_trait;
3
4/// Core trait for text embedding engines
5///
6/// Provides async interface for embedding generation while allowing
7/// synchronous implementations (e.g., ONNX) to be wrapped via spawn_blocking.
8///
9/// All returned embeddings are L2-normalized to unit vectors for cosine similarity.
10#[async_trait]
11pub trait EmbeddingEngine: Send + Sync {
12    /// Embed a batch of text strings into normalized vectors
13    ///
14    /// # Arguments
15    /// * `texts` - Slice of strings to embed
16    ///
17    /// # Returns
18    /// * Vector of embeddings, one per input text. Each embedding is L2-normalized.
19    ///
20    /// # Errors
21    /// * Returns error if tokenization, inference, or normalization fails
22    ///
23    /// # Example
24    /// ```ignore
25    /// let texts = vec!["Hello world".to_string()];
26    /// let embeddings = engine.embed(&texts).await?;
27    /// assert_eq!(embeddings.len(), 1);
28    /// assert_eq!(embeddings[0].len(), engine.dimension());
29    /// ```
30    async fn embed(&self, texts: &[&str]) -> EmbeddingResult<Vec<Vec<f32>>>;
31
32    /// Get the dimensionality of embeddings produced by this engine
33    ///
34    /// # Returns
35    /// * Number of dimensions in output vectors (e.g., 384 for BGE-Small)
36    fn dimension(&self) -> usize;
37
38    /// Get the optimal batch size for this engine
39    ///
40    /// Batches larger than this should be chunked by the caller.
41    ///
42    /// # Returns
43    /// * Maximum number of texts to process in a single embed() call
44    fn batch_size(&self) -> usize;
45
46    /// Get the maximum sequence length (in tokens) supported
47    ///
48    /// Input texts will be truncated to this length during tokenization.
49    ///
50    /// # Returns
51    /// * Maximum token count (e.g., 512 for BGE-Small-v1.5)
52    fn max_sequence_length(&self) -> usize;
53}
cognee_embedding/engine.rs

cognee_embedding/
engine.rs
