1pub static KNOWN_MODELS: &[(&str, u64, bool)] = &[
12 ("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 ("gpt-5.5", 1_050_000, true),
19 ("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 ("o4-mini", 200_000, true),
25 ("o3", 200_000, true),
26 ("o3-pro", 200_000, true),
27 ("gpt-4o", 128_000, false),
29 ("gpt-4o-mini", 128_000, false),
30 ("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 ("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-pro", 1_000_000, true),
41 ("deepseek-v4-flash", 1_000_000, false),
42 ("deepseek-chat", 1_000_000, false),
44 ("deepseek-reasoner", 1_000_000, true),
45 ("glm-5.2", 1_000_000, true),
47 ("glm-4-plus", 128_000, false),
48 ("glm-4-flash", 128_000, false),
49 ("qwen2.5-72b-instruct", 131_072, false),
51 ("qwen2.5-coder-32b-instruct", 131_072, false),
52 ("llama3.3-70b-instruct", 131_072, false),
54 ("llama3.1-8b-instruct", 131_072, false),
55];
56
57pub 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
78fn 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 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}