1use std::collections::HashMap;
2
3use deepseek_config::ProviderKind;
4use serde::{Deserialize, Serialize};
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct ModelInfo {
8 pub id: String,
9 pub provider: ProviderKind,
10 pub aliases: Vec<String>,
11 pub supports_tools: bool,
12 pub supports_reasoning: bool,
13}
14
15#[derive(Debug, Clone, Serialize, Deserialize)]
16pub struct ModelResolution {
17 pub requested: Option<String>,
18 pub resolved: ModelInfo,
19 pub used_fallback: bool,
20 pub fallback_chain: Vec<String>,
21}
22
23#[derive(Debug, Clone)]
24pub struct ModelRegistry {
25 models: Vec<ModelInfo>,
26 alias_map: HashMap<String, usize>,
27}
28
29impl Default for ModelRegistry {
30 fn default() -> Self {
31 let models = vec![
32 ModelInfo {
33 id: "deepseek-v4-pro".to_string(),
34 provider: ProviderKind::Deepseek,
35 aliases: vec![],
36 supports_tools: true,
37 supports_reasoning: true,
38 },
39 ModelInfo {
40 id: "deepseek-v4-flash".to_string(),
41 provider: ProviderKind::Deepseek,
42 aliases: vec![
43 "deepseek-chat".to_string(),
44 "deepseek-reasoner".to_string(),
45 "deepseek-r1".to_string(),
46 "deepseek-v3".to_string(),
47 "deepseek-v3.2".to_string(),
48 ],
49 supports_tools: true,
50 supports_reasoning: true,
51 },
52 ModelInfo {
53 id: "deepseek-ai/deepseek-v4-pro".to_string(),
54 provider: ProviderKind::NvidiaNim,
55 aliases: vec![
56 "deepseek-v4-pro".to_string(),
57 "nvidia-deepseek-v4-pro".to_string(),
58 "nim-deepseek-v4-pro".to_string(),
59 ],
60 supports_tools: true,
61 supports_reasoning: true,
62 },
63 ModelInfo {
64 id: "deepseek-ai/deepseek-v4-flash".to_string(),
65 provider: ProviderKind::NvidiaNim,
66 aliases: vec![
67 "deepseek-v4-flash".to_string(),
68 "deepseek-chat".to_string(),
69 "deepseek-reasoner".to_string(),
70 "nvidia-deepseek-v4-flash".to_string(),
71 "nim-deepseek-v4-flash".to_string(),
72 ],
73 supports_tools: true,
74 supports_reasoning: true,
75 },
76 ModelInfo {
77 id: "gpt-4.1".to_string(),
78 provider: ProviderKind::Openai,
79 aliases: vec!["gpt4.1".to_string(), "gpt-4o".to_string()],
80 supports_tools: true,
81 supports_reasoning: true,
82 },
83 ModelInfo {
84 id: "gpt-4.1-mini".to_string(),
85 provider: ProviderKind::Openai,
86 aliases: vec!["gpt-4o-mini".to_string()],
87 supports_tools: true,
88 supports_reasoning: false,
89 },
90 ModelInfo {
91 id: "deepseek/deepseek-v4-pro".to_string(),
92 provider: ProviderKind::Openrouter,
93 aliases: vec![
94 "deepseek-v4-pro".to_string(),
95 "openrouter-deepseek-v4-pro".to_string(),
96 ],
97 supports_tools: true,
98 supports_reasoning: true,
99 },
100 ModelInfo {
101 id: "deepseek/deepseek-v4-flash".to_string(),
102 provider: ProviderKind::Openrouter,
103 aliases: vec![
104 "deepseek-v4-flash".to_string(),
105 "deepseek-chat".to_string(),
106 "deepseek-reasoner".to_string(),
107 "openrouter-deepseek-v4-flash".to_string(),
108 ],
109 supports_tools: true,
110 supports_reasoning: true,
111 },
112 ModelInfo {
113 id: "deepseek/deepseek-v4-pro".to_string(),
114 provider: ProviderKind::Novita,
115 aliases: vec![
116 "deepseek-v4-pro".to_string(),
117 "novita-deepseek-v4-pro".to_string(),
118 ],
119 supports_tools: true,
120 supports_reasoning: true,
121 },
122 ModelInfo {
123 id: "deepseek/deepseek-v4-flash".to_string(),
124 provider: ProviderKind::Novita,
125 aliases: vec![
126 "deepseek-v4-flash".to_string(),
127 "deepseek-chat".to_string(),
128 "deepseek-reasoner".to_string(),
129 "novita-deepseek-v4-flash".to_string(),
130 ],
131 supports_tools: true,
132 supports_reasoning: true,
133 },
134 ModelInfo {
135 id: "accounts/fireworks/models/deepseek-v4-pro".to_string(),
136 provider: ProviderKind::Fireworks,
137 aliases: vec![
138 "deepseek-v4-pro".to_string(),
139 "fireworks-deepseek-v4-pro".to_string(),
140 ],
141 supports_tools: true,
142 supports_reasoning: true,
143 },
144 ModelInfo {
145 id: "deepseek-ai/DeepSeek-V4-Pro".to_string(),
146 provider: ProviderKind::Sglang,
147 aliases: vec![
148 "deepseek-v4-pro".to_string(),
149 "sglang-deepseek-v4-pro".to_string(),
150 ],
151 supports_tools: true,
152 supports_reasoning: true,
153 },
154 ModelInfo {
155 id: "deepseek-ai/DeepSeek-V4-Flash".to_string(),
156 provider: ProviderKind::Sglang,
157 aliases: vec![
158 "deepseek-v4-flash".to_string(),
159 "deepseek-chat".to_string(),
160 "deepseek-reasoner".to_string(),
161 "sglang-deepseek-v4-flash".to_string(),
162 ],
163 supports_tools: true,
164 supports_reasoning: true,
165 },
166 ];
167 Self::new(models)
168 }
169}
170
171impl ModelRegistry {
172 #[must_use]
173 pub fn new(models: Vec<ModelInfo>) -> Self {
174 let mut alias_map = HashMap::new();
175 for (idx, model) in models.iter().enumerate() {
176 alias_map.entry(normalize(&model.id)).or_insert(idx);
177 for alias in &model.aliases {
178 alias_map.entry(normalize(alias)).or_insert(idx);
179 }
180 }
181 Self { models, alias_map }
182 }
183
184 #[must_use]
185 pub fn list(&self) -> Vec<ModelInfo> {
186 self.models.clone()
187 }
188
189 #[must_use]
190 pub fn resolve(
191 &self,
192 requested: Option<&str>,
193 provider_hint: Option<ProviderKind>,
194 ) -> ModelResolution {
195 let mut fallback_chain = Vec::new();
196
197 if let Some(name) = requested {
198 fallback_chain.push(format!("requested:{name}"));
199 if let Some(provider) = provider_hint
200 && let Some(model) = self
201 .models
202 .iter()
203 .find(|m| m.provider == provider && model_matches(m, name))
204 .cloned()
205 {
206 return ModelResolution {
207 requested: Some(name.to_string()),
208 resolved: model,
209 used_fallback: false,
210 fallback_chain,
211 };
212 }
213 if let Some(idx) = self.alias_map.get(&normalize(name)) {
214 return ModelResolution {
215 requested: Some(name.to_string()),
216 resolved: self.models[*idx].clone(),
217 used_fallback: false,
218 fallback_chain,
219 };
220 }
221 }
222
223 let provider = provider_hint.unwrap_or(ProviderKind::Deepseek);
224 fallback_chain.push(format!("provider_default:{}", provider.as_str()));
225 if let Some(model) = self.models.iter().find(|m| m.provider == provider).cloned() {
226 return ModelResolution {
227 requested: requested.map(ToOwned::to_owned),
228 resolved: model,
229 used_fallback: true,
230 fallback_chain,
231 };
232 }
233
234 let final_fallback = self.models.first().cloned().unwrap_or(ModelInfo {
235 id: "deepseek-v4-pro".to_string(),
236 provider: ProviderKind::Deepseek,
237 aliases: Vec::new(),
238 supports_tools: true,
239 supports_reasoning: true,
240 });
241 fallback_chain.push("global_default:deepseek-v4-pro".to_string());
242 ModelResolution {
243 requested: requested.map(ToOwned::to_owned),
244 resolved: final_fallback,
245 used_fallback: true,
246 fallback_chain,
247 }
248 }
249}
250
251fn normalize(value: &str) -> String {
252 value.trim().to_ascii_lowercase()
253}
254
255fn model_matches(model: &ModelInfo, requested: &str) -> bool {
256 let requested = normalize(requested);
257 normalize(&model.id) == requested
258 || model
259 .aliases
260 .iter()
261 .any(|alias| normalize(alias) == requested)
262}
263
264#[cfg(test)]
265mod tests {
266 use super::*;
267
268 #[test]
269 fn deepseek_v4_pro_alias_stays_deepseek_by_default() {
270 let registry = ModelRegistry::default();
271 let resolved = registry.resolve(Some("deepseek-v4-pro"), None);
272
273 assert_eq!(resolved.resolved.provider, ProviderKind::Deepseek);
274 assert_eq!(resolved.resolved.id, "deepseek-v4-pro");
275 }
276
277 #[test]
278 fn deepseek_v4_pro_alias_resolves_to_nvidia_nim_when_provider_hinted() {
279 let registry = ModelRegistry::default();
280 let resolved = registry.resolve(Some("deepseek-v4-pro"), Some(ProviderKind::NvidiaNim));
281
282 assert_eq!(resolved.resolved.provider, ProviderKind::NvidiaNim);
283 assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-pro");
284 }
285
286 #[test]
287 fn nvidia_nim_default_uses_catalog_model_id() {
288 let registry = ModelRegistry::default();
289 let resolved = registry.resolve(None, Some(ProviderKind::NvidiaNim));
290
291 assert_eq!(resolved.resolved.provider, ProviderKind::NvidiaNim);
292 assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-pro");
293 }
294
295 #[test]
296 fn deepseek_v4_flash_alias_resolves_to_nvidia_nim_when_provider_hinted() {
297 let registry = ModelRegistry::default();
298 let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::NvidiaNim));
299
300 assert_eq!(resolved.resolved.provider, ProviderKind::NvidiaNim);
301 assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-flash");
302 }
303
304 #[test]
305 fn openrouter_default_uses_namespaced_model_id() {
306 let registry = ModelRegistry::default();
307 let resolved = registry.resolve(None, Some(ProviderKind::Openrouter));
308
309 assert_eq!(resolved.resolved.provider, ProviderKind::Openrouter);
310 assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-pro");
311 }
312
313 #[test]
314 fn novita_default_uses_namespaced_model_id() {
315 let registry = ModelRegistry::default();
316 let resolved = registry.resolve(None, Some(ProviderKind::Novita));
317
318 assert_eq!(resolved.resolved.provider, ProviderKind::Novita);
319 assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-pro");
320 }
321
322 #[test]
323 fn fireworks_default_uses_canonical_model_id() {
324 let registry = ModelRegistry::default();
325 let resolved = registry.resolve(None, Some(ProviderKind::Fireworks));
326
327 assert_eq!(resolved.resolved.provider, ProviderKind::Fireworks);
328 assert_eq!(
329 resolved.resolved.id,
330 "accounts/fireworks/models/deepseek-v4-pro"
331 );
332 }
333
334 #[test]
335 fn sglang_default_uses_canonical_model_id() {
336 let registry = ModelRegistry::default();
337 let resolved = registry.resolve(None, Some(ProviderKind::Sglang));
338
339 assert_eq!(resolved.resolved.provider, ProviderKind::Sglang);
340 assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Pro");
341 }
342
343 #[test]
344 fn deepseek_v4_flash_alias_resolves_to_openrouter_when_provider_hinted() {
345 let registry = ModelRegistry::default();
346 let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Openrouter));
347
348 assert_eq!(resolved.resolved.provider, ProviderKind::Openrouter);
349 assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-flash");
350 }
351
352 #[test]
353 fn deepseek_v4_flash_alias_resolves_to_novita_when_provider_hinted() {
354 let registry = ModelRegistry::default();
355 let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Novita));
356
357 assert_eq!(resolved.resolved.provider, ProviderKind::Novita);
358 assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-flash");
359 }
360
361 #[test]
362 fn deepseek_v4_flash_alias_resolves_to_sglang_when_provider_hinted() {
363 let registry = ModelRegistry::default();
364 let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Sglang));
365
366 assert_eq!(resolved.resolved.provider, ProviderKind::Sglang);
367 assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Flash");
368 }
369}