uni_xervo/
api.rs

1//! Public API types for configuring models, catalogs, and runtime behavior.
2
3use crate::error::{Result, RuntimeError};
4use serde::{Deserialize, Serialize};
5use std::path::Path;
6
7/// The kind of inference task a model performs.
8///
9/// Marked `#[non_exhaustive]` so adding new variants is non-breaking. Downstream
10/// pattern matches must include a wildcard `_ => ...` arm.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
12#[serde(rename_all = "snake_case")]
13#[non_exhaustive]
14pub enum ModelTask {
15    /// Produce dense vector embeddings from text.
16    Embed,
17    /// Re-score a set of documents against a query.
18    Rerank,
19    /// Generate text (chat completions, summarization, etc.).
20    Generate,
21    /// Resolve a provider-specific raw runtime without task interpretation.
22    Raw,
23    /// Produce dense vector embeddings from images.
24    EmbedImage,
25    /// Produce dense vector embeddings from audio.
26    EmbedAudio,
27    /// Produce dense vector embeddings from heterogeneous (text + image + audio) inputs.
28    EmbedMultimodal,
29    /// Produce learned-sparse term-weight vectors from text (SPLADE / BGE-M3 sparse).
30    EmbedSparse,
31    /// Produce per-token (multi-vector / ColBERT late-interaction) embeddings from text.
32    EmbedMultiVector,
33    /// Produce dense + sparse + multi-vector heads from a single forward pass on a
34    /// multi-output graph (e.g. BGE-M3 `aapot/bge-m3-onnx`).
35    EmbedHybrid,
36    /// Structured natural-language analysis (POS / NER / DEP / SRL / dialog-act).
37    Nlp,
38    /// Extract structured blocks (text / heading / table / figure) from document page images.
39    DocumentExtract,
40    /// Transcribe speech audio into text with timing information.
41    Transcribe,
42    /// Recognize text in images (optical character recognition).
43    Ocr,
44}
45
46/// Controls when a model or provider is initialized during runtime startup.
47#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
48#[serde(rename_all = "snake_case")]
49pub enum WarmupPolicy {
50    /// Load immediately during [`ModelRuntime::builder().build()`](crate::runtime::ModelRuntimeBuilder::build).
51    /// Startup blocks until the load completes (or fails).
52    Eager,
53    /// Defer loading until the first inference request. This is the default.
54    #[default]
55    Lazy,
56    /// Spawn loading in a background task at startup. Inference calls that arrive
57    /// before loading finishes will trigger a blocking wait.
58    Background,
59}
60
61impl std::fmt::Display for WarmupPolicy {
62    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
63        match self {
64            Self::Eager => write!(f, "eager"),
65            Self::Lazy => write!(f, "lazy"),
66            Self::Background => write!(f, "background"),
67        }
68    }
69}
70
71/// Declarative specification that maps a human-readable alias to a concrete
72/// provider and model.
73///
74/// A model catalog is a `Vec<ModelAliasSpec>` — either built programmatically or
75/// parsed from JSON with [`catalog_from_str`] / [`catalog_from_file`].
76///
77/// # Example JSON
78///
79/// ```json
80/// {
81///   "alias": "embed/default",
82///   "task": "embed",
83///   "provider_id": "local/candle",
84///   "model_id": "sentence-transformers/all-MiniLM-L6-v2",
85///   "warmup": "lazy"
86/// }
87/// ```
88#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
89pub struct ModelAliasSpec {
90    /// Human-readable name used to request this model (e.g. `"embed/default"`).
91    /// Must contain a `/` separator.
92    pub alias: String,
93    /// The inference task this alias targets.
94    pub task: ModelTask,
95    /// Identifier of the provider that will load this model (e.g. `"local/candle"`,
96    /// `"remote/openai"`).
97    pub provider_id: String,
98    /// Model identifier understood by the provider — typically a HuggingFace repo ID
99    /// for local providers or an API model name for remote providers.
100    pub model_id: String,
101    /// Optional HuggingFace revision (branch, tag, or commit hash).
102    #[serde(skip_serializing_if = "Option::is_none")]
103    pub revision: Option<String>,
104    /// When this model should be initialized. Defaults to [`WarmupPolicy::Lazy`].
105    #[serde(default)]
106    pub warmup: WarmupPolicy,
107    /// If `true`, a failed eager warmup aborts runtime startup. Defaults to `false`.
108    #[serde(default)]
109    pub required: bool,
110    /// Per-inference timeout in seconds. `None` means no timeout.
111    #[serde(skip_serializing_if = "Option::is_none")]
112    pub timeout: Option<u64>,
113    /// Model load timeout in seconds. Defaults to 600 s if unset.
114    #[serde(skip_serializing_if = "Option::is_none")]
115    pub load_timeout: Option<u64>,
116    /// Retry configuration for transient inference failures.
117    #[serde(skip_serializing_if = "Option::is_none")]
118    pub retry: Option<RetryConfig>,
119    /// Provider-specific options (e.g. `{"isq": "Q4K"}` for mistral.rs,
120    /// `{"api_key_env": "MY_KEY"}` for remote providers). Defaults to `{}`.
121    #[serde(default)]
122    pub options: serde_json::Value,
123}
124
125/// Configuration for exponential-backoff retries on transient inference errors.
126#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
127pub struct RetryConfig {
128    /// Maximum number of attempts (including the initial call).
129    pub max_attempts: u32,
130    /// Base delay in milliseconds; doubled on each subsequent attempt.
131    pub initial_backoff_ms: u64,
132}
133
134impl RetryConfig {
135    /// Compute the backoff duration for the given 1-based `attempt` number.
136    ///
137    /// Uses `initial_backoff_ms * 2^(attempt - 1)` with saturating arithmetic.
138    pub fn get_backoff(&self, attempt: u32) -> std::time::Duration {
139        std::time::Duration::from_millis(
140            self.initial_backoff_ms * 2u64.pow(attempt.saturating_sub(1)),
141        )
142    }
143}
144
145impl Default for RetryConfig {
146    fn default() -> Self {
147        Self {
148            max_attempts: 3,
149            initial_backoff_ms: 100,
150        }
151    }
152}
153
154/// Deduplication key used by the runtime to share a single loaded model instance
155/// across multiple aliases that point to the same provider, model, revision, and
156/// options.
157#[derive(Debug, Clone, PartialEq, Eq, Hash)]
158pub struct ModelRuntimeKey {
159    /// The task type (embed, rerank, generate, raw).
160    pub task: ModelTask,
161    /// Provider that owns this model instance.
162    pub provider_id: String,
163    /// Model identifier within the provider.
164    pub model_id: String,
165    /// Optional HuggingFace revision.
166    pub revision: Option<String>,
167    /// Hash of the provider-specific options JSON. Two specs with semantically
168    /// equivalent options (same keys/values, any object-key order) produce the
169    /// same hash.
170    pub variant_hash: u64,
171}
172
173impl ModelRuntimeKey {
174    /// Derive a runtime key from a [`ModelAliasSpec`], hashing the options JSON
175    /// in a key-order-independent manner.
176    pub fn new(spec: &ModelAliasSpec) -> Self {
177        let mut hasher = std::collections::hash_map::DefaultHasher::new();
178        use std::hash::Hasher;
179
180        // Hash all JSON option shapes with deterministic key ordering.
181        // This avoids collisions for non-object values while preserving
182        // object-order independence for semantically equivalent JSON.
183        hash_json_value(&spec.options, &mut hasher);
184
185        Self {
186            task: spec.task,
187            provider_id: spec.provider_id.clone(),
188            model_id: spec.model_id.clone(),
189            revision: spec.revision.clone(),
190            variant_hash: hasher.finish(),
191        }
192    }
193}
194
195/// Recursively hash a JSON value in a deterministic, key-order-independent way.
196///
197/// Each JSON variant is prefixed with a unique discriminant byte to avoid
198/// collisions between structurally different values (e.g. `null` vs `false`).
199/// Object keys are sorted before hashing so that `{"a":1,"b":2}` and
200/// `{"b":2,"a":1}` produce the same hash.
201fn hash_json_value<H: std::hash::Hasher>(value: &serde_json::Value, hasher: &mut H) {
202    use std::hash::Hash;
203
204    match value {
205        serde_json::Value::Null => {
206            0u8.hash(hasher);
207        }
208        serde_json::Value::Bool(v) => {
209            1u8.hash(hasher);
210            v.hash(hasher);
211        }
212        serde_json::Value::Number(v) => {
213            2u8.hash(hasher);
214            v.to_string().hash(hasher);
215        }
216        serde_json::Value::String(v) => {
217            3u8.hash(hasher);
218            v.hash(hasher);
219        }
220        serde_json::Value::Array(values) => {
221            4u8.hash(hasher);
222            values.len().hash(hasher);
223            for v in values {
224                hash_json_value(v, hasher);
225            }
226        }
227        serde_json::Value::Object(map) => {
228            5u8.hash(hasher);
229            map.len().hash(hasher);
230
231            let mut entries: Vec<_> = map.iter().collect();
232            entries.sort_by_key(|(k, _)| *k);
233            for (k, v) in entries {
234                k.hash(hasher);
235                hash_json_value(v, hasher);
236            }
237        }
238    }
239}
240
241impl ModelAliasSpec {
242    /// Validate invariants: alias must be non-empty and contain a `'/'`, timeouts
243    /// must be non-zero when set.
244    pub fn validate(&self) -> Result<()> {
245        if self.alias.is_empty() {
246            return Err(RuntimeError::Config("Alias cannot be empty".to_string()));
247        }
248        if !self.alias.contains('/') {
249            return Err(RuntimeError::Config(format!(
250                "Alias '{}' must be in 'task/name' format",
251                self.alias
252            )));
253        }
254        if self.timeout == Some(0) {
255            return Err(RuntimeError::Config(
256                "Inference timeout must be greater than 0".to_string(),
257            ));
258        }
259        if self.load_timeout == Some(0) {
260            return Err(RuntimeError::Config(
261                "Load timeout must be greater than 0".to_string(),
262            ));
263        }
264        Ok(())
265    }
266
267    /// Parse a single `ModelAliasSpec` from a JSON value.
268    pub fn from_json(value: serde_json::Value) -> Result<Self> {
269        let spec: Self = serde_json::from_value(value)
270            .map_err(|e| RuntimeError::Config(format!("Invalid ModelAliasSpec JSON: {}", e)))?;
271        spec.validate()?;
272        Ok(spec)
273    }
274
275    /// Parse a single `ModelAliasSpec` from a JSON string.
276    pub fn from_json_str(s: &str) -> Result<Self> {
277        let spec: Self = serde_json::from_str(s)
278            .map_err(|e| RuntimeError::Config(format!("Invalid ModelAliasSpec JSON: {}", e)))?;
279        spec.validate()?;
280        Ok(spec)
281    }
282}
283
284/// Parse a catalog (array) of `ModelAliasSpec` from a JSON string.
285pub fn catalog_from_str(s: &str) -> Result<Vec<ModelAliasSpec>> {
286    let specs: Vec<ModelAliasSpec> = serde_json::from_str(s)
287        .map_err(|e| RuntimeError::Config(format!("Invalid catalog JSON: {}", e)))?;
288    for spec in &specs {
289        spec.validate()?;
290    }
291    Ok(specs)
292}
293
294/// Read and parse a catalog from a JSON file.
295///
296/// The file must contain a JSON array of model alias specs.
297pub fn catalog_from_file(path: impl AsRef<Path>) -> Result<Vec<ModelAliasSpec>> {
298    let path = path.as_ref();
299    let contents = std::fs::read_to_string(path).map_err(|e| {
300        RuntimeError::Config(format!(
301            "Failed to read catalog file '{}': {}",
302            path.display(),
303            e
304        ))
305    })?;
306    catalog_from_str(&contents)
307}
308
309#[cfg(test)]
310mod tests {
311    use super::*;
312    use serde_json::json;
313
314    const VALID_JSON: &str = r#"{
315        "alias": "embed/default",
316        "task": "embed",
317        "provider_id": "local/candle",
318        "model_id": "sentence-transformers/all-MiniLM-L6-v2"
319    }"#;
320
321    const VALID_CATALOG_JSON: &str = r#"[
322        {
323            "alias": "embed/default",
324            "task": "embed",
325            "provider_id": "local/candle",
326            "model_id": "sentence-transformers/all-MiniLM-L6-v2"
327        },
328        {
329            "alias": "chat/fast",
330            "task": "generate",
331            "provider_id": "local/mistralrs",
332            "model_id": "mistralai/Mistral-7B-v0.1",
333            "warmup": "background",
334            "required": false,
335            "options": { "isq": "Q4K" }
336        }
337    ]"#;
338
339    #[test]
340    fn from_json_str_parses_valid_spec() {
341        let spec = ModelAliasSpec::from_json_str(VALID_JSON).unwrap();
342        assert_eq!(spec.alias, "embed/default");
343        assert_eq!(spec.task, ModelTask::Embed);
344        assert_eq!(spec.provider_id, "local/candle");
345        assert_eq!(spec.warmup, WarmupPolicy::Lazy); // default
346        assert!(!spec.required); // default
347    }
348
349    #[test]
350    fn from_json_value_parses_valid_spec() {
351        let value = json!({
352            "alias": "embed/fast",
353            "task": "embed",
354            "provider_id": "local/onnx",
355            "model_id": "BAAI/bge-small-en-v1.5",
356            "required": true,
357            "warmup": "eager"
358        });
359        let spec = ModelAliasSpec::from_json(value).unwrap();
360        assert_eq!(spec.alias, "embed/fast");
361        assert_eq!(spec.warmup, WarmupPolicy::Eager);
362        assert!(spec.required);
363    }
364
365    #[test]
366    fn from_json_str_rejects_missing_slash_in_alias() {
367        let json = r#"{"alias":"noSlash","task":"embed","provider_id":"x","model_id":"y"}"#;
368        assert!(ModelAliasSpec::from_json_str(json).is_err());
369    }
370
371    #[test]
372    fn from_json_str_rejects_invalid_json() {
373        assert!(ModelAliasSpec::from_json_str("{not valid}").is_err());
374    }
375
376    #[test]
377    fn catalog_from_str_parses_array() {
378        let specs = catalog_from_str(VALID_CATALOG_JSON).unwrap();
379        assert_eq!(specs.len(), 2);
380        assert_eq!(specs[0].alias, "embed/default");
381        assert_eq!(specs[1].alias, "chat/fast");
382        assert_eq!(specs[1].options["isq"], "Q4K");
383    }
384
385    #[test]
386    fn catalog_from_str_rejects_invalid_spec() {
387        let json = r#"[{"alias":"bad","task":"embed","provider_id":"x","model_id":"y"}]"#;
388        assert!(catalog_from_str(json).is_err()); // alias has no '/'
389    }
390
391    #[test]
392    fn catalog_from_file_reads_and_parses() {
393        let dir = std::env::temp_dir();
394        let path = dir.join("test_catalog.json");
395        std::fs::write(&path, VALID_CATALOG_JSON).unwrap();
396        let specs = catalog_from_file(&path).unwrap();
397        assert_eq!(specs.len(), 2);
398        std::fs::remove_file(&path).unwrap();
399    }
400
401    #[test]
402    fn catalog_from_file_errors_on_missing_file() {
403        assert!(catalog_from_file("/nonexistent/path/catalog.json").is_err());
404    }
405
406    #[test]
407    fn runtime_key_distinguishes_non_object_options() {
408        let mut spec_null = ModelAliasSpec::from_json_str(VALID_JSON).unwrap();
409        spec_null.options = serde_json::Value::Null;
410
411        let mut spec_bool = spec_null.clone();
412        spec_bool.options = json!(true);
413
414        let mut spec_array = spec_null.clone();
415        spec_array.options = json!(["a", 1]);
416
417        let key_null = ModelRuntimeKey::new(&spec_null);
418        let key_bool = ModelRuntimeKey::new(&spec_bool);
419        let key_array = ModelRuntimeKey::new(&spec_array);
420
421        assert_ne!(key_null, key_bool);
422        assert_ne!(key_null, key_array);
423        assert_ne!(key_bool, key_array);
424    }
425
426    #[test]
427    fn model_task_serde_round_trip_for_all_variants() {
428        // Covers existing variants plus the seven multimodal additions.
429        // Wire format is documented (snake_case) and downstream catalog
430        // files depend on it — regression-protect every variant.
431        let cases: &[(ModelTask, &str)] = &[
432            (ModelTask::Embed, "\"embed\""),
433            (ModelTask::Rerank, "\"rerank\""),
434            (ModelTask::Generate, "\"generate\""),
435            (ModelTask::Raw, "\"raw\""),
436            (ModelTask::EmbedImage, "\"embed_image\""),
437            (ModelTask::EmbedAudio, "\"embed_audio\""),
438            (ModelTask::EmbedMultimodal, "\"embed_multimodal\""),
439            (ModelTask::EmbedSparse, "\"embed_sparse\""),
440            (ModelTask::EmbedMultiVector, "\"embed_multi_vector\""),
441            (ModelTask::EmbedHybrid, "\"embed_hybrid\""),
442            (ModelTask::Nlp, "\"nlp\""),
443            (ModelTask::DocumentExtract, "\"document_extract\""),
444            (ModelTask::Transcribe, "\"transcribe\""),
445            (ModelTask::Ocr, "\"ocr\""),
446        ];
447        for (variant, wire) in cases {
448            let serialized = serde_json::to_string(variant).unwrap();
449            assert_eq!(&serialized, wire, "serialize {:?}", variant);
450            let parsed: ModelTask = serde_json::from_str(wire).unwrap();
451            assert_eq!(&parsed, variant, "deserialize {}", wire);
452        }
453    }
454
455    #[test]
456    fn runtime_key_nested_option_order_independence() {
457        let mut spec1 = ModelAliasSpec::from_json_str(VALID_JSON).unwrap();
458        spec1.options = json!({
459            "outer": {
460                "b": [3, 2, 1],
461                "a": {"y": 2, "x": 1}
462            }
463        });
464
465        let mut spec2 = ModelAliasSpec::from_json_str(VALID_JSON).unwrap();
466        spec2.options = json!({
467            "outer": {
468                "a": {"x": 1, "y": 2},
469                "b": [3, 2, 1]
470            }
471        });
472
473        let key1 = ModelRuntimeKey::new(&spec1);
474        let key2 = ModelRuntimeKey::new(&spec2);
475        assert_eq!(key1, key2);
476    }
477}
uni_xervo/api.rs

uni_xervo/
api.rs
