uni_xervo/traits/
multimodal.rs

1//! Multimodal embedding types and traits (image / audio / mixed-modality).
2//!
3//! These types accompany the trait definitions for embedders that consume
4//! non-text inputs. Each trait sits alongside
5//! [`EmbeddingModel`](crate::traits::EmbeddingModel) and returns an
6//! [`EmbedResult`] so remote providers can attach per-call usage counts.
7
8use crate::error::Result;
9use crate::traits::{EmbedResult, ImageInput, ModelInfo};
10use async_trait::async_trait;
11
12/// A model that produces dense vector embeddings from images.
13///
14/// One vector is returned per input image, each of dimension
15/// [`dimensions`](ImageEmbeddingModel::dimensions). Implementations
16/// typically batch internally for GPU efficiency; callers should pass as
17/// large a batch as fits their latency budget.
18#[async_trait]
19pub trait ImageEmbeddingModel: ModelInfo {
20    /// Embed a batch of images.
21    ///
22    /// # Errors
23    /// Returns an error if the provider fails to load any input, runs out
24    /// of memory, or the upstream API rejects the request.
25    async fn embed(&self, images: Vec<ImageInput>) -> Result<EmbedResult>;
26
27    /// The dimensionality of vectors produced by this model.
28    fn dimensions(&self) -> u32;
29
30    /// Optional warmup hook (e.g. load weights into memory on first access).
31    /// The default is a no-op.
32    ///
33    /// # Errors
34    /// Returns an error if the underlying model fails to initialize.
35    async fn warmup(&self) -> Result<()> {
36        Ok(())
37    }
38}
39
40/// A model that produces dense vector embeddings from audio.
41///
42/// One vector is returned per input audio, each of dimension
43/// [`dimensions`](AudioEmbeddingModel::dimensions). Use cases include
44/// CLAP-style audio similarity, music tagging, and audio-text retrieval.
45#[async_trait]
46pub trait AudioEmbeddingModel: ModelInfo {
47    /// Embed a batch of audio inputs.
48    ///
49    /// # Errors
50    /// Returns an error if the provider cannot decode an input or runs out
51    /// of memory.
52    async fn embed(&self, audios: Vec<AudioInput>) -> Result<EmbedResult>;
53
54    /// The dimensionality of vectors produced by this model.
55    fn dimensions(&self) -> u32;
56
57    /// Optional warmup hook. The default is a no-op.
58    ///
59    /// # Errors
60    /// Returns an error if the underlying model fails to initialize.
61    async fn warmup(&self) -> Result<()> {
62        Ok(())
63    }
64}
65
66/// A model that produces a single dense vector from heterogeneous content
67/// (text + image + audio together).
68///
69/// Supports models like Cohere Embed v4, Jina Embeddings v4, and Gemini
70/// Embedding 2 that map mixed-modality inputs into a single shared embedding
71/// space.
72#[async_trait]
73pub trait MultimodalEmbeddingModel: ModelInfo {
74    /// Embed a batch of mixed-modality inputs.
75    ///
76    /// Each [`MultimodalInput`] produces one vector spanning all its blocks.
77    ///
78    /// # Errors
79    /// Returns an error if the provider rejects an unsupported modality
80    /// combination or runs out of memory.
81    async fn embed(&self, inputs: Vec<MultimodalInput>) -> Result<EmbedResult>;
82
83    /// The dimensionality of vectors produced by this model.
84    fn dimensions(&self) -> u32;
85
86    /// Modalities accepted by this model (e.g. `&[Text, Image]` for Cohere
87    /// Embed v4 vs `&[Text, Image, Audio, Video]` for Gemini Embedding 2).
88    fn supported_modalities(&self) -> &[Modality];
89
90    /// Optional warmup hook. The default is a no-op.
91    ///
92    /// # Errors
93    /// Returns an error if the underlying model fails to initialize.
94    async fn warmup(&self) -> Result<()> {
95        Ok(())
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use super::*;
102    use std::sync::Arc;
103
104    // Dyn-safety smoke checks. These will fail to compile if any of these
105    // traits accidentally becomes non-object-safe.
106    #[test]
107    fn image_embedding_model_is_dyn_safe() {
108        fn _accept(_: Arc<dyn ImageEmbeddingModel>) {}
109    }
110
111    #[test]
112    fn audio_embedding_model_is_dyn_safe() {
113        fn _accept(_: Arc<dyn AudioEmbeddingModel>) {}
114    }
115
116    #[test]
117    fn multimodal_embedding_model_is_dyn_safe() {
118        fn _accept(_: Arc<dyn MultimodalEmbeddingModel>) {}
119    }
120}
121
122/// An audio input to a transcription or audio embedding model.
123///
124/// Two variants reflect the realistic input shapes:
125/// raw container bytes for HTTP handlers, and pre-decoded PCM for in-process
126/// audio capture. Local-file ingest is the caller's responsibility — uni-xervo
127/// providers do not perform file I/O, matching the [`ImageInput`] design.
128#[derive(Debug, Clone)]
129pub enum AudioInput {
130    /// Raw container bytes (e.g. WAV, MP3, FLAC). The provider decides how to
131    /// decode based on `media_type`.
132    Bytes {
133        /// Audio payload bytes in container format.
134        data: Vec<u8>,
135        /// MIME type indicating the container format (e.g. `"audio/wav"`).
136        media_type: String,
137    },
138    /// Pre-decoded PCM samples. The provider skips the decode step.
139    Pcm {
140        /// Samples per second (Hz).
141        sample_rate: u32,
142        /// Number of interleaved channels (1 = mono, 2 = stereo, …).
143        channels: u16,
144        /// Sample data, length `sample_rate * channels * duration_seconds`.
145        samples: Vec<f32>,
146    },
147}
148
149/// One block of input to a [`MultimodalEmbeddingModel`].
150///
151/// Distinct from [`ContentBlock`](crate::traits::ContentBlock) — that type
152/// is for conversational generation contexts and only carries text + image.
153/// This type extends to audio for use cases like Cohere Embed v4 and Gemini
154/// Embedding 2 that produce a single embedding spanning multiple modalities.
155#[derive(Debug, Clone)]
156pub enum MultimodalBlock {
157    /// Plain text content.
158    Text(String),
159    /// An image input (URL or raw bytes — see [`ImageInput`]).
160    Image(ImageInput),
161    /// An audio input.
162    Audio(AudioInput),
163}
164
165/// A heterogeneous input to a [`MultimodalEmbeddingModel`].
166///
167/// Each input produces one embedding spanning all of its blocks. Multiple
168/// inputs are batched together by passing a `Vec<MultimodalInput>`.
169#[derive(Debug, Clone)]
170pub struct MultimodalInput {
171    /// Ordered sequence of content blocks that share a single embedding.
172    pub blocks: Vec<MultimodalBlock>,
173}
174
175/// Modality labels reported by [`MultimodalEmbeddingModel::supported_modalities`].
176///
177/// The exact set varies per model: Cohere Embed v4 supports `Text + Image`,
178/// Gemini Embedding 2 supports `Text + Image + Audio + Video`, etc.
179#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
180pub enum Modality {
181    /// Plain text.
182    Text,
183    /// Still images.
184    Image,
185    /// Audio (speech, music, ambient).
186    Audio,
187    /// Video (frame sequences with optional audio).
188    Video,
189}
uni_xervo/traits/multimodal.rs

uni_xervo/traits/
multimodal.rs
