codewhale_agent/
lib.rs

1use std::collections::HashMap;
2
3use codewhale_config::ProviderKind;
4use serde::{Deserialize, Serialize};
5
6/// Metadata for a single model entry in the registry.
7///
8/// Each model has a canonical `id` used by the provider, a list of `aliases`
9/// that users may reference, and capability flags indicating whether the model
10/// supports tool use and reasoning.
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct ModelInfo {
13    /// The canonical model identifier used by the provider (e.g. `"deepseek-v4-pro"`).
14    pub id: String,
15    /// The provider that serves this model.
16    pub provider: ProviderKind,
17    /// Alternative names that users can use to reference this model (case-insensitive).
18    pub aliases: Vec<String>,
19    /// Whether this model supports tool/function calling.
20    pub supports_tools: bool,
21    /// Whether this model supports extended reasoning.
22    pub supports_reasoning: bool,
23}
24
25/// The result of resolving a user-requested model name to a concrete model entry.
26///
27/// Contains the resolved [`ModelInfo`], whether a fallback was used, and the
28/// chain of resolution strategies that were attempted.
29#[derive(Debug, Clone, Serialize, Deserialize)]
30pub struct ModelResolution {
31    /// The original model name requested by the user, if any.
32    pub requested: Option<String>,
33    /// The concrete model that was resolved.
34    pub resolved: ModelInfo,
35    /// Whether a fallback was used because the requested model was not found.
36    pub used_fallback: bool,
37    /// The ordered list of resolution strategies that were attempted.
38    pub fallback_chain: Vec<String>,
39}
40
41/// A registry of supported models and their aliases, used to resolve user-facing
42/// model names to concrete provider-specific model entries.
43///
44/// The default registry is populated with all built-in models across supported
45/// providers (DeepSeek, NVIDIA NIM, OpenAI-compatible, and others).
46#[derive(Debug, Clone)]
47pub struct ModelRegistry {
48    models: Vec<ModelInfo>,
49    alias_map: HashMap<String, usize>,
50}
51
52/// Creates a registry pre-populated with all built-in models and their aliases.
53impl Default for ModelRegistry {
54    fn default() -> Self {
55        let models = vec![
56            ModelInfo {
57                id: "deepseek-v4-pro".to_string(),
58                provider: ProviderKind::Deepseek,
59                aliases: vec![],
60                supports_tools: true,
61                supports_reasoning: true,
62            },
63            ModelInfo {
64                id: "deepseek-v4-flash".to_string(),
65                provider: ProviderKind::Deepseek,
66                aliases: vec![
67                    "deepseek-chat".to_string(),
68                    "deepseek-reasoner".to_string(),
69                    "deepseek-r1".to_string(),
70                    "deepseek-v3".to_string(),
71                    "deepseek-v3.2".to_string(),
72                ],
73                supports_tools: true,
74                supports_reasoning: true,
75            },
76            ModelInfo {
77                id: "deepseek-ai/deepseek-v4-pro".to_string(),
78                provider: ProviderKind::NvidiaNim,
79                aliases: vec![
80                    "deepseek-v4-pro".to_string(),
81                    "nvidia-deepseek-v4-pro".to_string(),
82                    "nim-deepseek-v4-pro".to_string(),
83                ],
84                supports_tools: true,
85                supports_reasoning: true,
86            },
87            ModelInfo {
88                id: "deepseek-ai/deepseek-v4-flash".to_string(),
89                provider: ProviderKind::NvidiaNim,
90                aliases: vec![
91                    "deepseek-v4-flash".to_string(),
92                    "deepseek-chat".to_string(),
93                    "deepseek-reasoner".to_string(),
94                    "nvidia-deepseek-v4-flash".to_string(),
95                    "nim-deepseek-v4-flash".to_string(),
96                ],
97                supports_tools: true,
98                supports_reasoning: true,
99            },
100            ModelInfo {
101                id: "deepseek-v4-pro".to_string(),
102                provider: ProviderKind::Openai,
103                aliases: vec!["openai-compatible-deepseek-v4-pro".to_string()],
104                supports_tools: true,
105                supports_reasoning: true,
106            },
107            ModelInfo {
108                id: "deepseek-v4-flash".to_string(),
109                provider: ProviderKind::Openai,
110                aliases: vec!["openai-compatible-deepseek-v4-flash".to_string()],
111                supports_tools: true,
112                supports_reasoning: true,
113            },
114            ModelInfo {
115                id: "deepseek-ai/deepseek-v4-flash".to_string(),
116                provider: ProviderKind::Atlascloud,
117                aliases: vec![
118                    "deepseek-v4-flash".to_string(),
119                    "atlascloud-deepseek-v4-flash".to_string(),
120                ],
121                supports_tools: true,
122                supports_reasoning: true,
123            },
124            ModelInfo {
125                id: "deepseek-ai/deepseek-v4-pro".to_string(),
126                provider: ProviderKind::Atlascloud,
127                aliases: vec![
128                    "deepseek-v4-pro".to_string(),
129                    "atlascloud-deepseek-v4-pro".to_string(),
130                ],
131                supports_tools: true,
132                supports_reasoning: true,
133            },
134            ModelInfo {
135                id: "deepseek-reasoner".to_string(),
136                provider: ProviderKind::WanjieArk,
137                aliases: vec![
138                    "wanjie-deepseek-reasoner".to_string(),
139                    "ark-wanjie-deepseek-reasoner".to_string(),
140                ],
141                supports_tools: true,
142                supports_reasoning: true,
143            },
144            ModelInfo {
145                id: "DeepSeek-V4-Pro".to_string(),
146                provider: ProviderKind::Volcengine,
147                aliases: vec![
148                    "deepseek-v4-pro".to_string(),
149                    "volcengine-deepseek-v4-pro".to_string(),
150                    "ark-deepseek-v4-pro".to_string(),
151                ],
152                supports_tools: true,
153                supports_reasoning: true,
154            },
155            ModelInfo {
156                id: "DeepSeek-V4-Flash".to_string(),
157                provider: ProviderKind::Volcengine,
158                aliases: vec![
159                    "deepseek-v4-flash".to_string(),
160                    "deepseek-chat".to_string(),
161                    "volcengine-deepseek-v4-flash".to_string(),
162                    "ark-deepseek-v4-flash".to_string(),
163                ],
164                supports_tools: true,
165                supports_reasoning: true,
166            },
167            ModelInfo {
168                id: "trinity-large-thinking".to_string(),
169                provider: ProviderKind::Arcee,
170                aliases: vec![
171                    "trinity".to_string(),
172                    "arcee-trinity".to_string(),
173                    "arcee-trinity-large-thinking".to_string(),
174                ],
175                supports_tools: true,
176                supports_reasoning: true,
177            },
178            ModelInfo {
179                id: "deepseek/deepseek-v4-pro".to_string(),
180                provider: ProviderKind::Openrouter,
181                aliases: vec![
182                    "deepseek-v4-pro".to_string(),
183                    "openrouter-deepseek-v4-pro".to_string(),
184                ],
185                supports_tools: true,
186                supports_reasoning: true,
187            },
188            ModelInfo {
189                id: "deepseek/deepseek-v4-flash".to_string(),
190                provider: ProviderKind::Openrouter,
191                aliases: vec![
192                    "deepseek-v4-flash".to_string(),
193                    "deepseek-chat".to_string(),
194                    "deepseek-reasoner".to_string(),
195                    "openrouter-deepseek-v4-flash".to_string(),
196                ],
197                supports_tools: true,
198                supports_reasoning: true,
199            },
200            ModelInfo {
201                id: "arcee-ai/trinity-large-thinking".to_string(),
202                provider: ProviderKind::Openrouter,
203                aliases: vec![
204                    "trinity".to_string(),
205                    "trinity-large-thinking".to_string(),
206                    "arcee-trinity-large-thinking".to_string(),
207                ],
208                supports_tools: true,
209                supports_reasoning: true,
210            },
211            ModelInfo {
212                id: "xiaomi/mimo-v2.5-pro".to_string(),
213                provider: ProviderKind::Openrouter,
214                aliases: vec![
215                    "openrouter-mimo-v2.5-pro".to_string(),
216                    "openrouter-xiaomi-mimo-v2.5-pro".to_string(),
217                ],
218                supports_tools: true,
219                supports_reasoning: true,
220            },
221            ModelInfo {
222                id: "xiaomi/mimo-v2.5".to_string(),
223                provider: ProviderKind::Openrouter,
224                aliases: vec![
225                    "openrouter-mimo-v2.5".to_string(),
226                    "openrouter-xiaomi-mimo-v2.5".to_string(),
227                ],
228                supports_tools: true,
229                supports_reasoning: true,
230            },
231            ModelInfo {
232                id: "qwen/qwen3.6-flash".to_string(),
233                provider: ProviderKind::Openrouter,
234                aliases: vec!["qwen3.6-flash".to_string(), "qwen-3.6-flash".to_string()],
235                supports_tools: true,
236                supports_reasoning: true,
237            },
238            ModelInfo {
239                id: "qwen/qwen3.6-35b-a3b".to_string(),
240                provider: ProviderKind::Openrouter,
241                aliases: vec![
242                    "qwen3.6-35b-a3b".to_string(),
243                    "qwen-3.6-35b-a3b".to_string(),
244                ],
245                supports_tools: true,
246                supports_reasoning: true,
247            },
248            ModelInfo {
249                id: "qwen/qwen3.6-max-preview".to_string(),
250                provider: ProviderKind::Openrouter,
251                aliases: vec![
252                    "qwen3.6-max-preview".to_string(),
253                    "qwen-3.6-max-preview".to_string(),
254                    "qwen-max-preview".to_string(),
255                ],
256                supports_tools: true,
257                supports_reasoning: true,
258            },
259            ModelInfo {
260                id: "qwen/qwen3.6-27b".to_string(),
261                provider: ProviderKind::Openrouter,
262                aliases: vec!["qwen3.6-27b".to_string(), "qwen-3.6-27b".to_string()],
263                supports_tools: true,
264                supports_reasoning: true,
265            },
266            ModelInfo {
267                id: "qwen/qwen3.6-plus".to_string(),
268                provider: ProviderKind::Openrouter,
269                aliases: vec!["qwen3.6-plus".to_string(), "qwen-3.6-plus".to_string()],
270                supports_tools: true,
271                supports_reasoning: true,
272            },
273            ModelInfo {
274                id: "moonshotai/kimi-k2.6".to_string(),
275                provider: ProviderKind::Openrouter,
276                aliases: vec!["openrouter-kimi-k2.6".to_string()],
277                supports_tools: true,
278                supports_reasoning: true,
279            },
280            ModelInfo {
281                id: "minimax/minimax-m3".to_string(),
282                provider: ProviderKind::Openrouter,
283                aliases: vec![
284                    "minimax-m3".to_string(),
285                    "minimax-m-3".to_string(),
286                    "openrouter-minimax-m3".to_string(),
287                ],
288                supports_tools: true,
289                supports_reasoning: true,
290            },
291            ModelInfo {
292                id: "z-ai/glm-5.1".to_string(),
293                provider: ProviderKind::Openrouter,
294                aliases: vec!["glm-5.1".to_string(), "zai-glm-5.1".to_string()],
295                supports_tools: true,
296                supports_reasoning: true,
297            },
298            ModelInfo {
299                id: "tencent/hy3-preview".to_string(),
300                provider: ProviderKind::Openrouter,
301                aliases: vec!["hy3-preview".to_string(), "tencent-hy3-preview".to_string()],
302                supports_tools: true,
303                supports_reasoning: true,
304            },
305            ModelInfo {
306                id: "google/gemma-4-31b-it".to_string(),
307                provider: ProviderKind::Openrouter,
308                aliases: vec!["gemma-4-31b".to_string(), "gemma-4-31b-it".to_string()],
309                supports_tools: true,
310                supports_reasoning: true,
311            },
312            ModelInfo {
313                id: "google/gemma-4-26b-a4b-it".to_string(),
314                provider: ProviderKind::Openrouter,
315                aliases: vec![
316                    "gemma-4-26b-a4b".to_string(),
317                    "gemma-4-26b-a4b-it".to_string(),
318                ],
319                supports_tools: true,
320                supports_reasoning: true,
321            },
322            ModelInfo {
323                id: "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:free".to_string(),
324                provider: ProviderKind::Openrouter,
325                aliases: vec![
326                    "nemotron-3-nano-omni".to_string(),
327                    "nemotron-3-nano-omni-reasoning".to_string(),
328                ],
329                supports_tools: true,
330                supports_reasoning: true,
331            },
332            ModelInfo {
333                id: "mimo-v2.5-pro".to_string(),
334                provider: ProviderKind::XiaomiMimo,
335                aliases: vec![
336                    "mimo".to_string(),
337                    "pro".to_string(),
338                    "xiaomi-mimo-v2.5-pro".to_string(),
339                    "xiaomi-mimo-v2-5-pro".to_string(),
340                ],
341                supports_tools: true,
342                supports_reasoning: true,
343            },
344            ModelInfo {
345                id: "mimo-v2.5".to_string(),
346                provider: ProviderKind::XiaomiMimo,
347                aliases: vec![
348                    "omni".to_string(),
349                    "mimo-omni".to_string(),
350                    "v2.5-omni".to_string(),
351                    "mimo-v2.5-omni".to_string(),
352                    "xiaomi-mimo-v2.5".to_string(),
353                    "xiaomi-mimo-v2.5-omni".to_string(),
354                ],
355                supports_tools: true,
356                supports_reasoning: true,
357            },
358            ModelInfo {
359                id: "mimo-v2.5-asr".to_string(),
360                provider: ProviderKind::XiaomiMimo,
361                aliases: vec![
362                    "asr".to_string(),
363                    "speech-to-text".to_string(),
364                    "transcribe".to_string(),
365                ],
366                supports_tools: false,
367                supports_reasoning: false,
368            },
369            ModelInfo {
370                id: "mimo-v2.5-tts".to_string(),
371                provider: ProviderKind::XiaomiMimo,
372                aliases: vec![
373                    "tts".to_string(),
374                    "speech".to_string(),
375                    "mimo-tts".to_string(),
376                ],
377                supports_tools: false,
378                supports_reasoning: false,
379            },
380            ModelInfo {
381                id: "mimo-v2.5-tts-voicedesign".to_string(),
382                provider: ProviderKind::XiaomiMimo,
383                aliases: vec![
384                    "voicedesign".to_string(),
385                    "voice-design".to_string(),
386                    "mimo-voice-design".to_string(),
387                ],
388                supports_tools: false,
389                supports_reasoning: false,
390            },
391            ModelInfo {
392                id: "mimo-v2.5-tts-voiceclone".to_string(),
393                provider: ProviderKind::XiaomiMimo,
394                aliases: vec![
395                    "voiceclone".to_string(),
396                    "voice-clone".to_string(),
397                    "mimo-voice-clone".to_string(),
398                ],
399                supports_tools: false,
400                supports_reasoning: false,
401            },
402            ModelInfo {
403                id: "mimo-v2-tts".to_string(),
404                provider: ProviderKind::XiaomiMimo,
405                aliases: vec!["mimo-v2-speech".to_string()],
406                supports_tools: false,
407                supports_reasoning: false,
408            },
409            ModelInfo {
410                id: "deepseek/deepseek-v4-pro".to_string(),
411                provider: ProviderKind::Novita,
412                aliases: vec![
413                    "deepseek-v4-pro".to_string(),
414                    "novita-deepseek-v4-pro".to_string(),
415                ],
416                supports_tools: true,
417                supports_reasoning: true,
418            },
419            ModelInfo {
420                id: "deepseek/deepseek-v4-flash".to_string(),
421                provider: ProviderKind::Novita,
422                aliases: vec![
423                    "deepseek-v4-flash".to_string(),
424                    "deepseek-chat".to_string(),
425                    "deepseek-reasoner".to_string(),
426                    "novita-deepseek-v4-flash".to_string(),
427                ],
428                supports_tools: true,
429                supports_reasoning: true,
430            },
431            ModelInfo {
432                id: "accounts/fireworks/models/deepseek-v4-pro".to_string(),
433                provider: ProviderKind::Fireworks,
434                aliases: vec![
435                    "deepseek-v4-pro".to_string(),
436                    "fireworks-deepseek-v4-pro".to_string(),
437                ],
438                supports_tools: true,
439                supports_reasoning: true,
440            },
441            ModelInfo {
442                id: "deepseek-ai/DeepSeek-V4-Pro".to_string(),
443                provider: ProviderKind::Siliconflow,
444                aliases: vec![
445                    "deepseek-v4-pro".to_string(),
446                    "deepseek-reasoner".to_string(),
447                    "deepseek-r1".to_string(),
448                    "siliconflow-deepseek-v4-pro".to_string(),
449                ],
450                supports_tools: true,
451                supports_reasoning: true,
452            },
453            ModelInfo {
454                id: "deepseek-ai/DeepSeek-V4-Flash".to_string(),
455                provider: ProviderKind::Siliconflow,
456                aliases: vec![
457                    "deepseek-v4-flash".to_string(),
458                    "deepseek-chat".to_string(),
459                    "deepseek-v3".to_string(),
460                    "siliconflow-deepseek-v4-flash".to_string(),
461                ],
462                supports_tools: true,
463                supports_reasoning: true,
464            },
465            ModelInfo {
466                id: "trinity-large-preview".to_string(),
467                provider: ProviderKind::Arcee,
468                aliases: vec!["arcee-trinity-large-preview".to_string()],
469                supports_tools: true,
470                supports_reasoning: false,
471            },
472            ModelInfo {
473                id: "kimi-k2.6".to_string(),
474                provider: ProviderKind::Moonshot,
475                aliases: vec![
476                    "kimi".to_string(),
477                    "kimi-k2".to_string(),
478                    "moonshot-kimi-k2.6".to_string(),
479                ],
480                supports_tools: true,
481                supports_reasoning: true,
482            },
483            ModelInfo {
484                id: "deepseek-ai/DeepSeek-V4-Pro".to_string(),
485                provider: ProviderKind::Sglang,
486                aliases: vec![
487                    "deepseek-v4-pro".to_string(),
488                    "sglang-deepseek-v4-pro".to_string(),
489                ],
490                supports_tools: true,
491                supports_reasoning: true,
492            },
493            ModelInfo {
494                id: "deepseek-ai/DeepSeek-V4-Flash".to_string(),
495                provider: ProviderKind::Sglang,
496                aliases: vec![
497                    "deepseek-v4-flash".to_string(),
498                    "deepseek-chat".to_string(),
499                    "deepseek-reasoner".to_string(),
500                    "sglang-deepseek-v4-flash".to_string(),
501                ],
502                supports_tools: true,
503                supports_reasoning: true,
504            },
505            ModelInfo {
506                id: "deepseek-ai/DeepSeek-V4-Pro".to_string(),
507                provider: ProviderKind::Vllm,
508                aliases: vec![
509                    "deepseek-v4-pro".to_string(),
510                    "vllm-deepseek-v4-pro".to_string(),
511                ],
512                supports_tools: true,
513                supports_reasoning: true,
514            },
515            ModelInfo {
516                id: "deepseek-ai/DeepSeek-V4-Flash".to_string(),
517                provider: ProviderKind::Vllm,
518                aliases: vec![
519                    "deepseek-v4-flash".to_string(),
520                    "deepseek-chat".to_string(),
521                    "deepseek-reasoner".to_string(),
522                    "vllm-deepseek-v4-flash".to_string(),
523                ],
524                supports_tools: true,
525                supports_reasoning: true,
526            },
527            ModelInfo {
528                id: "deepseek-coder:1.3b".to_string(),
529                provider: ProviderKind::Ollama,
530                aliases: vec![],
531                supports_tools: true,
532                supports_reasoning: false,
533            },
534        ];
535        Self::new(models)
536    }
537}
538
539impl ModelRegistry {
540    /// Creates a new registry from a list of [`ModelInfo`] entries.
541    ///
542    /// Builds an internal alias map for fast lookup by model id or alias.
543    /// If multiple models share the same id or alias, the first one registered
544    /// takes priority.
545    #[must_use]
546    pub fn new(models: Vec<ModelInfo>) -> Self {
547        let mut alias_map = HashMap::new();
548        for (idx, model) in models.iter().enumerate() {
549            alias_map.entry(normalize(&model.id)).or_insert(idx);
550            for alias in &model.aliases {
551                alias_map.entry(normalize(alias)).or_insert(idx);
552            }
553        }
554        Self { models, alias_map }
555    }
556
557    /// Returns a clone of all models in the registry.
558    #[must_use]
559    pub fn list(&self) -> Vec<ModelInfo> {
560        self.models.clone()
561    }
562
563    /// Resolves a user-requested model name to a concrete [`ModelInfo`].
564    ///
565    /// Resolution follows this priority order:
566    /// 1. If the provider is Ollama, the requested name is used as-is (to
567    ///    support arbitrary local model tags like `qwen2.5-coder:7b`).
568    /// 2. If a `provider_hint` is given, search for a model matching that
569    ///    provider whose id or alias matches the request (case-insensitive).
570    /// 3. Look up the alias map for a case-insensitive match.
571    /// 4. Fall back to the first model belonging to the hinted provider
572    ///    (or DeepSeek if no hint was given).
573    /// 5. As a last resort, fall back to the first model in the registry.
574    #[must_use]
575    pub fn resolve(
576        &self,
577        requested: Option<&str>,
578        provider_hint: Option<ProviderKind>,
579    ) -> ModelResolution {
580        let mut fallback_chain = Vec::new();
581
582        if let Some(name) = requested {
583            fallback_chain.push(format!("requested:{name}"));
584            if provider_hint == Some(ProviderKind::Ollama) {
585                return ModelResolution {
586                    requested: Some(name.to_string()),
587                    resolved: ModelInfo {
588                        id: name.trim().to_string(),
589                        provider: ProviderKind::Ollama,
590                        aliases: Vec::new(),
591                        supports_tools: true,
592                        supports_reasoning: false,
593                    },
594                    used_fallback: false,
595                    fallback_chain,
596                };
597            }
598            if let Some(provider) = provider_hint
599                && let Some(model) = self
600                    .models
601                    .iter()
602                    .find(|m| m.provider == provider && model_matches(m, name))
603                    .cloned()
604            {
605                return ModelResolution {
606                    requested: Some(name.to_string()),
607                    resolved: model,
608                    used_fallback: false,
609                    fallback_chain,
610                };
611            }
612            if provider_hint == Some(ProviderKind::Atlascloud)
613                && let Some(model) = atlascloud_passthrough_model(name)
614            {
615                return ModelResolution {
616                    requested: Some(name.to_string()),
617                    resolved: model,
618                    used_fallback: false,
619                    fallback_chain,
620                };
621            }
622            if provider_hint == Some(ProviderKind::Arcee)
623                && let Some(model) = arcee_passthrough_model(name)
624            {
625                return ModelResolution {
626                    requested: Some(name.to_string()),
627                    resolved: model,
628                    used_fallback: false,
629                    fallback_chain,
630                };
631            }
632            if provider_hint == Some(ProviderKind::XiaomiMimo)
633                && let Some(model) = xiaomi_mimo_passthrough_model(name)
634            {
635                return ModelResolution {
636                    requested: Some(name.to_string()),
637                    resolved: model,
638                    used_fallback: false,
639                    fallback_chain,
640                };
641            }
642            if let Some(idx) = self.alias_map.get(&normalize(name)) {
643                return ModelResolution {
644                    requested: Some(name.to_string()),
645                    resolved: preserve_requested_model_id_case(self.models[*idx].clone(), name),
646                    used_fallback: false,
647                    fallback_chain,
648                };
649            }
650        }
651
652        let provider = provider_hint.unwrap_or(ProviderKind::Deepseek);
653        fallback_chain.push(format!("provider_default:{}", provider.as_str()));
654        if let Some(model) = self.models.iter().find(|m| m.provider == provider).cloned() {
655            return ModelResolution {
656                requested: requested.map(ToOwned::to_owned),
657                resolved: model,
658                used_fallback: true,
659                fallback_chain,
660            };
661        }
662
663        let final_fallback = self.models.first().cloned().unwrap_or(ModelInfo {
664            id: "deepseek-v4-pro".to_string(),
665            provider: ProviderKind::Deepseek,
666            aliases: Vec::new(),
667            supports_tools: true,
668            supports_reasoning: true,
669        });
670        fallback_chain.push("global_default:deepseek-v4-pro".to_string());
671        ModelResolution {
672            requested: requested.map(ToOwned::to_owned),
673            resolved: final_fallback,
674            used_fallback: true,
675            fallback_chain,
676        }
677    }
678}
679
680fn normalize(value: &str) -> String {
681    value.trim().to_ascii_lowercase()
682}
683
684fn model_matches(model: &ModelInfo, requested: &str) -> bool {
685    let requested = normalize(requested);
686    normalize(&model.id) == requested
687        || model
688            .aliases
689            .iter()
690            .any(|alias| normalize(alias) == requested)
691}
692
693fn preserve_requested_model_id_case(mut model: ModelInfo, requested: &str) -> ModelInfo {
694    let requested = requested.trim();
695    if model.id.eq_ignore_ascii_case(requested) {
696        model.id = requested.to_string();
697    }
698    model
699}
700
701fn atlascloud_passthrough_model(requested: &str) -> Option<ModelInfo> {
702    let requested = requested.trim();
703    if requested.is_empty() || !requested.contains('/') {
704        return None;
705    }
706
707    Some(ModelInfo {
708        id: requested.to_string(),
709        provider: ProviderKind::Atlascloud,
710        aliases: Vec::new(),
711        supports_tools: true,
712        supports_reasoning: true,
713    })
714}
715
716fn arcee_passthrough_model(requested: &str) -> Option<ModelInfo> {
717    let requested = requested.trim();
718    if requested.is_empty() {
719        return None;
720    }
721    let supports_reasoning = requested.to_ascii_lowercase().contains("thinking");
722
723    Some(ModelInfo {
724        id: requested.to_string(),
725        provider: ProviderKind::Arcee,
726        aliases: Vec::new(),
727        supports_tools: true,
728        supports_reasoning,
729    })
730}
731
732fn xiaomi_mimo_passthrough_model(requested: &str) -> Option<ModelInfo> {
733    let requested = requested.trim();
734    if requested.is_empty() || requested.chars().any(char::is_control) {
735        return None;
736    }
737
738    Some(ModelInfo {
739        id: requested.to_string(),
740        provider: ProviderKind::XiaomiMimo,
741        aliases: Vec::new(),
742        supports_tools: true,
743        supports_reasoning: true,
744    })
745}
746
747#[cfg(test)]
748mod tests {
749    use super::*;
750
751    #[test]
752    fn deepseek_v4_pro_alias_stays_deepseek_by_default() {
753        let registry = ModelRegistry::default();
754        let resolved = registry.resolve(Some("deepseek-v4-pro"), None);
755
756        assert_eq!(resolved.resolved.provider, ProviderKind::Deepseek);
757        assert_eq!(resolved.resolved.id, "deepseek-v4-pro");
758    }
759
760    #[test]
761    fn deepseek_v4_pro_alias_resolves_to_nvidia_nim_when_provider_hinted() {
762        let registry = ModelRegistry::default();
763        let resolved = registry.resolve(Some("deepseek-v4-pro"), Some(ProviderKind::NvidiaNim));
764
765        assert_eq!(resolved.resolved.provider, ProviderKind::NvidiaNim);
766        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-pro");
767    }
768
769    #[test]
770    fn nvidia_nim_default_uses_catalog_model_id() {
771        let registry = ModelRegistry::default();
772        let resolved = registry.resolve(None, Some(ProviderKind::NvidiaNim));
773
774        assert_eq!(resolved.resolved.provider, ProviderKind::NvidiaNim);
775        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-pro");
776    }
777
778    #[test]
779    fn deepseek_v4_flash_alias_resolves_to_nvidia_nim_when_provider_hinted() {
780        let registry = ModelRegistry::default();
781        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::NvidiaNim));
782
783        assert_eq!(resolved.resolved.provider, ProviderKind::NvidiaNim);
784        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-flash");
785    }
786
787    #[test]
788    fn atlascloud_default_uses_namespaced_model_id() {
789        let registry = ModelRegistry::default();
790        let resolved = registry.resolve(None, Some(ProviderKind::Atlascloud));
791
792        assert_eq!(resolved.resolved.provider, ProviderKind::Atlascloud);
793        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-flash");
794        assert!(resolved.resolved.supports_reasoning);
795    }
796
797    #[test]
798    fn deepseek_v4_flash_alias_resolves_to_atlascloud_when_provider_hinted() {
799        let registry = ModelRegistry::default();
800        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Atlascloud));
801
802        assert_eq!(resolved.resolved.provider, ProviderKind::Atlascloud);
803        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-flash");
804    }
805
806    #[test]
807    fn deepseek_v4_pro_alias_resolves_to_atlascloud_when_provider_hinted() {
808        let registry = ModelRegistry::default();
809        let resolved = registry.resolve(Some("deepseek-v4-pro"), Some(ProviderKind::Atlascloud));
810
811        assert_eq!(resolved.resolved.provider, ProviderKind::Atlascloud);
812        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-pro");
813    }
814
815    #[test]
816    fn atlascloud_provider_hint_passes_through_explicit_model_id() {
817        let registry = ModelRegistry::default();
818        let resolved =
819            registry.resolve(Some("openai/gpt-5.2-chat"), Some(ProviderKind::Atlascloud));
820
821        assert_eq!(resolved.resolved.provider, ProviderKind::Atlascloud);
822        assert_eq!(resolved.resolved.id, "openai/gpt-5.2-chat");
823        assert!(resolved.resolved.supports_tools);
824        assert!(resolved.resolved.supports_reasoning);
825        assert!(!resolved.used_fallback);
826    }
827
828    #[test]
829    fn atlascloud_provider_hint_preserves_explicit_model_id_case() {
830        let registry = ModelRegistry::default();
831        let resolved = registry.resolve(Some("Qwen/Qwen3-Coder"), Some(ProviderKind::Atlascloud));
832
833        assert_eq!(resolved.resolved.provider, ProviderKind::Atlascloud);
834        assert_eq!(resolved.resolved.id, "Qwen/Qwen3-Coder");
835        assert!(!resolved.used_fallback);
836    }
837
838    #[test]
839    fn atlascloud_plain_unknown_model_still_uses_provider_default() {
840        let registry = ModelRegistry::default();
841        let resolved = registry.resolve(Some("not-in-atlas"), Some(ProviderKind::Atlascloud));
842
843        assert_eq!(resolved.resolved.provider, ProviderKind::Atlascloud);
844        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-flash");
845        assert!(resolved.used_fallback);
846    }
847
848    #[test]
849    fn openrouter_default_uses_namespaced_model_id() {
850        let registry = ModelRegistry::default();
851        let resolved = registry.resolve(None, Some(ProviderKind::Openrouter));
852
853        assert_eq!(resolved.resolved.provider, ProviderKind::Openrouter);
854        assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-pro");
855    }
856
857    #[test]
858    fn xiaomi_mimo_default_uses_canonical_model_id() {
859        let registry = ModelRegistry::default();
860        let resolved = registry.resolve(None, Some(ProviderKind::XiaomiMimo));
861
862        assert_eq!(resolved.resolved.provider, ProviderKind::XiaomiMimo);
863        assert_eq!(resolved.resolved.id, "mimo-v2.5-pro");
864        assert!(resolved.resolved.supports_reasoning);
865    }
866
867    #[test]
868    fn xiaomi_mimo_tts_aliases_resolve_when_provider_hinted() {
869        let registry = ModelRegistry::default();
870        let resolved = registry.resolve(Some("tts"), Some(ProviderKind::XiaomiMimo));
871        assert_eq!(resolved.resolved.provider, ProviderKind::XiaomiMimo);
872        assert_eq!(resolved.resolved.id, "mimo-v2.5-tts");
873        assert!(!resolved.resolved.supports_tools);
874        assert!(!resolved.resolved.supports_reasoning);
875
876        let resolved = registry.resolve(Some("voice-design"), Some(ProviderKind::XiaomiMimo));
877        assert_eq!(resolved.resolved.id, "mimo-v2.5-tts-voicedesign");
878
879        let resolved = registry.resolve(Some("voiceclone"), Some(ProviderKind::XiaomiMimo));
880        assert_eq!(resolved.resolved.id, "mimo-v2.5-tts-voiceclone");
881    }
882
883    #[test]
884    fn xiaomi_mimo_chat_aliases_resolve_when_provider_hinted() {
885        let registry = ModelRegistry::default();
886
887        let resolved = registry.resolve(Some("omni"), Some(ProviderKind::XiaomiMimo));
888        assert_eq!(resolved.resolved.provider, ProviderKind::XiaomiMimo);
889        assert_eq!(resolved.resolved.id, "mimo-v2.5");
890        assert!(resolved.resolved.supports_tools);
891    }
892
893    #[test]
894    fn xiaomi_mimo_provider_hint_preserves_custom_model_id() {
895        let registry = ModelRegistry::default();
896        let resolved =
897            registry.resolve(Some("account-custom-mimo"), Some(ProviderKind::XiaomiMimo));
898
899        assert_eq!(resolved.resolved.provider, ProviderKind::XiaomiMimo);
900        assert_eq!(resolved.resolved.id, "account-custom-mimo");
901        assert!(!resolved.used_fallback);
902    }
903
904    #[test]
905    fn xiaomi_mimo_provider_hint_does_not_reclassify_openrouter_model_id() {
906        let registry = ModelRegistry::default();
907        let resolved = registry.resolve(
908            Some("deepseek/deepseek-v4-pro"),
909            Some(ProviderKind::XiaomiMimo),
910        );
911
912        assert_eq!(resolved.resolved.provider, ProviderKind::XiaomiMimo);
913        assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-pro");
914        assert!(!resolved.used_fallback);
915    }
916
917    #[test]
918    fn wanjie_ark_default_uses_reasoner_model_id() {
919        let registry = ModelRegistry::default();
920        let resolved = registry.resolve(None, Some(ProviderKind::WanjieArk));
921
922        assert_eq!(resolved.resolved.provider, ProviderKind::WanjieArk);
923        assert_eq!(resolved.resolved.id, "deepseek-reasoner");
924        assert!(resolved.resolved.supports_reasoning);
925    }
926
927    #[test]
928    fn novita_default_uses_namespaced_model_id() {
929        let registry = ModelRegistry::default();
930        let resolved = registry.resolve(None, Some(ProviderKind::Novita));
931
932        assert_eq!(resolved.resolved.provider, ProviderKind::Novita);
933        assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-pro");
934    }
935
936    #[test]
937    fn fireworks_default_uses_canonical_model_id() {
938        let registry = ModelRegistry::default();
939        let resolved = registry.resolve(None, Some(ProviderKind::Fireworks));
940
941        assert_eq!(resolved.resolved.provider, ProviderKind::Fireworks);
942        assert_eq!(
943            resolved.resolved.id,
944            "accounts/fireworks/models/deepseek-v4-pro"
945        );
946    }
947
948    #[test]
949    fn siliconflow_default_uses_canonical_pro_model_id() {
950        let registry = ModelRegistry::default();
951        let resolved = registry.resolve(None, Some(ProviderKind::Siliconflow));
952
953        assert_eq!(resolved.resolved.provider, ProviderKind::Siliconflow);
954        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Pro");
955        assert!(resolved.resolved.supports_reasoning);
956    }
957
958    #[test]
959    fn arcee_default_uses_direct_trinity_large_thinking_model_id() {
960        let registry = ModelRegistry::default();
961        let resolved = registry.resolve(None, Some(ProviderKind::Arcee));
962
963        assert_eq!(resolved.resolved.provider, ProviderKind::Arcee);
964        assert_eq!(resolved.resolved.id, "trinity-large-thinking");
965        assert!(resolved.resolved.supports_reasoning);
966    }
967
968    #[test]
969    fn arcee_trinity_alias_resolves_to_direct_large_thinking_not_openrouter() {
970        let registry = ModelRegistry::default();
971        let resolved = registry.resolve(Some("trinity"), Some(ProviderKind::Arcee));
972
973        assert_eq!(resolved.resolved.provider, ProviderKind::Arcee);
974        assert_eq!(resolved.resolved.id, "trinity-large-thinking");
975        assert!(resolved.resolved.supports_reasoning);
976    }
977
978    #[test]
979    fn arcee_trinity_mini_remains_explicit_compatibility_model() {
980        let registry = ModelRegistry::default();
981        let resolved = registry.resolve(Some("trinity-mini"), Some(ProviderKind::Arcee));
982
983        assert_eq!(resolved.resolved.provider, ProviderKind::Arcee);
984        assert_eq!(resolved.resolved.id, "trinity-mini");
985        assert!(!resolved.resolved.supports_reasoning);
986    }
987
988    #[test]
989    fn arcee_provider_hint_preserves_explicit_future_model_id() {
990        let registry = ModelRegistry::default();
991        let resolved = registry.resolve(Some("trinity-large-next"), Some(ProviderKind::Arcee));
992
993        assert_eq!(resolved.resolved.provider, ProviderKind::Arcee);
994        assert_eq!(resolved.resolved.id, "trinity-large-next");
995        assert!(!resolved.resolved.supports_reasoning);
996        assert!(!resolved.used_fallback);
997    }
998
999    #[test]
1000    fn deepseek_reasoner_alias_resolves_to_siliconflow_pro_when_provider_hinted() {
1001        let registry = ModelRegistry::default();
1002        let resolved = registry.resolve(Some("deepseek-reasoner"), Some(ProviderKind::Siliconflow));
1003
1004        assert_eq!(resolved.resolved.provider, ProviderKind::Siliconflow);
1005        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Pro");
1006    }
1007
1008    #[test]
1009    fn deepseek_v4_flash_alias_resolves_to_siliconflow_flash_when_provider_hinted() {
1010        let registry = ModelRegistry::default();
1011        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Siliconflow));
1012
1013        assert_eq!(resolved.resolved.provider, ProviderKind::Siliconflow);
1014        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Flash");
1015    }
1016
1017    #[test]
1018    fn sglang_default_uses_canonical_model_id() {
1019        let registry = ModelRegistry::default();
1020        let resolved = registry.resolve(None, Some(ProviderKind::Sglang));
1021
1022        assert_eq!(resolved.resolved.provider, ProviderKind::Sglang);
1023        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Pro");
1024    }
1025
1026    #[test]
1027    fn deepseek_v4_flash_alias_resolves_to_openrouter_when_provider_hinted() {
1028        let registry = ModelRegistry::default();
1029        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Openrouter));
1030
1031        assert_eq!(resolved.resolved.provider, ProviderKind::Openrouter);
1032        assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-flash");
1033    }
1034
1035    #[test]
1036    fn recent_openrouter_large_model_aliases_resolve_when_provider_hinted() {
1037        let registry = ModelRegistry::default();
1038
1039        for (alias, expected) in [
1040            ("trinity-large-thinking", "arcee-ai/trinity-large-thinking"),
1041            ("qwen3.6-flash", "qwen/qwen3.6-flash"),
1042            ("qwen3.6-35b-a3b", "qwen/qwen3.6-35b-a3b"),
1043            ("qwen3.6-max-preview", "qwen/qwen3.6-max-preview"),
1044            ("qwen3.6-plus", "qwen/qwen3.6-plus"),
1045            ("gemma-4-31b-it", "google/gemma-4-31b-it"),
1046            ("glm-5.1", "z-ai/glm-5.1"),
1047            ("minimax-m3", "minimax/minimax-m3"),
1048            ("openrouter-mimo-v2.5-pro", "xiaomi/mimo-v2.5-pro"),
1049            ("openrouter-kimi-k2.6", "moonshotai/kimi-k2.6"),
1050        ] {
1051            let resolved = registry.resolve(Some(alias), Some(ProviderKind::Openrouter));
1052
1053            assert_eq!(resolved.resolved.provider, ProviderKind::Openrouter);
1054            assert_eq!(resolved.resolved.id, expected);
1055            assert!(resolved.resolved.supports_tools);
1056            assert!(resolved.resolved.supports_reasoning);
1057        }
1058    }
1059
1060    #[test]
1061    fn deepseek_v4_flash_alias_resolves_to_novita_when_provider_hinted() {
1062        let registry = ModelRegistry::default();
1063        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Novita));
1064
1065        assert_eq!(resolved.resolved.provider, ProviderKind::Novita);
1066        assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-flash");
1067    }
1068
1069    #[test]
1070    fn deepseek_v4_flash_alias_resolves_to_sglang_when_provider_hinted() {
1071        let registry = ModelRegistry::default();
1072        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Sglang));
1073
1074        assert_eq!(resolved.resolved.provider, ProviderKind::Sglang);
1075        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Flash");
1076    }
1077
1078    #[test]
1079    fn vllm_default_uses_canonical_model_id() {
1080        let registry = ModelRegistry::default();
1081        let resolved = registry.resolve(None, Some(ProviderKind::Vllm));
1082
1083        assert_eq!(resolved.resolved.provider, ProviderKind::Vllm);
1084        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Pro");
1085    }
1086
1087    #[test]
1088    fn ollama_default_uses_small_local_model_id() {
1089        let registry = ModelRegistry::default();
1090        let resolved = registry.resolve(None, Some(ProviderKind::Ollama));
1091
1092        assert_eq!(resolved.resolved.provider, ProviderKind::Ollama);
1093        assert_eq!(resolved.resolved.id, "deepseek-coder:1.3b");
1094        assert!(!resolved.resolved.supports_reasoning);
1095    }
1096
1097    #[test]
1098    fn ollama_requested_model_tag_is_preserved() {
1099        let registry = ModelRegistry::default();
1100        let resolved = registry.resolve(Some("qwen2.5-coder:7b"), Some(ProviderKind::Ollama));
1101
1102        assert_eq!(resolved.resolved.provider, ProviderKind::Ollama);
1103        assert_eq!(resolved.resolved.id, "qwen2.5-coder:7b");
1104        assert!(!resolved.used_fallback);
1105    }
1106
1107    #[test]
1108    fn deepseek_v4_flash_alias_resolves_to_vllm_when_provider_hinted() {
1109        let registry = ModelRegistry::default();
1110        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Vllm));
1111
1112        assert_eq!(resolved.resolved.provider, ProviderKind::Vllm);
1113        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Flash");
1114    }
1115
1116    #[test]
1117    fn preserves_requested_model_casing_for_third_party_providers() {
1118        let registry = ModelRegistry::default();
1119        let resolved = registry.resolve(Some("DeepSeek-V4-Pro"), None);
1120
1121        assert_eq!(resolved.resolved.provider, ProviderKind::Deepseek);
1122        assert_eq!(resolved.resolved.id, "DeepSeek-V4-Pro");
1123    }
1124
1125    #[test]
1126    fn registry_casing_takes_priority_over_requested_casing_with_provider_hint() {
1127        let registry = ModelRegistry::default();
1128        let resolved = registry.resolve(Some("DeepSeek-V4-Pro"), Some(ProviderKind::Deepseek));
1129
1130        assert_eq!(resolved.resolved.provider, ProviderKind::Deepseek);
1131        // Registry's canonical id is used even when user provides different casing
1132        assert_eq!(resolved.resolved.id, "deepseek-v4-pro");
1133    }
1134
1135    #[test]
1136    fn preserves_requested_model_casing_without_surrounding_whitespace() {
1137        let registry = ModelRegistry::default();
1138        let resolved = registry.resolve(Some("  DeepSeek-V4-Pro  "), None);
1139
1140        assert_eq!(resolved.resolved.provider, ProviderKind::Deepseek);
1141        assert_eq!(resolved.resolved.id, "DeepSeek-V4-Pro");
1142    }
1143
1144    #[test]
1145    fn alias_match_does_not_override_requested_casing() {
1146        let registry = ModelRegistry::default();
1147        let resolved = registry.resolve(Some("deepseek-reasoner"), None);
1148
1149        assert_eq!(resolved.resolved.provider, ProviderKind::Deepseek);
1150        assert_eq!(resolved.resolved.id, "deepseek-v4-flash");
1151    }
1152}
codewhale_agent/lib.rs

codewhale_agent/
lib.rs
