lc_embeddings/
vision.rs

1// lc-embeddings/src/vision.rs
2//! Vision (multimodal) embeddings — image **and** text in one shared vector
3//! space (B7, v0.22.4).
4//!
5//! Text-only [`Embeddings`] models map strings to vectors. Multimodal models
6//! (Cohere Embed v4, Alibaba DashScope `multimodal-embedding-v1`) map images
7//! *and* text into the **same** space, which is what makes cross-modal
8//! retrieval possible: store vectors of product photos, query them with a
9//! plain sentence such as “红色运动鞋”. This module defines:
10//!
11//! - [`ImageInput`] — provider-neutral image reference (URL / data URI / raw base64);
12//! - [`VisionEmbeddings`] — the cross-modal embedding trait;
13//! - [`MockVisionEmbeddings`] — deterministic offline backend for tests.
14//!
15//! As with [`Embeddings`], every returned vector is L2-normalized by the
16//! concrete backends so cosine/dot-product results stay comparable across
17//! providers.
18
19use std::collections::HashMap;
20use std::sync::Mutex;
21
22use async_trait::async_trait;
23use serde::{Deserialize, Serialize};
24
25use crate::{l2_normalize, EmbeddingError};
26
27/// Provider-neutral image input for vision embedding models.
28#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
29#[non_exhaustive]
30pub enum ImageInput {
31    /// A complete data URI (`data:image/png;base64,<data>`). Sent inline.
32    DataUri(String),
33    /// A provider-fetchable image URL (support depends on the backend —
34    /// DashScope fetches public URLs; the Cohere API requires inline bytes).
35    Url(String),
36    /// Raw base64-encoded image bytes with an explicit MIME type
37    /// (e.g. `image/png`).
38    Base64 {
39        /// Base64-encoded image bytes (no data-URI header).
40        data: String,
41        /// Image MIME type.
42        mime_type: String,
43    },
44}
45
46impl ImageInput {
47    /// Creates an image input from a provider-fetchable URL.
48    pub fn from_url(url: impl Into<String>) -> Self {
49        ImageInput::Url(url.into())
50    }
51
52    /// Creates an image input from a complete data URI.
53    pub fn from_data_uri(uri: impl Into<String>) -> Self {
54        ImageInput::DataUri(uri.into())
55    }
56
57    /// Creates an image input from raw base64 bytes and a MIME type.
58    pub fn from_base64(data: impl Into<String>, mime_type: impl Into<String>) -> Self {
59        ImageInput::Base64 {
60            data: data.into(),
61            mime_type: mime_type.into(),
62        }
63    }
64
65    /// Validates the input and returns its `(raw base64 bytes, MIME type)` for
66    /// backends that require inline bytes (Cohere).
67    ///
68    /// Plain [`ImageInput::Url`] values are rejected: such backends cannot ask
69    /// the model provider to perform a server-side fetch, and fetching here
70    /// would route caller-supplied URLs through embedding hosts without the
71    /// SSRF guard. Callers that need URL ingestion should resolve the image
72    /// themselves (or use a backend with a reference field, like DashScope).
73    pub(crate) fn inline_parts(&self) -> Result<(String, String), EmbeddingError> {
74        match self {
75            ImageInput::Base64 { data, mime_type } => {
76                Self::validate_inline(data, mime_type)?;
77                Ok((data.clone(), mime_type.clone()))
78            }
79            ImageInput::DataUri(uri) => {
80                let (mime, data) = parse_data_uri(uri)?;
81                Ok((data.to_string(), mime.to_string()))
82            }
83            ImageInput::Url(url) => Err(EmbeddingError::Config(format!(
84                "this vision embedding backend requires inline image bytes, got a URL: {url}"
85            ))),
86        }
87    }
88
89    /// Returns the reference string for a backend `image` field (DashScope):
90    /// URLs pass through, base64 inputs are rendered as data URIs.
91    pub(crate) fn reference(&self) -> Result<String, EmbeddingError> {
92        match self {
93            ImageInput::Url(url) => {
94                if url.trim().is_empty() {
95                    return Err(EmbeddingError::EmptyInput);
96                }
97                Ok(url.clone())
98            }
99            ImageInput::DataUri(uri) => {
100                let _ = parse_data_uri(uri)?;
101                Ok(uri.clone())
102            }
103            ImageInput::Base64 { data, mime_type } => {
104                Self::validate_inline(data, mime_type)?;
105                Ok(format!("data:{mime_type};base64,{data}"))
106            }
107        }
108    }
109
110    /// Stable key for mock-vector lookup: the full data URI for inline inputs.
111    pub(crate) fn mock_key(&self) -> Result<String, EmbeddingError> {
112        match self {
113            ImageInput::Url(url) => Ok(url.clone()),
114            ImageInput::DataUri(uri) => Ok(uri.clone()),
115            ImageInput::Base64 { data, mime_type } => Ok(format!("data:{mime_type};base64,{data}")),
116        }
117    }
118
119    fn validate_inline(data: &str, mime: &str) -> Result<(), EmbeddingError> {
120        if data.trim().is_empty() {
121            return Err(EmbeddingError::EmptyInput);
122        }
123        if !mime.starts_with("image/") {
124            return Err(EmbeddingError::Config(format!(
125                "vision embeddings require an image/* MIME type, got {mime:?}"
126            )));
127        }
128        Ok(())
129    }
130}
131
132/// Splits `data:<mime>;base64,<data>` into `(mime, raw base64)`.
133pub(crate) fn parse_data_uri(uri: &str) -> Result<(&str, &str), EmbeddingError> {
134    let rest = uri
135        .strip_prefix("data:")
136        .ok_or_else(|| EmbeddingError::Config(format!("not a data URI: {uri}")))?;
137    let comma = rest.find(',').ok_or_else(|| {
138        EmbeddingError::Config(format!("malformed data URI (missing comma): {uri}"))
139    })?;
140    let meta = &rest[..comma];
141    let data = &rest[comma + 1..];
142    if !meta.contains("base64") {
143        return Err(EmbeddingError::Config(
144            "data URI must carry base64-encoded image bytes".to_string(),
145        ));
146    }
147    let mime = meta
148        .split(';')
149        .next()
150        .map(str::trim)
151        .filter(|m| m.starts_with("image/") && !m.is_empty())
152        .ok_or_else(|| {
153            EmbeddingError::Config(format!("data URI is missing an image/* MIME type: {uri}"))
154        })?;
155    if data.trim().is_empty() {
156        return Err(EmbeddingError::EmptyInput);
157    }
158    Ok((mime, data))
159}
160
161/// Cross-modal embedding model: images and text mapped to one shared space.
162///
163/// # Contract
164///
165/// - `embed_text` and `embed_image` return vectors of identical dimension
166///   ([`VisionEmbeddings::dimension`]) that are directly comparable
167///   (cosine similarity) across modalities;
168/// - all vectors are L2-normalized;
169/// - empty/whitespace text or image payloads raise [`EmbeddingError::EmptyInput`];
170/// - an empty image slice is `Ok(vec![])` (nothing to do is not an error).
171#[async_trait]
172pub trait VisionEmbeddings: Send + Sync {
173    /// Embeds a single image.
174    async fn embed_image(&self, image: &ImageInput) -> Result<Vec<f32>, EmbeddingError>;
175
176    /// Embeds multiple images. The default loops one-by-one; HTTP backends
177    /// override it with a single batched request and must enforce batch
178    /// alignment (requested count == returned count).
179    async fn embed_images(&self, images: &[ImageInput]) -> Result<Vec<Vec<f32>>, EmbeddingError> {
180        if images.is_empty() {
181            return Ok(Vec::new());
182        }
183        let mut out = Vec::with_capacity(images.len());
184        for image in images {
185            out.push(self.embed_image(image).await?);
186        }
187        Ok(out)
188    }
189
190    /// Embeds text **into the same vector space as the images**.
191    ///
192    /// Use this (not a text-only [`crate::Embeddings`] model) when building
193    /// text-to-image retrieval: the query vector and the stored image
194    /// vectors must come from one model family.
195    async fn embed_text(&self, text: &str) -> Result<Vec<f32>, EmbeddingError>;
196
197    /// Embedding dimension (identical for both modalities).
198    fn dimension(&self) -> usize;
199
200    /// Model name.
201    fn model_name(&self) -> &str;
202}
203
204/// Deterministic in-memory [`VisionEmbeddings`] for offline tests.
205///
206/// Without overrides, vectors are derived deterministically from the input
207/// key (same key → same normalized vector; different keys → different
208/// vectors). Tests demonstrating cross-modal retrieval register aligned
209/// vectors explicitly via [`MockVisionEmbeddings::with_text_vector`] and
210/// [`MockVisionEmbeddings::with_image_vector`].
211#[derive(Debug)]
212pub struct MockVisionEmbeddings {
213    model: String,
214    dimension: usize,
215    text_vectors: Mutex<HashMap<String, Vec<f32>>>,
216    image_vectors: Mutex<HashMap<String, Vec<f32>>>,
217}
218
219impl MockVisionEmbeddings {
220    /// Creates a mock producing `dimension`-wide vectors.
221    pub fn new(dimension: usize) -> Self {
222        Self {
223            model: "mock-vision-embeddings".to_string(),
224            dimension,
225            text_vectors: Mutex::new(HashMap::new()),
226            image_vectors: Mutex::new(HashMap::new()),
227        }
228    }
229
230    /// Overrides the vector returned for a given text query/label.
231    ///
232    /// The vector is L2-normalized before being stored.
233    pub fn with_text_vector(&self, text: impl Into<String>, mut vector: Vec<f32>) -> &Self {
234        l2_normalize(&mut vector);
235        self.text_vectors
236            .lock()
237            .unwrap_or_else(|e| e.into_inner())
238            .insert(text.into(), vector);
239        self
240    }
241
242    /// Overrides the vector returned for an image, keyed exactly by its URL,
243    /// data URI, or the synthesized `data:<mime>;base64,<data>` reference.
244    ///
245    /// The vector is L2-normalized before being stored.
246    pub fn with_image_vector(&self, image: &ImageInput, mut vector: Vec<f32>) -> &Self {
247        if let Ok(key) = image.mock_key() {
248            l2_normalize(&mut vector);
249            self.image_vectors
250                .lock()
251                .unwrap_or_else(|e| e.into_inner())
252                .insert(key, vector);
253        }
254        self
255    }
256}
257
258#[async_trait]
259impl VisionEmbeddings for MockVisionEmbeddings {
260    async fn embed_image(&self, image: &ImageInput) -> Result<Vec<f32>, EmbeddingError> {
261        // Validate through the same path as real backends: empty URLs/data
262        // payloads must raise EmptyInput before mock-key lookup.
263        let _ = image.reference()?;
264        let key = image.mock_key()?;
265        let map = self.image_vectors.lock().unwrap_or_else(|e| e.into_inner());
266        Ok(map
267            .get(&key)
268            .cloned()
269            .unwrap_or_else(|| deterministic_vector(&key, self.dimension)))
270    }
271
272    async fn embed_text(&self, text: &str) -> Result<Vec<f32>, EmbeddingError> {
273        if text.trim().is_empty() {
274            return Err(EmbeddingError::EmptyInput);
275        }
276        let map = self.text_vectors.lock().unwrap_or_else(|e| e.into_inner());
277        Ok(map
278            .get(text)
279            .cloned()
280            .unwrap_or_else(|| deterministic_vector(text, self.dimension)))
281    }
282
283    fn dimension(&self) -> usize {
284        self.dimension
285    }
286
287    fn model_name(&self) -> &str {
288        &self.model
289    }
290}
291
292/// Stable pseudo-random normalized vector derived from a string key.
293pub(crate) fn deterministic_vector(key: &str, dimension: usize) -> Vec<f32> {
294    use std::collections::hash_map::DefaultHasher;
295    use std::hash::{Hash, Hasher};
296
297    let mut hasher = DefaultHasher::new();
298    key.hash(&mut hasher);
299    let mut state = hasher.finish();
300    let mut vector = Vec::with_capacity(dimension);
301    for _ in 0..dimension {
302        // LCG over u64, projected into [-1, 1).
303        state = state
304            .wrapping_mul(6_364_136_223_846_793_005)
305            .wrapping_add(1_442_695_040_888_963_407);
306        let unit = (state >> 33) as f32 / (1u64 << 30) as f32 - 1.0;
307        vector.push(unit);
308    }
309    l2_normalize(&mut vector);
310    vector
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316
317    #[test]
318    fn parses_data_uri_into_mime_and_bytes() {
319        let (mime, data) = parse_data_uri("data:image/png;base64,aW1n").unwrap();
320        assert_eq!(mime, "image/png");
321        assert_eq!(data, "aW1n");
322    }
323
324    #[test]
325    fn rejects_non_image_and_malformed_data_uris() {
326        assert!(parse_data_uri("data:application/pdf;base64,ZG9j").is_err());
327        assert!(parse_data_uri("image/png;base64,aW1n").is_err());
328        assert!(parse_data_uri("data:image/png,aW1n").is_err());
329        assert!(parse_data_uri("data:image/png;base64,").is_err());
330    }
331
332    #[test]
333    fn inline_parts_rejects_plain_url_but_accepts_base64() {
334        assert!(ImageInput::from_url("https://example.com/a.png")
335            .inline_parts()
336            .is_err());
337        let (data, mime) = ImageInput::from_base64("aW1n", "image/png")
338            .inline_parts()
339            .unwrap();
340        assert_eq!(data, "aW1n");
341        assert_eq!(mime, "image/png");
342    }
343
344    #[test]
345    fn reference_renders_base64_as_data_uri() {
346        let reference = ImageInput::from_base64("aW1n", "image/png")
347            .reference()
348            .unwrap();
349        assert_eq!(reference, "data:image/png;base64,aW1n");
350        assert_eq!(
351            ImageInput::from_url("https://example.com/a.png")
352                .reference()
353                .unwrap(),
354            "https://example.com/a.png"
355        );
356    }
357
358    #[tokio::test]
359    async fn mock_is_deterministic_and_normalized() {
360        let mock = MockVisionEmbeddings::new(16);
361        let a = mock.embed_text("cat").await.unwrap();
362        let b = mock.embed_text("cat").await.unwrap();
363        assert_eq!(a, b);
364        let norm: f32 = a.iter().map(|x| x * x).sum::<f32>().sqrt();
365        assert!((norm - 1.0).abs() < 1e-5);
366
367        let img = ImageInput::from_url("https://example.com/cat.png");
368        let va = mock.embed_image(&img).await.unwrap();
369        let vb = mock.embed_image(&img).await.unwrap();
370        assert_eq!(va, vb);
371        assert_ne!(va.len(), 0);
372    }
373
374    #[tokio::test]
375    async fn mock_overrides_align_text_and_image_in_shared_space() {
376        let mock = MockVisionEmbeddings::new(4);
377        mock.with_text_vector("cat", vec![1.0, 0.0, 0.0, 0.0]);
378        let img = ImageInput::from_url("https://example.com/cat.png");
379        mock.with_image_vector(&img, vec![1.0, 0.0, 0.0, 0.0]);
380
381        let text_v = mock.embed_text("cat").await.unwrap();
382        let image_v = mock.embed_image(&img).await.unwrap();
383        let sim = crate::cosine_similarity(&text_v, &image_v).unwrap();
384        assert!((sim - 1.0).abs() < 1e-5, "aligned pair sim = {sim}");
385    }
386
387    #[tokio::test]
388    async fn empty_inputs_are_rejected() {
389        let mock = MockVisionEmbeddings::new(4);
390        assert!(matches!(
391            mock.embed_text("  ").await,
392            Err(EmbeddingError::EmptyInput)
393        ));
394        assert!(matches!(
395            mock.embed_image(&ImageInput::from_base64(" ", "image/png"))
396                .await,
397            Err(EmbeddingError::EmptyInput)
398        ));
399        assert_eq!(
400            mock.embed_images(&[]).await.unwrap(),
401            Vec::<Vec<f32>>::new()
402        );
403    }
404}
lc_embeddings/vision.rs

lc_embeddings/
vision.rs
