lc_embeddings/
cohere.rs

1// lc-embeddings/src/cohere.rs
2//! Cohere Embeddings — embed-english-v3.0 / embed-multilingual-v3.0.
3//!
4//! Uses Cohere's v2/embed endpoint for generating text embeddings.
5
6use async_trait::async_trait;
7use serde::Deserialize;
8use serde_json::json;
9
10use crate::{EmbeddingError, Embeddings};
11
12/// Default Cohere embedding model.
13pub const COHERE_EMBED_MODEL: &str = "embed-english-v3.0";
14
15/// Cohere API base URL.
16pub const COHERE_EMBED_BASE_URL: &str = "https://api.cohere.com/v2";
17
18/// Cohere embedding input type.
19#[derive(Debug, Clone, Copy)]
20pub enum CohereEmbedInputType {
21    /// Search query embedding.
22    SearchQuery,
23    /// Search document embedding.
24    SearchDocument,
25    /// Classification embedding.
26    Classification,
27    /// Clustering embedding.
28    Clustering,
29}
30
31impl CohereEmbedInputType {
32    fn as_str(&self) -> &'static str {
33        match self {
34            CohereEmbedInputType::SearchQuery => "search_query",
35            CohereEmbedInputType::SearchDocument => "search_document",
36            CohereEmbedInputType::Classification => "classification",
37            CohereEmbedInputType::Clustering => "clustering",
38        }
39    }
40}
41
42/// Cohere embedding configuration.
43#[derive(Debug, Clone)]
44pub struct CohereEmbeddingsConfig {
45    pub api_key: String,
46    pub base_url: String,
47    pub model: String,
48    pub input_type: CohereEmbedInputType,
49}
50
51impl Default for CohereEmbeddingsConfig {
52    fn default() -> Self {
53        Self {
54            api_key: String::new(),
55            base_url: COHERE_EMBED_BASE_URL.to_string(),
56            model: COHERE_EMBED_MODEL.to_string(),
57            input_type: CohereEmbedInputType::SearchQuery,
58        }
59    }
60}
61
62impl CohereEmbeddingsConfig {
63    /// Creates a new config with the given API key.
64    pub fn new(api_key: impl Into<String>) -> Self {
65        Self {
66            api_key: api_key.into(),
67            ..Default::default()
68        }
69    }
70
71    /// Creates config from environment variables.
72    pub fn from_env_result() -> Result<Self, String> {
73        let api_key = std::env::var("COHERE_API_KEY")
74            .map_err(|_| "COHERE_API_KEY environment variable not set".to_string())?;
75        let base_url =
76            std::env::var("COHERE_BASE_URL").unwrap_or_else(|_| COHERE_EMBED_BASE_URL.to_string());
77        let model = std::env::var("COHERE_EMBED_MODEL")
78            .unwrap_or_else(|_| COHERE_EMBED_MODEL.to_string());
79        Ok(Self {
80            api_key,
81            base_url,
82            model,
83            ..Default::default()
84        })
85    }
86
87    /// Sets the model name.
88    pub fn with_model(mut self, model: impl Into<String>) -> Self {
89        self.model = model.into();
90        self
91    }
92
93    /// Sets the input type.
94    pub fn with_input_type(mut self, input_type: CohereEmbedInputType) -> Self {
95        self.input_type = input_type;
96        self
97    }
98}
99
100/// Cohere embedding response.
101#[derive(Debug, Deserialize)]
102struct CohereEmbedResponse {
103    data: Vec<CohereEmbedData>,
104}
105
106#[derive(Debug, Deserialize)]
107struct CohereEmbedData {
108    embedding: Vec<f32>,
109}
110
111/// Cohere embedding provider.
112pub struct CohereEmbeddings {
113    config: CohereEmbeddingsConfig,
114    client: reqwest::Client,
115    dimension: usize,
116}
117
118impl std::fmt::Debug for CohereEmbeddings {
119    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120        f.debug_struct("CohereEmbeddings")
121            .field("model", &self.config.model)
122            .finish()
123    }
124}
125
126impl CohereEmbeddings {
127    /// Creates a new CohereEmbeddings with the given configuration.
128    pub fn new(config: CohereEmbeddingsConfig) -> Self {
129        let dimension = 1024;
130        Self {
131            config,
132            client: reqwest::Client::new(),
133            dimension,
134        }
135    }
136
137    /// Creates from environment variables.
138    pub fn from_env_result() -> Result<Self, String> {
139        Ok(Self::new(CohereEmbeddingsConfig::from_env_result()?))
140    }
141}
142
143#[async_trait]
144impl Embeddings for CohereEmbeddings {
145    async fn embed_query(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
146        let url = format!("{}/embed", self.config.base_url);
147        let body = json!({
148            "model": self.config.model,
149            "input_type": self.config.input_type.as_str(),
150            "texts": [text],
151            "embedding_types": ["float"],
152        });
153
154        let response = self
155            .client
156            .post(&url)
157            .header("Authorization", format!("Bearer {}", self.config.api_key))
158            .header("Content-Type", "application/json")
159            .json(&body)
160            .send()
161            .await
162            .map_err(|e| EmbeddingError::HttpError(e.to_string()))?;
163
164        let status = response.status();
165        if !status.is_success() {
166            let error_text = response.text().await.unwrap_or_default();
167            return Err(EmbeddingError::ApiError(format!(
168                "HTTP {}: {}",
169                status, error_text
170            )));
171        }
172
173        let embed_response: CohereEmbedResponse = response
174            .json()
175            .await
176            .map_err(|e| EmbeddingError::ParseError(e.to_string()))?;
177
178        embed_response
179            .data
180            .first()
181            .map(|d| d.embedding.clone())
182            .ok_or_else(|| EmbeddingError::ApiError("No embedding in response".to_string()))
183    }
184
185    async fn embed_documents(&self, texts: &[&str]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
186        if texts.is_empty() {
187            return Err(EmbeddingError::EmptyInput);
188        }
189
190        let url = format!("{}/embed", self.config.base_url);
191        let body = json!({
192            "model": self.config.model,
193            "input_type": CohereEmbedInputType::SearchDocument.as_str(),
194            "texts": texts,
195            "embedding_types": ["float"],
196        });
197
198        let response = self
199            .client
200            .post(&url)
201            .header("Authorization", format!("Bearer {}", self.config.api_key))
202            .header("Content-Type", "application/json")
203            .json(&body)
204            .send()
205            .await
206            .map_err(|e| EmbeddingError::HttpError(e.to_string()))?;
207
208        let status = response.status();
209        if !status.is_success() {
210            let error_text = response.text().await.unwrap_or_default();
211            return Err(EmbeddingError::ApiError(format!(
212                "HTTP {}: {}",
213                status, error_text
214            )));
215        }
216
217        let embed_response: CohereEmbedResponse = response
218            .json()
219            .await
220            .map_err(|e| EmbeddingError::ParseError(e.to_string()))?;
221
222        Ok(embed_response
223            .data
224            .into_iter()
225            .map(|d| d.embedding)
226            .collect())
227    }
228
229    fn dimension(&self) -> usize {
230        self.dimension
231    }
232
233    fn model_name(&self) -> &str {
234        &self.config.model
235    }
236}
237
238#[cfg(test)]
239mod tests {
240    use super::*;
241
242    #[test]
243    fn test_config_new() {
244        let config = CohereEmbeddingsConfig::new("test-key");
245        assert_eq!(config.api_key, "test-key");
246        assert_eq!(config.model, COHERE_EMBED_MODEL);
247    }
248
249    #[test]
250    fn test_config_builder() {
251        let config = CohereEmbeddingsConfig::new("key")
252            .with_model("embed-multilingual-v3.0")
253            .with_input_type(CohereEmbedInputType::SearchDocument);
254        assert_eq!(config.model, "embed-multilingual-v3.0");
255        assert!(matches!(
256            config.input_type,
257            CohereEmbedInputType::SearchDocument
258        ));
259    }
260
261    #[test]
262    fn test_input_type_str() {
263        assert_eq!(CohereEmbedInputType::SearchQuery.as_str(), "search_query");
264        assert_eq!(
265            CohereEmbedInputType::SearchDocument.as_str(),
266            "search_document"
267        );
268        assert_eq!(CohereEmbedInputType::Classification.as_str(), "classification");
269        assert_eq!(CohereEmbedInputType::Clustering.as_str(), "clustering");
270    }
271
272    #[test]
273    fn test_embeddings_new() {
274        let config = CohereEmbeddingsConfig::new("key");
275        let embeddings = CohereEmbeddings::new(config);
276        assert_eq!(embeddings.model_name(), COHERE_EMBED_MODEL);
277        assert_eq!(embeddings.dimension(), 1024);
278    }
279
280    #[test]
281    fn test_embeddings_multilingual_dimension() {
282        let config = CohereEmbeddingsConfig::new("key").with_model("embed-multilingual-v3.0");
283        let embeddings = CohereEmbeddings::new(config);
284        assert_eq!(embeddings.dimension(), 1024);
285    }
286}
lc_embeddings/cohere.rs

lc_embeddings/
cohere.rs
