Skip to main content

lc_embeddings/
lib.rs

1// lc-embeddings/src/lib.rs
2//! Embedding model implementations for LangChainRust.
3//!
4//! Provides embedding generation via multiple backends:
5//! - OpenAI (`text-embedding-ada-002`, `text-embedding-3-small/large`)
6//! - DeepSeek
7//! - Qwen (Alibaba Cloud / DashScope)
8//! - Local: `BagOfWordsEmbeddings` (always available) and `LocalEmbeddings` (ONNX, feature-gated)
9//! - `MockEmbeddings` for testing
10
11mod deepseek;
12mod cohere;
13mod local;
14mod mock;
15mod openai;
16mod qwen;
17
18#[cfg(feature = "fastembed")]
19mod fastembed_emb;
20
21pub use cohere::{
22    CohereEmbeddings, CohereEmbeddingsConfig, CohereEmbedInputType, COHERE_EMBED_BASE_URL,
23    COHERE_EMBED_MODEL,
24};
25pub use deepseek::{DeepSeekEmbeddings, DeepSeekEmbeddingsConfig, DEEPSEEK_EMBED_MODEL};
26pub use local::{BagOfWordsEmbeddings, LocalEmbeddings};
27pub use mock::MockEmbeddings;
28pub use openai::{OpenAIEmbeddings, OpenAIEmbeddingsConfig};
29pub use qwen::{QwenEmbeddings, QwenEmbeddingsConfig, QWEN_EMBED_MODEL};
30
31#[cfg(feature = "fastembed")]
32pub use fastembed_emb::FastEmbedEmbeddings;
33
34use async_trait::async_trait;
35
36/// Embedding error type
37#[derive(Debug, thiserror::Error)]
38pub enum EmbeddingError {
39    /// HTTP request error
40    #[error("HTTP error: {0}")]
41    HttpError(String),
42
43    /// API error
44    #[error("API error: {0}")]
45    ApiError(String),
46
47    /// Parse error
48    #[error("Parse error: {0}")]
49    ParseError(String),
50
51    /// Empty input
52    #[error("Input is empty")]
53    EmptyInput,
54}
55
56/// Embedding model trait
57///
58/// Defines the interface for generating text embedding vectors.
59#[async_trait]
60pub trait Embeddings: Send + Sync {
61    /// Generate an embedding vector for a single text.
62    ///
63    /// # Arguments
64    /// * `text` - Input text
65    ///
66    /// # Returns
67    /// Embedding vector (typically 1536 dimensions or higher)
68    async fn embed_query(&self, text: &str) -> Result<Vec<f32>, EmbeddingError>;
69
70    /// Generate embedding vectors for multiple documents.
71    ///
72    /// # Arguments
73    /// * `texts` - List of input texts
74    ///
75    /// # Returns
76    /// List of embedding vectors
77    async fn embed_documents(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
78        let mut embeddings = Vec::new();
79        for text in texts {
80            embeddings.push(self.embed_query(text).await?);
81        }
82        Ok(embeddings)
83    }
84
85    /// Get the embedding vector dimension.
86    fn dimension(&self) -> usize;
87
88    /// Get the model name.
89    fn model_name(&self) -> &str;
90}
91
92/// Compute cosine similarity between two vectors.
93///
94/// Re-exported from [`lc_core::math::cosine_similarity`].
95pub use lc_core::math::cosine_similarity;
96
97/// Mutex for synchronizing environment-variable mutations in tests.
98///
99/// Tests that set/remove env vars must acquire this lock to avoid data races
100/// when running tests in parallel (the default for `cargo test`).
101#[cfg(test)]
102static ENV_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
103
104#[cfg(test)]
105mod tests {
106    use super::*;
107
108    #[test]
109    fn test_cosine_similarity() {
110        // Identical vectors
111        let a = vec![1.0, 0.0, 0.0];
112        let b = vec![1.0, 0.0, 0.0];
113        assert!((cosine_similarity(&a, &b).unwrap() - 1.0).abs() < 0.0001);
114
115        // Orthogonal vectors
116        let a = vec![1.0, 0.0, 0.0];
117        let b = vec![0.0, 1.0, 0.0];
118        assert!((cosine_similarity(&a, &b).unwrap() - 0.0).abs() < 0.0001);
119
120        // Opposite vectors
121        let a = vec![1.0, 0.0, 0.0];
122        let b = vec![-1.0, 0.0, 0.0];
123        assert!((cosine_similarity(&a, &b).unwrap() - (-1.0)).abs() < 0.0001);
124    }
125
126    #[test]
127    fn test_cosine_similarity_different_lengths() {
128        let a = vec![1.0, 0.0];
129        let b = vec![1.0, 0.0, 0.0];
130        assert!(cosine_similarity(&a, &b).is_err());
131    }
132}