uni_xervo/
traits.rs

1//! Core traits that every provider and model implementation must satisfy.
2
3pub mod asr;
4pub mod docs;
5pub mod hybrid;
6pub mod multimodal;
7pub mod multivector;
8pub mod nlp;
9pub mod raw_tensor_model;
10pub mod sparse;
11
12use crate::api::{ModelAliasSpec, ModelTask};
13use crate::error::Result;
14use async_trait::async_trait;
15use std::any::Any;
16
17pub use raw_tensor_model::{
18    DimSize, RawTensorModel, TensorBatch, TensorDtype, TensorSpec, TensorValue,
19};
20
21pub use asr::{
22    TranscribeOptions, TranscribeResult, TranscribeSegment, TranscribeWord, TranscriptionModel,
23};
24pub use docs::{
25    DocBlock, DocBlockKind, DocExtractOptions, DocExtractResult, DocOutputFormat,
26    DocumentExtractionModel, OcrBlock, OcrModel, OcrResult,
27};
28pub use hybrid::{HeadSet, HybridEmbedResult, HybridEmbeddingModel};
29pub use multimodal::{
30    AudioEmbeddingModel, AudioInput, ImageEmbeddingModel, Modality, MultimodalBlock,
31    MultimodalEmbeddingModel, MultimodalInput,
32};
33pub use multivector::{MultiVectorEmbedResult, MultiVectorEmbeddingModel};
34pub use nlp::{
35    DepLink, NerEntity, NlpLabelMaps, NlpModel, NlpRequest, NlpResult, NlpSentence, NlpTasks,
36    NlpToken, SpeechAct, SrlFrame, SrlRole,
37};
38pub use sparse::{SparseEmbedResult, SparseEmbeddingModel, SparseVector};
39
40/// Advertised capabilities of a [`ModelProvider`].
41#[derive(Debug, Clone)]
42pub struct ProviderCapabilities {
43    /// The set of [`ModelTask`] variants this provider can handle.
44    pub supported_tasks: Vec<ModelTask>,
45}
46
47/// Health status reported by a provider.
48#[derive(Debug, Clone)]
49pub enum ProviderHealth {
50    /// The provider is fully operational.
51    Healthy,
52    /// The provider is operational but experiencing partial issues.
53    Degraded(String),
54    /// The provider cannot serve requests.
55    Unhealthy(String),
56}
57
58/// A pluggable backend that knows how to load models for one or more
59/// [`ModelTask`] types.
60///
61/// Providers are registered with [`ModelRuntimeBuilder::register_provider`](crate::runtime::ModelRuntimeBuilder::register_provider)
62/// and are identified by their [`provider_id`](ModelProvider::provider_id)
63/// (e.g. `"local/candle"`, `"remote/openai"`).
64#[async_trait]
65pub trait ModelProvider: Send + Sync {
66    /// Unique identifier for this provider (e.g. `"local/candle"`, `"remote/openai"`).
67    fn provider_id(&self) -> &'static str;
68
69    /// Return the set of tasks this provider supports.
70    fn capabilities(&self) -> ProviderCapabilities;
71
72    /// Load (or connect to) a model described by `spec` and return a type-erased
73    /// handle.
74    ///
75    /// The returned [`LoadedModelHandle`] is expected to contain an
76    /// `Arc<dyn EmbeddingModel>`, `Arc<dyn RerankerModel>`,
77    /// `Arc<dyn GeneratorModel>`, or `Arc<dyn RawTensorModel>` depending on the task.
78    async fn load(&self, spec: &ModelAliasSpec) -> Result<LoadedModelHandle>;
79
80    /// Report the current health of this provider.
81    async fn health(&self) -> ProviderHealth;
82
83    /// Optional one-time warmup hook called during runtime startup.
84    ///
85    /// Use this for provider-wide initialization such as setting up API clients
86    /// or pre-caching shared resources. The default implementation is a no-op.
87    async fn warmup(&self) -> Result<()> {
88        Ok(())
89    }
90}
91
92/// A type-erased, reference-counted handle to a loaded model instance.
93///
94/// Providers wrap their concrete model (e.g. `Arc<dyn EmbeddingModel>`) inside
95/// this `Arc<dyn Any + Send + Sync>` so the runtime can store them uniformly.
96/// The runtime later downcasts the handle back to the expected trait object.
97pub type LoadedModelHandle = std::sync::Arc<dyn Any + Send + Sync>;
98
99/// Metadata common to every loaded model handle.
100///
101/// Supertrait of every task trait ([`EmbeddingModel`], [`RerankerModel`], …) so a
102/// caller holding any typed handle can identify the model and inspect its
103/// requested execution backends uniformly.
104pub trait ModelInfo: Send + Sync {
105    /// The underlying model identifier (e.g. a HuggingFace repo ID or API model name).
106    fn model_id(&self) -> &str;
107
108    /// Names of the ONNX Runtime execution providers requested for the underlying
109    /// session, in priority order. Empty for remote and non-ONNX models.
110    ///
111    /// "Requested" is not "attached": a provider may fall back if a backend is
112    /// unavailable at session-build time. This reports what was asked for.
113    fn active_execution_providers(&self) -> Vec<String> {
114        Vec::new()
115    }
116}
117
118/// A model that produces dense vector embeddings from text.
119#[async_trait]
120pub trait EmbeddingModel: ModelInfo {
121    /// Embed a batch of text strings into dense vectors.
122    ///
123    /// Returns an [`EmbedResult`] whose `vectors` holds one `Vec<f32>` per input
124    /// text, each with [`dimensions()`](EmbeddingModel::dimensions) elements.
125    async fn embed(&self, texts: &[&str]) -> Result<EmbedResult>;
126
127    /// The dimensionality of the embedding vectors produced by this model.
128    fn dimensions(&self) -> u32;
129
130    /// Optional warmup hook (e.g. load weights into memory on first access).
131    /// The default is a no-op.
132    async fn warmup(&self) -> Result<()> {
133        Ok(())
134    }
135}
136
137/// A single scored document returned by a [`RerankerModel`].
138#[derive(Debug, Clone)]
139pub struct ScoredDoc {
140    /// Zero-based index into the original `docs` slice passed to
141    /// [`RerankerModel::rerank`].
142    pub index: usize,
143    /// Relevance score assigned by the reranker (higher is more relevant).
144    pub score: f32,
145    /// The document text, if the provider returns it. May be `None`.
146    pub text: Option<String>,
147}
148
149/// A model that re-scores documents against a query for relevance ranking.
150#[async_trait]
151pub trait RerankerModel: ModelInfo {
152    /// Rerank `docs` by relevance to `query`, returning scored results
153    /// (typically sorted by descending score).
154    async fn rerank(&self, query: &str, docs: &[&str]) -> Result<Vec<ScoredDoc>>;
155
156    /// Optional warmup hook. The default is a no-op.
157    async fn warmup(&self) -> Result<()> {
158        Ok(())
159    }
160}
161
162// ---------------------------------------------------------------------------
163// Multimodal message types
164// ---------------------------------------------------------------------------
165
166/// The role of a message in a conversation.
167#[derive(Debug, Clone, PartialEq, Eq)]
168pub enum MessageRole {
169    /// System-level instructions.
170    System,
171    /// A user turn.
172    User,
173    /// An assistant (model) turn.
174    Assistant,
175}
176
177/// Image data that can be passed as part of a [`ContentBlock`].
178#[derive(Debug, Clone)]
179pub enum ImageInput {
180    /// Raw image bytes with a MIME type (e.g. `"image/png"`).
181    Bytes { data: Vec<u8>, media_type: String },
182    /// A URL pointing to an image.
183    Url(String),
184}
185
186/// A single block of content within a [`Message`].
187#[derive(Debug, Clone)]
188pub enum ContentBlock {
189    /// Plain text content.
190    Text(String),
191    /// An image (for vision models).
192    Image(ImageInput),
193}
194
195/// A single message in a conversation, containing one or more content blocks.
196#[derive(Debug, Clone)]
197pub struct Message {
198    /// The role of the message sender.
199    pub role: MessageRole,
200    /// The content blocks that make up this message.
201    pub content: Vec<ContentBlock>,
202}
203
204impl Message {
205    /// Create a user message with a single text block.
206    pub fn user(text: impl Into<String>) -> Self {
207        Self {
208            role: MessageRole::User,
209            content: vec![ContentBlock::Text(text.into())],
210        }
211    }
212
213    /// Create an assistant message with a single text block.
214    pub fn assistant(text: impl Into<String>) -> Self {
215        Self {
216            role: MessageRole::Assistant,
217            content: vec![ContentBlock::Text(text.into())],
218        }
219    }
220
221    /// Create a system message with a single text block.
222    pub fn system(text: impl Into<String>) -> Self {
223        Self {
224            role: MessageRole::System,
225            content: vec![ContentBlock::Text(text.into())],
226        }
227    }
228
229    /// Extract the concatenated text from all [`ContentBlock::Text`] blocks.
230    pub fn text(&self) -> String {
231        self.content
232            .iter()
233            .filter_map(|b| match b {
234                ContentBlock::Text(t) => Some(t.as_str()),
235                _ => None,
236            })
237            .collect::<Vec<_>>()
238            .join(" ")
239    }
240}
241
242// ---------------------------------------------------------------------------
243// Generation options and results
244// ---------------------------------------------------------------------------
245
246/// Sampling and length parameters for text generation.
247#[derive(Debug, Clone, Default)]
248pub struct GenerationOptions {
249    /// Maximum number of tokens to generate. Provider default if `None`.
250    pub max_tokens: Option<usize>,
251    /// Sampling temperature (0.0 = greedy, higher = more random).
252    pub temperature: Option<f32>,
253    /// Nucleus sampling threshold.
254    pub top_p: Option<f32>,
255    /// Desired image width (for diffusion models; ignored by text/vision).
256    pub width: Option<u32>,
257    /// Desired image height (for diffusion models; ignored by text/vision).
258    pub height: Option<u32>,
259}
260
261/// An image produced by a generation call (e.g. from a diffusion model).
262#[derive(Debug, Clone)]
263pub struct GeneratedImage {
264    /// Raw image bytes (e.g. PNG).
265    pub data: Vec<u8>,
266    /// MIME type (e.g. `"image/png"`).
267    pub media_type: String,
268}
269
270/// Audio output produced by a speech model.
271#[derive(Debug, Clone)]
272pub struct AudioOutput {
273    /// PCM sample data.
274    pub pcm_data: Vec<f32>,
275    /// Sample rate in Hz.
276    pub sample_rate: usize,
277    /// Number of audio channels.
278    pub channels: usize,
279}
280
281/// The output of a generation call.
282#[derive(Debug, Clone)]
283pub struct GenerationResult {
284    /// The generated text (may be empty for image/audio-only results).
285    pub text: String,
286    /// Token usage statistics, if reported by the provider.
287    pub usage: Option<TokenUsage>,
288    /// Generated images (non-empty for diffusion models).
289    pub images: Vec<GeneratedImage>,
290    /// Generated audio (present for speech models).
291    pub audio: Option<AudioOutput>,
292}
293
294/// Token counts for a generation request.
295#[derive(Debug, Clone)]
296pub struct TokenUsage {
297    /// Number of tokens in the prompt / input.
298    pub prompt_tokens: usize,
299    /// Number of tokens generated.
300    pub completion_tokens: usize,
301    /// Sum of prompt and completion tokens.
302    pub total_tokens: usize,
303}
304
305/// Result of a dense-embedding call, carrying the vectors plus optional usage.
306///
307/// Returned by [`EmbeddingModel`], [`ImageEmbeddingModel`],
308/// [`AudioEmbeddingModel`], and [`MultimodalEmbeddingModel`] — every dense
309/// embedder shares this shape so remote providers (OpenAI / Cohere / Gemini /
310/// etc.) can report per-call token usage.
311///
312/// `usage` is `Some` when the provider reports it (remote APIs), and `None`
313/// for local providers where the concept does not apply.
314#[derive(Debug, Clone)]
315pub struct EmbedResult {
316    /// One dense vector per input, in input order.
317    pub vectors: Vec<Vec<f32>>,
318    /// Token usage reported by the provider, if any.
319    pub usage: Option<TokenUsage>,
320}
321
322/// A model that generates text, images, or audio from a conversational
323/// message history.
324///
325/// Messages carry explicit roles via [`Message`] and may contain multimodal
326/// content (text and images). The output [`GenerationResult`] is a union:
327/// text, images, and audio fields — consumers check what is populated.
328#[async_trait]
329pub trait GeneratorModel: ModelInfo {
330    /// Generate a response given a conversation history and sampling options.
331    async fn generate(
332        &self,
333        messages: &[Message],
334        options: GenerationOptions,
335    ) -> Result<GenerationResult>;
336
337    /// Optional warmup hook. The default is a no-op.
338    async fn warmup(&self) -> Result<()> {
339        Ok(())
340    }
341}
uni_xervo/traits.rs

uni_xervo/
traits.rs
