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: "deepseek/deepseek-v4-pro".to_string(),
169                provider: ProviderKind::Openrouter,
170                aliases: vec![
171                    "deepseek-v4-pro".to_string(),
172                    "openrouter-deepseek-v4-pro".to_string(),
173                ],
174                supports_tools: true,
175                supports_reasoning: true,
176            },
177            ModelInfo {
178                id: "deepseek/deepseek-v4-flash".to_string(),
179                provider: ProviderKind::Openrouter,
180                aliases: vec![
181                    "deepseek-v4-flash".to_string(),
182                    "deepseek-chat".to_string(),
183                    "deepseek-reasoner".to_string(),
184                    "openrouter-deepseek-v4-flash".to_string(),
185                ],
186                supports_tools: true,
187                supports_reasoning: true,
188            },
189            ModelInfo {
190                id: "arcee-ai/trinity-large-thinking".to_string(),
191                provider: ProviderKind::Openrouter,
192                aliases: vec![
193                    "trinity".to_string(),
194                    "trinity-large-thinking".to_string(),
195                    "arcee-trinity-large-thinking".to_string(),
196                ],
197                supports_tools: true,
198                supports_reasoning: true,
199            },
200            ModelInfo {
201                id: "xiaomi/mimo-v2.5-pro".to_string(),
202                provider: ProviderKind::Openrouter,
203                aliases: vec![
204                    "openrouter-mimo-v2.5-pro".to_string(),
205                    "openrouter-xiaomi-mimo-v2.5-pro".to_string(),
206                ],
207                supports_tools: true,
208                supports_reasoning: true,
209            },
210            ModelInfo {
211                id: "xiaomi/mimo-v2.5".to_string(),
212                provider: ProviderKind::Openrouter,
213                aliases: vec![
214                    "openrouter-mimo-v2.5".to_string(),
215                    "openrouter-xiaomi-mimo-v2.5".to_string(),
216                ],
217                supports_tools: true,
218                supports_reasoning: true,
219            },
220            ModelInfo {
221                id: "qwen/qwen3.6-35b-a3b".to_string(),
222                provider: ProviderKind::Openrouter,
223                aliases: vec![
224                    "qwen3.6-35b-a3b".to_string(),
225                    "qwen-3.6-35b-a3b".to_string(),
226                ],
227                supports_tools: true,
228                supports_reasoning: true,
229            },
230            ModelInfo {
231                id: "qwen/qwen3.6-27b".to_string(),
232                provider: ProviderKind::Openrouter,
233                aliases: vec!["qwen3.6-27b".to_string(), "qwen-3.6-27b".to_string()],
234                supports_tools: true,
235                supports_reasoning: true,
236            },
237            ModelInfo {
238                id: "moonshotai/kimi-k2.6".to_string(),
239                provider: ProviderKind::Openrouter,
240                aliases: vec!["openrouter-kimi-k2.6".to_string()],
241                supports_tools: true,
242                supports_reasoning: true,
243            },
244            ModelInfo {
245                id: "minimax/minimax-m3".to_string(),
246                provider: ProviderKind::Openrouter,
247                aliases: vec![
248                    "minimax-m3".to_string(),
249                    "minimax-m-3".to_string(),
250                    "openrouter-minimax-m3".to_string(),
251                ],
252                supports_tools: true,
253                supports_reasoning: true,
254            },
255            ModelInfo {
256                id: "z-ai/glm-5.1".to_string(),
257                provider: ProviderKind::Openrouter,
258                aliases: vec!["glm-5.1".to_string(), "zai-glm-5.1".to_string()],
259                supports_tools: true,
260                supports_reasoning: true,
261            },
262            ModelInfo {
263                id: "tencent/hy3-preview".to_string(),
264                provider: ProviderKind::Openrouter,
265                aliases: vec!["hy3-preview".to_string(), "tencent-hy3-preview".to_string()],
266                supports_tools: true,
267                supports_reasoning: true,
268            },
269            ModelInfo {
270                id: "google/gemma-4-31b-it".to_string(),
271                provider: ProviderKind::Openrouter,
272                aliases: vec!["gemma-4-31b".to_string(), "gemma-4-31b-it".to_string()],
273                supports_tools: true,
274                supports_reasoning: true,
275            },
276            ModelInfo {
277                id: "google/gemma-4-26b-a4b-it".to_string(),
278                provider: ProviderKind::Openrouter,
279                aliases: vec![
280                    "gemma-4-26b-a4b".to_string(),
281                    "gemma-4-26b-a4b-it".to_string(),
282                ],
283                supports_tools: true,
284                supports_reasoning: true,
285            },
286            ModelInfo {
287                id: "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning:free".to_string(),
288                provider: ProviderKind::Openrouter,
289                aliases: vec![
290                    "nemotron-3-nano-omni".to_string(),
291                    "nemotron-3-nano-omni-reasoning".to_string(),
292                ],
293                supports_tools: true,
294                supports_reasoning: true,
295            },
296            ModelInfo {
297                id: "mimo-v2.5-pro".to_string(),
298                provider: ProviderKind::XiaomiMimo,
299                aliases: vec!["mimo".to_string()],
300                supports_tools: true,
301                supports_reasoning: true,
302            },
303            ModelInfo {
304                id: "mimo-v2.5".to_string(),
305                provider: ProviderKind::XiaomiMimo,
306                aliases: vec!["xiaomi-mimo-v2.5".to_string()],
307                supports_tools: true,
308                supports_reasoning: true,
309            },
310            ModelInfo {
311                id: "deepseek/deepseek-v4-pro".to_string(),
312                provider: ProviderKind::Novita,
313                aliases: vec![
314                    "deepseek-v4-pro".to_string(),
315                    "novita-deepseek-v4-pro".to_string(),
316                ],
317                supports_tools: true,
318                supports_reasoning: true,
319            },
320            ModelInfo {
321                id: "deepseek/deepseek-v4-flash".to_string(),
322                provider: ProviderKind::Novita,
323                aliases: vec![
324                    "deepseek-v4-flash".to_string(),
325                    "deepseek-chat".to_string(),
326                    "deepseek-reasoner".to_string(),
327                    "novita-deepseek-v4-flash".to_string(),
328                ],
329                supports_tools: true,
330                supports_reasoning: true,
331            },
332            ModelInfo {
333                id: "accounts/fireworks/models/deepseek-v4-pro".to_string(),
334                provider: ProviderKind::Fireworks,
335                aliases: vec![
336                    "deepseek-v4-pro".to_string(),
337                    "fireworks-deepseek-v4-pro".to_string(),
338                ],
339                supports_tools: true,
340                supports_reasoning: true,
341            },
342            ModelInfo {
343                id: "deepseek-ai/DeepSeek-V4-Pro".to_string(),
344                provider: ProviderKind::Siliconflow,
345                aliases: vec![
346                    "deepseek-v4-pro".to_string(),
347                    "deepseek-reasoner".to_string(),
348                    "deepseek-r1".to_string(),
349                    "siliconflow-deepseek-v4-pro".to_string(),
350                ],
351                supports_tools: true,
352                supports_reasoning: true,
353            },
354            ModelInfo {
355                id: "deepseek-ai/DeepSeek-V4-Flash".to_string(),
356                provider: ProviderKind::Siliconflow,
357                aliases: vec![
358                    "deepseek-v4-flash".to_string(),
359                    "deepseek-chat".to_string(),
360                    "deepseek-v3".to_string(),
361                    "siliconflow-deepseek-v4-flash".to_string(),
362                ],
363                supports_tools: true,
364                supports_reasoning: true,
365            },
366            ModelInfo {
367                id: "kimi-k2.6".to_string(),
368                provider: ProviderKind::Moonshot,
369                aliases: vec![
370                    "kimi".to_string(),
371                    "kimi-k2".to_string(),
372                    "moonshot-kimi-k2.6".to_string(),
373                ],
374                supports_tools: true,
375                supports_reasoning: true,
376            },
377            ModelInfo {
378                id: "deepseek-ai/DeepSeek-V4-Pro".to_string(),
379                provider: ProviderKind::Sglang,
380                aliases: vec![
381                    "deepseek-v4-pro".to_string(),
382                    "sglang-deepseek-v4-pro".to_string(),
383                ],
384                supports_tools: true,
385                supports_reasoning: true,
386            },
387            ModelInfo {
388                id: "deepseek-ai/DeepSeek-V4-Flash".to_string(),
389                provider: ProviderKind::Sglang,
390                aliases: vec![
391                    "deepseek-v4-flash".to_string(),
392                    "deepseek-chat".to_string(),
393                    "deepseek-reasoner".to_string(),
394                    "sglang-deepseek-v4-flash".to_string(),
395                ],
396                supports_tools: true,
397                supports_reasoning: true,
398            },
399            ModelInfo {
400                id: "deepseek-ai/DeepSeek-V4-Pro".to_string(),
401                provider: ProviderKind::Vllm,
402                aliases: vec![
403                    "deepseek-v4-pro".to_string(),
404                    "vllm-deepseek-v4-pro".to_string(),
405                ],
406                supports_tools: true,
407                supports_reasoning: true,
408            },
409            ModelInfo {
410                id: "deepseek-ai/DeepSeek-V4-Flash".to_string(),
411                provider: ProviderKind::Vllm,
412                aliases: vec![
413                    "deepseek-v4-flash".to_string(),
414                    "deepseek-chat".to_string(),
415                    "deepseek-reasoner".to_string(),
416                    "vllm-deepseek-v4-flash".to_string(),
417                ],
418                supports_tools: true,
419                supports_reasoning: true,
420            },
421            ModelInfo {
422                id: "deepseek-coder:1.3b".to_string(),
423                provider: ProviderKind::Ollama,
424                aliases: vec![],
425                supports_tools: true,
426                supports_reasoning: false,
427            },
428        ];
429        Self::new(models)
430    }
431}
432
433impl ModelRegistry {
434    /// Creates a new registry from a list of [`ModelInfo`] entries.
435    ///
436    /// Builds an internal alias map for fast lookup by model id or alias.
437    /// If multiple models share the same id or alias, the first one registered
438    /// takes priority.
439    #[must_use]
440    pub fn new(models: Vec<ModelInfo>) -> Self {
441        let mut alias_map = HashMap::new();
442        for (idx, model) in models.iter().enumerate() {
443            alias_map.entry(normalize(&model.id)).or_insert(idx);
444            for alias in &model.aliases {
445                alias_map.entry(normalize(alias)).or_insert(idx);
446            }
447        }
448        Self { models, alias_map }
449    }
450
451    /// Returns a clone of all models in the registry.
452    #[must_use]
453    pub fn list(&self) -> Vec<ModelInfo> {
454        self.models.clone()
455    }
456
457    /// Resolves a user-requested model name to a concrete [`ModelInfo`].
458    ///
459    /// Resolution follows this priority order:
460    /// 1. If the provider is Ollama, the requested name is used as-is (to
461    ///    support arbitrary local model tags like `qwen2.5-coder:7b`).
462    /// 2. If a `provider_hint` is given, search for a model matching that
463    ///    provider whose id or alias matches the request (case-insensitive).
464    /// 3. Look up the alias map for a case-insensitive match.
465    /// 4. Fall back to the first model belonging to the hinted provider
466    ///    (or DeepSeek if no hint was given).
467    /// 5. As a last resort, fall back to the first model in the registry.
468    #[must_use]
469    pub fn resolve(
470        &self,
471        requested: Option<&str>,
472        provider_hint: Option<ProviderKind>,
473    ) -> ModelResolution {
474        let mut fallback_chain = Vec::new();
475
476        if let Some(name) = requested {
477            fallback_chain.push(format!("requested:{name}"));
478            if provider_hint == Some(ProviderKind::Ollama) {
479                return ModelResolution {
480                    requested: Some(name.to_string()),
481                    resolved: ModelInfo {
482                        id: name.trim().to_string(),
483                        provider: ProviderKind::Ollama,
484                        aliases: Vec::new(),
485                        supports_tools: true,
486                        supports_reasoning: false,
487                    },
488                    used_fallback: false,
489                    fallback_chain,
490                };
491            }
492            if let Some(provider) = provider_hint
493                && let Some(model) = self
494                    .models
495                    .iter()
496                    .find(|m| m.provider == provider && model_matches(m, name))
497                    .cloned()
498            {
499                return ModelResolution {
500                    requested: Some(name.to_string()),
501                    resolved: model,
502                    used_fallback: false,
503                    fallback_chain,
504                };
505            }
506            if let Some(idx) = self.alias_map.get(&normalize(name)) {
507                return ModelResolution {
508                    requested: Some(name.to_string()),
509                    resolved: preserve_requested_model_id_case(self.models[*idx].clone(), name),
510                    used_fallback: false,
511                    fallback_chain,
512                };
513            }
514        }
515
516        let provider = provider_hint.unwrap_or(ProviderKind::Deepseek);
517        fallback_chain.push(format!("provider_default:{}", provider.as_str()));
518        if let Some(model) = self.models.iter().find(|m| m.provider == provider).cloned() {
519            return ModelResolution {
520                requested: requested.map(ToOwned::to_owned),
521                resolved: model,
522                used_fallback: true,
523                fallback_chain,
524            };
525        }
526
527        let final_fallback = self.models.first().cloned().unwrap_or(ModelInfo {
528            id: "deepseek-v4-pro".to_string(),
529            provider: ProviderKind::Deepseek,
530            aliases: Vec::new(),
531            supports_tools: true,
532            supports_reasoning: true,
533        });
534        fallback_chain.push("global_default:deepseek-v4-pro".to_string());
535        ModelResolution {
536            requested: requested.map(ToOwned::to_owned),
537            resolved: final_fallback,
538            used_fallback: true,
539            fallback_chain,
540        }
541    }
542}
543
544fn normalize(value: &str) -> String {
545    value.trim().to_ascii_lowercase()
546}
547
548fn model_matches(model: &ModelInfo, requested: &str) -> bool {
549    let requested = normalize(requested);
550    normalize(&model.id) == requested
551        || model
552            .aliases
553            .iter()
554            .any(|alias| normalize(alias) == requested)
555}
556
557fn preserve_requested_model_id_case(mut model: ModelInfo, requested: &str) -> ModelInfo {
558    let requested = requested.trim();
559    if model.id.eq_ignore_ascii_case(requested) {
560        model.id = requested.to_string();
561    }
562    model
563}
564
565#[cfg(test)]
566mod tests {
567    use super::*;
568
569    #[test]
570    fn deepseek_v4_pro_alias_stays_deepseek_by_default() {
571        let registry = ModelRegistry::default();
572        let resolved = registry.resolve(Some("deepseek-v4-pro"), None);
573
574        assert_eq!(resolved.resolved.provider, ProviderKind::Deepseek);
575        assert_eq!(resolved.resolved.id, "deepseek-v4-pro");
576    }
577
578    #[test]
579    fn deepseek_v4_pro_alias_resolves_to_nvidia_nim_when_provider_hinted() {
580        let registry = ModelRegistry::default();
581        let resolved = registry.resolve(Some("deepseek-v4-pro"), Some(ProviderKind::NvidiaNim));
582
583        assert_eq!(resolved.resolved.provider, ProviderKind::NvidiaNim);
584        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-pro");
585    }
586
587    #[test]
588    fn nvidia_nim_default_uses_catalog_model_id() {
589        let registry = ModelRegistry::default();
590        let resolved = registry.resolve(None, Some(ProviderKind::NvidiaNim));
591
592        assert_eq!(resolved.resolved.provider, ProviderKind::NvidiaNim);
593        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-pro");
594    }
595
596    #[test]
597    fn deepseek_v4_flash_alias_resolves_to_nvidia_nim_when_provider_hinted() {
598        let registry = ModelRegistry::default();
599        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::NvidiaNim));
600
601        assert_eq!(resolved.resolved.provider, ProviderKind::NvidiaNim);
602        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-flash");
603    }
604
605    #[test]
606    fn atlascloud_default_uses_namespaced_model_id() {
607        let registry = ModelRegistry::default();
608        let resolved = registry.resolve(None, Some(ProviderKind::Atlascloud));
609
610        assert_eq!(resolved.resolved.provider, ProviderKind::Atlascloud);
611        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-flash");
612        assert!(resolved.resolved.supports_reasoning);
613    }
614
615    #[test]
616    fn deepseek_v4_flash_alias_resolves_to_atlascloud_when_provider_hinted() {
617        let registry = ModelRegistry::default();
618        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Atlascloud));
619
620        assert_eq!(resolved.resolved.provider, ProviderKind::Atlascloud);
621        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-flash");
622    }
623
624    #[test]
625    fn deepseek_v4_pro_alias_resolves_to_atlascloud_when_provider_hinted() {
626        let registry = ModelRegistry::default();
627        let resolved = registry.resolve(Some("deepseek-v4-pro"), Some(ProviderKind::Atlascloud));
628
629        assert_eq!(resolved.resolved.provider, ProviderKind::Atlascloud);
630        assert_eq!(resolved.resolved.id, "deepseek-ai/deepseek-v4-pro");
631    }
632
633    #[test]
634    fn openrouter_default_uses_namespaced_model_id() {
635        let registry = ModelRegistry::default();
636        let resolved = registry.resolve(None, Some(ProviderKind::Openrouter));
637
638        assert_eq!(resolved.resolved.provider, ProviderKind::Openrouter);
639        assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-pro");
640    }
641
642    #[test]
643    fn xiaomi_mimo_default_uses_canonical_model_id() {
644        let registry = ModelRegistry::default();
645        let resolved = registry.resolve(None, Some(ProviderKind::XiaomiMimo));
646
647        assert_eq!(resolved.resolved.provider, ProviderKind::XiaomiMimo);
648        assert_eq!(resolved.resolved.id, "mimo-v2.5-pro");
649        assert!(resolved.resolved.supports_reasoning);
650    }
651
652    #[test]
653    fn wanjie_ark_default_uses_reasoner_model_id() {
654        let registry = ModelRegistry::default();
655        let resolved = registry.resolve(None, Some(ProviderKind::WanjieArk));
656
657        assert_eq!(resolved.resolved.provider, ProviderKind::WanjieArk);
658        assert_eq!(resolved.resolved.id, "deepseek-reasoner");
659        assert!(resolved.resolved.supports_reasoning);
660    }
661
662    #[test]
663    fn novita_default_uses_namespaced_model_id() {
664        let registry = ModelRegistry::default();
665        let resolved = registry.resolve(None, Some(ProviderKind::Novita));
666
667        assert_eq!(resolved.resolved.provider, ProviderKind::Novita);
668        assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-pro");
669    }
670
671    #[test]
672    fn fireworks_default_uses_canonical_model_id() {
673        let registry = ModelRegistry::default();
674        let resolved = registry.resolve(None, Some(ProviderKind::Fireworks));
675
676        assert_eq!(resolved.resolved.provider, ProviderKind::Fireworks);
677        assert_eq!(
678            resolved.resolved.id,
679            "accounts/fireworks/models/deepseek-v4-pro"
680        );
681    }
682
683    #[test]
684    fn siliconflow_default_uses_canonical_pro_model_id() {
685        let registry = ModelRegistry::default();
686        let resolved = registry.resolve(None, Some(ProviderKind::Siliconflow));
687
688        assert_eq!(resolved.resolved.provider, ProviderKind::Siliconflow);
689        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Pro");
690        assert!(resolved.resolved.supports_reasoning);
691    }
692
693    #[test]
694    fn deepseek_reasoner_alias_resolves_to_siliconflow_pro_when_provider_hinted() {
695        let registry = ModelRegistry::default();
696        let resolved = registry.resolve(Some("deepseek-reasoner"), Some(ProviderKind::Siliconflow));
697
698        assert_eq!(resolved.resolved.provider, ProviderKind::Siliconflow);
699        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Pro");
700    }
701
702    #[test]
703    fn deepseek_v4_flash_alias_resolves_to_siliconflow_flash_when_provider_hinted() {
704        let registry = ModelRegistry::default();
705        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Siliconflow));
706
707        assert_eq!(resolved.resolved.provider, ProviderKind::Siliconflow);
708        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Flash");
709    }
710
711    #[test]
712    fn sglang_default_uses_canonical_model_id() {
713        let registry = ModelRegistry::default();
714        let resolved = registry.resolve(None, Some(ProviderKind::Sglang));
715
716        assert_eq!(resolved.resolved.provider, ProviderKind::Sglang);
717        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Pro");
718    }
719
720    #[test]
721    fn deepseek_v4_flash_alias_resolves_to_openrouter_when_provider_hinted() {
722        let registry = ModelRegistry::default();
723        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Openrouter));
724
725        assert_eq!(resolved.resolved.provider, ProviderKind::Openrouter);
726        assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-flash");
727    }
728
729    #[test]
730    fn recent_openrouter_large_model_aliases_resolve_when_provider_hinted() {
731        let registry = ModelRegistry::default();
732
733        for (alias, expected) in [
734            ("trinity-large-thinking", "arcee-ai/trinity-large-thinking"),
735            ("qwen3.6-35b-a3b", "qwen/qwen3.6-35b-a3b"),
736            ("gemma-4-31b-it", "google/gemma-4-31b-it"),
737            ("glm-5.1", "z-ai/glm-5.1"),
738            ("minimax-m3", "minimax/minimax-m3"),
739            ("openrouter-mimo-v2.5-pro", "xiaomi/mimo-v2.5-pro"),
740            ("openrouter-kimi-k2.6", "moonshotai/kimi-k2.6"),
741        ] {
742            let resolved = registry.resolve(Some(alias), Some(ProviderKind::Openrouter));
743
744            assert_eq!(resolved.resolved.provider, ProviderKind::Openrouter);
745            assert_eq!(resolved.resolved.id, expected);
746            assert!(resolved.resolved.supports_tools);
747            assert!(resolved.resolved.supports_reasoning);
748        }
749    }
750
751    #[test]
752    fn deepseek_v4_flash_alias_resolves_to_novita_when_provider_hinted() {
753        let registry = ModelRegistry::default();
754        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Novita));
755
756        assert_eq!(resolved.resolved.provider, ProviderKind::Novita);
757        assert_eq!(resolved.resolved.id, "deepseek/deepseek-v4-flash");
758    }
759
760    #[test]
761    fn deepseek_v4_flash_alias_resolves_to_sglang_when_provider_hinted() {
762        let registry = ModelRegistry::default();
763        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Sglang));
764
765        assert_eq!(resolved.resolved.provider, ProviderKind::Sglang);
766        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Flash");
767    }
768
769    #[test]
770    fn vllm_default_uses_canonical_model_id() {
771        let registry = ModelRegistry::default();
772        let resolved = registry.resolve(None, Some(ProviderKind::Vllm));
773
774        assert_eq!(resolved.resolved.provider, ProviderKind::Vllm);
775        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Pro");
776    }
777
778    #[test]
779    fn ollama_default_uses_small_local_model_id() {
780        let registry = ModelRegistry::default();
781        let resolved = registry.resolve(None, Some(ProviderKind::Ollama));
782
783        assert_eq!(resolved.resolved.provider, ProviderKind::Ollama);
784        assert_eq!(resolved.resolved.id, "deepseek-coder:1.3b");
785        assert!(!resolved.resolved.supports_reasoning);
786    }
787
788    #[test]
789    fn ollama_requested_model_tag_is_preserved() {
790        let registry = ModelRegistry::default();
791        let resolved = registry.resolve(Some("qwen2.5-coder:7b"), Some(ProviderKind::Ollama));
792
793        assert_eq!(resolved.resolved.provider, ProviderKind::Ollama);
794        assert_eq!(resolved.resolved.id, "qwen2.5-coder:7b");
795        assert!(!resolved.used_fallback);
796    }
797
798    #[test]
799    fn deepseek_v4_flash_alias_resolves_to_vllm_when_provider_hinted() {
800        let registry = ModelRegistry::default();
801        let resolved = registry.resolve(Some("deepseek-v4-flash"), Some(ProviderKind::Vllm));
802
803        assert_eq!(resolved.resolved.provider, ProviderKind::Vllm);
804        assert_eq!(resolved.resolved.id, "deepseek-ai/DeepSeek-V4-Flash");
805    }
806
807    #[test]
808    fn preserves_requested_model_casing_for_third_party_providers() {
809        let registry = ModelRegistry::default();
810        let resolved = registry.resolve(Some("DeepSeek-V4-Pro"), None);
811
812        assert_eq!(resolved.resolved.provider, ProviderKind::Deepseek);
813        assert_eq!(resolved.resolved.id, "DeepSeek-V4-Pro");
814    }
815
816    #[test]
817    fn registry_casing_takes_priority_over_requested_casing_with_provider_hint() {
818        let registry = ModelRegistry::default();
819        let resolved = registry.resolve(Some("DeepSeek-V4-Pro"), Some(ProviderKind::Deepseek));
820
821        assert_eq!(resolved.resolved.provider, ProviderKind::Deepseek);
822        // Registry's canonical id is used even when user provides different casing
823        assert_eq!(resolved.resolved.id, "deepseek-v4-pro");
824    }
825
826    #[test]
827    fn preserves_requested_model_casing_without_surrounding_whitespace() {
828        let registry = ModelRegistry::default();
829        let resolved = registry.resolve(Some("  DeepSeek-V4-Pro  "), None);
830
831        assert_eq!(resolved.resolved.provider, ProviderKind::Deepseek);
832        assert_eq!(resolved.resolved.id, "DeepSeek-V4-Pro");
833    }
834
835    #[test]
836    fn alias_match_does_not_override_requested_casing() {
837        let registry = ModelRegistry::default();
838        let resolved = registry.resolve(Some("deepseek-reasoner"), None);
839
840        assert_eq!(resolved.resolved.provider, ProviderKind::Deepseek);
841        assert_eq!(resolved.resolved.id, "deepseek-v4-flash");
842    }
843}
codewhale_agent/lib.rs

codewhale_agent/
lib.rs
