| 1 | //! Concise model-capability badges for Fleet configuration UI (#5038). |
| 2 | //! |
| 3 | //! One resolver answers "what can this model do" for both the Fleet setup |
| 4 | //! model-selection step and the roster detail pane. Facts come from the |
| 5 | //! existing owners — the merged Models.dev catalog (via |
| 6 | //! [`crate::provider_lake::catalog_offering_for_model`], provider-aware, with |
| 7 | //! bundled/live/override layers) first, then the seeded |
| 8 | //! [`crate::model_registry`] facts for ids no catalog row covers (custom |
| 9 | //! providers, local models). No second model catalog is introduced here. |
| 10 | //! |
| 11 | //! Honesty rules: unknown facts are omitted rather than guessed, an explicit |
| 12 | //! catalog "unsupported" renders as `no <badge>`, provenance is always named, |
| 13 | //! and a completely unknown model resolves to `None` so callers can skip the |
| 14 | //! line entirely instead of blocking selection or fabricating capabilities. |
| 15 | |
| 16 | use codewhale_config::catalog::CatalogSource; |
| 17 | use codewhale_config::route::{CapabilityState, RouteCapabilities, RouteLimits}; |
| 18 | |
| 19 | use crate::config::ApiProvider; |
| 20 | use crate::model_registry::{self, ModelMetadata}; |
| 21 | use crate::tui::model_picker::format_picker_context_window; |
| 22 | |
| 23 | /// Resolved capability badges for one Fleet route. |
| 24 | #[derive(Debug, Clone, PartialEq, Eq)] |
| 25 | pub struct RouteCapabilityBadges { |
| 26 | /// Ordered, concise badges, e.g. `["1M ctx", "384K out", "tools", |
| 27 | /// "reasoning", "no vision"]`. Never empty when resolution succeeds. |
| 28 | pub badges: Vec<String>, |
| 29 | /// Where the facts came from: `bundled catalog`, `live catalog`, |
| 30 | /// `override`, or `registry`. |
| 31 | pub provenance: &'static str, |
| 32 | } |
| 33 | |
| 34 | impl RouteCapabilityBadges { |
| 35 | /// One-line summary with provenance, sized for narrow detail panes: |
| 36 | /// `1M ctx · 384K out · tools · reasoning · no vision (bundled catalog)`. |
| 37 | #[must_use] |
| 38 | pub fn summary(&self) -> String { |
| 39 | format!("{} ({})", self.badges.join(" · "), self.provenance) |
| 40 | } |
| 41 | } |
| 42 | |
| 43 | /// Resolve capability badges for one `(provider, model)` Fleet route. |
| 44 | /// |
| 45 | /// `provider_id` is the exact configured route key when known (canonical |
| 46 | /// built-in id or named custom table key); pass `None` when the route has no |
| 47 | /// resolvable provider (e.g. a pinned model that inherits the provider). Only |
| 48 | /// exact built-in ids reach the provider-scoped catalog; every other id falls |
| 49 | /// back to provider-agnostic registry facts rather than guessing a route. |
| 50 | /// |
| 51 | /// The session `auto -> model` display form is accepted and resolved against |
| 52 | /// the effective model. Returns `None` when nothing is known about the model, |
| 53 | /// so absence renders as absence. |
| 54 | #[must_use] |
| 55 | pub fn resolve_route_capability_badges( |
| 56 | provider_id: Option<&str>, |
| 57 | model: &str, |
| 58 | ) -> Option<RouteCapabilityBadges> { |
| 59 | let model = effective_model(model)?; |
| 60 | if let Some(provider) = provider_id.and_then(exact_builtin_provider) |
| 61 | && let Some(offering) = crate::provider_lake::catalog_offering_for_model(provider, model) |
| 62 | { |
| 63 | let route = offering.to_offering(); |
| 64 | let badges = badges_from_route(&route.limits, &route.capabilities); |
| 65 | if !badges.is_empty() { |
| 66 | return Some(RouteCapabilityBadges { |
| 67 | badges, |
| 68 | provenance: catalog_provenance(&offering.source), |
| 69 | }); |
| 70 | } |
| 71 | } |
| 72 | let meta = model_registry::lookup(model)?; |
| 73 | let badges = badges_from_registry(&meta); |
| 74 | (!badges.is_empty()).then_some(RouteCapabilityBadges { |
| 75 | badges, |
| 76 | provenance: "registry", |
| 77 | }) |
| 78 | } |
| 79 | |
| 80 | /// Strip the session `auto -> model` display form down to the effective model. |
| 81 | /// A bare `auto` (nothing resolved yet) or empty id yields `None`. |
| 82 | fn effective_model(model: &str) -> Option<&str> { |
| 83 | let model = model |
| 84 | .split_once("->") |
| 85 | .map_or(model, |(_, effective)| effective) |
| 86 | .trim(); |
| 87 | (!model.is_empty() && !model.eq_ignore_ascii_case("auto")).then_some(model) |
| 88 | } |
| 89 | |
| 90 | /// Accept only exact canonical built-in provider ids. Display labels and named |
| 91 | /// custom table keys must not inherit a built-in catalog by similarity. |
| 92 | fn exact_builtin_provider(provider_id: &str) -> Option<ApiProvider> { |
| 93 | ApiProvider::parse(provider_id).filter(|provider| provider.as_str() == provider_id) |
| 94 | } |
| 95 | |
| 96 | const fn catalog_provenance(source: &CatalogSource) -> &'static str { |
| 97 | match source { |
| 98 | CatalogSource::Bundled => "bundled catalog", |
| 99 | CatalogSource::Live { .. } => "live catalog", |
| 100 | CatalogSource::UserOverride => "override", |
| 101 | } |
| 102 | } |
| 103 | |
| 104 | /// Badges from exact provider-offering facts. Three-state facts keep their |
| 105 | /// explicit `Unsupported` (`no tools` / `no vision`); `Unknown` is omitted. |
| 106 | fn badges_from_route(limits: &RouteLimits, capabilities: &RouteCapabilities) -> Vec<String> { |
| 107 | let mut badges = Vec::new(); |
| 108 | if let Some(context) = limits.context_tokens { |
| 109 | badges.push(format!("{} ctx", format_picker_context_window(context))); |
| 110 | } |
| 111 | if let Some(output) = limits.output_tokens { |
| 112 | badges.push(format!("{} out", format_picker_context_window(output))); |
| 113 | } |
| 114 | push_state_badge(&mut badges, capabilities.native_tool_calls, "tools"); |
| 115 | push_state_badge(&mut badges, capabilities.reasoning, "reasoning"); |
| 116 | push_state_badge(&mut badges, capabilities.image_input, "vision"); |
| 117 | badges |
| 118 | } |
| 119 | |
| 120 | /// Badges from seeded registry facts. The registry has no tool/vision facts, |
| 121 | /// and its `supports_reasoning: false` is a heuristic default rather than a |
| 122 | /// sourced denial, so only a positive reasoning fact is shown. |
| 123 | fn badges_from_registry(meta: &ModelMetadata) -> Vec<String> { |
| 124 | let mut badges = Vec::new(); |
| 125 | if let Some(context) = meta.context_window { |
| 126 | badges.push(format!( |
| 127 | "{} ctx", |
| 128 | format_picker_context_window(u64::from(context)) |
| 129 | )); |
| 130 | } |
| 131 | if let Some(output) = meta.max_output { |
| 132 | badges.push(format!( |
| 133 | "{} out", |
| 134 | format_picker_context_window(u64::from(output)) |
| 135 | )); |
| 136 | } |
| 137 | if meta.supports_reasoning { |
| 138 | badges.push("reasoning".to_string()); |
| 139 | } |
| 140 | badges |
| 141 | } |
| 142 | |
| 143 | fn push_state_badge(badges: &mut Vec<String>, state: CapabilityState, name: &str) { |
| 144 | match state { |
| 145 | CapabilityState::Supported => badges.push(name.to_string()), |
| 146 | CapabilityState::Unsupported => badges.push(format!("no {name}")), |
| 147 | CapabilityState::Unknown => {} |
| 148 | } |
| 149 | } |
| 150 | |
| 151 | #[cfg(test)] |
| 152 | mod tests { |
| 153 | use super::*; |
| 154 | use codewhale_config::catalog::CatalogOffering; |
| 155 | use codewhale_config::models_dev::{ModelsDevLimit, ModelsDevModalities}; |
| 156 | |
| 157 | #[test] |
| 158 | fn known_catalog_model_resolves_provider_aware_badges() { |
| 159 | let badges = resolve_route_capability_badges(Some("deepseek"), "deepseek-v4-pro") |
| 160 | .expect("bundled catalog knows deepseek-v4-pro"); |
| 161 | assert!( |
| 162 | badges.badges.contains(&"1M ctx".to_string()), |
| 163 | "missing context badge: {:?}", |
| 164 | badges.badges |
| 165 | ); |
| 166 | assert!(badges.badges.contains(&"384K out".to_string())); |
| 167 | assert!(badges.badges.contains(&"tools".to_string())); |
| 168 | assert!(badges.badges.contains(&"reasoning".to_string())); |
| 169 | // Text-only modalities are an explicit sourced fact, not an unknown. |
| 170 | assert!(badges.badges.contains(&"no vision".to_string())); |
| 171 | assert!( |
| 172 | badges.provenance.contains("catalog"), |
| 173 | "catalog facts must carry catalog provenance, got {}", |
| 174 | badges.provenance |
| 175 | ); |
| 176 | assert!(badges.summary().contains(" · ")); |
| 177 | } |
| 178 | |
| 179 | #[test] |
| 180 | fn unknown_model_resolves_to_graceful_absence() { |
| 181 | assert_eq!( |
| 182 | resolve_route_capability_badges(Some("deepseek"), "totally-made-up-model-xyz"), |
| 183 | None |
| 184 | ); |
| 185 | assert_eq!(resolve_route_capability_badges(None, ""), None); |
| 186 | assert_eq!(resolve_route_capability_badges(None, "auto"), None); |
| 187 | } |
| 188 | |
| 189 | #[test] |
| 190 | fn registry_fallback_covers_routes_without_catalog_rows() { |
| 191 | // A named custom route key never inherits a built-in catalog; the |
| 192 | // provider-agnostic registry still answers for the model id. |
| 193 | let badges = resolve_route_capability_badges(Some("my-custom-endpoint"), "claude-fable-5") |
| 194 | .expect("registry knows claude-fable-5"); |
| 195 | assert_eq!(badges.provenance, "registry"); |
| 196 | assert!(badges.badges.contains(&"1M ctx".to_string())); |
| 197 | assert!(badges.badges.contains(&"reasoning".to_string())); |
| 198 | // Tool/vision facts are unknown here — omitted, never fabricated. |
| 199 | assert!( |
| 200 | !badges |
| 201 | .badges |
| 202 | .iter() |
| 203 | .any(|badge| badge.contains("tools") || badge.contains("vision")), |
| 204 | "unsourced facts must not appear: {:?}", |
| 205 | badges.badges |
| 206 | ); |
| 207 | } |
| 208 | |
| 209 | #[test] |
| 210 | fn auto_display_route_resolves_the_effective_model() { |
| 211 | let badges = resolve_route_capability_badges(None, "auto -> deepseek-v4-pro") |
| 212 | .expect("effective model resolves"); |
| 213 | assert!(badges.badges.contains(&"1M ctx".to_string())); |
| 214 | } |
| 215 | |
| 216 | #[test] |
| 217 | fn explicit_unsupported_catalog_facts_render_as_no_badges() { |
| 218 | let offering = CatalogOffering { |
| 219 | provider: "deepseek".to_string(), |
| 220 | wire_model_id: "fixture-model".to_string(), |
| 221 | endpoint_key: "chat".to_string(), |
| 222 | limit: Some(ModelsDevLimit { |
| 223 | context: Some(131_072), |
| 224 | input: None, |
| 225 | output: None, |
| 226 | }), |
| 227 | reasoning: Some(false), |
| 228 | tool_call: Some(false), |
| 229 | modalities: Some(ModelsDevModalities { |
| 230 | input: vec!["text".to_string(), "image".to_string()], |
| 231 | output: vec!["text".to_string()], |
| 232 | }), |
| 233 | ..CatalogOffering::default() |
| 234 | }; |
| 235 | let route = offering.to_offering(); |
| 236 | let badges = badges_from_route(&route.limits, &route.capabilities); |
| 237 | assert_eq!( |
| 238 | badges, |
| 239 | vec![ |
| 240 | "131K ctx".to_string(), |
| 241 | "no tools".to_string(), |
| 242 | "no reasoning".to_string(), |
| 243 | "vision".to_string(), |
| 244 | ] |
| 245 | ); |
| 246 | } |
| 247 | |
| 248 | #[test] |
| 249 | fn token_labels_match_picker_vocabulary() { |
| 250 | assert_eq!(format_picker_context_window(1_000_000), "1M"); |
| 251 | assert_eq!(format_picker_context_window(1_050_000), "1.05M"); |
| 252 | assert_eq!(format_picker_context_window(262_144), "262K"); |
| 253 | assert_eq!(format_picker_context_window(500), "500"); |
| 254 | } |
| 255 | } |
| 256 |