rig_core/providers/
voyageai.rs

1use crate::client::{self, BearerAuth, DebugExt, Provider};
2use crate::embeddings;
3use crate::embeddings::EmbeddingError;
4use crate::http_client::HttpClientExt;
5use crate::rerank;
6use crate::rerank::RerankError;
7use bytes::Bytes;
8use serde::Deserialize;
9use serde_json::json;
10
11// ================================================================
12// Main Voyage AI Client
13// ================================================================
14const VOYAGEAI_API_BASE_URL: &str = "https://api.voyageai.com/v1";
15
16#[derive(Debug, Default, Clone, Copy)]
17pub struct VoyageExt;
18
19#[derive(Debug, Default, Clone, Copy)]
20pub struct VoyageBuilder;
21
22type VoyageApiKey = BearerAuth;
23
24impl Provider for VoyageExt {
25    type Builder = VoyageBuilder;
26
27    /// There is currently no way to verify a Voyage api key without consuming tokens
28    const VERIFY_PATH: &'static str = "";
29}
30
31client::impl_capabilities!(
32    VoyageExt,
33    embeddings = EmbeddingModel<H>,
34    rerank = RerankModel<H>,
35);
36
37impl DebugExt for VoyageExt {}
38
39client::impl_default_provider_builder!(
40    VoyageBuilder => VoyageExt,
41    api_key = VoyageApiKey,
42    base_url = VOYAGEAI_API_BASE_URL,
43);
44
45pub type Client<H = reqwest::Client> = client::Client<VoyageExt, H>;
46pub type ClientBuilder<H = crate::markers::Missing> =
47    client::ClientBuilder<VoyageBuilder, VoyageApiKey, H>;
48
49client::impl_provider_client!(Client, input = String, api_key_env = "VOYAGE_API_KEY");
50
51impl<T> EmbeddingModel<T> {
52    pub fn new(client: Client<T>, model: impl Into<String>, ndims: usize) -> Self {
53        Self {
54            client,
55            model: model.into(),
56            ndims,
57            options: EmbeddingOptions::default(),
58        }
59    }
60
61    pub fn with_model(client: Client<T>, model: &str, ndims: usize) -> Self {
62        Self {
63            client,
64            model: model.into(),
65            ndims,
66            options: EmbeddingOptions::default(),
67        }
68    }
69
70    /// Set optional request parameters for every embedding call made through
71    /// this model. Defaults to [`EmbeddingOptions::default()`] (all `None`).
72    pub fn with_options(mut self, options: EmbeddingOptions) -> Self {
73        self.options = options;
74        self
75    }
76}
77
78// ================================================================
79// Voyage AI Embedding API
80// ================================================================
81
82/// `voyage-3-large` embedding model (Voyage AI)
83pub const VOYAGE_3_LARGE: &str = "voyage-3-large";
84/// `voyage-3.5` embedding model (Voyage AI)
85pub const VOYAGE_3_5: &str = "voyage-3.5";
86/// `voyage-3.5-lite` embedding model (Voyage AI)
87pub const VOYAGE_3_5_LITE: &str = "voyage.3-5.lite";
88/// `voyage-code-3` embedding model (Voyage AI)
89pub const VOYAGE_CODE_3: &str = "voyage-code-3";
90/// `voyage-finance-2` embedding model (Voyage AI)
91pub const VOYAGE_FINANCE_2: &str = "voyage-finance-2";
92/// `voyage-law-2` embedding model (Voyage AI)
93pub const VOYAGE_LAW_2: &str = "voyage-law-2";
94/// `voyage-code-2` embedding model (Voyage AI)
95pub const VOYAGE_CODE_2: &str = "voyage-code-2";
96
97pub fn model_dimensions_from_identifier(model_identifier: &str) -> Option<usize> {
98    match model_identifier {
99        "voyage-code-2" => Some(1536),
100        "voyage-3-large" | "voyage-3.5" | "voyage.3-5.lite" | "voyage-code-3"
101        | "voyage-finance-2" | "voyage-law-2" => Some(1024),
102        _ => None,
103    }
104}
105
106#[derive(Debug, Deserialize)]
107pub struct EmbeddingResponse {
108    pub object: String,
109    pub data: Vec<EmbeddingData>,
110    pub model: String,
111    pub usage: Usage,
112}
113
114#[derive(Clone, Debug, Deserialize)]
115pub struct Usage {
116    pub total_tokens: usize,
117}
118
119#[derive(Debug)]
120pub struct ApiErrorResponse {
121    /// Provider error message; tolerant of `{"message": "..."}`,
122    /// `{"error": "..."}`, nested `{"error": {"message": ...}}`, and bodies
123    /// carrying both keys. Used for logging only — the raw body is preserved
124    /// on the returned error.
125    pub(crate) message: String,
126}
127
128impl<'de> Deserialize<'de> for ApiErrorResponse {
129    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
130    where
131        D: serde::Deserializer<'de>,
132    {
133        Ok(Self {
134            message: crate::providers::internal::envelope::error_message(deserializer)?,
135        })
136    }
137}
138
139#[derive(Debug, Deserialize)]
140#[serde(untagged)]
141pub(crate) enum ApiResponse<T> {
142    Ok(T),
143    Err(ApiErrorResponse),
144}
145
146#[derive(Debug, Deserialize)]
147pub struct EmbeddingData {
148    pub object: String,
149    pub embedding: Vec<f64>,
150    pub index: usize,
151}
152
153/// Optional request parameters for Voyage AI embedding calls.
154///
155/// All fields default to `None`, which matches Voyage's own server defaults:
156/// no `input_type`, `truncation` enabled, and the model's default output
157/// dimension.
158///
159/// TODO: `output_dtype` (`float` | `int8` | `uint8` | `binary` | `ubinary`) is
160/// intentionally not implemented yet. Quantized dtypes return integer vectors
161/// instead of floats, so they need a separate response-parsing change before
162/// they can be supported here.
163#[derive(Clone, Debug, Default, PartialEq, Eq)]
164pub struct EmbeddingOptions {
165    /// Prepends a retrieval prompt to the input text. Use `"document"` when
166    /// embedding stored chunks and `"query"` when embedding search queries.
167    ///
168    /// Embeddings produced with and without `input_type` are compatible.
169    pub input_type: Option<String>,
170    /// Whether to truncate inputs that exceed the model's maximum context
171    /// length. Voyage's server default is `true`.
172    pub truncation: Option<bool>,
173    /// Dimensionality of the returned embeddings. Defaults to the model's
174    /// default output dimension when unset.
175    pub output_dimension: Option<usize>,
176}
177
178#[derive(Clone)]
179pub struct EmbeddingModel<T> {
180    client: Client<T>,
181    pub model: String,
182    ndims: usize,
183    options: EmbeddingOptions,
184}
185
186impl<T> embeddings::EmbeddingModel for EmbeddingModel<T>
187where
188    T: HttpClientExt + Clone + std::fmt::Debug + Default + 'static,
189{
190    const MAX_DOCUMENTS: usize = 1024;
191
192    type Client = Client<T>;
193
194    fn make(client: &Self::Client, model: impl Into<String>, dims: Option<usize>) -> Self {
195        let model = model.into();
196        let dims = dims
197            .or(model_dimensions_from_identifier(&model))
198            .unwrap_or_default();
199
200        Self::new(client.clone(), model, dims)
201    }
202
203    fn ndims(&self) -> usize {
204        self.ndims
205    }
206
207    async fn embed_texts(
208        &self,
209        documents: impl IntoIterator<Item = String>,
210    ) -> Result<Vec<embeddings::Embedding>, EmbeddingError> {
211        let documents: Vec<String> = documents.into_iter().collect();
212        let response = self.embed_texts_with_usage(documents).await?;
213        Ok(response.embeddings)
214    }
215
216    async fn embed_texts_with_usage(
217        &self,
218        documents: impl IntoIterator<Item = String>,
219    ) -> Result<embeddings::EmbeddingResponse, EmbeddingError> {
220        let documents: Vec<String> = documents.into_iter().collect();
221        let mut request = json!({
222            "model": self.model,
223            "input": documents,
224        });
225
226        let request_obj = request.as_object_mut().ok_or_else(|| {
227            EmbeddingError::ResponseError("embedding request body must be a JSON object".into())
228        })?;
229
230        if let Some(input_type) = &self.options.input_type {
231            request_obj.insert("input_type".to_owned(), json!(input_type));
232        }
233        if let Some(truncation) = self.options.truncation {
234            request_obj.insert("truncation".to_owned(), json!(truncation));
235        }
236        if let Some(output_dimension) = self.options.output_dimension {
237            request_obj.insert("output_dimension".to_owned(), json!(output_dimension));
238        }
239
240        let body = serde_json::to_vec(&request)?;
241
242        let req = self
243            .client
244            .post("/embeddings")?
245            .body(body)
246            .map_err(|x| EmbeddingError::HttpError(x.into()))?;
247
248        let response = self.client.send::<_, Bytes>(req).await?;
249        let status = response.status();
250        let response_body = response.into_body().into_future().await?.to_vec();
251
252        if status.is_success() {
253            match serde_json::from_slice::<ApiResponse<EmbeddingResponse>>(&response_body)? {
254                ApiResponse::Ok(response) => {
255                    tracing::info!(target: "rig",
256                        "VoyageAI embedding token usage: {}",
257                        response.usage.total_tokens
258                    );
259
260                    if response.data.len() != documents.len() {
261                        return Err(EmbeddingError::ResponseError(
262                            "Response data length does not match input length".into(),
263                        ));
264                    }
265
266                    let usage = crate::completion::Usage {
267                        input_tokens: response.usage.total_tokens as u64,
268                        output_tokens: 0,
269                        total_tokens: response.usage.total_tokens as u64,
270                        cached_input_tokens: 0,
271                        cache_creation_input_tokens: 0,
272                        tool_use_prompt_tokens: 0,
273                        reasoning_tokens: 0,
274                    };
275
276                    let embeddings = response
277                        .data
278                        .into_iter()
279                        .zip(documents.into_iter())
280                        .map(|(embedding, document)| embeddings::Embedding {
281                            document,
282                            vec: embedding.embedding,
283                        })
284                        .collect();
285
286                    Ok(embeddings::EmbeddingResponse { embeddings, usage })
287                }
288                ApiResponse::Err(err) => {
289                    tracing::warn!(message = %err.message, "provider returned an error response");
290                    Err(EmbeddingError::from_http_response(
291                        status,
292                        String::from_utf8_lossy(&response_body),
293                    ))
294                }
295            }
296        } else {
297            Err(EmbeddingError::from_http_response(
298                status,
299                String::from_utf8_lossy(&response_body),
300            ))
301        }
302    }
303}
304
305// ================================================================
306// Voyage AI Rerank API
307// ================================================================
308
309/// `rerank-2.5` reranker model (Voyage AI)
310pub const RERANK_2_5: &str = "rerank-2.5";
311/// `rerank-2.5-lite` reranker model (Voyage AI)
312pub const RERANK_2_5_LITE: &str = "rerank-2.5-lite";
313/// `rerank-2` reranker model (Voyage AI)
314pub const RERANK_2: &str = "rerank-2";
315/// `rerank-2-lite` reranker model (Voyage AI)
316pub const RERANK_2_LITE: &str = "rerank-2-lite";
317/// `rerank-1` reranker model (Voyage AI)
318pub const RERANK_1: &str = "rerank-1";
319/// `rerank-lite-1` reranker model (Voyage AI)
320pub const RERANK_LITE_1: &str = "rerank-lite-1";
321
322#[derive(Debug, Deserialize)]
323pub struct RerankApiResponse {
324    pub data: Vec<RerankApiData>,
325    pub model: String,
326    pub usage: RerankApiUsage,
327}
328
329#[derive(Debug, Deserialize)]
330pub struct RerankApiUsage {
331    pub total_tokens: usize,
332}
333
334#[derive(Debug, Deserialize)]
335pub struct RerankApiData {
336    pub index: usize,
337    pub relevance_score: f64,
338    #[serde(default)]
339    pub document: Option<String>,
340}
341
342#[derive(Clone)]
343pub struct RerankModel<T = reqwest::Client> {
344    client: Client<T>,
345    pub model: String,
346    pub top_k: Option<usize>,
347    pub return_documents: bool,
348    pub truncation: Option<bool>,
349}
350
351impl<T> RerankModel<T> {
352    pub fn new(client: Client<T>, model: impl Into<String>) -> Self {
353        Self {
354            client,
355            model: model.into(),
356            top_k: None,
357            return_documents: false,
358            truncation: None,
359        }
360    }
361
362    pub fn top_k(mut self, top_k: usize) -> Self {
363        self.top_k = Some(top_k);
364        self
365    }
366
367    pub fn return_documents(mut self, return_documents: bool) -> Self {
368        self.return_documents = return_documents;
369        self
370    }
371
372    pub fn truncation(mut self, truncation: bool) -> Self {
373        self.truncation = Some(truncation);
374        self
375    }
376}
377
378impl<T> rerank::RerankModel for RerankModel<T>
379where
380    T: HttpClientExt + Clone + std::fmt::Debug + Default + 'static,
381{
382    const MAX_DOCUMENTS: usize = 1000;
383
384    type Client = Client<T>;
385
386    fn make(client: &Self::Client, model: impl Into<String>) -> Self {
387        Self::new(client.clone(), model)
388    }
389
390    async fn rerank(
391        &self,
392        query: &str,
393        documents: Vec<String>,
394    ) -> Result<rerank::RerankResponse, RerankError> {
395        let mut body = json!({
396            "query": query,
397            "documents": documents,
398            "model": self.model,
399        });
400
401        let body_obj = body.as_object_mut().ok_or_else(|| {
402            RerankError::ResponseError("rerank request body must be a JSON object".into())
403        })?;
404
405        if let Some(top_k) = self.top_k {
406            body_obj.insert("top_k".to_owned(), json!(top_k));
407        }
408
409        body_obj.insert("return_documents".to_owned(), json!(self.return_documents));
410
411        if let Some(truncation) = self.truncation {
412            body_obj.insert("truncation".to_owned(), json!(truncation));
413        }
414
415        let body = serde_json::to_vec(&body)?;
416
417        let req = self
418            .client
419            .post("/rerank")?
420            .body(body)
421            .map_err(|x| RerankError::HttpError(x.into()))?;
422
423        let response = self.client.send::<_, Bytes>(req).await?;
424        let status = response.status();
425        let response_body = response.into_body().into_future().await?.to_vec();
426
427        if status.is_success() {
428            match serde_json::from_slice::<ApiResponse<RerankApiResponse>>(&response_body)? {
429                ApiResponse::Ok(response) => {
430                    tracing::info!(target: "rig",
431                        "VoyageAI rerank token usage: {}",
432                        response.usage.total_tokens
433                    );
434
435                    let usage = crate::completion::Usage {
436                        input_tokens: response.usage.total_tokens as u64,
437                        output_tokens: 0,
438                        total_tokens: response.usage.total_tokens as u64,
439                        cached_input_tokens: 0,
440                        cache_creation_input_tokens: 0,
441                        reasoning_tokens: 0,
442                        tool_use_prompt_tokens: 0,
443                    };
444
445                    let results = response
446                        .data
447                        .into_iter()
448                        .map(|d| rerank::RerankResult {
449                            index: d.index,
450                            document: d.document,
451                            relevance_score: d.relevance_score,
452                        })
453                        .collect();
454
455                    Ok(rerank::RerankResponse {
456                        results,
457                        model: response.model,
458                        usage,
459                    })
460                }
461                ApiResponse::Err(err) => {
462                    tracing::warn!(message = %err.message, "provider returned an error response");
463                    Err(RerankError::from_http_response(
464                        status,
465                        String::from_utf8_lossy(&response_body),
466                    ))
467                }
468            }
469        } else {
470            Err(RerankError::from_http_response(
471                status,
472                String::from_utf8_lossy(&response_body),
473            ))
474        }
475    }
476}
477
478#[cfg(test)]
479mod tests {
480    #[test]
481    fn test_client_initialization() {
482        let _client =
483            crate::providers::voyageai::Client::new("dummy-key").expect("Client::new() failed");
484        let _client_from_builder = crate::providers::voyageai::Client::builder()
485            .api_key("dummy-key")
486            .build()
487            .expect("Client::builder() failed");
488    }
489
490    #[tokio::test]
491    async fn rerank_non_success_preserves_status_and_body() {
492        use crate::client::RerankingClient;
493        use crate::rerank::{RerankError, RerankModel as _};
494        use crate::test_utils::RecordingHttpClient;
495
496        let body = r#"{"error":{"message":"boom"}}"#;
497        let http_client =
498            RecordingHttpClient::with_error_response(http::StatusCode::SERVICE_UNAVAILABLE, body);
499        let client = super::Client::builder()
500            .api_key("test-key")
501            .http_client(http_client)
502            .build()
503            .expect("build client");
504        let model = client.rerank_model(super::RERANK_2_5);
505
506        let error = model
507            .rerank("query", vec!["doc one".to_string(), "doc two".to_string()])
508            .await
509            .expect_err("rerank should fail with non-success status");
510
511        assert!(matches!(error, RerankError::HttpError(_)));
512        assert_eq!(
513            error.provider_response_status(),
514            Some(http::StatusCode::SERVICE_UNAVAILABLE)
515        );
516        assert_eq!(error.provider_response_body(), Some(body));
517    }
518
519    #[tokio::test]
520    async fn rerank_2xx_error_envelope_preserves_status_and_body() {
521        use crate::client::RerankingClient;
522        use crate::rerank::{RerankError, RerankModel as _};
523        use crate::test_utils::RecordingHttpClient;
524
525        let body = r#"{"message":"boom"}"#;
526        let http_client = RecordingHttpClient::new(body); // 200 OK
527        let client = super::Client::builder()
528            .api_key("test-key")
529            .http_client(http_client)
530            .build()
531            .expect("build client");
532        let model = client.rerank_model(super::RERANK_2_5);
533
534        let error = model
535            .rerank("query", vec!["doc one".to_string(), "doc two".to_string()])
536            .await
537            .expect_err("rerank should fail with provider error envelope");
538
539        match &error {
540            RerankError::ProviderResponse(stored) => {
541                assert_eq!(stored.body, body);
542                assert_eq!(stored.status, Some(http::StatusCode::OK));
543            }
544            other => panic!("expected ProviderResponse, got {other:?}"),
545        }
546    }
547
548    #[tokio::test]
549    async fn embedding_request_includes_options_when_set() {
550        use crate::client::EmbeddingsClient;
551        use crate::embeddings::EmbeddingModel as _;
552        use crate::test_utils::RecordingHttpClient;
553
554        let response_body = r#"{
555            "object": "list",
556            "data": [{"object": "embedding", "embedding": [0.1, 0.2, 0.3], "index": 0}],
557            "model": "voyage-3-large",
558            "usage": {"total_tokens": 7}
559        }"#;
560        let http_client = RecordingHttpClient::new(response_body);
561        let client = super::Client::builder()
562            .api_key("test-key")
563            .http_client(http_client.clone())
564            .build()
565            .expect("build client");
566        let model = client.embedding_model(super::VOYAGE_3_LARGE);
567
568        model
569            .with_options(super::EmbeddingOptions {
570                input_type: Some("document".to_string()),
571                truncation: Some(true),
572                output_dimension: Some(256),
573            })
574            .embed_texts_with_usage(vec!["doc".to_string()])
575            .await
576            .expect("embed should succeed");
577
578        let captured = http_client.requests();
579        assert_eq!(captured.len(), 1);
580        let body: serde_json::Value =
581            serde_json::from_slice(&captured[0].body).expect("request body is valid JSON");
582        assert_eq!(body["model"], super::VOYAGE_3_LARGE);
583        assert_eq!(body["input_type"], "document");
584        assert_eq!(body["truncation"], true);
585        assert_eq!(body["output_dimension"], serde_json::json!(256));
586        assert_eq!(body.get("output_dtype"), None);
587    }
588
589    #[tokio::test]
590    async fn embedding_request_omits_options_when_unset() {
591        use crate::client::EmbeddingsClient;
592        use crate::embeddings::EmbeddingModel as _;
593        use crate::test_utils::RecordingHttpClient;
594
595        let response_body = r#"{
596            "object": "list",
597            "data": [{"object": "embedding", "embedding": [0.1, 0.2, 0.3], "index": 0}],
598            "model": "voyage-3-large",
599            "usage": {"total_tokens": 7}
600        }"#;
601        let http_client = RecordingHttpClient::new(response_body);
602        let client = super::Client::builder()
603            .api_key("test-key")
604            .http_client(http_client.clone())
605            .build()
606            .expect("build client");
607        let model = client.embedding_model(super::VOYAGE_3_LARGE);
608
609        model
610            .embed_texts_with_usage(vec!["doc".to_string()])
611            .await
612            .expect("embed should succeed");
613
614        let captured = http_client.requests();
615        assert_eq!(captured.len(), 1);
616        let body: serde_json::Value =
617            serde_json::from_slice(&captured[0].body).expect("request body is valid JSON");
618        assert_eq!(body["model"], super::VOYAGE_3_LARGE);
619        assert_eq!(body.get("input_type"), None);
620        assert_eq!(body.get("truncation"), None);
621        assert_eq!(body.get("output_dimension"), None);
622    }
623}
rig_core/providers/voyageai.rs

rig_core/providers/
voyageai.rs
