Skip to main content

atman_runtime/
known_models.rs

1//! Known model metadata table.
2//!
3//! Used to fill `context_budget` and `thinking` for models discovered via
4//! API (`GET /v1/models` returns IDs only, not metadata).
5//!
6//! Update this table when new models are released. Entries are matched by
7//! exact match first, then longest-prefix match (so `gpt-4` doesn't shadow
8//! `gpt-4o`).
9
10/// (model_id_prefix, context_budget, thinking_enabled)
11pub static KNOWN_MODELS: &[(&str, u64, bool)] = &[
12    // OpenAI GPT-5.6 (current flagship, 2026-08)
13    ("gpt-5.6-sol", 1_050_000, true),
14    ("gpt-5.6-terra", 1_050_000, false),
15    ("gpt-5.6-luna", 1_050_000, false),
16    ("gpt-5.6", 1_050_000, true),
17    // OpenAI GPT-5.5 (superseded)
18    ("gpt-5.5", 1_050_000, true),
19    // OpenAI GPT-5.4 (superseded)
20    ("gpt-5.4", 1_000_000, false),
21    ("gpt-5.4-mini", 1_000_000, false),
22    ("gpt-5.4-nano", 1_000_000, false),
23    // OpenAI o-series (reasoning)
24    ("o4-mini", 200_000, true),
25    ("o3", 200_000, true),
26    ("o3-pro", 200_000, true),
27    // OpenAI legacy
28    ("gpt-4o", 128_000, false),
29    ("gpt-4o-mini", 128_000, false),
30    // Anthropic Claude (current, 2026-08)
31    ("claude-opus-5", 1_000_000, true),
32    ("claude-sonnet-5", 1_000_000, true),
33    ("claude-fable-5", 1_000_000, false),
34    ("claude-haiku-4-5", 200_000, false),
35    // Anthropic legacy
36    ("claude-opus-4-8", 1_000_000, true),
37    ("claude-opus-4-7", 1_000_000, true),
38    ("claude-sonnet-4-6", 1_000_000, true),
39    // DeepSeek V4 (current, 2026-08)
40    ("deepseek-v4-pro", 1_000_000, true),
41    ("deepseek-v4-flash", 1_000_000, false),
42    // DeepSeek legacy (deprecated Jul 24 2026, aliased to v4-flash)
43    ("deepseek-chat", 1_000_000, false),
44    ("deepseek-reasoner", 1_000_000, true),
45    // ZhipuAI GLM
46    ("glm-5.2", 1_000_000, true),
47    ("glm-4-plus", 128_000, false),
48    ("glm-4-flash", 128_000, false),
49    // Qwen (Ollama / Alibaba)
50    ("qwen2.5-72b-instruct", 131_072, false),
51    ("qwen2.5-coder-32b-instruct", 131_072, false),
52    // Llama (Ollama / Meta)
53    ("llama3.3-70b-instruct", 131_072, false),
54    ("llama3.1-8b-instruct", 131_072, false),
55];
56
57/// Look up a model ID in [`KNOWN_MODELS`].
58///
59/// Strips a `-YYYY-MM-DD` date suffix, tries exact match, then falls back to
60/// longest-prefix match (so `gpt-4` doesn't shadow `gpt-4o`).
61///
62/// Returns `(context_budget, thinking_enabled)` on match, or `None` if the
63/// model ID is not in the table.
64pub fn lookup_known_model(model_id: &str) -> Option<(u64, bool)> {
65    let stripped = strip_date_suffix(model_id);
66    if let Some((_, budget, thinking)) = KNOWN_MODELS.iter().find(|(k, _, _)| *k == stripped) {
67        return Some((*budget, *thinking));
68    }
69    let mut best: Option<(&str, u64, bool)> = None;
70    for (key, budget, thinking) in KNOWN_MODELS {
71        if stripped.starts_with(key) && best.is_none_or(|(k, _, _)| key.len() > k.len()) {
72            best = Some((key, *budget, *thinking));
73        }
74    }
75    best.map(|(_, b, t)| (b, t))
76}
77
78/// Strip a `-YYYY-MM-DD` date suffix from a model ID.
79fn strip_date_suffix(s: &str) -> &str {
80    if s.len() >= 11 {
81        let suffix = &s[s.len() - 11..];
82        if let Some(rest) = suffix.strip_prefix('-') {
83            let parts: Vec<&str> = rest.split('-').collect();
84            if parts.len() == 3 && parts.iter().all(|p| p.chars().all(|c| c.is_ascii_digit())) {
85                return &s[..s.len() - 11];
86            }
87        }
88    }
89    s
90}
91
92#[cfg(test)]
93mod tests {
94    use super::*;
95
96    #[test]
97    fn exact_match() {
98        assert_eq!(lookup_known_model("gpt-4o"), Some((128_000, false)));
99        assert_eq!(lookup_known_model("claude-opus-5"), Some((1_000_000, true)));
100    }
101
102    #[test]
103    fn date_suffix_stripped() {
104        assert_eq!(
105            lookup_known_model("gpt-4o-2024-08-06"),
106            Some((128_000, false))
107        );
108    }
109
110    #[test]
111    fn longest_prefix_wins() {
112        // gpt-4o-2024-08-06 should match gpt-4o (128K), not gpt-4 (would be 8K)
113        let result = lookup_known_model("gpt-4o-2024-08-06");
114        assert_eq!(result, Some((128_000, false)));
115    }
116
117    #[test]
118    fn unknown_model_returns_none() {
119        assert_eq!(lookup_known_model("mystery-model"), None);
120    }
121}