1mod 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#[derive(Debug, thiserror::Error)]
38pub enum EmbeddingError {
39 #[error("HTTP error: {0}")]
41 HttpError(String),
42
43 #[error("API error: {0}")]
45 ApiError(String),
46
47 #[error("Parse error: {0}")]
49 ParseError(String),
50
51 #[error("Input is empty")]
53 EmptyInput,
54}
55
56#[async_trait]
60pub trait Embeddings: Send + Sync {
61 async fn embed_query(&self, text: &str) -> Result<Vec<f32>, EmbeddingError>;
69
70 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 fn dimension(&self) -> usize;
87
88 fn model_name(&self) -> &str;
90}
91
92pub use lc_core::math::cosine_similarity;
96
97#[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 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 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 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}