1use crate::{Api, CompatSettings, Cost, InputModality, MaxTokensField, Model, ThinkingFormat};
8use parking_lot::RwLock;
9use std::collections::HashMap;
10use std::sync::LazyLock;
11
12fn extract_model_name(id: &str) -> &str {
14 id.rsplit_once('/').map(|(_, name)| name).unwrap_or(id)
15}
16
17fn 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 _ => None,
45 }
46}
47
48static STATIC_MODELS: LazyLock<HashMap<String, Model>> = LazyLock::new(|| {
50 let mut map = HashMap::new();
51
52 add_openai_models(&mut map);
54
55 add_anthropic_models(&mut map);
57
58 add_google_models(&mut map);
60
61 add_deepseek_models(&mut map);
63
64 add_mistral_models(&mut map);
66
67 add_groq_models(&mut map);
69
70 add_cerebras_models(&mut map);
72
73 add_xai_models(&mut map);
75
76 add_openrouter_models(&mut map);
78
79 add_azure_models(&mut map);
81
82 add_zai_models(&mut map);
84
85 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 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 let v4_models = [
348 (
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 (
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 cache_read: if input_cost < 0.2 { 0.0028 } else { 0.003625 },
386 cache_write: 0.0, },
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#[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 pub fn new() -> Self {
829 Self {
830 static_models: HashMap::new(),
831 dynamic_models: RwLock::new(HashMap::new()),
832 }
833 }
834
835 pub fn from_static() -> Self {
839 Self {
840 static_models: STATIC_MODELS.clone(),
841 dynamic_models: RwLock::new(HashMap::new()),
842 }
843 }
844
845 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 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 pub fn lookup(&self, provider: &str, model_id: &str) -> Option<Model> {
864 let key = format!("{}/{}", provider, model_id);
865 if let Some(m) = self.dynamic_models.read().get(&key) {
867 return Some(m.clone());
868 }
869 self.static_models.get(&key).cloned()
871 }
872
873 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 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 pub fn all() -> Vec<&'static Model> {
889 STATIC_MODELS.values().collect()
890 }
891
892 pub fn dynamic_models(&self) -> Vec<Model> {
894 self.dynamic_models.read().values().cloned().collect()
895 }
896
897 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 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
917static GLOBAL_REGISTRY: LazyLock<ModelRegistry> = LazyLock::new(ModelRegistry::from_static);
921
922pub fn register_model(model: Model) {
930 GLOBAL_REGISTRY.register(model);
931}
932
933pub fn unregister_model(provider: &str, model_id: &str) {
935 GLOBAL_REGISTRY.unregister(provider, model_id);
936}
937
938pub fn lookup_model(provider: &str, model_id: &str) -> Option<Model> {
942 GLOBAL_REGISTRY.lookup(provider, model_id)
943}
944
945pub fn get_model(provider: &str, model_id: &str) -> Option<&'static Model> {
947 ModelRegistry::get(provider, model_id)
948}
949
950pub 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
961pub fn get_models(provider: &str) -> Vec<&'static Model> {
963 ModelRegistry::get_by_provider(provider)
964}
965
966pub 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 }
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 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}