mentedb_embedding/
http_provider.rs

1//! Generic HTTP-based embedding provider for OpenAI, Cohere, Voyage, and other APIs.
2
3use std::collections::HashMap;
4
5use mentedb_core::MenteError;
6use mentedb_core::error::MenteResult;
7use serde::{Deserialize, Serialize};
8
9use crate::provider::{AsyncEmbeddingProvider, EmbeddingProvider};
10
11/// Configuration for an HTTP-based embedding API.
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct HttpEmbeddingConfig {
14    /// The API endpoint URL.
15    pub api_url: String,
16    /// The API key for authentication.
17    pub api_key: String,
18    /// The model name to request.
19    pub model_name: String,
20    /// The dimensionality of the returned embeddings.
21    pub dimensions: usize,
22    /// Additional headers to include in requests.
23    pub headers: HashMap<String, String>,
24}
25
26impl HttpEmbeddingConfig {
27    /// Create a configuration for OpenAI's embedding API.
28    ///
29    /// Default dimensions: 1536 for text-embedding-ada-002, 3072 for text-embedding-3-large.
30    pub fn openai(api_key: impl Into<String>, model: impl Into<String>) -> Self {
31        let model = model.into();
32        let dimensions = match model.as_str() {
33            "text-embedding-3-small" => 1536,
34            "text-embedding-3-large" => 3072,
35            "text-embedding-ada-002" => 1536,
36            _ => 1536,
37        };
38
39        let mut headers = HashMap::new();
40        headers.insert("Content-Type".to_string(), "application/json".to_string());
41
42        Self {
43            api_url: "https://api.openai.com/v1/embeddings".to_string(),
44            api_key: api_key.into(),
45            model_name: model,
46            dimensions,
47            headers,
48        }
49    }
50
51    /// Create a configuration for Cohere's embedding API.
52    ///
53    /// Default dimensions: 1024 for embed-english-v3.0.
54    pub fn cohere(api_key: impl Into<String>, model: impl Into<String>) -> Self {
55        let model = model.into();
56        let dimensions = match model.as_str() {
57            "embed-english-v3.0" => 1024,
58            "embed-multilingual-v3.0" => 1024,
59            "embed-english-light-v3.0" => 384,
60            "embed-multilingual-light-v3.0" => 384,
61            _ => 1024,
62        };
63
64        let mut headers = HashMap::new();
65        headers.insert("Content-Type".to_string(), "application/json".to_string());
66
67        Self {
68            api_url: "https://api.cohere.ai/v1/embed".to_string(),
69            api_key: api_key.into(),
70            model_name: model,
71            dimensions,
72            headers,
73        }
74    }
75
76    /// Create a configuration for Voyage AI's embedding API.
77    ///
78    /// Default dimensions: 1024 for voyage-2.
79    pub fn voyage(api_key: impl Into<String>, model: impl Into<String>) -> Self {
80        let model = model.into();
81        let dimensions = match model.as_str() {
82            "voyage-2" => 1024,
83            "voyage-large-2" => 1536,
84            "voyage-code-2" => 1536,
85            "voyage-lite-02-instruct" => 1024,
86            _ => 1024,
87        };
88
89        let mut headers = HashMap::new();
90        headers.insert("Content-Type".to_string(), "application/json".to_string());
91
92        Self {
93            api_url: "https://api.voyageai.com/v1/embeddings".to_string(),
94            api_key: api_key.into(),
95            model_name: model,
96            dimensions,
97            headers,
98        }
99    }
100
101    /// Override the embedding dimensions.
102    pub fn with_dimensions(mut self, dimensions: usize) -> Self {
103        self.dimensions = dimensions;
104        self
105    }
106
107    /// Add a custom header.
108    pub fn with_header(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
109        self.headers.insert(key.into(), value.into());
110        self
111    }
112}
113
114/// HTTP-based embedding provider.
115///
116/// Currently requires an external HTTP client feature to function.
117/// The structure and configuration are fully usable for setup and validation.
118pub struct HttpEmbeddingProvider {
119    config: HttpEmbeddingConfig,
120}
121
122impl HttpEmbeddingProvider {
123    /// Create a new HTTP embedding provider with the given configuration.
124    pub fn new(config: HttpEmbeddingConfig) -> Self {
125        Self { config }
126    }
127
128    /// Get a reference to the provider's configuration.
129    pub fn config(&self) -> &HttpEmbeddingConfig {
130        &self.config
131    }
132}
133
134impl AsyncEmbeddingProvider for HttpEmbeddingProvider {
135    async fn embed(&self, _text: &str) -> MenteResult<Vec<f32>> {
136        Err(MenteError::Storage(
137            "HTTP embedding requires the 'http' feature for async, use sync EmbeddingProvider instead".to_string(),
138        ))
139    }
140
141    async fn embed_batch(&self, _texts: &[&str]) -> MenteResult<Vec<Vec<f32>>> {
142        Err(MenteError::Storage(
143            "HTTP embedding requires the 'http' feature for async, use sync EmbeddingProvider instead".to_string(),
144        ))
145    }
146
147    fn dimensions(&self) -> usize {
148        self.config.dimensions
149    }
150
151    fn model_name(&self) -> &str {
152        &self.config.model_name
153    }
154}
155
156#[cfg(feature = "http")]
157mod http_impl {
158    use super::*;
159    use serde_json::json;
160    use std::time::Duration;
161    use ureq::config::Config;
162
163    #[derive(Deserialize)]
164    struct OpenAIEmbeddingResponse {
165        data: Vec<OpenAIEmbeddingData>,
166    }
167
168    #[derive(Deserialize)]
169    struct OpenAIEmbeddingData {
170        embedding: Vec<f32>,
171    }
172
173    impl HttpEmbeddingProvider {
174        /// Create a ureq agent with a 60-second global timeout to prevent hangs.
175        fn agent(&self) -> ureq::Agent {
176            Config::builder()
177                .timeout_global(Some(Duration::from_secs(60)))
178                .build()
179                .new_agent()
180        }
181
182        /// Retry-aware single embedding call with exponential backoff.
183        fn embed_with_retry(&self, text: &str, max_attempts: u32) -> MenteResult<Vec<f32>> {
184            let agent = self.agent();
185            let mut last_err = None;
186            for attempt in 0..max_attempts {
187                if attempt > 0 {
188                    std::thread::sleep(std::time::Duration::from_millis(500 * (1 << attempt)));
189                }
190
191                let body = json!({
192                    "model": self.config.model_name,
193                    "input": text,
194                });
195
196                let mut req = agent
197                    .post(&self.config.api_url)
198                    .header("Authorization", &format!("Bearer {}", self.config.api_key));
199
200                for (k, v) in &self.config.headers {
201                    if k.to_lowercase() != "content-type" {
202                        req = req.header(k, v);
203                    }
204                }
205
206                match req.send_json(&body) {
207                    Ok(mut resp) => match resp.body_mut().read_json::<OpenAIEmbeddingResponse>() {
208                        Ok(parsed) => {
209                            return parsed
210                                .data
211                                .into_iter()
212                                .next()
213                                .map(|d| d.embedding)
214                                .ok_or_else(|| {
215                                    MenteError::Storage("Empty embedding response".to_string())
216                                });
217                        }
218                        Err(e) => {
219                            last_err = Some(format!("Failed to parse embedding response: {}", e));
220                        }
221                    },
222                    Err(e) => {
223                        last_err = Some(format!("HTTP embedding request failed: {}", e));
224                    }
225                }
226            }
227            Err(MenteError::Storage(last_err.unwrap_or_else(|| {
228                "embedding failed after retries".to_string()
229            })))
230        }
231
232        /// Retry-aware batch embedding call with exponential backoff.
233        fn embed_batch_with_retry(
234            &self,
235            texts: &[&str],
236            max_attempts: u32,
237        ) -> MenteResult<Vec<Vec<f32>>> {
238            let agent = self.agent();
239            let mut last_err = None;
240            for attempt in 0..max_attempts {
241                if attempt > 0 {
242                    std::thread::sleep(std::time::Duration::from_millis(500 * (1 << attempt)));
243                }
244
245                let body = json!({
246                    "model": self.config.model_name,
247                    "input": texts,
248                });
249
250                let mut req = agent
251                    .post(&self.config.api_url)
252                    .header("Authorization", &format!("Bearer {}", self.config.api_key));
253
254                for (k, v) in &self.config.headers {
255                    if k.to_lowercase() != "content-type" {
256                        req = req.header(k, v);
257                    }
258                }
259
260                match req.send_json(&body) {
261                    Ok(mut resp) => match resp.body_mut().read_json::<OpenAIEmbeddingResponse>() {
262                        Ok(parsed) => {
263                            return Ok(parsed.data.into_iter().map(|d| d.embedding).collect());
264                        }
265                        Err(e) => {
266                            last_err = Some(format!("Failed to parse embedding response: {}", e));
267                        }
268                    },
269                    Err(e) => {
270                        last_err = Some(format!("HTTP embedding request failed: {}", e));
271                    }
272                }
273            }
274            Err(MenteError::Storage(last_err.unwrap_or_else(|| {
275                "batch embedding failed after retries".to_string()
276            })))
277        }
278    }
279
280    impl EmbeddingProvider for HttpEmbeddingProvider {
281        fn embed(&self, text: &str) -> MenteResult<Vec<f32>> {
282            self.embed_with_retry(text, 3)
283        }
284
285        fn embed_batch(&self, texts: &[&str]) -> MenteResult<Vec<Vec<f32>>> {
286            self.embed_batch_with_retry(texts, 3)
287        }
288
289        fn dimensions(&self) -> usize {
290            self.config.dimensions
291        }
292
293        fn model_name(&self) -> &str {
294            &self.config.model_name
295        }
296    }
297}
298
299#[cfg(not(feature = "http"))]
300impl EmbeddingProvider for HttpEmbeddingProvider {
301    fn embed(&self, _text: &str) -> MenteResult<Vec<f32>> {
302        Err(MenteError::Storage(
303            "HTTP embedding requires the 'http' feature. Enable it in Cargo.toml.".to_string(),
304        ))
305    }
306
307    fn embed_batch(&self, _texts: &[&str]) -> MenteResult<Vec<Vec<f32>>> {
308        Err(MenteError::Storage(
309            "HTTP embedding requires the 'http' feature. Enable it in Cargo.toml.".to_string(),
310        ))
311    }
312
313    fn dimensions(&self) -> usize {
314        self.config.dimensions
315    }
316
317    fn model_name(&self) -> &str {
318        &self.config.model_name
319    }
320}
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325
326    #[test]
327    fn test_openai_config() {
328        let config = HttpEmbeddingConfig::openai("sk-test", "text-embedding-3-small");
329        assert_eq!(config.api_url, "https://api.openai.com/v1/embeddings");
330        assert_eq!(config.dimensions, 1536);
331        assert_eq!(config.model_name, "text-embedding-3-small");
332    }
333
334    #[test]
335    fn test_cohere_config() {
336        let config = HttpEmbeddingConfig::cohere("key", "embed-english-v3.0");
337        assert_eq!(config.api_url, "https://api.cohere.ai/v1/embed");
338        assert_eq!(config.dimensions, 1024);
339    }
340
341    #[test]
342    fn test_voyage_config() {
343        let config = HttpEmbeddingConfig::voyage("key", "voyage-2");
344        assert_eq!(config.api_url, "https://api.voyageai.com/v1/embeddings");
345        assert_eq!(config.dimensions, 1024);
346    }
347
348    #[test]
349    fn test_with_dimensions_override() {
350        let config =
351            HttpEmbeddingConfig::openai("key", "text-embedding-3-small").with_dimensions(256);
352        assert_eq!(config.dimensions, 256);
353    }
354}
mentedb_embedding/http_provider.rs

mentedb_embedding/
http_provider.rs
