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