Skip to main content

mur_common/muragent/
model_class.rs

1//! Static model-class table: map a (provider, name) binding to a `ModelHint`
2//! (tier, RAM estimate, local-capability). Heuristic + conservative on
3//! unknown providers. See spec §7.1.
4
5use crate::model::provider_is_local;
6use crate::muragent::manifest::{ModelHint, ModelTier};
7
8const CLOUD_PROVIDERS: &[&str] = &[
9    "anthropic",
10    "openai",
11    "google",
12    "gemini",
13    "mistral",
14    "groq",
15    "cohere",
16    "deepseek",
17    "xai",
18    "openrouter",
19];
20const SMALL_MARKERS: &[&str] = &[
21    "1b", "1.5b", "2b", "3b", "mini", "nano", "haiku", "flash", "small", "tiny",
22];
23const LARGE_LOCAL_MARKERS: &[&str] = &["70b", "72b", "405b", "mixtral", "command-r-plus"];
24
25/// Map a model binding to a `ModelHint`. Local providers classify by size
26/// markers in the name; known cloud providers are frontier unless a "small"
27/// variant marker is present; unknown providers fall back to a conservative
28/// local-capable mid tier (the wizard still offers all options).
29pub fn classify(provider: &str, name: &str) -> ModelHint {
30    let p = provider.to_ascii_lowercase();
31    let n = name.to_ascii_lowercase();
32    let has = |markers: &[&str]| markers.iter().any(|m| n.contains(m));
33
34    if provider_is_local(&p) {
35        let (tier, min_ram_gb) = if has(SMALL_MARKERS) {
36            (ModelTier::Small, 8)
37        } else if has(LARGE_LOCAL_MARKERS) {
38            (ModelTier::Frontier, 64)
39        } else {
40            (ModelTier::Mid, 16)
41        };
42        return ModelHint {
43            provider: provider.to_string(),
44            name: name.to_string(),
45            tier,
46            min_ram_gb,
47            local_capable: true,
48        };
49    }
50
51    if CLOUD_PROVIDERS.contains(&p.as_str()) {
52        let tier = if has(SMALL_MARKERS) {
53            ModelTier::Mid
54        } else {
55            ModelTier::Frontier
56        };
57        return ModelHint {
58            provider: provider.to_string(),
59            name: name.to_string(),
60            tier,
61            min_ram_gb: 0,
62            local_capable: false,
63        };
64    }
65
66    // Unknown provider: conservative — assume a local-capable mid model.
67    ModelHint {
68        provider: provider.to_string(),
69        name: name.to_string(),
70        tier: ModelTier::Mid,
71        min_ram_gb: 16,
72        local_capable: true,
73    }
74}
75
76#[cfg(test)]
77mod tests {
78    use super::*;
79
80    #[test]
81    fn local_small_model() {
82        let h = classify("ollama", "llama3.2:3b");
83        assert_eq!(h.tier, ModelTier::Small);
84        assert!(h.local_capable);
85        assert_eq!(h.min_ram_gb, 8);
86    }
87
88    #[test]
89    fn cloud_frontier_model() {
90        let h = classify("anthropic", "claude-opus-4-7");
91        assert_eq!(h.tier, ModelTier::Frontier);
92        assert!(!h.local_capable);
93        assert_eq!(h.min_ram_gb, 0);
94    }
95
96    #[test]
97    fn cloud_small_variant_is_mid() {
98        let h = classify("openai", "gpt-4o-mini");
99        assert_eq!(h.tier, ModelTier::Mid);
100        assert!(!h.local_capable);
101    }
102
103    #[test]
104    fn unknown_provider_is_conservative_local_mid() {
105        let h = classify("acme-llm", "whatever-v2");
106        assert_eq!(h.tier, ModelTier::Mid);
107        assert!(h.local_capable);
108        assert_eq!(h.min_ram_gb, 16);
109    }
110}