lc_embeddings/
vision_cohere.rs

1// lc-embeddings/src/vision_cohere.rs
2//! Cohere **vision** embeddings — Embed v4 (multimodal) over `v2/embed`.
3//!
4//! Embed v4 maps images and text into one shared 1536-dimensional space.
5//! Images must be sent inline (base64 bytes + format); the API does not
6//! fetch caller-supplied URLs, so [`ImageInput::Url`] is rejected at the
7//! boundary rather than fetched server-side.
8//!
9//! Request (images):
10//! ```json
11//! {
12//!   "model": "embed-v4.0",
13//!   "input_type": "image",
14//!   "images": [{"image_bytes": {"bytes": "<base64>"}, "format": "png"}],
15//!   "embedding_types": ["float"]
16//! }
17//! ```
18//! Response: `{"embeddings": {"float": [[...], ...]}}`.
19
20use async_trait::async_trait;
21use serde::Deserialize;
22use serde_json::json;
23
24use crate::retry::{post_json_with_retry, DEFAULT_RETRY};
25use crate::vision::{ImageInput, VisionEmbeddings};
26use crate::{l2_normalize, CohereEmbedInputType, EmbeddingError};
27
28/// Default Cohere multimodal embedding model.
29pub const COHERE_VISION_EMBED_MODEL: &str = "embed-v4.0";
30
31/// Output dimension of Embed v4 (English and multilingual variants share it).
32pub const COHERE_VISION_DIMENSION: usize = 1536;
33
34/// Cohere vision embedding configuration.
35#[derive(Debug, Clone)]
36pub struct CohereVisionEmbeddingsConfig {
37    /// Cohere API key.
38    pub api_key: String,
39    /// Base URL for the Cohere API (v2).
40    pub base_url: String,
41    /// Multimodal model name.
42    pub model: String,
43    /// Input type used for [`VisionEmbeddings::embed_text`] (image batches
44    /// always use `"image"` per the API).
45    pub text_input_type: CohereEmbedInputType,
46}
47
48impl Default for CohereVisionEmbeddingsConfig {
49    fn default() -> Self {
50        Self {
51            api_key: String::new(),
52            base_url: crate::cohere::COHERE_EMBED_BASE_URL.to_string(),
53            model: COHERE_VISION_EMBED_MODEL.to_string(),
54            text_input_type: CohereEmbedInputType::SearchQuery,
55        }
56    }
57}
58
59impl CohereVisionEmbeddingsConfig {
60    /// Creates a new config with the given API key and defaults
61    /// (`embed-v4.0`, `search_query` for text queries).
62    pub fn new(api_key: impl Into<String>) -> Self {
63        Self {
64            api_key: api_key.into(),
65            ..Default::default()
66        }
67    }
68
69    /// Creates config from `COHERE_API_KEY` (and optional `COHERE_BASE_URL` /
70    /// `COHERE_VISION_EMBED_MODEL`) environment variables.
71    pub fn from_env_result() -> Result<Self, EmbeddingError> {
72        let api_key = std::env::var("COHERE_API_KEY").map_err(|_| {
73            EmbeddingError::Config("COHERE_API_KEY environment variable not set".to_string())
74        })?;
75        let base_url = std::env::var("COHERE_BASE_URL")
76            .unwrap_or_else(|_| crate::cohere::COHERE_EMBED_BASE_URL.to_string());
77        let model = std::env::var("COHERE_VISION_EMBED_MODEL")
78            .unwrap_or_else(|_| COHERE_VISION_EMBED_MODEL.to_string());
79        Ok(Self {
80            api_key,
81            base_url,
82            model,
83            ..Default::default()
84        })
85    }
86
87    /// Sets the model name.
88    pub fn with_model(mut self, model: impl Into<String>) -> Self {
89        self.model = model.into();
90        self
91    }
92
93    /// Sets the input type used for text queries/documents.
94    pub fn with_text_input_type(mut self, input_type: CohereEmbedInputType) -> Self {
95        self.text_input_type = input_type;
96        self
97    }
98}
99
100/// Cohere Embed v4 multimodal embedding provider.
101pub struct CohereVisionEmbeddings {
102    config: CohereVisionEmbeddingsConfig,
103    client: reqwest::Client,
104    dimension: usize,
105}
106
107impl std::fmt::Debug for CohereVisionEmbeddings {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        f.debug_struct("CohereVisionEmbeddings")
110            .field("model", &self.config.model)
111            .finish()
112    }
113}
114
115/// v2/embed response (with `embedding_types: ["float"]`).
116#[derive(Debug, Deserialize)]
117struct CohereV2EmbedResponse {
118    embeddings: CohereV2FloatEmbeddings,
119}
120
121#[derive(Debug, Deserialize)]
122struct CohereV2FloatEmbeddings {
123    float: Vec<Vec<f32>>,
124}
125
126impl CohereVisionEmbeddings {
127    /// Creates the provider, failing fast on an empty API key or an unknown
128    /// model dimension.
129    pub fn new(config: CohereVisionEmbeddingsConfig) -> Result<Self, EmbeddingError> {
130        if config.api_key.trim().is_empty() {
131            return Err(EmbeddingError::Config(
132                "COHERE_API_KEY is empty".to_string(),
133            ));
134        }
135        let dimension = Self::dimension_for(&config.model)?;
136        Ok(Self {
137            config,
138            client: reqwest::Client::new(),
139            dimension,
140        })
141    }
142
143    /// Creates from environment variables.
144    pub fn from_env_result() -> Result<Self, EmbeddingError> {
145        Self::new(CohereVisionEmbeddingsConfig::from_env_result()?)
146    }
147
148    fn dimension_for(model: &str) -> Result<usize, EmbeddingError> {
149        match model {
150            // Embed v4 family is 1536-dimensional (multimodal, multilingual).
151            "embed-v4.0" => Ok(COHERE_VISION_DIMENSION),
152            other => Err(EmbeddingError::Config(format!(
153                "unknown embedding dimension for Cohere vision model '{other}' \
154                 (supported: 'embed-v4.0')"
155            ))),
156        }
157    }
158
159    /// Maps an image MIME type to Cohere's `format` token.
160    fn image_format(mime: &str) -> Result<&'static str, EmbeddingError> {
161        match mime {
162            "image/png" => Ok("png"),
163            "image/jpeg" | "image/jpg" => Ok("jpeg"),
164            "image/webp" => Ok("webp"),
165            "image/gif" => Ok("gif"),
166            other => Err(EmbeddingError::Config(format!(
167                "Cohere vision embeddings support png/jpeg/webp/gif images, got {other}"
168            ))),
169        }
170    }
171
172    fn endpoint(&self) -> String {
173        format!("{}/embed", self.config.base_url.trim_end_matches('/'))
174    }
175
176    /// POST helper shared by text and image calls: applies retry, surfaces
177    /// non-2xx bodies (never silently defaulting), and extracts the float
178    /// embedding matrix.
179    async fn post_float_embeddings(
180        &self,
181        body: &serde_json::Value,
182    ) -> Result<Vec<Vec<f32>>, EmbeddingError> {
183        let response = post_json_with_retry(
184            &self.client,
185            &self.endpoint(),
186            &self.config.api_key,
187            body,
188            &DEFAULT_RETRY,
189        )
190        .await
191        .map_err(|e| EmbeddingError::HttpError(e.to_string()))?;
192
193        let status = response.status();
194        if !status.is_success() {
195            let error_text = response.text().await.map_err(|e| {
196                EmbeddingError::HttpError(format!("failed to read error response body: {e}"))
197            })?;
198            return Err(EmbeddingError::ApiError(format!(
199                "HTTP {}: {}",
200                status, error_text
201            )));
202        }
203
204        let parsed: CohereV2EmbedResponse = response
205            .json()
206            .await
207            .map_err(|e| EmbeddingError::ParseError(e.to_string()))?;
208        Ok(parsed.embeddings.float)
209    }
210}
211
212#[async_trait]
213impl VisionEmbeddings for CohereVisionEmbeddings {
214    async fn embed_image(&self, image: &ImageInput) -> Result<Vec<f32>, EmbeddingError> {
215        let mut batch = self.embed_images(std::slice::from_ref(image)).await?;
216        batch
217            .pop()
218            .ok_or_else(|| EmbeddingError::ApiError("No embedding in response".to_string()))
219    }
220
221    async fn embed_images(&self, images: &[ImageInput]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
222        if images.is_empty() {
223            return Ok(Vec::new());
224        }
225
226        // Inline-only: plain URLs are rejected before any request is made.
227        let mut payload = Vec::with_capacity(images.len());
228        for image in images {
229            let (bytes, mime) = image.inline_parts()?;
230            let format = Self::image_format(&mime)?;
231            payload.push(json!({
232                "image_bytes": {"bytes": bytes},
233                "format": format,
234            }));
235        }
236
237        let body = json!({
238            "model": self.config.model,
239            "input_type": "image",
240            "images": payload,
241            "embedding_types": ["float"],
242        });
243
244        let mut embeddings = self.post_float_embeddings(&body).await?;
245        if embeddings.len() != images.len() {
246            return Err(EmbeddingError::BatchMismatch {
247                expected: images.len(),
248                actual: embeddings.len(),
249            });
250        }
251        for vector in embeddings.iter_mut() {
252            l2_normalize(vector);
253        }
254        Ok(embeddings)
255    }
256
257    async fn embed_text(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
258        if text.trim().is_empty() {
259            return Err(EmbeddingError::EmptyInput);
260        }
261
262        let body = json!({
263            "model": self.config.model,
264            "input_type": self.config.text_input_type.as_str(),
265            "texts": [text],
266            "embedding_types": ["float"],
267        });
268
269        let mut embeddings = self.post_float_embeddings(&body).await?;
270        let mut embedding = embeddings
271            .drain(..)
272            .next()
273            .ok_or_else(|| EmbeddingError::ApiError("No embedding in response".to_string()))?;
274        l2_normalize(&mut embedding);
275        Ok(embedding)
276    }
277
278    fn dimension(&self) -> usize {
279        self.dimension
280    }
281
282    fn model_name(&self) -> &str {
283        &self.config.model
284    }
285}
286
287#[cfg(test)]
288mod tests {
289    use super::*;
290    use crate::test_support::{spawn_json_recording_stub, spawn_status_stub};
291    use std::sync::atomic::Ordering;
292
293    fn config_with(base_url: String) -> CohereVisionEmbeddingsConfig {
294        CohereVisionEmbeddingsConfig {
295            api_key: "test-key".into(),
296            base_url,
297            model: COHERE_VISION_EMBED_MODEL.into(),
298            text_input_type: CohereEmbedInputType::SearchQuery,
299        }
300    }
301
302    #[test]
303    fn construction_validates_key_and_model() {
304        let mut bad_key = config_with("http://127.0.0.1:1".into());
305        bad_key.api_key = String::new();
306        assert!(matches!(
307            CohereVisionEmbeddings::new(bad_key),
308            Err(EmbeddingError::Config(_))
309        ));
310
311        let mut bad_model = config_with("http://127.0.0.1:1".into());
312        bad_model.model = "unknown-vision".into();
313        assert!(matches!(
314            CohereVisionEmbeddings::new(bad_model),
315            Err(EmbeddingError::Config(_))
316        ));
317
318        let ok = CohereVisionEmbeddings::new(config_with("http://127.0.0.1:1".into())).unwrap();
319        assert_eq!(ok.dimension(), COHERE_VISION_DIMENSION);
320        assert_eq!(ok.model_name(), COHERE_VISION_EMBED_MODEL);
321    }
322
323    #[test]
324    fn image_format_maps_known_mimes() {
325        assert_eq!(
326            CohereVisionEmbeddings::image_format("image/png").unwrap(),
327            "png"
328        );
329        assert_eq!(
330            CohereVisionEmbeddings::image_format("image/jpeg").unwrap(),
331            "jpeg"
332        );
333        assert_eq!(
334            CohereVisionEmbeddings::image_format("image/webp").unwrap(),
335            "webp"
336        );
337        assert!(CohereVisionEmbeddings::image_format("image/bmp").is_err());
338    }
339
340    /// B7 snapshot: image requests use the v2 image shape with inline bytes.
341    #[tokio::test]
342    async fn image_request_body_matches_v2_snapshot() {
343        let response = json!({"embeddings": {"float": [[0.6, 0.8], [0.8, -0.6]]}});
344        let (base_url, bodies) = spawn_json_recording_stub(response).await;
345        let embeddings = CohereVisionEmbeddings::new(config_with(base_url)).unwrap();
346
347        let images = vec![
348            ImageInput::from_base64("aW1n", "image/png"),
349            ImageInput::from_data_uri("data:image/jpeg;base64,amZlZw"),
350        ];
351        let vectors = embeddings.embed_images(&images).await.unwrap();
352        assert_eq!(vectors.len(), 2);
353        // L2-normalized.
354        for v in &vectors {
355            let norm: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
356            assert!((norm - 1.0).abs() < 1e-5);
357        }
358
359        let recorded = bodies.lock().unwrap();
360        assert_eq!(recorded.len(), 1);
361        let body = &recorded[0];
362        assert_eq!(body["model"], COHERE_VISION_EMBED_MODEL);
363        assert_eq!(body["input_type"], "image");
364        assert_eq!(body["embedding_types"][0], "float");
365        assert_eq!(
366            body["images"],
367            json!([
368                {"image_bytes": {"bytes": "aW1n"}, "format": "png"},
369                {"image_bytes": {"bytes": "amZlZw"}, "format": "jpeg"},
370            ])
371        );
372    }
373
374    #[tokio::test]
375    async fn text_request_uses_configured_input_type() {
376        let response = json!({"embeddings": {"float": [[0.6, 0.8]]}});
377        let (base_url, bodies) = spawn_json_recording_stub(response).await;
378        let cfg = config_with(base_url).with_text_input_type(CohereEmbedInputType::SearchDocument);
379        let embeddings = CohereVisionEmbeddings::new(cfg).unwrap();
380
381        let v = embeddings.embed_text("a red shoe").await.unwrap();
382        assert_eq!(v.len(), 2);
383
384        let recorded = bodies.lock().unwrap();
385        assert_eq!(recorded[0]["input_type"], "search_document");
386        assert_eq!(recorded[0]["texts"][0], "a red shoe");
387    }
388
389    /// Plain URLs are rejected (Cohere has no URL fetch field; fetching here
390    /// would bypass the SSRF guard).
391    #[tokio::test]
392    async fn plain_url_images_are_rejected() {
393        let response = json!({"embeddings": {"float": [[0.6, 0.8]]}});
394        let (base_url, _bodies) = spawn_json_recording_stub(response).await;
395        let embeddings = CohereVisionEmbeddings::new(config_with(base_url)).unwrap();
396
397        let err = embeddings
398            .embed_image(&ImageInput::from_url("https://example.com/a.png"))
399            .await
400            .unwrap_err();
401        assert!(matches!(err, EmbeddingError::Config(_)));
402    }
403
404    #[tokio::test]
405    async fn batch_mismatch_is_reported() {
406        // Two images sent, stub returns one vector.
407        let response = json!({"embeddings": {"float": [[0.6, 0.8]]}});
408        let (base_url, _bodies) = spawn_json_recording_stub(response).await;
409        let embeddings = CohereVisionEmbeddings::new(config_with(base_url)).unwrap();
410
411        let images = vec![
412            ImageInput::from_base64("aW1n", "image/png"),
413            ImageInput::from_base64("amZlZw", "image/jpeg"),
414        ];
415        let err = embeddings.embed_images(&images).await.unwrap_err();
416        assert!(matches!(
417            err,
418            EmbeddingError::BatchMismatch {
419                expected: 2,
420                actual: 1
421            }
422        ));
423    }
424
425    #[tokio::test]
426    async fn retries_on_429_and_rejects_empty_inputs() {
427        let success_body = r#"{"embeddings":{"float":[[0.6,0.8]]}}"#;
428        let (base_url, requests) = spawn_status_stub(429, 2, 200, success_body).await;
429        let embeddings = CohereVisionEmbeddings::new(config_with(base_url)).unwrap();
430
431        let v = embeddings
432            .embed_text("hello")
433            .await
434            .expect("should retry after two 429s");
435        assert_eq!(v.len(), 2);
436        assert_eq!(requests.load(Ordering::SeqCst), 3, "1 initial + 2 retries");
437
438        assert!(matches!(
439            embeddings.embed_text("  ").await,
440            Err(EmbeddingError::EmptyInput)
441        ));
442        assert_eq!(
443            embeddings.embed_images(&[]).await.unwrap(),
444            Vec::<Vec<f32>>::new()
445        );
446    }
447}
lc_embeddings/vision_cohere.rs

lc_embeddings/
vision_cohere.rs
