uni_xervo/provider/
voyageai.rs

1use crate::api::{ModelAliasSpec, ModelTask};
2use crate::error::{Result, RuntimeError};
3use crate::provider::remote_common::{RemoteProviderBase, check_http_status, resolve_api_key};
4use crate::traits::{
5    EmbedResult, EmbeddingModel, LoadedModelHandle, ModelProvider, ProviderCapabilities,
6    ProviderHealth, RerankerModel, ScoredDoc, TokenUsage,
7};
8use async_trait::async_trait;
9use reqwest::Client;
10use serde_json::json;
11use std::sync::Arc;
12
13/// Remote provider that calls the [Voyage AI API](https://docs.voyageai.com/reference/embeddings-api)
14/// for embedding and reranking. Does not support text generation.
15///
16/// Requires the `VOYAGE_API_KEY` environment variable (or a custom env var
17/// name via the `api_key_env` option).
18pub struct RemoteVoyageAIProvider {
19    base: RemoteProviderBase,
20}
21
22impl Default for RemoteVoyageAIProvider {
23    fn default() -> Self {
24        Self {
25            base: RemoteProviderBase::new(),
26        }
27    }
28}
29
30impl RemoteVoyageAIProvider {
31    pub fn new() -> Self {
32        Self::default()
33    }
34
35    #[cfg(test)]
36    fn insert_test_breaker(&self, key: crate::api::ModelRuntimeKey, age: std::time::Duration) {
37        self.base.insert_test_breaker(key, age);
38    }
39
40    #[cfg(test)]
41    fn breaker_count(&self) -> usize {
42        self.base.breaker_count()
43    }
44
45    #[cfg(test)]
46    fn force_cleanup_now_for_test(&self) {
47        self.base.force_cleanup_now_for_test();
48    }
49}
50
51#[async_trait]
52impl ModelProvider for RemoteVoyageAIProvider {
53    fn provider_id(&self) -> &'static str {
54        "remote/voyageai"
55    }
56
57    fn capabilities(&self) -> ProviderCapabilities {
58        ProviderCapabilities {
59            supported_tasks: vec![ModelTask::Embed, ModelTask::Rerank],
60        }
61    }
62
63    async fn load(&self, spec: &ModelAliasSpec) -> Result<LoadedModelHandle> {
64        let cb = self.base.circuit_breaker_for(spec);
65        let api_key = resolve_api_key(&spec.options, "api_key_env", "VOYAGE_API_KEY")?;
66
67        let embedding_dimensions = spec
68            .options
69            .get("embedding_dimensions")
70            .and_then(|v| v.as_u64())
71            .map(|v| v as u32);
72
73        match spec.task {
74            ModelTask::Embed => {
75                let default_dims = match spec.model_id.as_str() {
76                    "voyage-large-2" => 1536,
77                    "voyage-3-lite" => 512,
78                    _ => 1024,
79                };
80                let model = VoyageAIEmbeddingModel {
81                    client: self.base.client.clone(),
82                    cb: cb.clone(),
83                    model_id: spec.model_id.clone(),
84                    api_key,
85                    dimensions: embedding_dimensions.unwrap_or(default_dims),
86                };
87                let handle: Arc<dyn EmbeddingModel> = Arc::new(model);
88                Ok(Arc::new(handle) as LoadedModelHandle)
89            }
90            ModelTask::Rerank => {
91                let model = VoyageAIRerankerModel {
92                    client: self.base.client.clone(),
93                    cb,
94                    model_id: spec.model_id.clone(),
95                    api_key,
96                };
97                let handle: Arc<dyn RerankerModel> = Arc::new(model);
98                Ok(Arc::new(handle) as LoadedModelHandle)
99            }
100            ModelTask::Raw => Err(RuntimeError::CapabilityMismatch(
101                "Voyage AI provider does not support task Raw".to_string(),
102            )),
103            _ => Err(RuntimeError::CapabilityMismatch(format!(
104                "Voyage AI provider does not support task {:?}",
105                spec.task
106            ))),
107        }
108    }
109
110    async fn health(&self) -> ProviderHealth {
111        ProviderHealth::Healthy
112    }
113}
114
115struct VoyageAIEmbeddingModel {
116    client: Client,
117    cb: crate::reliability::CircuitBreakerWrapper,
118    model_id: String,
119    api_key: String,
120    dimensions: u32,
121}
122
123#[async_trait]
124impl EmbeddingModel for VoyageAIEmbeddingModel {
125    async fn embed(&self, texts: &[&str]) -> Result<EmbedResult> {
126        let texts: Vec<String> = texts.iter().map(|s| s.to_string()).collect();
127
128        self.cb
129            .call(move || async move {
130                let response = self
131                    .client
132                    .post("https://api.voyageai.com/v1/embeddings")
133                    .header("Authorization", format!("Bearer {}", self.api_key))
134                    .json(&json!({
135                        "input": texts,
136                        "model": self.model_id
137                    }))
138                    .send()
139                    .await
140                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
141
142                let body: serde_json::Value = check_http_status("Voyage AI", response)?
143                    .json()
144                    .await
145                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
146
147                let mut vectors = Vec::new();
148                if let Some(data) = body.get("data").and_then(|d| d.as_array()) {
149                    for item in data {
150                        if let Some(embedding) = item.get("embedding").and_then(|e| e.as_array()) {
151                            let vec: Vec<f32> = embedding
152                                .iter()
153                                .filter_map(|v| v.as_f64().map(|f| f as f32))
154                                .collect();
155                            vectors.push(vec);
156                        }
157                    }
158                }
159
160                // Voyage AI embeddings response carries a `usage` object that
161                // only reports `total_tokens` (no prompt/completion split).
162                let usage = body
163                    .get("usage")
164                    .and_then(|u| u.get("total_tokens"))
165                    .and_then(|t| t.as_u64())
166                    .map(|tokens| TokenUsage {
167                        prompt_tokens: tokens as usize,
168                        completion_tokens: 0,
169                        total_tokens: tokens as usize,
170                    });
171
172                Ok(EmbedResult { vectors, usage })
173            })
174            .await
175    }
176
177    fn dimensions(&self) -> u32 {
178        self.dimensions
179    }
180}
181
182impl crate::traits::ModelInfo for VoyageAIEmbeddingModel {
183    fn model_id(&self) -> &str {
184        &self.model_id
185    }
186}
187
188struct VoyageAIRerankerModel {
189    client: Client,
190    cb: crate::reliability::CircuitBreakerWrapper,
191    model_id: String,
192    api_key: String,
193}
194
195impl crate::traits::ModelInfo for VoyageAIRerankerModel {
196    fn model_id(&self) -> &str {
197        &self.model_id
198    }
199}
200
201#[async_trait]
202impl RerankerModel for VoyageAIRerankerModel {
203    async fn rerank(&self, query: &str, docs: &[&str]) -> Result<Vec<ScoredDoc>> {
204        let query = query.to_string();
205        let docs: Vec<String> = docs.iter().map(|s| s.to_string()).collect();
206
207        self.cb
208            .call(move || async move {
209                let response = self
210                    .client
211                    .post("https://api.voyageai.com/v1/rerank")
212                    .header("Authorization", format!("Bearer {}", self.api_key))
213                    .json(&json!({
214                        "query": query,
215                        "documents": docs,
216                        "model": self.model_id,
217                    }))
218                    .send()
219                    .await
220                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
221
222                let body: serde_json::Value = check_http_status("Voyage AI", response)?
223                    .json()
224                    .await
225                    .map_err(|e| RuntimeError::ApiError(e.to_string()))?;
226
227                let data = body.get("data").and_then(|d| d.as_array()).ok_or_else(|| {
228                    RuntimeError::ApiError("Invalid rerank response format".to_string())
229                })?;
230
231                let mut results = Vec::new();
232                for item in data {
233                    let index = item.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as usize;
234                    let score = item
235                        .get("relevance_score")
236                        .and_then(|s| s.as_f64())
237                        .unwrap_or(0.0) as f32;
238                    results.push(ScoredDoc {
239                        index,
240                        score,
241                        text: None,
242                    });
243                }
244                Ok(results)
245            })
246            .await
247    }
248}
249
250#[cfg(test)]
251mod tests {
252    use super::*;
253    use crate::api::ModelRuntimeKey;
254    use crate::provider::remote_common::RemoteProviderBase;
255    use crate::traits::ModelProvider;
256    use std::time::Duration;
257
258    static ENV_LOCK: tokio::sync::Mutex<()> = tokio::sync::Mutex::const_new(());
259
260    fn spec(alias: &str, task: ModelTask, model_id: &str) -> ModelAliasSpec {
261        ModelAliasSpec {
262            alias: alias.to_string(),
263            task,
264            provider_id: "remote/voyageai".to_string(),
265            model_id: model_id.to_string(),
266            revision: None,
267            warmup: crate::api::WarmupPolicy::Lazy,
268            required: false,
269            timeout: None,
270            load_timeout: None,
271            retry: None,
272            options: serde_json::Value::Null,
273        }
274    }
275
276    #[tokio::test]
277    async fn breaker_reused_for_same_runtime_key() {
278        let _lock = ENV_LOCK.lock().await;
279        unsafe { std::env::set_var("VOYAGE_API_KEY", "test-key") };
280
281        let provider = RemoteVoyageAIProvider::new();
282        let s1 = spec("embed/a", ModelTask::Embed, "voyage-3");
283        let s2 = spec("embed/b", ModelTask::Embed, "voyage-3");
284
285        let _ = provider.load(&s1).await.unwrap();
286        let _ = provider.load(&s2).await.unwrap();
287
288        assert_eq!(provider.breaker_count(), 1);
289
290        unsafe { std::env::remove_var("VOYAGE_API_KEY") };
291    }
292
293    #[tokio::test]
294    async fn breaker_isolated_by_task_and_model() {
295        let _lock = ENV_LOCK.lock().await;
296        unsafe { std::env::set_var("VOYAGE_API_KEY", "test-key") };
297
298        let provider = RemoteVoyageAIProvider::new();
299        let embed = spec("embed/a", ModelTask::Embed, "voyage-3");
300        let rerank = spec("rerank/a", ModelTask::Rerank, "rerank-2");
301
302        let _ = provider.load(&embed).await.unwrap();
303        let _ = provider.load(&rerank).await.unwrap();
304
305        assert_eq!(provider.breaker_count(), 2);
306
307        unsafe { std::env::remove_var("VOYAGE_API_KEY") };
308    }
309
310    #[tokio::test]
311    async fn breaker_cleanup_evicts_stale_entries() {
312        let _lock = ENV_LOCK.lock().await;
313        unsafe { std::env::set_var("VOYAGE_API_KEY", "test-key") };
314
315        let provider = RemoteVoyageAIProvider::new();
316        let stale = spec("embed/stale", ModelTask::Embed, "voyage-3");
317        let fresh = spec("rerank/fresh", ModelTask::Rerank, "rerank-2");
318        provider.insert_test_breaker(
319            ModelRuntimeKey::new(&stale),
320            RemoteProviderBase::BREAKER_TTL + Duration::from_secs(5),
321        );
322        provider.insert_test_breaker(ModelRuntimeKey::new(&fresh), Duration::from_secs(1));
323        assert_eq!(provider.breaker_count(), 2);
324
325        provider.force_cleanup_now_for_test();
326        let _ = provider.load(&fresh).await.unwrap();
327
328        assert_eq!(provider.breaker_count(), 1);
329
330        unsafe { std::env::remove_var("VOYAGE_API_KEY") };
331    }
332
333    #[tokio::test]
334    async fn generate_capability_mismatch() {
335        let _lock = ENV_LOCK.lock().await;
336        unsafe { std::env::set_var("VOYAGE_API_KEY", "test-key") };
337
338        let provider = RemoteVoyageAIProvider::new();
339        let s = spec("gen/a", ModelTask::Generate, "voyage-3");
340        let result = provider.load(&s).await;
341        assert!(result.is_err());
342        assert!(
343            result
344                .unwrap_err()
345                .to_string()
346                .contains("does not support task")
347        );
348
349        unsafe { std::env::remove_var("VOYAGE_API_KEY") };
350    }
351
352    fn spec_with_opts(
353        alias: &str,
354        task: ModelTask,
355        model_id: &str,
356        options: serde_json::Value,
357    ) -> ModelAliasSpec {
358        ModelAliasSpec {
359            alias: alias.to_string(),
360            task,
361            provider_id: "remote/voyageai".to_string(),
362            model_id: model_id.to_string(),
363            revision: None,
364            warmup: crate::api::WarmupPolicy::Lazy,
365            required: false,
366            timeout: None,
367            load_timeout: None,
368            retry: None,
369            options,
370        }
371    }
372
373    #[tokio::test]
374    async fn default_embedding_dimensions() {
375        let _lock = ENV_LOCK.lock().await;
376        unsafe { std::env::set_var("VOYAGE_API_KEY", "test-key") };
377
378        let provider = RemoteVoyageAIProvider::new();
379        let s = spec("embed/dim", ModelTask::Embed, "voyage-3");
380        let handle = provider.load(&s).await.unwrap();
381        let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
382        assert_eq!(model.dimensions(), 1024);
383
384        unsafe { std::env::remove_var("VOYAGE_API_KEY") };
385    }
386
387    #[tokio::test]
388    async fn custom_embedding_dimensions() {
389        let _lock = ENV_LOCK.lock().await;
390        unsafe { std::env::set_var("VOYAGE_API_KEY", "test-key") };
391
392        let provider = RemoteVoyageAIProvider::new();
393        let s = spec_with_opts(
394            "embed/dim-custom",
395            ModelTask::Embed,
396            "voyage-3",
397            json!({ "embedding_dimensions": 256 }),
398        );
399        let handle = provider.load(&s).await.unwrap();
400        let model = handle.downcast_ref::<Arc<dyn EmbeddingModel>>().unwrap();
401        assert_eq!(model.dimensions(), 256);
402
403        unsafe { std::env::remove_var("VOYAGE_API_KEY") };
404    }
405}
uni_xervo/provider/voyageai.rs

uni_xervo/provider/
voyageai.rs
