rig_core/providers/voyageai/
wire.rs

1//! Voyage AI configuration and embedding and reranking endpoint wires.
2//!
3//! ```no_run
4//! use rig_core::providers::voyageai::{VoyageAi, VOYAGE_3_5};
5//! let wire = VoyageAi::from_env()?.embedding(VOYAGE_3_5, None);
6//! # Ok::<(), Box<dyn std::error::Error>>(())
7//! ```
8
9use crate::client::env::{self, EnvError};
10use crate::error::EncodeError;
11use crate::error::ProviderError;
12use crate::operation::{Embedding, Rerank as RerankOp, RerankRequest};
13use crate::rerank::{RerankResponse, RerankResult};
14use crate::wire::{
15    Body, Capabilities, Descriptor, DocumentDecoder, Encoded, Framing, Mode, Secret, Wire,
16};
17use serde::{Deserialize, Serialize};
18
19use super::{
20    EmbeddingResponse as VoyageEmbeddingResponse, RerankApiResponse, VOYAGEAI_API_BASE_URL,
21    model_dimensions_from_identifier,
22};
23
24/// The provider descriptor name, as records and telemetry spell it.
25const PROVIDER_NAME: &str = "voyageai";
26
27/// The environment variable carrying the API key.
28const API_KEY_ENV: &str = "VOYAGE_API_KEY";
29
30/// The most texts `POST /embeddings` accepts in one call.
31const MAX_DOCUMENTS: usize = 1024;
32
33/// The most documents `POST /rerank` orders in one call.
34const MAX_RERANK_DOCUMENTS: usize = 1000;
35
36/// The settings of a Voyage AI provider: serializable, and the credential
37/// is never serialized. [`connect`](Self::connect) puts it on a transport as
38/// a [`VoyageAi`](super::VoyageAi) client.
39#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
40pub struct VoyageAiConfig {
41    /// The API key, sent as `Authorization: Bearer`.
42    pub api_key: Secret,
43    /// The API root, without a trailing slash.
44    pub base_url: String,
45}
46
47impl VoyageAiConfig {
48    /// Voyage AI with default settings.
49    pub fn new(api_key: impl Into<Secret>) -> Self {
50        Self {
51            api_key: api_key.into(),
52            base_url: VOYAGEAI_API_BASE_URL.to_owned(),
53        }
54    }
55
56    /// Voyage AI from `VOYAGE_API_KEY`.
57    pub fn from_env() -> Result<Self, EnvError> {
58        Ok(Self::new(env::required(API_KEY_ENV)?))
59    }
60
61    /// Point the wires at another API root.
62    pub fn with_base_url(mut self, base_url: impl AsRef<str>) -> Self {
63        self.base_url = base_url.as_ref().trim_end_matches('/').to_owned();
64        self
65    }
66
67    /// Build an embedding wire reporting `ndims`, or the known model width,
68    /// or zero if unknown. This does not send an output-dimension override;
69    /// use [`Embeddings::with_output_dimension`] for that.
70    pub(crate) fn embedding(&self, model: impl Into<String>, ndims: Option<usize>) -> Embeddings {
71        let model = model.into();
72        let ndims = ndims
73            .or_else(|| model_dimensions_from_identifier(&model))
74            .unwrap_or_default();
75        Embeddings {
76            provider: self.clone(),
77            model,
78            ndims,
79            input_type: None,
80            truncation: None,
81            output_dimension: None,
82        }
83    }
84
85    /// The rerank wire for `model`.
86    pub(crate) fn rerank(&self, model: impl Into<String>) -> Rerank {
87        Rerank {
88            provider: self.clone(),
89            model: model.into(),
90            top_k: None,
91            return_documents: false,
92            truncation: None,
93        }
94    }
95
96    /// One request, authenticated and typed as JSON.
97    fn post(&self, path: &str) -> http::request::Builder {
98        http::Request::post(format!("{}{path}", self.base_url))
99            .header(http::header::CONTENT_TYPE, "application/json")
100            .header(
101                http::header::AUTHORIZATION,
102                format!("Bearer {}", self.api_key.expose()),
103            )
104    }
105}
106
107/// The embedding wire: `POST /embeddings`.
108///
109/// Every option defaults to `None`, which is Voyage's own server default:
110/// no retrieval prompt, truncation enabled, and the model's default output
111/// dimension.
112#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
113pub struct Embeddings {
114    /// The provider this wire speaks to.
115    pub provider: VoyageAiConfig,
116    /// The model to address.
117    pub model: String,
118    /// The width this wire reports, from the caller or the model's published
119    /// dimensions. `0` means neither named one.
120    pub ndims: usize,
121    /// Prepends a retrieval prompt to the input text: `"document"` when
122    /// embedding stored chunks, `"query"` when embedding search queries.
123    /// Embeddings produced with and without it are compatible.
124    pub input_type: Option<String>,
125    /// Whether to truncate inputs longer than the model's context. Voyage's
126    /// server default is `true`.
127    pub truncation: Option<bool>,
128    /// Dimensionality of the returned vectors, when overriding the model's
129    /// default output dimension.
130    pub output_dimension: Option<usize>,
131}
132
133impl Embeddings {
134    /// Prepend Voyage's retrieval prompt for this input's role.
135    pub fn with_input_type(mut self, input_type: impl Into<String>) -> Self {
136        self.input_type = Some(input_type.into());
137        self
138    }
139
140    /// Set whether overlong inputs are truncated (`true`) or rejected (`false`).
141    pub fn with_truncation(mut self, truncation: bool) -> Self {
142        self.truncation = Some(truncation);
143        self
144    }
145
146    /// Ask Voyage for vectors of this width.
147    ///
148    /// This is the width the request carries; [`Self::ndims`] is the width
149    /// the wire reports to a vector store, so they are set together.
150    pub fn with_output_dimension(mut self, output_dimension: usize) -> Self {
151        self.output_dimension = Some(output_dimension);
152        self.ndims = output_dimension;
153        self
154    }
155}
156
157impl Wire for Embeddings {
158    type Op = Embedding;
159    type Payload = crate::wire::Encoded;
160    type Frame = crate::wire::WireFrame;
161    type Event = VoyageEmbeddingResponse;
162    type Decoder<'id> = DocumentDecoder<Embedding, Self::Event>;
163    type Reassembler = crate::wire::document::Unreassembled;
164
165    fn describe(&self) -> Descriptor<'_> {
166        Descriptor::new(PROVIDER_NAME)
167            .model(self.model.as_str())
168            .capabilities(Capabilities::embedding(MAX_DOCUMENTS, self.ndims))
169    }
170
171    fn encode(&self, texts: Vec<String>, _mode: Mode) -> Result<Encoded, EncodeError> {
172        let mut body = serde_json::Map::new();
173        body.insert("model".to_owned(), serde_json::json!(self.model));
174        body.insert("input".to_owned(), serde_json::Value::from(texts));
175        if let Some(input_type) = &self.input_type {
176            body.insert("input_type".to_owned(), serde_json::json!(input_type));
177        }
178        if let Some(truncation) = self.truncation {
179            body.insert("truncation".to_owned(), serde_json::json!(truncation));
180        }
181        if let Some(output_dimension) = self.output_dimension {
182            body.insert(
183                "output_dimension".to_owned(),
184                serde_json::json!(output_dimension),
185            );
186        }
187        let request = self
188            .provider
189            .post("/embeddings")
190            .body(Body::Bytes(serde_json::to_vec(&body)?))?;
191        Ok(Encoded::new(request, Framing::Whole))
192    }
193
194    /// Decodes one `/embeddings` reply.
195    fn decoder<'id>(&self) -> Self::Decoder<'id> {
196        DocumentDecoder::new(
197            |body| crate::providers::internal::wire::classify_marker_keyed_frame(body, &["data"]),
198            |reply| {
199                // Voyage reports one count; every token of an embedding is input.
200                let usage = crate::completion::Usage {
201                    input_tokens: Some(reply.usage.total_tokens as u64),
202                    total_tokens: Some(reply.usage.total_tokens as u64),
203                    ..Default::default()
204                };
205                let vectors = reply.data.into_iter().map(|embedding| embedding.embedding);
206                Ok(crate::embeddings::EmbeddingResponse {
207                    model: Some(reply.model),
208                    usage,
209                    ..crate::embeddings::EmbeddingResponse::from_vectors(vectors)
210                })
211            },
212        )
213    }
214}
215
216/// The rerank wire: `POST /rerank`.
217#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
218pub struct Rerank {
219    /// The provider this wire speaks to.
220    pub provider: VoyageAiConfig,
221    /// The model to address.
222    pub model: String,
223    /// Return only the `top_k` most relevant documents. `None` returns them
224    /// all.
225    pub top_k: Option<usize>,
226    /// Whether the reply echoes each document's text back.
227    pub return_documents: bool,
228    /// Whether to truncate documents longer than the model's context.
229    /// Voyage's server default is `true`.
230    pub truncation: Option<bool>,
231}
232
233impl Rerank {
234    /// Order only the `top_k` most relevant documents.
235    pub fn with_top_k(mut self, top_k: usize) -> Self {
236        self.top_k = Some(top_k);
237        self
238    }
239
240    /// Ask Voyage to echo each ordered document's text back.
241    pub fn with_return_documents(mut self, return_documents: bool) -> Self {
242        self.return_documents = return_documents;
243        self
244    }
245
246    /// Set whether overlong documents are truncated (`true`) or rejected (`false`).
247    pub fn with_truncation(mut self, truncation: bool) -> Self {
248        self.truncation = Some(truncation);
249        self
250    }
251}
252
253impl Wire for Rerank {
254    type Op = RerankOp;
255    type Payload = crate::wire::Encoded;
256    type Frame = crate::wire::WireFrame;
257    type Event = Result<RerankApiResponse, String>;
258    type Decoder<'id> = DocumentDecoder<RerankOp, Self::Event>;
259    type Reassembler = crate::wire::document::Unreassembled;
260
261    fn describe(&self) -> Descriptor<'_> {
262        Descriptor::new(PROVIDER_NAME)
263            .model(self.model.as_str())
264            .capabilities(Capabilities::rerank(MAX_RERANK_DOCUMENTS))
265    }
266
267    fn encode(&self, request: RerankRequest, _mode: Mode) -> Result<Encoded, EncodeError> {
268        let mut body = serde_json::Map::new();
269        body.insert("query".to_owned(), serde_json::Value::from(request.query));
270        body.insert(
271            "documents".to_owned(),
272            serde_json::Value::from(request.documents),
273        );
274        body.insert("model".to_owned(), serde_json::json!(self.model));
275        if let Some(top_k) = self.top_k {
276            body.insert("top_k".to_owned(), serde_json::json!(top_k));
277        }
278        body.insert(
279            "return_documents".to_owned(),
280            serde_json::json!(self.return_documents),
281        );
282        if let Some(truncation) = self.truncation {
283            body.insert("truncation".to_owned(), serde_json::json!(truncation));
284        }
285        let request = self
286            .provider
287            .post("/rerank")
288            .body(Body::Bytes(serde_json::to_vec(&body)?))?;
289        Ok(Encoded::new(request, Framing::Whole))
290    }
291
292    /// Decodes one `/rerank` reply.
293    /// The ordering, or the error envelope Voyage can answer a **200** with.
294    fn decoder<'id>(&self) -> Self::Decoder<'id> {
295        DocumentDecoder::new(
296            |body| {
297                crate::providers::internal::wire::classify_reply_or_message_envelope(body, "data")
298            },
299            |reply| {
300                let reply = reply.map_err(ProviderError::from_provider_body)?;
301                // Voyage reports one count; every token of a rerank is input.
302                let usage = crate::completion::Usage {
303                    input_tokens: Some(reply.usage.total_tokens as u64),
304                    total_tokens: Some(reply.usage.total_tokens as u64),
305                    ..Default::default()
306                };
307                let results = reply
308                    .data
309                    .into_iter()
310                    .map(|result| RerankResult {
311                        index: result.index,
312                        document: result.document,
313                        relevance_score: result.relevance_score,
314                    })
315                    .collect();
316                Ok(RerankResponse {
317                    model: Some(reply.model),
318                    usage,
319                    ..RerankResponse::new(results)
320                })
321            },
322        )
323    }
324}
325
326#[cfg(test)]
327mod tests;
rig_core/providers/voyageai/wire.rs

rig_core/providers/voyageai/
wire.rs
