uni_xervo/
runtime.rs

1//! The core runtime that manages providers, catalogs, and loaded model instances.
2
3mod hybrid_adapter;
4
5use crate::api::{ModelAliasSpec, ModelRuntimeKey};
6use crate::error::{Result, RuntimeError};
7use crate::options_validation::validate_provider_options;
8use crate::reliability::{
9    InstrumentedAudioEmbeddingModel, InstrumentedDocumentExtractionModel,
10    InstrumentedEmbeddingModel, InstrumentedGeneratorModel, InstrumentedHybridEmbeddingModel,
11    InstrumentedImageEmbeddingModel, InstrumentedMultiVectorEmbeddingModel,
12    InstrumentedMultimodalEmbeddingModel, InstrumentedNlpModel, InstrumentedOcrModel,
13    InstrumentedRawTensorModel, InstrumentedRerankerModel, InstrumentedSparseEmbeddingModel,
14    InstrumentedTranscriptionModel,
15};
16use crate::traits::{
17    AudioEmbeddingModel, DocumentExtractionModel, EmbeddingModel, GeneratorModel,
18    HybridEmbeddingModel, ImageEmbeddingModel, LoadedModelHandle, ModelProvider,
19    MultiVectorEmbeddingModel, MultimodalEmbeddingModel, NlpModel, OcrModel, RawTensorModel,
20    RerankerModel, SparseEmbeddingModel, TranscriptionModel,
21};
22use dashmap::DashMap;
23use std::any::Any;
24use std::collections::HashMap;
25use std::sync::Arc;
26use tokio::sync::{Mutex, RwLock};
27
28/// Per-alias cache of fully-instrumented typed handles.
29///
30/// Eliminates repeated spec lookups, key hashing, registry reads, and
31/// wrapper allocations on hot paths where the same alias is resolved
32/// many times (e.g. per-turn inference in a pipeline).
33///
34/// Uses [`DashMap`] for lock-free concurrent reads on the hot path.
35/// Entries are populated on first successful resolve via
36/// `entry().or_insert_with(...)`, so concurrent first-callers all observe
37/// the same `Arc` and the wrapper is allocated exactly once per alias.
38/// Once populated, an entry is never mutated.
39#[derive(Default)]
40struct HandleCache {
41    embeddings: DashMap<String, Arc<dyn EmbeddingModel>>,
42    rerankers: DashMap<String, Arc<dyn RerankerModel>>,
43    generators: DashMap<String, Arc<dyn GeneratorModel>>,
44    raw_tensor_models: DashMap<String, Arc<dyn RawTensorModel>>,
45    // Multimodal extension surface — added in Phase 6, instrumentation
46    // filled in by Phase 7.
47    image_embedders: DashMap<String, Arc<dyn ImageEmbeddingModel>>,
48    audio_embedders: DashMap<String, Arc<dyn AudioEmbeddingModel>>,
49    multimodal_embedders: DashMap<String, Arc<dyn MultimodalEmbeddingModel>>,
50    sparse_embedders: DashMap<String, Arc<dyn SparseEmbeddingModel>>,
51    multi_vector_embedders: DashMap<String, Arc<dyn MultiVectorEmbeddingModel>>,
52    hybrid_embedders: DashMap<String, Arc<dyn HybridEmbeddingModel>>,
53    nlp_models: DashMap<String, Arc<dyn NlpModel>>,
54    document_extractors: DashMap<String, Arc<dyn DocumentExtractionModel>>,
55    transcribers: DashMap<String, Arc<dyn TranscriptionModel>>,
56    ocr_models: DashMap<String, Arc<dyn OcrModel>>,
57}
58
59/// Default load timeout applied when [`ModelAliasSpec::load_timeout`] is `None`.
60const DEFAULT_LOAD_TIMEOUT_SECS: u64 = 600;
61
62/// The central runtime that owns registered providers and a catalog of model
63/// aliases.
64///
65/// Obtain an instance via [`ModelRuntime::builder()`] and the
66/// [`ModelRuntimeBuilder`].  Once built, use [`embedding`](Self::embedding),
67/// [`reranker`](Self::reranker), or [`generator`](Self::generator) to obtain
68/// typed, instrumented model handles.
69///
70/// Models are loaded lazily on first access (unless configured for eager or
71/// background warmup) and cached in an internal registry so that subsequent
72/// requests for the same model are served instantly.
73pub struct ModelRuntime {
74    providers: HashMap<String, Box<dyn ModelProvider>>,
75    registry: Arc<ModelRegistry>,
76    catalog: RwLock<HashMap<String, ModelAliasSpec>>,
77    handle_cache: HandleCache,
78}
79
80/// Internal registry that caches loaded model instances and coordinates
81/// concurrent load requests to prevent duplicate work.
82#[derive(Default)]
83pub struct ModelRegistry {
84    instances: RwLock<HashMap<ModelRuntimeKey, LoadedModelHandle>>,
85    /// Per-key mutexes to prevent concurrent loads of the same model.
86    loader_locks: Mutex<HashMap<ModelRuntimeKey, Arc<Mutex<()>>>>,
87}
88
89impl ModelRuntime {
90    /// Create a new [`ModelRuntimeBuilder`] for configuring and constructing a
91    /// runtime.
92    pub fn builder() -> ModelRuntimeBuilder {
93        ModelRuntimeBuilder::default()
94    }
95
96    /// Register a new model alias at runtime.
97    pub async fn register(&self, spec: ModelAliasSpec) -> Result<()> {
98        spec.validate()?;
99        if !self.providers.contains_key(&spec.provider_id) {
100            return Err(RuntimeError::Config(format!(
101                "Unknown provider '{}' for alias '{}'",
102                spec.provider_id, spec.alias
103            )));
104        }
105        validate_provider_options(&spec.provider_id, spec.task, &spec.options)?;
106        let mut catalog = self.catalog.write().await;
107        if catalog.contains_key(&spec.alias) {
108            return Err(RuntimeError::Config(format!(
109                "Alias '{}' already exists",
110                spec.alias
111            )));
112        }
113        catalog.insert(spec.alias.clone(), spec);
114        Ok(())
115    }
116
117    /// Check if an alias exists in the catalog.
118    pub async fn contains_alias(&self, alias: &str) -> bool {
119        let catalog = self.catalog.read().await;
120        catalog.contains_key(alias)
121    }
122
123    /// Look up a spec by alias, returning an error if not found.
124    async fn lookup_spec(&self, alias: &str) -> Result<ModelAliasSpec> {
125        let catalog = self.catalog.read().await;
126        catalog
127            .get(alias)
128            .cloned()
129            .ok_or_else(|| RuntimeError::AliasNotFound {
130                alias: alias.to_string(),
131            })
132    }
133
134    /// Pre-load and cache every model in the catalog.
135    ///
136    /// Models already loaded are skipped. Fails fast on the first error.
137    /// Call this during application startup to avoid cold-start latency on
138    /// first inference.
139    pub async fn prefetch_all(&self) -> Result<()> {
140        let specs: Vec<ModelAliasSpec> = {
141            let catalog = self.catalog.read().await;
142            catalog.values().cloned().collect()
143        };
144        for spec in specs {
145            tracing::info!(alias = %spec.alias, "Prefetching model");
146            self.resolve_and_load_internal(&spec).await?;
147        }
148        Ok(())
149    }
150
151    /// Pre-load and cache specific aliases.
152    ///
153    /// Returns an error immediately if an alias is not found in the catalog
154    /// or if any model fails to load. Models already loaded are skipped.
155    pub async fn prefetch(&self, aliases: &[&str]) -> Result<()> {
156        for alias in aliases {
157            let spec = self.lookup_spec(alias).await?;
158            tracing::info!(alias = %alias, "Prefetching model");
159            self.resolve_and_load_internal(&spec).await?;
160        }
161        Ok(())
162    }
163
164    /// Resolve, load (if necessary), and return an instrumented [`EmbeddingModel`]
165    /// handle for the given alias.
166    ///
167    /// The returned handle is cached per alias so that repeated calls skip
168    /// spec lookup, key hashing, and wrapper allocation.
169    pub async fn embedding(&self, alias: &str) -> Result<Arc<dyn EmbeddingModel>> {
170        if let Some(cached) = self.handle_cache.embeddings.get(alias) {
171            return Ok(cached.clone());
172        }
173
174        let spec = self.lookup_spec(alias).await?;
175        let handle = self.resolve_and_load_internal(&spec).await?;
176        if let Some(model) = handle.downcast_ref::<Arc<dyn EmbeddingModel>>() {
177            let cached = self
178                .handle_cache
179                .embeddings
180                .entry(alias.to_string())
181                .or_insert_with(|| {
182                    let wrapper: Arc<dyn EmbeddingModel> = Arc::new(InstrumentedEmbeddingModel {
183                        inner: model.clone(),
184                        alias: alias.to_string(),
185                        provider_id: spec.provider_id.clone(),
186                        timeout: spec.timeout.map(std::time::Duration::from_secs),
187                        retry: spec.retry.clone(),
188                    });
189                    wrapper
190                })
191                .clone();
192            return Ok(cached);
193        }
194
195        // As in `sparse_embedder`: an `embed_hybrid` alias can serve a dense
196        // query from its dense head.
197        if let Some(model) = handle.downcast_ref::<Arc<dyn HybridEmbeddingModel>>() {
198            let dimensions = hybrid_adapter::head_width_of(
199                model,
200                crate::traits::HeadSet::DENSE,
201                alias,
202                &spec.provider_id,
203                "EmbeddingModel",
204            )?;
205            let view: Arc<dyn EmbeddingModel> = Arc::new(hybrid_adapter::HybridAsDense {
206                inner: model.clone(),
207                dimensions,
208            });
209            let cached = self
210                .handle_cache
211                .embeddings
212                .entry(alias.to_string())
213                .or_insert_with(|| {
214                    let wrapper: Arc<dyn EmbeddingModel> = Arc::new(InstrumentedEmbeddingModel {
215                        inner: view,
216                        alias: alias.to_string(),
217                        provider_id: spec.provider_id.clone(),
218                        timeout: spec.timeout.map(std::time::Duration::from_secs),
219                        retry: spec.retry.clone(),
220                    });
221                    wrapper
222                })
223                .clone();
224            return Ok(cached);
225        }
226
227        Err(RuntimeError::CapabilityMismatch(format!(
228            "Model for alias '{}' does not implement EmbeddingModel",
229            alias
230        )))
231    }
232
233    /// Resolve a dense text [`EmbeddingModel`] by alias.
234    ///
235    /// Agent-noun alias for [`embedding`](Self::embedding), matching the
236    /// `image_embedder` / `sparse_embedder` / `multi_vector_embedder` naming.
237    pub async fn embedder(&self, alias: &str) -> Result<Arc<dyn EmbeddingModel>> {
238        self.embedding(alias).await
239    }
240
241    /// Resolve, load (if necessary), and return an instrumented [`RerankerModel`]
242    /// handle for the given alias.
243    ///
244    /// The returned handle is cached per alias so that repeated calls skip
245    /// spec lookup, key hashing, and wrapper allocation.
246    pub async fn reranker(&self, alias: &str) -> Result<Arc<dyn RerankerModel>> {
247        if let Some(cached) = self.handle_cache.rerankers.get(alias) {
248            return Ok(cached.clone());
249        }
250
251        let spec = self.lookup_spec(alias).await?;
252        let handle = self.resolve_and_load_internal(&spec).await?;
253        if let Some(model) = handle.downcast_ref::<Arc<dyn RerankerModel>>() {
254            let cached = self
255                .handle_cache
256                .rerankers
257                .entry(alias.to_string())
258                .or_insert_with(|| {
259                    let wrapper: Arc<dyn RerankerModel> = Arc::new(InstrumentedRerankerModel {
260                        inner: model.clone(),
261                        alias: alias.to_string(),
262                        provider_id: spec.provider_id.clone(),
263                        timeout: spec.timeout.map(std::time::Duration::from_secs),
264                        retry: spec.retry.clone(),
265                    });
266                    wrapper
267                })
268                .clone();
269            return Ok(cached);
270        }
271        Err(RuntimeError::CapabilityMismatch(format!(
272            "Model for alias '{}' does not implement RerankerModel",
273            alias
274        )))
275    }
276
277    /// Resolve, load (if necessary), and return an instrumented [`GeneratorModel`]
278    /// handle for the given alias.
279    ///
280    /// The returned handle is cached per alias so that repeated calls skip
281    /// spec lookup, key hashing, and wrapper allocation.
282    pub async fn generator(&self, alias: &str) -> Result<Arc<dyn GeneratorModel>> {
283        if let Some(cached) = self.handle_cache.generators.get(alias) {
284            return Ok(cached.clone());
285        }
286
287        let spec = self.lookup_spec(alias).await?;
288        let handle = self.resolve_and_load_internal(&spec).await?;
289        if let Some(model) = handle.downcast_ref::<Arc<dyn GeneratorModel>>() {
290            let cached = self
291                .handle_cache
292                .generators
293                .entry(alias.to_string())
294                .or_insert_with(|| {
295                    let wrapper: Arc<dyn GeneratorModel> = Arc::new(InstrumentedGeneratorModel {
296                        inner: model.clone(),
297                        alias: alias.to_string(),
298                        provider_id: spec.provider_id.clone(),
299                        timeout: spec.timeout.map(std::time::Duration::from_secs),
300                        retry: spec.retry.clone(),
301                    });
302                    wrapper
303                })
304                .clone();
305            return Ok(cached);
306        }
307        Err(RuntimeError::CapabilityMismatch(format!(
308            "Model for alias '{}' does not implement GeneratorModel",
309            alias
310        )))
311    }
312
313    /// Resolve, load (if necessary), and return an instrumented [`RawTensorModel`]
314    /// handle for the given alias.
315    ///
316    /// The returned handle is cached per alias so that repeated calls skip
317    /// spec lookup, key hashing, and wrapper allocation.
318    pub async fn raw_tensor_model(&self, alias: &str) -> Result<Arc<dyn RawTensorModel>> {
319        if let Some(cached) = self.handle_cache.raw_tensor_models.get(alias) {
320            return Ok(cached.clone());
321        }
322
323        let spec = self.lookup_spec(alias).await?;
324        let handle = self.resolve_and_load_internal(&spec).await?;
325        if let Some(model) = handle.downcast_ref::<Arc<dyn RawTensorModel>>() {
326            let cached = self
327                .handle_cache
328                .raw_tensor_models
329                .entry(alias.to_string())
330                .or_insert_with(|| {
331                    let wrapper: Arc<dyn RawTensorModel> = Arc::new(InstrumentedRawTensorModel {
332                        inner: model.clone(),
333                        alias: alias.to_string(),
334                        provider_id: spec.provider_id.clone(),
335                        timeout: spec.timeout.map(std::time::Duration::from_secs),
336                        retry: spec.retry.clone(),
337                    });
338                    wrapper
339                })
340                .clone();
341            return Ok(cached);
342        }
343
344        Err(RuntimeError::ProviderCapabilityMissing {
345            alias: alias.to_string(),
346            provider_id: spec.provider_id,
347            capability: "RawTensorModel".to_string(),
348        })
349    }
350
351    /// Resolve, load (if necessary), and return an instrumented
352    /// [`ImageEmbeddingModel`] handle for the given alias.
353    ///
354    /// # Examples
355    /// ```no_run
356    /// # use std::sync::Arc;
357    /// # use uni_xervo::runtime::ModelRuntime;
358    /// # use uni_xervo::traits::ImageInput;
359    /// # async fn run(runtime: Arc<ModelRuntime>) -> uni_xervo::error::Result<()> {
360    /// let embedder = runtime.image_embedder("embed/siglip").await?;
361    /// let image = ImageInput::Bytes {
362    ///     data: std::fs::read("photo.png").unwrap(),
363    ///     media_type: "image/png".to_string(),
364    /// };
365    /// let result = embedder.embed(vec![image]).await?;
366    /// println!("dimension: {}", result.vectors[0].len());
367    /// # Ok(())
368    /// # }
369    /// ```
370    ///
371    /// # Errors
372    /// Returns an error if the alias is unknown, the model fails to load, or
373    /// the provider does not implement image embedding.
374    pub async fn image_embedder(&self, alias: &str) -> Result<Arc<dyn ImageEmbeddingModel>> {
375        if let Some(cached) = self.handle_cache.image_embedders.get(alias) {
376            return Ok(cached.clone());
377        }
378        let spec = self.lookup_spec(alias).await?;
379        let handle = self.resolve_and_load_internal(&spec).await?;
380        if let Some(model) = handle.downcast_ref::<Arc<dyn ImageEmbeddingModel>>() {
381            let cached = self
382                .handle_cache
383                .image_embedders
384                .entry(alias.to_string())
385                .or_insert_with(|| {
386                    let wrapper: Arc<dyn ImageEmbeddingModel> =
387                        Arc::new(InstrumentedImageEmbeddingModel {
388                            inner: model.clone(),
389                            alias: alias.to_string(),
390                            provider_id: spec.provider_id.clone(),
391                            timeout: spec.timeout.map(std::time::Duration::from_secs),
392                            retry: spec.retry.clone(),
393                        });
394                    wrapper
395                })
396                .clone();
397            return Ok(cached);
398        }
399        Err(RuntimeError::ProviderCapabilityMissing {
400            alias: alias.to_string(),
401            provider_id: spec.provider_id,
402            capability: "ImageEmbeddingModel".to_string(),
403        })
404    }
405
406    /// Resolve, load (if necessary), and return an instrumented
407    /// [`AudioEmbeddingModel`] handle for the given alias.
408    pub async fn audio_embedder(&self, alias: &str) -> Result<Arc<dyn AudioEmbeddingModel>> {
409        if let Some(cached) = self.handle_cache.audio_embedders.get(alias) {
410            return Ok(cached.clone());
411        }
412        let spec = self.lookup_spec(alias).await?;
413        let handle = self.resolve_and_load_internal(&spec).await?;
414        if let Some(model) = handle.downcast_ref::<Arc<dyn AudioEmbeddingModel>>() {
415            let cached = self
416                .handle_cache
417                .audio_embedders
418                .entry(alias.to_string())
419                .or_insert_with(|| {
420                    let wrapper: Arc<dyn AudioEmbeddingModel> =
421                        Arc::new(InstrumentedAudioEmbeddingModel {
422                            inner: model.clone(),
423                            alias: alias.to_string(),
424                            provider_id: spec.provider_id.clone(),
425                            timeout: spec.timeout.map(std::time::Duration::from_secs),
426                            retry: spec.retry.clone(),
427                        });
428                    wrapper
429                })
430                .clone();
431            return Ok(cached);
432        }
433        Err(RuntimeError::ProviderCapabilityMissing {
434            alias: alias.to_string(),
435            provider_id: spec.provider_id,
436            capability: "AudioEmbeddingModel".to_string(),
437        })
438    }
439
440    /// Resolve, load (if necessary), and return an instrumented
441    /// [`MultimodalEmbeddingModel`] handle for the given alias.
442    pub async fn multimodal_embedder(
443        &self,
444        alias: &str,
445    ) -> Result<Arc<dyn MultimodalEmbeddingModel>> {
446        if let Some(cached) = self.handle_cache.multimodal_embedders.get(alias) {
447            return Ok(cached.clone());
448        }
449        let spec = self.lookup_spec(alias).await?;
450        let handle = self.resolve_and_load_internal(&spec).await?;
451        if let Some(model) = handle.downcast_ref::<Arc<dyn MultimodalEmbeddingModel>>() {
452            let cached = self
453                .handle_cache
454                .multimodal_embedders
455                .entry(alias.to_string())
456                .or_insert_with(|| {
457                    let wrapper: Arc<dyn MultimodalEmbeddingModel> =
458                        Arc::new(InstrumentedMultimodalEmbeddingModel {
459                            inner: model.clone(),
460                            alias: alias.to_string(),
461                            provider_id: spec.provider_id.clone(),
462                            timeout: spec.timeout.map(std::time::Duration::from_secs),
463                            retry: spec.retry.clone(),
464                        });
465                    wrapper
466                })
467                .clone();
468            return Ok(cached);
469        }
470        Err(RuntimeError::ProviderCapabilityMissing {
471            alias: alias.to_string(),
472            provider_id: spec.provider_id,
473            capability: "MultimodalEmbeddingModel".to_string(),
474        })
475    }
476
477    /// Resolve, load (if necessary), and return an instrumented
478    /// [`SparseEmbeddingModel`] handle for the given alias.
479    ///
480    /// # Errors
481    /// Returns an error if the alias is unknown, the model fails to load, or
482    /// the provider does not implement sparse embedding.
483    pub async fn sparse_embedder(&self, alias: &str) -> Result<Arc<dyn SparseEmbeddingModel>> {
484        if let Some(cached) = self.handle_cache.sparse_embedders.get(alias) {
485            return Ok(cached.clone());
486        }
487        let spec = self.lookup_spec(alias).await?;
488        let handle = self.resolve_and_load_internal(&spec).await?;
489        if let Some(model) = handle.downcast_ref::<Arc<dyn SparseEmbeddingModel>>() {
490            let cached = self
491                .handle_cache
492                .sparse_embedders
493                .entry(alias.to_string())
494                .or_insert_with(|| {
495                    let wrapper: Arc<dyn SparseEmbeddingModel> =
496                        Arc::new(InstrumentedSparseEmbeddingModel {
497                            inner: model.clone(),
498                            alias: alias.to_string(),
499                            provider_id: spec.provider_id.clone(),
500                            timeout: spec.timeout.map(std::time::Duration::from_secs),
501                            retry: spec.retry.clone(),
502                        });
503                    wrapper
504                })
505                .clone();
506            return Ok(cached);
507        }
508        // An `embed_hybrid` alias stores `Arc<dyn HybridEmbeddingModel>`, so the
509        // downcast above misses even though the graph has a sparse head. Serve
510        // it through a single-head view rather than making the caller register a
511        // second alias for the same weights.
512        if let Some(model) = handle.downcast_ref::<Arc<dyn HybridEmbeddingModel>>() {
513            let vocab_size = hybrid_adapter::head_width_of(
514                model,
515                crate::traits::HeadSet::SPARSE,
516                alias,
517                &spec.provider_id,
518                "SparseEmbeddingModel",
519            )?;
520            let view: Arc<dyn SparseEmbeddingModel> = Arc::new(hybrid_adapter::HybridAsSparse {
521                inner: model.clone(),
522                vocab_size,
523            });
524            let cached = self
525                .handle_cache
526                .sparse_embedders
527                .entry(alias.to_string())
528                .or_insert_with(|| {
529                    // Instrument outside the view so timeouts, retries and
530                    // metrics are attributed to this alias and task.
531                    let wrapper: Arc<dyn SparseEmbeddingModel> =
532                        Arc::new(InstrumentedSparseEmbeddingModel {
533                            inner: view,
534                            alias: alias.to_string(),
535                            provider_id: spec.provider_id.clone(),
536                            timeout: spec.timeout.map(std::time::Duration::from_secs),
537                            retry: spec.retry.clone(),
538                        });
539                    wrapper
540                })
541                .clone();
542            return Ok(cached);
543        }
544        Err(RuntimeError::ProviderCapabilityMissing {
545            alias: alias.to_string(),
546            provider_id: spec.provider_id,
547            capability: "SparseEmbeddingModel".to_string(),
548        })
549    }
550
551    /// Resolve, load (if necessary), and return an instrumented
552    /// [`MultiVectorEmbeddingModel`] handle for the given alias.
553    ///
554    /// # Errors
555    /// Returns an error if the alias is unknown, the model fails to load, or
556    /// the provider does not implement multi-vector embedding.
557    pub async fn multi_vector_embedder(
558        &self,
559        alias: &str,
560    ) -> Result<Arc<dyn MultiVectorEmbeddingModel>> {
561        if let Some(cached) = self.handle_cache.multi_vector_embedders.get(alias) {
562            return Ok(cached.clone());
563        }
564        let spec = self.lookup_spec(alias).await?;
565        let handle = self.resolve_and_load_internal(&spec).await?;
566        if let Some(model) = handle.downcast_ref::<Arc<dyn MultiVectorEmbeddingModel>>() {
567            let cached = self
568                .handle_cache
569                .multi_vector_embedders
570                .entry(alias.to_string())
571                .or_insert_with(|| {
572                    let wrapper: Arc<dyn MultiVectorEmbeddingModel> =
573                        Arc::new(InstrumentedMultiVectorEmbeddingModel {
574                            inner: model.clone(),
575                            alias: alias.to_string(),
576                            provider_id: spec.provider_id.clone(),
577                            timeout: spec.timeout.map(std::time::Duration::from_secs),
578                            retry: spec.retry.clone(),
579                        });
580                    wrapper
581                })
582                .clone();
583            return Ok(cached);
584        }
585        // Same single-head view as `sparse_embedder`, for the ColBERT head.
586        if let Some(model) = handle.downcast_ref::<Arc<dyn HybridEmbeddingModel>>() {
587            let dimensions = hybrid_adapter::head_width_of(
588                model,
589                crate::traits::HeadSet::MULTI_VECTOR,
590                alias,
591                &spec.provider_id,
592                "MultiVectorEmbeddingModel",
593            )?;
594            let view: Arc<dyn MultiVectorEmbeddingModel> =
595                Arc::new(hybrid_adapter::HybridAsMultiVector {
596                    inner: model.clone(),
597                    dimensions,
598                });
599            let cached = self
600                .handle_cache
601                .multi_vector_embedders
602                .entry(alias.to_string())
603                .or_insert_with(|| {
604                    let wrapper: Arc<dyn MultiVectorEmbeddingModel> =
605                        Arc::new(InstrumentedMultiVectorEmbeddingModel {
606                            inner: view,
607                            alias: alias.to_string(),
608                            provider_id: spec.provider_id.clone(),
609                            timeout: spec.timeout.map(std::time::Duration::from_secs),
610                            retry: spec.retry.clone(),
611                        });
612                    wrapper
613                })
614                .clone();
615            return Ok(cached);
616        }
617        Err(RuntimeError::ProviderCapabilityMissing {
618            alias: alias.to_string(),
619            provider_id: spec.provider_id,
620            capability: "MultiVectorEmbeddingModel".to_string(),
621        })
622    }
623
624    /// Resolve, load (if necessary), and return an instrumented
625    /// [`HybridEmbeddingModel`] handle for the given alias.
626    ///
627    /// The handle serves dense, sparse, and multi-vector heads from a single
628    /// forward pass; select which to materialize with a
629    /// [`HeadSet`](crate::traits::HeadSet). Only multi-output graphs with a
630    /// hybrid preset (e.g. `BGEM3Hybrid`) resolve here — single-head models use
631    /// the per-task resolvers.
632    ///
633    /// # Errors
634    /// Returns an error if the alias is unknown, the model has no hybrid preset,
635    /// the model fails to load, or the loaded handle lacks the hybrid capability.
636    pub async fn hybrid_embedder(&self, alias: &str) -> Result<Arc<dyn HybridEmbeddingModel>> {
637        if let Some(cached) = self.handle_cache.hybrid_embedders.get(alias) {
638            return Ok(cached.clone());
639        }
640        let spec = self.lookup_spec(alias).await?;
641        let handle = self.resolve_and_load_internal(&spec).await?;
642        if let Some(model) = handle.downcast_ref::<Arc<dyn HybridEmbeddingModel>>() {
643            let cached = self
644                .handle_cache
645                .hybrid_embedders
646                .entry(alias.to_string())
647                .or_insert_with(|| {
648                    let wrapper: Arc<dyn HybridEmbeddingModel> =
649                        Arc::new(InstrumentedHybridEmbeddingModel {
650                            inner: model.clone(),
651                            alias: alias.to_string(),
652                            provider_id: spec.provider_id.clone(),
653                            timeout: spec.timeout.map(std::time::Duration::from_secs),
654                            retry: spec.retry.clone(),
655                        });
656                    wrapper
657                })
658                .clone();
659            return Ok(cached);
660        }
661        Err(RuntimeError::ProviderCapabilityMissing {
662            alias: alias.to_string(),
663            provider_id: spec.provider_id,
664            capability: "HybridEmbeddingModel".to_string(),
665        })
666    }
667
668    /// Resolve, load (if necessary), and return an instrumented [`NlpModel`]
669    /// handle for the given alias.
670    pub async fn nlp_model(&self, alias: &str) -> Result<Arc<dyn NlpModel>> {
671        if let Some(cached) = self.handle_cache.nlp_models.get(alias) {
672            return Ok(cached.clone());
673        }
674        let spec = self.lookup_spec(alias).await?;
675        let handle = self.resolve_and_load_internal(&spec).await?;
676        if let Some(model) = handle.downcast_ref::<Arc<dyn NlpModel>>() {
677            let cached = self
678                .handle_cache
679                .nlp_models
680                .entry(alias.to_string())
681                .or_insert_with(|| {
682                    let wrapper: Arc<dyn NlpModel> = Arc::new(InstrumentedNlpModel {
683                        inner: model.clone(),
684                        alias: alias.to_string(),
685                        provider_id: spec.provider_id.clone(),
686                        timeout: spec.timeout.map(std::time::Duration::from_secs),
687                        retry: spec.retry.clone(),
688                    });
689                    wrapper
690                })
691                .clone();
692            return Ok(cached);
693        }
694        Err(RuntimeError::ProviderCapabilityMissing {
695            alias: alias.to_string(),
696            provider_id: spec.provider_id,
697            capability: "NlpModel".to_string(),
698        })
699    }
700
701    /// Resolve, load (if necessary), and return an instrumented
702    /// [`DocumentExtractionModel`] handle for the given alias.
703    ///
704    /// # Examples
705    /// ```no_run
706    /// # use std::sync::Arc;
707    /// # use uni_xervo::runtime::ModelRuntime;
708    /// # use uni_xervo::traits::{DocExtractOptions, DocOutputFormat, ImageInput};
709    /// # async fn run(runtime: Arc<ModelRuntime>) -> uni_xervo::error::Result<()> {
710    /// let extractor = runtime.document_extractor("docext/olmocr").await?;
711    /// let page = ImageInput::Bytes {
712    ///     data: std::fs::read("page.png").unwrap(),
713    ///     media_type: "image/png".to_string(),
714    /// };
715    /// let options = DocExtractOptions {
716    ///     output: DocOutputFormat::Markdown,
717    ///     include_tables: true,
718    ///     include_formulas: true,
719    ///     include_bboxes: false,
720    /// };
721    /// let pages = extractor.extract(vec![page], options).await?;
722    /// println!("{}", pages[0].plain_markdown);
723    /// # Ok(())
724    /// # }
725    /// ```
726    ///
727    /// # Errors
728    /// Returns an error if the alias is unknown, the model fails to load, or
729    /// the provider does not implement document extraction.
730    pub async fn document_extractor(
731        &self,
732        alias: &str,
733    ) -> Result<Arc<dyn DocumentExtractionModel>> {
734        if let Some(cached) = self.handle_cache.document_extractors.get(alias) {
735            return Ok(cached.clone());
736        }
737        let spec = self.lookup_spec(alias).await?;
738        let handle = self.resolve_and_load_internal(&spec).await?;
739        if let Some(model) = handle.downcast_ref::<Arc<dyn DocumentExtractionModel>>() {
740            let cached = self
741                .handle_cache
742                .document_extractors
743                .entry(alias.to_string())
744                .or_insert_with(|| {
745                    let wrapper: Arc<dyn DocumentExtractionModel> =
746                        Arc::new(InstrumentedDocumentExtractionModel {
747                            inner: model.clone(),
748                            alias: alias.to_string(),
749                            provider_id: spec.provider_id.clone(),
750                            timeout: spec.timeout.map(std::time::Duration::from_secs),
751                            retry: spec.retry.clone(),
752                        });
753                    wrapper
754                })
755                .clone();
756            return Ok(cached);
757        }
758        Err(RuntimeError::ProviderCapabilityMissing {
759            alias: alias.to_string(),
760            provider_id: spec.provider_id,
761            capability: "DocumentExtractionModel".to_string(),
762        })
763    }
764
765    /// Resolve, load (if necessary), and return an instrumented
766    /// [`TranscriptionModel`] handle for the given alias.
767    pub async fn transcriber(&self, alias: &str) -> Result<Arc<dyn TranscriptionModel>> {
768        if let Some(cached) = self.handle_cache.transcribers.get(alias) {
769            return Ok(cached.clone());
770        }
771        let spec = self.lookup_spec(alias).await?;
772        let handle = self.resolve_and_load_internal(&spec).await?;
773        if let Some(model) = handle.downcast_ref::<Arc<dyn TranscriptionModel>>() {
774            let cached = self
775                .handle_cache
776                .transcribers
777                .entry(alias.to_string())
778                .or_insert_with(|| {
779                    let wrapper: Arc<dyn TranscriptionModel> =
780                        Arc::new(InstrumentedTranscriptionModel {
781                            inner: model.clone(),
782                            alias: alias.to_string(),
783                            provider_id: spec.provider_id.clone(),
784                            timeout: spec.timeout.map(std::time::Duration::from_secs),
785                            retry: spec.retry.clone(),
786                        });
787                    wrapper
788                })
789                .clone();
790            return Ok(cached);
791        }
792        Err(RuntimeError::ProviderCapabilityMissing {
793            alias: alias.to_string(),
794            provider_id: spec.provider_id,
795            capability: "TranscriptionModel".to_string(),
796        })
797    }
798
799    /// Resolve, load (if necessary), and return an instrumented [`OcrModel`]
800    /// handle for the given alias.
801    ///
802    /// # Examples
803    /// ```no_run
804    /// # use std::sync::Arc;
805    /// # use uni_xervo::runtime::ModelRuntime;
806    /// # use uni_xervo::traits::ImageInput;
807    /// # async fn run(runtime: Arc<ModelRuntime>) -> uni_xervo::error::Result<()> {
808    /// let ocr = runtime.ocr_model("ocr/ppocr-en").await?;
809    /// let image = ImageInput::Bytes {
810    ///     data: std::fs::read("scan.png").unwrap(),
811    ///     media_type: "image/png".to_string(),
812    /// };
813    /// let results = ocr.recognize(vec![image]).await?;
814    /// println!("{}", results[0].plain_text);
815    /// # Ok(())
816    /// # }
817    /// ```
818    ///
819    /// # Errors
820    /// Returns an error if the alias is unknown, the model fails to load, or
821    /// the provider does not implement OCR.
822    pub async fn ocr_model(&self, alias: &str) -> Result<Arc<dyn OcrModel>> {
823        if let Some(cached) = self.handle_cache.ocr_models.get(alias) {
824            return Ok(cached.clone());
825        }
826        let spec = self.lookup_spec(alias).await?;
827        let handle = self.resolve_and_load_internal(&spec).await?;
828        if let Some(model) = handle.downcast_ref::<Arc<dyn OcrModel>>() {
829            let cached = self
830                .handle_cache
831                .ocr_models
832                .entry(alias.to_string())
833                .or_insert_with(|| {
834                    let wrapper: Arc<dyn OcrModel> = Arc::new(InstrumentedOcrModel {
835                        inner: model.clone(),
836                        alias: alias.to_string(),
837                        provider_id: spec.provider_id.clone(),
838                        timeout: spec.timeout.map(std::time::Duration::from_secs),
839                        retry: spec.retry.clone(),
840                    });
841                    wrapper
842                })
843                .clone();
844            return Ok(cached);
845        }
846        Err(RuntimeError::ProviderCapabilityMissing {
847            alias: alias.to_string(),
848            provider_id: spec.provider_id,
849            capability: "OcrModel".to_string(),
850        })
851    }
852
853    #[tracing::instrument(skip(self, spec), fields(provider, model))]
854    async fn resolve_and_load_internal(
855        &self,
856        spec: &ModelAliasSpec,
857    ) -> Result<Arc<dyn Any + Send + Sync>> {
858        let key = ModelRuntimeKey::new(spec);
859
860        // Fast path: already loaded
861        {
862            let registry = self.registry.instances.read().await;
863            if let Some(handle) = registry.get(&key) {
864                return Ok(handle.clone());
865            }
866        }
867
868        // Slow path: coordinate loading
869        let lock = {
870            let mut locks = self.registry.loader_locks.lock().await;
871            locks
872                .entry(key.clone())
873                .or_insert_with(|| Arc::new(Mutex::new(())))
874                .clone()
875        };
876
877        // Acquire loader lock for this key
878        let _guard = lock.lock().await;
879
880        // Double-check after acquiring the loader lock
881        {
882            let registry = self.registry.instances.read().await;
883            if let Some(handle) = registry.get(&key) {
884                let result = Ok(handle.clone());
885                let mut locks = self.registry.loader_locks.lock().await;
886                locks.remove(&key);
887                return result;
888            }
889        }
890
891        let load_timeout =
892            std::time::Duration::from_secs(spec.load_timeout.unwrap_or(DEFAULT_LOAD_TIMEOUT_SECS));
893
894        let result = match tokio::time::timeout(load_timeout, async {
895            let provider = self.providers.get(&spec.provider_id).ok_or_else(|| {
896                RuntimeError::ProviderNotFound(format!("Provider '{}' not found", spec.provider_id))
897            })?;
898
899            tracing::info!(alias = %spec.alias, provider = %spec.provider_id, "Loading model instance");
900            let start = std::time::Instant::now();
901            let handle_result = provider.load(spec).await;
902            let duration = start.elapsed().as_secs_f64();
903
904            metrics::histogram!("model_load.duration_seconds").record(duration);
905
906            let handle = match handle_result {
907                Ok(h) => {
908                    metrics::counter!("model_load.total", "status" => "success").increment(1);
909                    h
910                }
911                Err(e) => {
912                    metrics::counter!("model_load.total", "status" => "failure").increment(1);
913                    tracing::error!(alias = %spec.alias, error = %e, "Model load failed");
914                    return Err(e);
915                }
916            };
917
918            // Model warmup. Dispatch covers every task trait — including the
919            // multimodal surface, which previously fell through to no warmup at
920            // all. The trait default is a no-op, so adding arms only enables
921            // eager weight loading for providers that override `warmup`.
922            if let Some(model) = handle.downcast_ref::<Arc<dyn EmbeddingModel>>() {
923                model.warmup().await?;
924            } else if let Some(model) = handle.downcast_ref::<Arc<dyn RerankerModel>>() {
925                model.warmup().await?;
926            } else if let Some(model) = handle.downcast_ref::<Arc<dyn GeneratorModel>>() {
927                model.warmup().await?;
928            } else if let Some(model) = handle.downcast_ref::<Arc<dyn RawTensorModel>>() {
929                model.warmup().await?;
930            } else if let Some(model) = handle.downcast_ref::<Arc<dyn ImageEmbeddingModel>>() {
931                model.warmup().await?;
932            } else if let Some(model) = handle.downcast_ref::<Arc<dyn AudioEmbeddingModel>>() {
933                model.warmup().await?;
934            } else if let Some(model) = handle.downcast_ref::<Arc<dyn MultimodalEmbeddingModel>>() {
935                model.warmup().await?;
936            } else if let Some(model) = handle.downcast_ref::<Arc<dyn SparseEmbeddingModel>>() {
937                model.warmup().await?;
938            } else if let Some(model) = handle.downcast_ref::<Arc<dyn MultiVectorEmbeddingModel>>() {
939                model.warmup().await?;
940            } else if let Some(model) = handle.downcast_ref::<Arc<dyn HybridEmbeddingModel>>() {
941                model.warmup().await?;
942            } else if let Some(model) = handle.downcast_ref::<Arc<dyn NlpModel>>() {
943                model.warmup().await?;
944            } else if let Some(model) = handle.downcast_ref::<Arc<dyn DocumentExtractionModel>>() {
945                model.warmup().await?;
946            } else if let Some(model) = handle.downcast_ref::<Arc<dyn TranscriptionModel>>() {
947                model.warmup().await?;
948            } else if let Some(model) = handle.downcast_ref::<Arc<dyn OcrModel>>() {
949                model.warmup().await?;
950            }
951
952            {
953                let mut registry = self.registry.instances.write().await;
954                registry.insert(key.clone(), handle.clone());
955            }
956
957            Ok(handle)
958        })
959        .await
960        {
961            Ok(res) => res,
962            Err(_) => {
963                metrics::counter!("model_load.total", "status" => "failure").increment(1);
964                tracing::error!(
965                    alias = %spec.alias,
966                    provider = %spec.provider_id,
967                    timeout_secs = load_timeout.as_secs(),
968                    "Model load timed out"
969                );
970                Err(RuntimeError::Timeout)
971            }
972        };
973
974        // Bound loader lock map growth by removing this key once the load path completes.
975        // Existing waiters hold cloned lock Arcs, so this is safe.
976        {
977            let mut locks = self.registry.loader_locks.lock().await;
978            locks.remove(&key);
979        }
980
981        result
982    }
983}
984
985/// Builder for constructing a [`ModelRuntime`] with registered providers,
986/// a model catalog, and a warmup policy.
987///
988/// ```rust,no_run
989/// # use uni_xervo::runtime::ModelRuntime;
990/// # async fn example() -> Result<(), Box<dyn std::error::Error>> {
991/// let runtime = ModelRuntime::builder()
992///     // .register_provider(...)
993///     // .catalog(...)
994///     .build()
995///     .await?;
996/// # Ok(())
997/// # }
998/// ```
999#[derive(Default)]
1000pub struct ModelRuntimeBuilder {
1001    providers: HashMap<String, Box<dyn ModelProvider>>,
1002    catalog: Vec<ModelAliasSpec>,
1003    warmup_policy: crate::api::WarmupPolicy,
1004}
1005
1006impl ModelRuntimeBuilder {
1007    /// Register a provider. The provider's
1008    /// [`provider_id`](crate::traits::ModelProvider::provider_id) is used as
1009    /// the lookup key; registering a second provider with the same ID
1010    /// replaces the first.
1011    pub fn register_provider<P: ModelProvider + 'static>(mut self, provider: P) -> Self {
1012        self.providers
1013            .insert(provider.provider_id().to_string(), Box::new(provider));
1014        self
1015    }
1016
1017    /// Set the model catalog from a pre-built vector of specs.
1018    pub fn catalog(mut self, catalog: Vec<ModelAliasSpec>) -> Self {
1019        self.catalog = catalog;
1020        self
1021    }
1022
1023    /// Load catalog from a JSON string (array of model alias specs).
1024    pub fn catalog_from_str(mut self, s: &str) -> Result<Self> {
1025        self.catalog = crate::api::catalog_from_str(s)?;
1026        Ok(self)
1027    }
1028
1029    /// Load catalog from a JSON file (array of model alias specs).
1030    pub fn catalog_from_file(mut self, path: impl AsRef<std::path::Path>) -> Result<Self> {
1031        self.catalog = crate::api::catalog_from_file(path)?;
1032        Ok(self)
1033    }
1034
1035    /// Set the global warmup policy applied to providers during
1036    /// [`build`](Self::build).
1037    pub fn warmup_policy(mut self, policy: crate::api::WarmupPolicy) -> Self {
1038        self.warmup_policy = policy;
1039        self
1040    }
1041
1042    /// Validate the catalog, execute the warmup policy, and return the
1043    /// constructed [`ModelRuntime`].
1044    ///
1045    /// Returns an error if any spec references an unknown provider, contains
1046    /// invalid options, or if a required eager warmup fails.
1047    pub async fn build(self) -> Result<Arc<ModelRuntime>> {
1048        let mut catalog_map = HashMap::new();
1049        for spec in self.catalog {
1050            spec.validate()?;
1051            if !self.providers.contains_key(&spec.provider_id) {
1052                return Err(RuntimeError::Config(format!(
1053                    "Unknown provider '{}' for alias '{}'",
1054                    spec.provider_id, spec.alias
1055                )));
1056            }
1057            validate_provider_options(&spec.provider_id, spec.task, &spec.options)?;
1058            if catalog_map.insert(spec.alias.clone(), spec).is_some() {
1059                return Err(RuntimeError::Config(
1060                    "Duplicate alias in catalog".to_string(),
1061                ));
1062            }
1063        }
1064
1065        let runtime = Arc::new(ModelRuntime {
1066            providers: self.providers,
1067            registry: Arc::new(ModelRegistry::default()),
1068            catalog: RwLock::new(catalog_map),
1069            handle_cache: HandleCache::default(),
1070        });
1071
1072        // Provider Warmup Phase
1073        match self.warmup_policy {
1074            crate::api::WarmupPolicy::Eager => {
1075                for (id, provider) in &runtime.providers {
1076                    tracing::info!(provider = %id, "Eagerly warming up provider");
1077                    provider.warmup().await.map_err(|e| {
1078                        RuntimeError::Load(format!("Failed to warmup provider {}: {}", id, e))
1079                    })?;
1080                }
1081            }
1082            crate::api::WarmupPolicy::Background => {
1083                for id in runtime.providers.keys() {
1084                    tracing::info!(provider = %id, "Scheduling background provider warmup");
1085                    // We have the Arc<ModelRuntime> already.
1086                    let rt = runtime.clone();
1087                    let provider_id = id.clone();
1088                    tokio::spawn(async move {
1089                        if let Some(provider) = rt.providers.get(&provider_id)
1090                            && let Err(e) = provider.warmup().await
1091                        {
1092                            tracing::error!(provider = %provider_id, error = %e, "Background provider warmup failed");
1093                        }
1094                    });
1095                }
1096            }
1097            crate::api::WarmupPolicy::Lazy => {
1098                tracing::debug!("Lazy provider warmup (no-op)");
1099            }
1100        }
1101
1102        // Model Warmup Phase
1103        let mut warmup_tasks = Vec::new();
1104
1105        let specs: Vec<ModelAliasSpec> = {
1106            let catalog = runtime.catalog.read().await;
1107            catalog.values().cloned().collect()
1108        };
1109
1110        for spec in specs {
1111            match spec.warmup {
1112                crate::api::WarmupPolicy::Eager => {
1113                    tracing::info!(alias = %spec.alias, "Eagerly warming up model");
1114                    if let Err(e) = runtime.resolve_and_load_internal(&spec).await {
1115                        if spec.required {
1116                            return Err(e);
1117                        }
1118                        tracing::error!(
1119                            alias = %spec.alias,
1120                            provider = %spec.provider_id,
1121                            error = %e,
1122                            "Optional eager model warmup failed; continuing startup"
1123                        );
1124                    }
1125                }
1126                crate::api::WarmupPolicy::Background => {
1127                    tracing::info!(alias = %spec.alias, "Scheduling background warmup");
1128                    let rt = runtime.clone();
1129                    let spec_clone = spec.clone();
1130                    // Spawn background task
1131                    warmup_tasks.push(tokio::spawn(async move {
1132                        if let Err(e) = rt.resolve_and_load_internal(&spec_clone).await {
1133                            tracing::error!(alias = %spec_clone.alias, error = %e, "Background warmup failed");
1134                        }
1135                    }));
1136                }
1137                crate::api::WarmupPolicy::Lazy => {
1138                    tracing::debug!(alias = %spec.alias, "Lazy warmup (no-op)");
1139                }
1140            }
1141        }
1142
1143        // We don't await background tasks here, they run detached.
1144        // Eager tasks are already awaited.
1145
1146        Ok(runtime)
1147    }
1148}
1149
1150#[cfg(test)]
1151mod tests {
1152    use super::*;
1153    use crate::api::ModelTask;
1154    use crate::mock::{MockProvider, make_spec};
1155
1156    #[tokio::test]
1157    async fn loader_lock_entries_cleaned_after_successful_load() {
1158        let spec = make_spec("embed/test", ModelTask::Embed, "mock/embed", "test-model");
1159        let runtime = ModelRuntime::builder()
1160            .register_provider(MockProvider::embed_only())
1161            .catalog(vec![spec])
1162            .build()
1163            .await
1164            .unwrap();
1165
1166        let _ = runtime.embedding("embed/test").await.unwrap();
1167
1168        let locks = runtime.registry.loader_locks.lock().await;
1169        assert!(
1170            locks.is_empty(),
1171            "loader lock map should be empty after load"
1172        );
1173    }
1174
1175    #[tokio::test]
1176    async fn loader_lock_entries_cleaned_after_failed_load() {
1177        let mut spec = make_spec("embed/test", ModelTask::Embed, "mock/failing", "test-model");
1178        spec.warmup = crate::api::WarmupPolicy::Lazy;
1179        let runtime = ModelRuntime::builder()
1180            .register_provider(MockProvider::failing())
1181            .catalog(vec![spec])
1182            .build()
1183            .await
1184            .unwrap();
1185
1186        let err = runtime.embedding("embed/test").await;
1187        assert!(err.is_err());
1188
1189        let locks = runtime.registry.loader_locks.lock().await;
1190        assert!(
1191            locks.is_empty(),
1192            "loader lock map should be empty after failure"
1193        );
1194    }
1195
1196    #[tokio::test]
1197    async fn loader_lock_entries_cleaned_after_load_timeout() {
1198        let mut spec = make_spec("embed/test", ModelTask::Embed, "mock/embed", "test-model");
1199        spec.warmup = crate::api::WarmupPolicy::Lazy;
1200        spec.load_timeout = Some(1);
1201
1202        let runtime = ModelRuntime::builder()
1203            .register_provider(MockProvider::embed_only().with_load_delay(2_000))
1204            .catalog(vec![spec])
1205            .build()
1206            .await
1207            .unwrap();
1208
1209        let err = runtime.embedding("embed/test").await;
1210        assert!(matches!(err, Err(RuntimeError::Timeout)));
1211
1212        let locks = runtime.registry.loader_locks.lock().await;
1213        assert!(
1214            locks.is_empty(),
1215            "loader lock map should be empty after load timeout"
1216        );
1217    }
1218
1219    // -----------------------------------------------------------------------
1220    // An EmbedHybrid alias serves the single-head accessors (issue #49)
1221    // -----------------------------------------------------------------------
1222
1223    use crate::traits::HeadSet;
1224
1225    /// Register one `EmbedHybrid` alias, the way a hybrid retrieval pipeline
1226    /// does, and resolve it through each single-head accessor.
1227    async fn hybrid_runtime(provider: MockProvider) -> Arc<ModelRuntime> {
1228        let spec = make_spec(
1229            "embed/hybrid",
1230            ModelTask::EmbedHybrid,
1231            "mock/hybrid",
1232            "aapot/bge-m3-onnx",
1233        );
1234        ModelRuntime::builder()
1235            .register_provider(provider)
1236            .catalog(vec![spec])
1237            .build()
1238            .await
1239            .unwrap()
1240    }
1241
1242    fn hybrid_provider() -> MockProvider {
1243        MockProvider::new("mock/hybrid", vec![ModelTask::EmbedHybrid])
1244    }
1245
1246    /// The reported bug: both query channels failed with
1247    /// `ProviderCapabilityMissing` while ingest wrote the columns fine.
1248    #[tokio::test]
1249    async fn hybrid_alias_serves_sparse_and_multi_vector() {
1250        let runtime = hybrid_runtime(hybrid_provider()).await;
1251
1252        let sparse = runtime
1253            .sparse_embedder("embed/hybrid")
1254            .await
1255            .expect("a hybrid alias must satisfy a sparse query");
1256        assert_eq!(sparse.vocab_size(), 250002);
1257        let out = sparse.embed(&["a", "b"]).await.unwrap();
1258        assert_eq!(out.vectors.len(), 2);
1259
1260        let mv = runtime
1261            .multi_vector_embedder("embed/hybrid")
1262            .await
1263            .expect("a hybrid alias must satisfy a multi-vector query");
1264        assert_eq!(mv.dimensions(), 1024);
1265        let out = mv.embed(&["a", "b"]).await.unwrap();
1266        assert_eq!(out.vectors.len(), 2);
1267    }
1268
1269    /// Not in the original report: the dense accessor has the same shape, so a
1270    /// hybrid alias must serve a dense query too.
1271    #[tokio::test]
1272    async fn hybrid_alias_serves_dense() {
1273        let runtime = hybrid_runtime(hybrid_provider()).await;
1274
1275        let dense = runtime
1276            .embedding("embed/hybrid")
1277            .await
1278            .expect("a hybrid alias must satisfy a dense query");
1279        assert_eq!(dense.dimensions(), 1024);
1280        let out = dense.embed(&["a"]).await.unwrap();
1281        assert_eq!(out.vectors.len(), 1);
1282        assert_eq!(out.vectors[0].len(), 1024);
1283    }
1284
1285    /// The hybrid handle itself must keep working unchanged.
1286    #[tokio::test]
1287    async fn hybrid_alias_still_serves_the_hybrid_accessor() {
1288        let runtime = hybrid_runtime(hybrid_provider()).await;
1289
1290        let hybrid = runtime.hybrid_embedder("embed/hybrid").await.unwrap();
1291        let out = hybrid.embed(&["a"], HeadSet::ALL).await.unwrap();
1292        assert!(out.dense.is_some() && out.sparse.is_some() && out.multi_vector.is_some());
1293    }
1294
1295    /// A head the graph does not expose must still fail — but say why, rather
1296    /// than claiming the provider has no sparse capability at all.
1297    #[tokio::test]
1298    async fn hybrid_alias_without_the_head_fails_with_a_specific_message() {
1299        let provider = hybrid_provider().with_hybrid_heads(HeadSet::DENSE);
1300        let runtime = hybrid_runtime(provider).await;
1301
1302        let err = runtime
1303            .sparse_embedder("embed/hybrid")
1304            .await
1305            .err()
1306            .expect("a hybrid model with no sparse head cannot serve sparse");
1307        let msg = err.to_string();
1308        assert!(msg.contains("does not expose that head"), "{msg}");
1309        assert!(msg.contains("embed/hybrid"), "{msg}");
1310    }
1311
1312    /// A hybrid implementation that doesn't report widths cannot be adapted,
1313    /// because the single-head traits must answer `vocab_size` / `dimensions`.
1314    #[tokio::test]
1315    async fn hybrid_alias_without_widths_is_rejected() {
1316        let runtime = hybrid_runtime(hybrid_provider().without_hybrid_widths()).await;
1317
1318        let err = runtime
1319            .sparse_embedder("embed/hybrid")
1320            .await
1321            .err()
1322            .expect("no reported width means no faithful vocab_size");
1323        assert!(err.to_string().contains("does not report its width"));
1324    }
1325
1326    /// The per-alias handle cache must hold the adapted view, so the second
1327    /// resolve is served from cache rather than rebuilt.
1328    #[tokio::test]
1329    async fn adapted_hybrid_handles_are_cached_per_alias() {
1330        let runtime = hybrid_runtime(hybrid_provider()).await;
1331
1332        let first = runtime.sparse_embedder("embed/hybrid").await.unwrap();
1333        let second = runtime.sparse_embedder("embed/hybrid").await.unwrap();
1334        assert!(
1335            Arc::ptr_eq(&first, &second),
1336            "expected the cached adapter, not a fresh one"
1337        );
1338    }
1339
1340    /// A non-hybrid alias that genuinely lacks the capability must keep its
1341    /// original error — the fallback must not mask real mismatches.
1342    #[tokio::test]
1343    async fn non_hybrid_alias_still_reports_capability_missing() {
1344        let spec = make_spec("embed/test", ModelTask::Embed, "mock/embed", "test-model");
1345        let runtime = ModelRuntime::builder()
1346            .register_provider(MockProvider::embed_only())
1347            .catalog(vec![spec])
1348            .build()
1349            .await
1350            .unwrap();
1351
1352        let err = runtime.sparse_embedder("embed/test").await.err().unwrap();
1353        assert!(matches!(
1354            err,
1355            RuntimeError::ProviderCapabilityMissing { .. }
1356        ));
1357        assert!(err.to_string().contains("SparseEmbeddingModel"));
1358    }
1359}
uni_xervo/runtime.rs

uni_xervo/
runtime.rs
