runifold_model/
capability_catalog.rs1use std::collections::BTreeMap;
4
5use serde::{Deserialize, Serialize};
6
7use crate::{ModelCapabilities, ModelRef};
8
9#[derive(Clone, Debug, Default, Deserialize, PartialEq, Serialize)]
15pub struct ModelCapabilityCatalog {
16 models: BTreeMap<ModelRef, ModelCapabilities>,
17}
18
19impl ModelCapabilityCatalog {
20 #[must_use]
22 pub const fn new() -> Self {
23 Self {
24 models: BTreeMap::new(),
25 }
26 }
27
28 pub fn insert(
30 &mut self,
31 model: ModelRef,
32 capabilities: ModelCapabilities,
33 ) -> Option<ModelCapabilities> {
34 self.models.insert(model, capabilities)
35 }
36
37 #[must_use]
39 pub fn get(&self, model: &ModelRef) -> Option<&ModelCapabilities> {
40 self.models.get(model)
41 }
42
43 #[must_use]
45 pub fn len(&self) -> usize {
46 self.models.len()
47 }
48
49 #[must_use]
51 pub fn is_empty(&self) -> bool {
52 self.models.is_empty()
53 }
54}
55
56#[cfg(test)]
57mod tests {
58 use crate::{FeatureSupport, SupportLevel};
59
60 use super::*;
61
62 #[test]
63 fn catalog_requires_exact_provider_and_model_identity() {
64 let model = ModelRef::new("openai", "model-a");
65 let capabilities = ModelCapabilities {
66 tools: FeatureSupport::new(SupportLevel::Native),
67 ..ModelCapabilities::default()
68 };
69 let mut catalog = ModelCapabilityCatalog::new();
70 catalog.insert(model.clone(), capabilities.clone());
71
72 assert_eq!(catalog.get(&model), Some(&capabilities));
73 assert!(catalog.get(&ModelRef::new("openai", "model-b")).is_none());
74 assert!(catalog.get(&ModelRef::new("other", "model-a")).is_none());
75 }
76}