Skip to main content

oxicode_ai/
model_registry.rs

1//! Model registry for oxicode-ai
2//!
3//! Provides a centralized registry of available LLM models.
4//! Supports both static built-in models and dynamic runtime registration
5//! for custom OpenAI-compatible providers.
6
7use crate::{Api, CompatSettings, Cost, InputModality, MaxTokensField, Model, ThinkingFormat};
8use parking_lot::RwLock;
9use std::collections::HashMap;
10use std::sync::LazyLock;
11
12/// Extract the model name after the last '/', or return the whole id if no '/' is present.
13fn extract_model_name(id: &str) -> &str {
14    id.rsplit_once('/').map(|(_, name)| name).unwrap_or(id)
15}
16
17/// Return provider-specific compatibility defaults.
18///
19/// Internal helper used by `add_*_models()` functions so that every model
20/// from the same provider gets the same `compat` baseline.
21fn default_compat_for_provider(provider: &str) -> Option<CompatSettings> {
22    match provider {
23        "openai" | "openai-responses" | "openai-completions" => Some(CompatSettings {
24            thinking_format: Some(ThinkingFormat::OpenAI),
25            max_tokens_field: Some(MaxTokensField::MaxCompletionTokens),
26            ..CompatSettings::default()
27        }),
28        "openrouter" => Some(CompatSettings {
29            thinking_format: Some(ThinkingFormat::OpenRouter),
30            requires_tool_result_name: true,
31            ..CompatSettings::default()
32        }),
33        "deepseek" => Some(CompatSettings {
34            thinking_format: Some(ThinkingFormat::DeepSeek),
35            max_tokens_field: Some(MaxTokensField::MaxTokens),
36            ..CompatSettings::default()
37        }),
38        "zai" => Some(CompatSettings {
39            thinking_format: Some(ThinkingFormat::Zai),
40            ..CompatSettings::default()
41        }),
42        // azure-openai already has explicit CompatSettings in add_azure_models()
43        // All other providers: use defaults (return None)
44        _ => None,
45    }
46}
47
48/// Global model registry (static built-in models)
49static STATIC_MODELS: LazyLock<HashMap<String, Model>> = LazyLock::new(|| {
50    let mut map = HashMap::new();
51
52    // OpenAI models
53    add_openai_models(&mut map);
54
55    // Anthropic models
56    add_anthropic_models(&mut map);
57
58    // Google models
59    add_google_models(&mut map);
60
61    // DeepSeek models
62    add_deepseek_models(&mut map);
63
64    // Mistral models
65    add_mistral_models(&mut map);
66
67    // Groq models
68    add_groq_models(&mut map);
69
70    // Cerebras models
71    add_cerebras_models(&mut map);
72
73    // xAI models
74    add_xai_models(&mut map);
75
76    // OpenRouter models
77    add_openrouter_models(&mut map);
78
79    // Azure OpenAI models
80    add_azure_models(&mut map);
81
82    // ZAI models
83    add_zai_models(&mut map);
84
85    // MiniMax models
86    add_minimax_models(&mut map);
87
88    map
89});
90
91fn add_openai_models(map: &mut HashMap<String, Model>) {
92    let models = [
93        ("openai/gpt-4o", "GPT-4o", true, 2.5, 10.0),
94        ("openai/gpt-4o-mini", "GPT-4o Mini", true, 0.15, 0.60),
95        ("openai/gpt-4-turbo", "GPT-4 Turbo", true, 10.0, 30.0),
96        ("openai/gpt-4", "GPT-4", false, 30.0, 60.0),
97        ("openai/gpt-3.5-turbo", "GPT-3.5 Turbo", false, 0.5, 1.5),
98        ("openai/o1-preview", "OpenAI o1 Preview", true, 15.0, 60.0),
99        ("openai/o1-mini", "OpenAI o1 Mini", true, 15.0, 60.0),
100        ("openai/o1", "OpenAI o1", true, 15.0, 60.0),
101        ("openai/o3", "OpenAI o3", true, 15.0, 60.0),
102        ("openai/o3-mini", "OpenAI o3 Mini", true, 15.0, 60.0),
103    ];
104
105    for (id, name, reasoning, input_cost, output_cost) in models {
106        map.insert(
107            id.to_string(),
108            Model {
109                id: extract_model_name(id).to_string(),
110                name: name.to_string(),
111                api: Api::OpenAiCompletions,
112                provider: "openai".to_string(),
113                base_url: "https://api.openai.com/v1".to_string(),
114                reasoning,
115                input: if reasoning {
116                    vec![InputModality::Text]
117                } else {
118                    vec![InputModality::Text, InputModality::Image]
119                },
120                cost: Cost {
121                    input: input_cost,
122                    output: output_cost,
123                    cache_read: input_cost * 0.5,
124                    cache_write: input_cost * 7.5,
125                },
126                context_window: 128_000,
127                max_tokens: 32_000,
128                headers: Default::default(),
129                compat: default_compat_for_provider("openai"),
130            },
131        );
132    }
133}
134
135fn add_anthropic_models(map: &mut HashMap<String, Model>) {
136    let models = [
137        (
138            "anthropic/claude-sonnet-4-20250514",
139            "Claude Sonnet 4",
140            true,
141            3.0,
142            15.0,
143        ),
144        (
145            "anthropic/claude-opus-4-20250514",
146            "Claude Opus 4",
147            true,
148            15.0,
149            75.0,
150        ),
151        (
152            "anthropic/claude-3-5-sonnet-20241022",
153            "Claude 3.5 Sonnet",
154            true,
155            3.0,
156            15.0,
157        ),
158        (
159            "anthropic/claude-3-5-haiku-20241022",
160            "Claude 3.5 Haiku",
161            false,
162            0.8,
163            4.0,
164        ),
165        (
166            "anthropic/claude-3-opus",
167            "Claude 3 Opus",
168            false,
169            15.0,
170            75.0,
171        ),
172        (
173            "anthropic/claude-3-sonnet",
174            "Claude 3 Sonnet",
175            false,
176            3.0,
177            15.0,
178        ),
179        (
180            "anthropic/claude-3-haiku",
181            "Claude 3 Haiku",
182            false,
183            0.25,
184            1.25,
185        ),
186    ];
187
188    for (id, name, reasoning, input_cost, output_cost) in models {
189        map.insert(
190            id.to_string(),
191            Model {
192                id: extract_model_name(id).to_string(),
193                name: name.to_string(),
194                api: Api::AnthropicMessages,
195                provider: "anthropic".to_string(),
196                base_url: "https://api.anthropic.com".to_string(),
197                reasoning,
198                input: vec![InputModality::Text, InputModality::Image],
199                cost: Cost {
200                    input: input_cost,
201                    output: output_cost,
202                    cache_read: input_cost * 0.1,
203                    cache_write: input_cost * 1.25,
204                },
205                context_window: 200_000,
206                max_tokens: 8192,
207                headers: Default::default(),
208                compat: default_compat_for_provider("anthropic"),
209            },
210        );
211    }
212}
213
214fn add_google_models(map: &mut HashMap<String, Model>) {
215    let models = [
216        (
217            "google/gemini-2.0-flash",
218            "Gemini 2.0 Flash",
219            0.0,
220            0.0,
221            1_000_000,
222        ),
223        (
224            "google/gemini-2.5-flash",
225            "Gemini 2.5 Flash",
226            0.0,
227            0.0,
228            1_000_000,
229        ),
230        (
231            "google/gemini-2.5-pro",
232            "Gemini 2.5 Pro",
233            1.25,
234            5.0,
235            2_000_000,
236        ),
237        (
238            "google/gemini-1.5-flash",
239            "Gemini 1.5 Flash",
240            0.0,
241            0.0,
242            1_000_000,
243        ),
244        (
245            "google/gemini-1.5-pro",
246            "Gemini 1.5 Pro",
247            1.25,
248            5.0,
249            2_000_000,
250        ),
251        ("google/gemini-pro", "Gemini Pro", 0.125, 0.5, 32_000),
252    ];
253
254    for (id, name, input_cost, output_cost, ctx) in models {
255        map.insert(
256            id.to_string(),
257            Model {
258                id: extract_model_name(id).to_string(),
259                name: name.to_string(),
260                api: Api::GoogleGenerativeAi,
261                provider: "google".to_string(),
262                base_url: "https://generativelanguage.googleapis.com".to_string(),
263                reasoning: false,
264                input: vec![InputModality::Text, InputModality::Image],
265                cost: Cost {
266                    input: input_cost,
267                    output: output_cost,
268                    cache_read: 0.0,
269                    cache_write: 0.0,
270                },
271                context_window: ctx,
272                max_tokens: 8192,
273                headers: Default::default(),
274                compat: default_compat_for_provider("google"),
275            },
276        );
277    }
278}
279
280fn add_deepseek_models(map: &mut HashMap<String, Model>) {
281    // Legacy models (to be retired 2026-07-24)
282    let legacy_models = [
283        (
284            "deepseek/deepseek-chat",
285            "DeepSeek Chat",
286            false,
287            0.27,
288            1.1,
289            64_000,
290            8192,
291        ),
292        (
293            "deepseek/deepseek-chat-v3",
294            "DeepSeek Chat V3",
295            false,
296            0.27,
297            1.1,
298            64_000,
299            8192,
300        ),
301        (
302            "deepseek/deepseek-reasoner",
303            "DeepSeek Reasoner",
304            true,
305            0.55,
306            2.19,
307            64_000,
308            8192,
309        ),
310        (
311            "deepseek/deepseek-coder",
312            "DeepSeek Coder",
313            false,
314            0.27,
315            1.1,
316            64_000,
317            8192,
318        ),
319    ];
320
321    for (id, name, reasoning, input_cost, output_cost, ctx, max_out) in legacy_models {
322        map.insert(
323            id.to_string(),
324            Model {
325                id: extract_model_name(id).to_string(),
326                name: name.to_string(),
327                api: Api::OpenAiCompletions,
328                provider: "deepseek".to_string(),
329                base_url: "https://api.deepseek.com".to_string(),
330                reasoning,
331                input: vec![InputModality::Text],
332                cost: Cost {
333                    input: input_cost,
334                    output: output_cost,
335                    cache_read: 0.1,
336                    cache_write: 1.0,
337                },
338                context_window: ctx,
339                max_tokens: max_out,
340                headers: Default::default(),
341                compat: default_compat_for_provider("deepseek"),
342            },
343        );
344    }
345
346    // V4 models (released 2026-04-24)
347    let v4_models = [
348        // deepseek-v4-flash: 284B total / 13B active, $0.14/M input, $0.28/M output
349        (
350            "deepseek/deepseek-v4-flash",
351            "DeepSeek V4 Flash",
352            true,
353            0.14,
354            0.28,
355            1_000_000,
356            384_000,
357        ),
358        // deepseek-v4-pro: 1.6T total / 49B active, $0.435/M input, $0.87/M output
359        (
360            "deepseek/deepseek-v4-pro",
361            "DeepSeek V4 Pro",
362            true,
363            0.435,
364            0.87,
365            1_000_000,
366            384_000,
367        ),
368    ];
369
370    for (id, name, reasoning, input_cost, output_cost, ctx, max_out) in v4_models {
371        map.insert(
372            id.to_string(),
373            Model {
374                id: extract_model_name(id).to_string(),
375                name: name.to_string(),
376                api: Api::OpenAiCompletions,
377                provider: "deepseek".to_string(),
378                base_url: "https://api.deepseek.com".to_string(),
379                reasoning,
380                input: vec![InputModality::Text],
381                cost: Cost {
382                    input: input_cost,
383                    output: output_cost,
384                    // V4 cache pricing: flash $0.0028, pro $0.003625 per 1M tokens
385                    cache_read: if input_cost < 0.2 { 0.0028 } else { 0.003625 },
386                    cache_write: 0.0, // DeepSeek does not charge extra for cache writes
387                },
388                context_window: ctx,
389                max_tokens: max_out,
390                headers: Default::default(),
391                compat: default_compat_for_provider("deepseek"),
392            },
393        );
394    }
395}
396
397fn add_mistral_models(map: &mut HashMap<String, Model>) {
398    let models = [
399        (
400            "mistral/mistral-large-latest",
401            "Mistral Large",
402            false,
403            2.0,
404            6.0,
405        ),
406        (
407            "mistral/mistral-medium-latest",
408            "Mistral Medium",
409            false,
410            0.5,
411            1.5,
412        ),
413        (
414            "mistral/mistral-small-latest",
415            "Mistral Small",
416            false,
417            0.2,
418            0.6,
419        ),
420        ("mistral/mistral-nemo", "Mistral Nemo", false, 0.15, 0.15),
421        ("mistral/codestral", "Codestral", false, 0.3, 0.9),
422        (
423            "mistral/codestral-mamba",
424            "Codestral Mamba",
425            false,
426            0.25,
427            0.25,
428        ),
429        (
430            "mistral/open-mixtral-8x22b",
431            "Mixtral 8x22B",
432            false,
433            0.45,
434            1.4,
435        ),
436        (
437            "mistral/open-mixtral-8x7b",
438            "Mixtral 8x7B",
439            false,
440            0.24,
441            0.24,
442        ),
443    ];
444
445    for (id, name, reasoning, input_cost, output_cost) in models {
446        map.insert(
447            id.to_string(),
448            Model {
449                id: extract_model_name(id).to_string(),
450                name: name.to_string(),
451                api: Api::OpenAiCompletions,
452                provider: "mistral".to_string(),
453                base_url: "https://api.mistral.ai".to_string(),
454                reasoning,
455                input: vec![InputModality::Text],
456                cost: Cost {
457                    input: input_cost,
458                    output: output_cost,
459                    cache_read: 0.0,
460                    cache_write: 0.0,
461                },
462                context_window: 128_000,
463                max_tokens: 32_000,
464                headers: Default::default(),
465                compat: default_compat_for_provider("mistral"),
466            },
467        );
468    }
469}
470
471fn add_groq_models(map: &mut HashMap<String, Model>) {
472    let models = [
473        (
474            "groq/llama-3.3-70b-versatile",
475            "Llama 3.3 70B Versatile",
476            false,
477            0.0,
478            0.0,
479        ),
480        (
481            "groq/llama-3.1-70b-versatile",
482            "Llama 3.1 70B Versatile",
483            false,
484            0.0,
485            0.0,
486        ),
487        (
488            "groq/llama-3.1-8b-instant",
489            "Llama 3.1 8B Instant",
490            false,
491            0.0,
492            0.0,
493        ),
494        (
495            "groq/llama-3-70b-versatile",
496            "Llama 3 70B Versatile",
497            false,
498            0.0,
499            0.0,
500        ),
501        (
502            "groq/llama-3-8b-versatile",
503            "Llama 3 8B Versatile",
504            false,
505            0.0,
506            0.0,
507        ),
508        ("groq/mixtral-8x7b-32768", "Mixtral 8x7B", false, 0.0, 0.0),
509        ("groq/gemma2-9b-it", "Gemma 2 9B", false, 0.0, 0.0),
510        ("groq/gemma-7b-it", "Gemma 7B", false, 0.0, 0.0),
511    ];
512
513    for (id, name, reasoning, input_cost, output_cost) in models {
514        map.insert(
515            id.to_string(),
516            Model {
517                id: extract_model_name(id).to_string(),
518                name: name.to_string(),
519                api: Api::OpenAiCompletions,
520                provider: "groq".to_string(),
521                base_url: "https://api.groq.com/openai/v1".to_string(),
522                reasoning,
523                input: vec![InputModality::Text],
524                cost: Cost {
525                    input: input_cost,
526                    output: output_cost,
527                    cache_read: 0.0,
528                    cache_write: 0.0,
529                },
530                context_window: 128_000,
531                max_tokens: 8192,
532                headers: Default::default(),
533                compat: default_compat_for_provider("groq"),
534            },
535        );
536    }
537}
538
539fn add_cerebras_models(map: &mut HashMap<String, Model>) {
540    let models = [
541        ("cerebras/llama-3.3-70b", "Llama 3.3 70B", false, 0.0, 0.0),
542        ("cerebras/llama-3.1-8b", "Llama 3.1 8B", false, 0.0, 0.0),
543        ("cerebras/qwen-2.5-32b", "Qwen 2.5 32B", false, 0.0, 0.0),
544        ("cerebras/qwen-2.5-7b", "Qwen 2.5 7B", false, 0.0, 0.0),
545    ];
546
547    for (id, name, reasoning, input_cost, output_cost) in models {
548        map.insert(
549            id.to_string(),
550            Model {
551                id: extract_model_name(id).to_string(),
552                name: name.to_string(),
553                api: Api::OpenAiCompletions,
554                provider: "cerebras".to_string(),
555                base_url: "https://api.cerebras.ai".to_string(),
556                reasoning,
557                input: vec![InputModality::Text],
558                cost: Cost {
559                    input: input_cost,
560                    output: output_cost,
561                    cache_read: 0.0,
562                    cache_write: 0.0,
563                },
564                context_window: 128_000,
565                max_tokens: 8192,
566                headers: Default::default(),
567                compat: default_compat_for_provider("cerebras"),
568            },
569        );
570    }
571}
572
573fn add_xai_models(map: &mut HashMap<String, Model>) {
574    let models = [
575        ("xai/grok-2", "Grok 2", false, 5.0, 15.0),
576        ("xai/grok-2-mini", "Grok 2 Mini", false, 0.3, 0.5),
577        ("xai/grok-1", "Grok 1", false, 5.0, 15.0),
578        ("xai/grok-1.5", "Grok 1.5", false, 5.0, 15.0),
579    ];
580
581    for (id, name, reasoning, input_cost, output_cost) in models {
582        map.insert(
583            id.to_string(),
584            Model {
585                id: extract_model_name(id).to_string(),
586                name: name.to_string(),
587                api: Api::OpenAiCompletions,
588                provider: "xai".to_string(),
589                base_url: "https://api.x.ai/v1".to_string(),
590                reasoning,
591                input: vec![InputModality::Text],
592                cost: Cost {
593                    input: input_cost,
594                    output: output_cost,
595                    cache_read: 0.0,
596                    cache_write: 0.0,
597                },
598                context_window: 131_072,
599                max_tokens: 8192,
600                headers: Default::default(),
601                compat: default_compat_for_provider("xai"),
602            },
603        );
604    }
605}
606
607fn add_openrouter_models(map: &mut HashMap<String, Model>) {
608    let models = [
609        (
610            "openrouter/anthropic/claude-3.5-sonnet",
611            "Claude 3.5 Sonnet",
612            false,
613            3.0,
614            15.0,
615        ),
616        (
617            "openrouter/anthropic/claude-3-opus",
618            "Claude 3 Opus",
619            false,
620            15.0,
621            75.0,
622        ),
623        (
624            "openrouter/google/gemini-pro-1.5",
625            "Gemini Pro 1.5",
626            false,
627            1.25,
628            5.0,
629        ),
630        (
631            "openrouter/meta-llama/llama-3-70b",
632            "Llama 3 70B",
633            false,
634            0.65,
635            2.75,
636        ),
637        (
638            "openrouter/meta-llama/llama-3-8b",
639            "Llama 3 8B",
640            false,
641            0.2,
642            0.2,
643        ),
644        (
645            "openrouter/mistralai/mistral-large",
646            "Mistral Large",
647            false,
648            2.0,
649            6.0,
650        ),
651        (
652            "openrouter/deepseek/deepseek-chat",
653            "DeepSeek Chat",
654            false,
655            0.27,
656            1.1,
657        ),
658        ("openrouter/qwen/qwen-2-72b", "Qwen 2 72B", false, 0.9, 0.9),
659        (
660            "openrouter/nousresearch/hermes-3-llama-3-70b",
661            "Hermes 3 70B",
662            false,
663            0.5,
664            1.5,
665        ),
666    ];
667
668    for (id, name, reasoning, input_cost, output_cost) in models {
669        map.insert(
670            id.to_string(),
671            Model {
672                id: extract_model_name(id).to_string(),
673                name: name.to_string(),
674                api: Api::OpenAiCompletions,
675                provider: "openrouter".to_string(),
676                base_url: "https://openrouter.ai/api/v1".to_string(),
677                reasoning,
678                input: vec![InputModality::Text],
679                cost: Cost {
680                    input: input_cost,
681                    output: output_cost,
682                    cache_read: 0.0,
683                    cache_write: 0.0,
684                },
685                context_window: 128_000,
686                max_tokens: 32_000,
687                headers: [
688                    ("HTTP-Referer".to_string(), "https://oxicode-ai".to_string()),
689                    ("X-Title".to_string(), "oxicode-ai".to_string()),
690                ]
691                .into_iter()
692                .collect(),
693                compat: default_compat_for_provider("openrouter"),
694            },
695        );
696    }
697}
698
699fn add_azure_models(map: &mut HashMap<String, Model>) {
700    let models = [
701        ("azure-openai/gpt-4o", "GPT-4o", false, 2.5, 10.0),
702        ("azure-openai/gpt-4o-mini", "GPT-4o Mini", false, 0.15, 0.60),
703        ("azure-openai/gpt-4-turbo", "GPT-4 Turbo", false, 10.0, 30.0),
704    ];
705
706    for (id, name, reasoning, input_cost, output_cost) in models {
707        map.insert(
708            id.to_string(),
709            Model {
710                id: extract_model_name(id).to_string(),
711                name: name.to_string(),
712                api: Api::AzureOpenAiResponses,
713                provider: "azure-openai".to_string(),
714                base_url: "https://{your-resource-name}.openai.azure.com".to_string(),
715                reasoning,
716                input: vec![InputModality::Text, InputModality::Image],
717                cost: Cost {
718                    input: input_cost,
719                    output: output_cost,
720                    cache_read: 0.0,
721                    cache_write: 0.0,
722                },
723                context_window: 128_000,
724                max_tokens: 32_000,
725                headers: Default::default(),
726                compat: Some(crate::CompatSettings {
727                    supports_store: false,
728                    supports_developer_role: false,
729                    supports_reasoning_effort: false,
730                    supports_usage_in_streaming: false,
731                    max_tokens_field: Some(crate::MaxTokensField::MaxCompletionTokens),
732                    requires_tool_result_name: true,
733                    requires_assistant_after_tool_result: false,
734                    requires_thinking_as_text: false,
735                    thinking_format: None,
736                }),
737            },
738        );
739    }
740}
741
742fn add_zai_models(map: &mut HashMap<String, Model>) {
743    let models = [
744        ("zai/glm-4.7", "GLM-4.7", true, 0.0, 0.0),
745        ("zai/glm-5-turbo", "GLM-5-Turbo", true, 0.0, 0.0),
746        ("zai/glm-5.1", "GLM-5.1", true, 0.0, 0.0),
747        ("zai/glm-5v-turbo", "GLM-5V-Turbo", true, 0.0, 0.0),
748        ("zai/glm-4.5-air", "GLM-4.5-Air", true, 0.0, 0.0),
749    ];
750
751    for (id, name, reasoning, input_cost, output_cost) in models {
752        map.insert(
753            id.to_string(),
754            Model {
755                id: extract_model_name(id).to_string(),
756                name: name.to_string(),
757                api: Api::OpenAiCompletions,
758                provider: "zai".to_string(),
759                base_url: "https://api.z.ai/api/coding/paas/v4".to_string(),
760                reasoning,
761                input: vec![InputModality::Text],
762                cost: Cost {
763                    input: input_cost,
764                    output: output_cost,
765                    cache_read: 0.0,
766                    cache_write: 0.0,
767                },
768                context_window: 200_000,
769                max_tokens: 131_072,
770                headers: Default::default(),
771                compat: default_compat_for_provider("zai"),
772            },
773        );
774    }
775}
776
777fn add_minimax_models(map: &mut HashMap<String, Model>) {
778    let models = [
779        ("minimax/MiniMax-M2.7", "MiniMax-M2.7", true, 0.0, 0.0),
780        (
781            "minimax/MiniMax-M2.7-highspeed",
782            "MiniMax-M2.7-highspeed",
783            true,
784            0.0,
785            0.0,
786        ),
787    ];
788
789    for (id, name, reasoning, input_cost, output_cost) in models {
790        map.insert(
791            id.to_string(),
792            Model {
793                id: extract_model_name(id).to_string(),
794                name: name.to_string(),
795                api: Api::AnthropicMessages,
796                provider: "minimax".to_string(),
797                base_url: "https://api.minimax.io".to_string(),
798                reasoning,
799                input: vec![InputModality::Text],
800                cost: Cost {
801                    input: input_cost,
802                    output: output_cost,
803                    cache_read: 0.06,
804                    cache_write: 0.375,
805                },
806                context_window: 204_800,
807                max_tokens: 131_072,
808                headers: Default::default(),
809                compat: default_compat_for_provider("minimax"),
810            },
811        );
812    }
813}
814
815/// Lightweight model registry for SDK/engine usage.
816///
817/// Stores model metadata (provider, base_url, API type, costs) without
818/// authentication details. For CLI usage with auth integration, see
819/// `oxicode_store::CliModelRegistry`.
820#[derive(Default)]
821pub struct ModelRegistry {
822    static_models: HashMap<String, Model>,
823    dynamic_models: parking_lot::RwLock<HashMap<String, Model>>,
824}
825
826impl ModelRegistry {
827    /// Create a new empty registry.
828    pub fn new() -> Self {
829        Self {
830            static_models: HashMap::new(),
831            dynamic_models: RwLock::new(HashMap::new()),
832        }
833    }
834
835    /// Create a registry pre-populated with all built-in static models.
836    ///
837    /// This loads models from the embedded static database.
838    pub fn from_static() -> Self {
839        Self {
840            static_models: STATIC_MODELS.clone(),
841            dynamic_models: RwLock::new(HashMap::new()),
842        }
843    }
844
845    /// Register a model at runtime.
846    ///
847    /// If a model with the same `provider/model_id` key already exists,
848    /// the new one replaces it.
849    pub fn register(&self, model: Model) {
850        let key = format!("{}/{}", model.provider, model.id);
851        self.dynamic_models.write().insert(key, model);
852    }
853
854    /// Unregister a previously registered dynamic model.
855    pub fn unregister(&self, provider: &str, model_id: &str) {
856        let key = format!("{}/{}", provider, model_id);
857        self.dynamic_models.write().remove(&key);
858    }
859
860    /// Look up a model by provider and model ID.
861    ///
862    /// Dynamic models take priority over static ones.
863    pub fn lookup(&self, provider: &str, model_id: &str) -> Option<Model> {
864        let key = format!("{}/{}", provider, model_id);
865        // Dynamic models take priority
866        if let Some(m) = self.dynamic_models.read().get(&key) {
867            return Some(m.clone());
868        }
869        // Then static models
870        self.static_models.get(&key).cloned()
871    }
872
873    /// Get a model by provider/model ID (static models only).
874    pub fn get(provider: &str, model_id: &str) -> Option<&'static Model> {
875        let key = format!("{}/{}", provider, model_id);
876        STATIC_MODELS.get(&key)
877    }
878
879    /// Get all models from a provider (static only).
880    pub fn get_by_provider(provider: &str) -> Vec<&'static Model> {
881        STATIC_MODELS
882            .values()
883            .filter(|m| m.provider == provider)
884            .collect()
885    }
886
887    /// Get all available models (static only).
888    pub fn all() -> Vec<&'static Model> {
889        STATIC_MODELS.values().collect()
890    }
891
892    /// Get all dynamically registered models.
893    pub fn dynamic_models(&self) -> Vec<Model> {
894        self.dynamic_models.read().values().cloned().collect()
895    }
896
897    /// Get all registered model IDs as `provider/model` strings.
898    pub fn model_ids(&self) -> Vec<String> {
899        let static_ids: Vec<String> = self.static_models.keys().cloned().collect();
900        let dynamic_ids: Vec<String> = self.dynamic_models.read().keys().cloned().collect();
901        static_ids.into_iter().chain(dynamic_ids).collect()
902    }
903
904    /// Search models by pattern (static only).
905    pub fn search(pattern: &str) -> Vec<&'static Model> {
906        let pattern_lower = pattern.to_lowercase();
907        STATIC_MODELS
908            .values()
909            .filter(|m| {
910                m.id.to_lowercase().contains(&pattern_lower)
911                    || m.name.to_lowercase().contains(&pattern_lower)
912            })
913            .collect()
914    }
915}
916
917// ── Global registry instance ────────────────────────────────────────
918
919/// Global model registry instance (for convenience functions).
920static GLOBAL_REGISTRY: LazyLock<ModelRegistry> = LazyLock::new(ModelRegistry::from_static);
921
922// ── Convenience functions using global registry ─────────────────────
923
924/// Register a model at runtime.
925///
926/// Call this during startup for each custom provider's model.
927/// If a model with the same `provider/model_id` key already exists,
928/// the new one replaces it.
929pub fn register_model(model: Model) {
930    GLOBAL_REGISTRY.register(model);
931}
932
933/// Unregister a previously registered dynamic model.
934pub fn unregister_model(provider: &str, model_id: &str) {
935    GLOBAL_REGISTRY.unregister(provider, model_id);
936}
937
938/// Look up a model by provider and model ID, checking both dynamic and static registries.
939///
940/// Dynamic models take priority over static ones.
941pub fn lookup_model(provider: &str, model_id: &str) -> Option<Model> {
942    GLOBAL_REGISTRY.lookup(provider, model_id)
943}
944
945/// Convenience function to get a model (static registry only – use [`lookup_model`] for dynamic too)
946pub fn get_model(provider: &str, model_id: &str) -> Option<&'static Model> {
947    ModelRegistry::get(provider, model_id)
948}
949
950/// Get all available providers
951pub fn get_providers() -> Vec<&'static str> {
952    let mut providers: Vec<&'static str> = STATIC_MODELS
953        .values()
954        .map(|m| m.provider.as_str())
955        .collect();
956    providers.sort();
957    providers.dedup();
958    providers
959}
960
961/// Get all models from a provider
962pub fn get_models(provider: &str) -> Vec<&'static Model> {
963    ModelRegistry::get_by_provider(provider)
964}
965
966/// Get all dynamically registered models.
967pub fn dynamic_models() -> Vec<Model> {
968    GLOBAL_REGISTRY.dynamic_models()
969}
970
971#[cfg(test)]
972mod tests {
973    use super::*;
974
975    #[test]
976    fn test_get_model() {
977        let model = get_model("openai", "gpt-4o");
978        assert!(model.is_some());
979        let model = model.unwrap();
980        assert_eq!(model.provider, "openai");
981        // Note: gpt-4o has reasoning enabled
982    }
983
984    #[test]
985    fn test_get_providers() {
986        let providers = get_providers();
987        assert!(providers.contains(&"openai"));
988        assert!(providers.contains(&"anthropic"));
989        assert!(providers.contains(&"google"));
990        assert!(providers.contains(&"deepseek"));
991        assert!(providers.contains(&"mistral"));
992        assert!(providers.contains(&"groq"));
993    }
994
995    #[test]
996    fn test_deepseek_model() {
997        let model = get_model("deepseek", "deepseek-chat");
998        assert!(model.is_some());
999        let model = model.unwrap();
1000        assert_eq!(model.provider, "deepseek");
1001        assert_eq!(model.base_url, "https://api.deepseek.com");
1002    }
1003
1004    #[test]
1005    fn test_deepseek_v4_models() {
1006        let flash = get_model("deepseek", "deepseek-v4-flash");
1007        assert!(flash.is_some(), "deepseek-v4-flash should be registered");
1008        let flash = flash.unwrap();
1009        assert_eq!(flash.provider, "deepseek");
1010        assert_eq!(flash.context_window, 1_000_000);
1011        assert_eq!(flash.max_tokens, 384_000);
1012        assert!(flash.reasoning);
1013
1014        let pro = get_model("deepseek", "deepseek-v4-pro");
1015        assert!(pro.is_some(), "deepseek-v4-pro should be registered");
1016        let pro = pro.unwrap();
1017        assert_eq!(pro.provider, "deepseek");
1018        assert_eq!(pro.context_window, 1_000_000);
1019        assert_eq!(pro.max_tokens, 384_000);
1020        assert!(pro.reasoning);
1021        // V4 Pro is more expensive than V4 Flash
1022        assert!(pro.cost.input > flash.cost.input);
1023    }
1024
1025    #[test]
1026    fn test_search_models() {
1027        let results = ModelRegistry::search("gpt");
1028        assert!(!results.is_empty());
1029        assert!(
1030            results
1031                .iter()
1032                .all(|m| m.name.to_lowercase().contains("gpt"))
1033        );
1034    }
1035
1036    #[test]
1037    fn test_model_registry_instance() {
1038        let registry = ModelRegistry::from_static();
1039        assert!(registry.lookup("openai", "gpt-4o").is_some());
1040        assert!(registry.lookup("fake", "fake-model").is_none());
1041    }
1042
1043    #[test]
1044    fn test_model_registry_register_dynamic() {
1045        let registry = ModelRegistry::new();
1046        let custom_model = Model {
1047            id: "custom-model".to_string(),
1048            name: "Custom Model".to_string(),
1049            api: Api::OpenAiCompletions,
1050            provider: "custom".to_string(),
1051            base_url: "https://custom.example.com".to_string(),
1052            reasoning: false,
1053            input: vec![InputModality::Text],
1054            cost: Cost {
1055                input: 1.0,
1056                output: 2.0,
1057                cache_read: 0.5,
1058                cache_write: 5.0,
1059            },
1060            context_window: 100_000,
1061            max_tokens: 8192,
1062            headers: Default::default(),
1063            compat: None,
1064        };
1065        registry.register(custom_model.clone());
1066        assert!(registry.lookup("custom", "custom-model").is_some());
1067    }
1068}