1mod 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#[derive(Debug, thiserror::Error)]
27pub enum EmbeddingError {
28 #[error("HTTP error: {0}")]
30 HttpError(String),
31
32 #[error("API error: {0}")]
34 ApiError(String),
35
36 #[error("Parse error: {0}")]
38 ParseError(String),
39
40 #[error("Input is empty")]
42 EmptyInput,
43}
44
45#[async_trait]
49pub trait Embeddings: Send + Sync {
50 async fn embed_query(&self, text: &str) -> Result<Vec<f32>, EmbeddingError>;
58
59 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 fn dimension(&self) -> usize;
76
77 fn model_name(&self) -> &str;
79}
80
81pub use lc_core::math::cosine_similarity;
85
86#[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 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 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 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}