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

uni_xervo/
traits.rs
