rig_core/providers/cohere/
embeddings.rs

1//! Cohere's text-embedding and image-embedding wires (`POST /v1/embed`),
2//! their reply types, and image-input validation.
3//!
4//! ```
5//! use rig_core::providers::cohere::embeddings::FloatEmbeddings;
6//! let vectors: FloatEmbeddings = serde_json::from_str(r#"{"float": [[0.5]]}"#)?;
7//! assert_eq!(vectors.values.len(), 1);
8//! # Ok::<(), serde_json::Error>(())
9//! ```
10
11use crate::error::{EncodeError, ProviderError};
12use crate::operation::{Embedding, ImageEmbedding};
13use crate::providers::internal::wire::classify_reply_or_message_envelope;
14use crate::wire::{Body, Capabilities, Descriptor, DocumentDecoder, Encoded, Framing, Mode, Wire};
15use base64::{Engine as _, engine::general_purpose::STANDARD};
16use serde::{Deserialize, Serialize};
17
18use super::{CohereConfig, PROVIDER_NAME};
19
20/// `embed-v4.0` embedding model
21pub const EMBED_V4: &str = "embed-v4.0";
22/// `embed-english-v3.0` embedding model
23pub const EMBED_ENGLISH_V3: &str = "embed-english-v3.0";
24/// `embed-english-light-v3.0` embedding model
25pub const EMBED_ENGLISH_LIGHT_V3: &str = "embed-english-light-v3.0";
26/// `embed-multilingual-v3.0` embedding model
27pub const EMBED_MULTILINGUAL_V3: &str = "embed-multilingual-v3.0";
28/// `embed-multilingual-light-v3.0` embedding model
29pub const EMBED_MULTILINGUAL_LIGHT_V3: &str = "embed-multilingual-light-v3.0";
30
31pub(crate) fn model_dimensions_from_identifier(identifier: &str) -> Option<usize> {
32    match identifier {
33        EMBED_V4 => Some(1_536),
34        EMBED_ENGLISH_V3 | EMBED_MULTILINGUAL_V3 => Some(1_024),
35        EMBED_ENGLISH_LIGHT_V3 | EMBED_MULTILINGUAL_LIGHT_V3 => Some(384),
36        _ => None,
37    }
38}
39
40impl CohereConfig {
41    /// Build a text-embedding wire reporting the supplied or known model width,
42    /// or zero if unknown. This width is metadata, not a request parameter.
43    pub(crate) fn embedding(&self, model: impl Into<String>, ndims: Option<usize>) -> Embeddings {
44        let model = model.into();
45        let ndims = ndims
46            .or_else(|| model_dimensions_from_identifier(&model))
47            .unwrap_or_default();
48        Embeddings {
49            provider: self.clone(),
50            model,
51            ndims,
52            input_type: DEFAULT_INPUT_TYPE.to_owned(),
53        }
54    }
55
56    /// The image-embedding wire.
57    ///
58    /// Cohere Embed v3 embeds images with one fixed model, so this wire
59    /// names no model.
60    pub(crate) fn image_embedding(&self) -> ImageEmbeddings {
61        ImageEmbeddings {
62            provider: self.clone(),
63        }
64    }
65
66    /// One request, authenticated and typed as JSON.
67    pub(super) fn post(&self, path: &str) -> http::request::Builder {
68        http::Request::post(format!("{}{path}", self.base_url))
69            .header(http::header::CONTENT_TYPE, "application/json")
70            .header(
71                http::header::AUTHORIZATION,
72                format!("Bearer {}", self.api_key.expose()),
73            )
74    }
75}
76
77/// Default retrieval role for embeddings of stored document chunks.
78const DEFAULT_INPUT_TYPE: &str = "search_document";
79
80/// The most texts Cohere embeds in one `/v1/embed` call.
81const MAX_DOCUMENTS: usize = 96;
82
83/// The width Cohere's image embeddings come back at.
84const IMAGE_NDIMS: usize = 1_024;
85
86/// The text-embedding wire: `POST /v1/embed`.
87#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
88pub struct Embeddings {
89    /// The provider this wire speaks to.
90    pub provider: CohereConfig,
91    /// The model to address.
92    pub model: String,
93    /// The width this wire reports, from the caller or the model's published
94    /// dimensions. `0` means neither named one.
95    pub ndims: usize,
96    /// Cohere's retrieval prompt: `search_document` for stored chunks,
97    /// `search_query` for queries, `classification`, `clustering`.
98    pub input_type: String,
99}
100
101impl Embeddings {
102    /// Embed for a different purpose than storing chunks.
103    pub fn with_input_type(mut self, input_type: impl Into<String>) -> Self {
104        self.input_type = input_type.into();
105        self
106    }
107}
108
109impl Wire for Embeddings {
110    type Op = Embedding;
111    type Payload = crate::wire::Encoded;
112    type Frame = crate::wire::WireFrame;
113    type Event = Result<EmbeddingResponse, String>;
114    type Decoder<'id> = DocumentDecoder<Embedding, Self::Event>;
115    type Reassembler = crate::wire::document::Unreassembled;
116
117    fn describe(&self) -> Descriptor<'_> {
118        Descriptor::new(PROVIDER_NAME)
119            .model(self.model.as_str())
120            .capabilities(Capabilities::embedding(MAX_DOCUMENTS, self.ndims))
121    }
122
123    fn encode(&self, texts: Vec<String>, _mode: Mode) -> Result<Encoded, EncodeError> {
124        let body = crate::providers::internal::object! {
125            "model": self.model.as_str(),
126            "texts": texts,
127            "input_type": self.input_type.as_str(),
128        };
129        let request = self
130            .provider
131            .post("/v1/embed")
132            .body(Body::Bytes(serde_json::to_vec(&body)?))?;
133        Ok(Encoded::new(request, Framing::Whole))
134    }
135
136    /// The vectors, or the error envelope Cohere can answer a **200** with.
137    fn decoder<'id>(&self) -> Self::Decoder<'id> {
138        DocumentDecoder::new(
139            |body| classify_reply_or_message_envelope(body, "embeddings"),
140            |reply| {
141                let reply = reply.map_err(ProviderError::from_provider_body)?;
142                // Cohere's `/v1/embed` reply names no model.
143                Ok(crate::embeddings::EmbeddingResponse {
144                    response_id: Some(reply.id),
145                    usage: usage_of(reply.meta.as_ref()),
146                    ..crate::embeddings::EmbeddingResponse::from_vectors(reply.embeddings)
147                })
148            },
149        )
150    }
151}
152
153/// The image-embedding wire: `POST /v1/embed`, one image per request.
154///
155/// Cohere Embed v3 accepts a single image per call, so the wire's batch limit
156/// is one image.
157#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
158pub struct ImageEmbeddings {
159    /// The provider this wire speaks to.
160    pub provider: CohereConfig,
161}
162
163impl Wire for ImageEmbeddings {
164    type Op = ImageEmbedding;
165    type Payload = crate::wire::Encoded;
166    type Frame = crate::wire::WireFrame;
167    type Event = Result<ImageEmbeddingResponse, String>;
168    type Decoder<'id> = DocumentDecoder<ImageEmbedding, Self::Event>;
169    type Reassembler = crate::wire::document::Unreassembled;
170
171    fn describe(&self) -> Descriptor<'_> {
172        Descriptor::new(PROVIDER_NAME)
173            .model(EMBED_ENGLISH_V3)
174            .capabilities(Capabilities::embedding(1, IMAGE_NDIMS))
175    }
176
177    fn encode(&self, images: Vec<Vec<u8>>, _mode: Mode) -> Result<Encoded, EncodeError> {
178        // The wire's batch limit is one image: a caller splits larger
179        // batches into one call per image.
180        let [image] = images.as_slice() else {
181            return Err(EncodeError::request(format!(
182                "Cohere embeds one image per request, not {}",
183                images.len()
184            )));
185        };
186        let media_type = validate_image(image)?;
187        let body = serde_json::json!({
188            "model": EMBED_ENGLISH_V3,
189            "images": [image_data_url(image, media_type)],
190            "input_type": "image",
191            "embedding_types": ["float"],
192        });
193        let request = self
194            .provider
195            .post("/v1/embed")
196            .body(Body::Bytes(serde_json::to_vec(&body)?))?;
197        Ok(Encoded::new(request, Framing::Whole))
198    }
199
200    /// The vector, or the error envelope Cohere can answer a **200** with.
201    fn decoder<'id>(&self) -> Self::Decoder<'id> {
202        DocumentDecoder::new(
203            |body| classify_reply_or_message_envelope(body, "embeddings"),
204            |reply| {
205                let reply = reply.map_err(ProviderError::from_provider_body)?;
206                // Each request carries one image, so any other vector count is invalid.
207                let [vector] = reply.embeddings.values.as_slice() else {
208                    return Err(ProviderError::Response(format!(
209                        "Expected 1 image embedding, got {}",
210                        reply.embeddings.values.len()
211                    )));
212                };
213                // The fold names the input: an image has no text, and its bytes
214                // must never travel back in a response.
215                Ok(crate::embeddings::EmbeddingResponse {
216                    usage: usage_of(reply.meta.as_ref()),
217                    response_id: reply.id,
218                    ..crate::embeddings::EmbeddingResponse::from_vectors([vector.clone()])
219                })
220            },
221        )
222    }
223}
224
225const MAX_IMAGE_BYTES: usize = 5_000_000;
226
227/// The usage a reply's `meta` bills, or none.
228fn usage_of(meta: Option<&Meta>) -> crate::completion::Usage {
229    meta.map(|meta| meta.billed_units.to_usage())
230        .unwrap_or_default()
231}
232
233/// One `/v1/embed` reply, as far as rig reads it.
234#[derive(Debug, Clone, Serialize, Deserialize)]
235pub struct EmbeddingResponse {
236    pub id: String,
237    pub embeddings: Vec<Vec<f64>>,
238    #[serde(default)]
239    pub meta: Option<Meta>,
240}
241
242/// A reply's `meta`: what the call billed.
243#[derive(Debug, Clone, Serialize, Deserialize)]
244pub struct Meta {
245    #[serde(default)]
246    pub billed_units: BilledUnits,
247}
248
249/// Cohere's `meta.billed_units`. Token counters are absent when the request
250/// was not billed in tokens (image embeds bill `images`), so they are
251/// `Option`; the non-token counters default to zero.
252#[derive(Debug, Clone, Default, Serialize, Deserialize)]
253pub struct BilledUnits {
254    #[serde(skip_serializing_if = "Option::is_none")]
255    pub input_tokens: Option<u32>,
256    #[serde(skip_serializing_if = "Option::is_none")]
257    pub output_tokens: Option<u32>,
258    #[serde(default)]
259    pub search_units: u32,
260    #[serde(default)]
261    pub classifications: u32,
262    #[serde(default)]
263    pub images: u32,
264}
265
266/// One Cohere `/v1/embed` answer for a single image, one per input image:
267/// Cohere Embed v3 accepts a single image per call.
268#[derive(Debug, Clone, Serialize, Deserialize)]
269pub struct ImageEmbeddingResponse {
270    #[serde(default)]
271    pub id: Option<String>,
272    pub embeddings: FloatEmbeddings,
273    #[serde(default)]
274    pub meta: Option<Meta>,
275}
276
277#[derive(Debug, Clone, Serialize, Deserialize)]
278pub struct FloatEmbeddings {
279    #[serde(rename = "float")]
280    pub values: Vec<Vec<f64>>,
281}
282
283impl BilledUnits {
284    /// Maps the billed token counters straight through; `total_tokens` is
285    /// the sum of whichever counters Cohere sent (an embed bills input only,
286    /// so it reports `input_tokens` and `total_tokens`, no `output_tokens`).
287    pub(super) fn to_usage(&self) -> crate::completion::Usage {
288        let input_tokens = self.input_tokens.map(u64::from);
289        let output_tokens = self.output_tokens.map(u64::from);
290        let total_tokens = match (input_tokens, output_tokens) {
291            (None, None) => None,
292            (input, output) => Some(input.unwrap_or(0) + output.unwrap_or(0)),
293        };
294        crate::completion::Usage {
295            input_tokens,
296            output_tokens,
297            total_tokens,
298            ..Default::default()
299        }
300    }
301}
302
303#[derive(Debug, thiserror::Error)]
304pub(super) enum ImageInputError {
305    #[error("Cohere image embeddings support PNG, JPEG, WebP, or GIF file bytes")]
306    UnsupportedFormat,
307    #[error("Cohere image embeddings accept at most 5 MB per image; received {actual_bytes} bytes")]
308    TooLarge { actual_bytes: usize },
309}
310
311/// Detect an accepted image media type. Returns a document error for unsupported
312/// formats or inputs exceeding 5,000,000 bytes.
313pub(super) fn validate_image(bytes: &[u8]) -> Result<&'static str, EncodeError> {
314    if bytes.len() > MAX_IMAGE_BYTES {
315        return Err(EncodeError::request(ImageInputError::TooLarge {
316            actual_bytes: bytes.len(),
317        }));
318    }
319
320    crate::embeddings::image_media_type(bytes)
321        .ok_or_else(|| EncodeError::request(ImageInputError::UnsupportedFormat))
322}
323
324pub(super) fn image_data_url(bytes: &[u8], media_type: &str) -> String {
325    format!("data:{media_type};base64,{}", STANDARD.encode(bytes))
326}
327
328#[cfg(test)]
329mod tests;
rig_core/providers/cohere/embeddings.rs

rig_core/providers/cohere/
embeddings.rs
