1use std::fmt;
27
28#[derive(Debug, Clone, PartialEq, Eq)]
30pub enum ModelError {
31 InvalidPrefix {
33 model: String,
34 expected_prefixes: &'static [&'static str],
35 },
36 UnknownShorthand(String),
38}
39
40impl fmt::Display for ModelError {
41 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
42 match self {
43 Self::InvalidPrefix {
44 model,
45 expected_prefixes,
46 } => {
47 write!(
48 f,
49 "Invalid model name '{}'. Expected prefix: {}",
50 model,
51 expected_prefixes.join(" or ")
52 )
53 }
54 Self::UnknownShorthand(s) => write!(f, "Unknown model shorthand: {}", s),
55 }
56 }
57}
58
59impl std::error::Error for ModelError {}
60
61#[derive(Debug, Clone, PartialEq, Eq, Hash)]
84pub enum ClaudeModel {
85 Opus46,
87 Sonnet46,
89 Haiku45,
91 Opus45,
93 Sonnet45,
95 Opus41,
97 Opus4,
99 Sonnet4,
101 Custom(String),
103}
104
105impl Default for ClaudeModel {
106 fn default() -> Self {
107 Self::Sonnet46
108 }
109}
110
111impl ClaudeModel {
112 pub fn as_api_id(&self) -> &str {
116 match self {
117 Self::Opus46 => "claude-opus-4-6",
118 Self::Sonnet46 => "claude-sonnet-4-6",
119 Self::Haiku45 => "claude-haiku-4-5-20251001",
120 Self::Opus45 => "claude-opus-4-5-20251101",
121 Self::Sonnet45 => "claude-sonnet-4-5-20250929",
122 Self::Opus41 => "claude-opus-4-1-20250805",
123 Self::Opus4 => "claude-opus-4-20250514",
124 Self::Sonnet4 => "claude-sonnet-4-20250514",
125 Self::Custom(s) => s,
126 }
127 }
128
129 pub fn as_cli_name(&self) -> &str {
133 match self {
134 Self::Opus46 => "claude-opus-4.6",
135 Self::Sonnet46 => "claude-sonnet-4.6",
136 Self::Haiku45 => "claude-haiku-4.5",
137 Self::Opus45 => "claude-opus-4.5",
138 Self::Sonnet45 => "claude-sonnet-4.5",
139 Self::Opus41 => "claude-opus-4.1",
140 Self::Opus4 => "claude-opus-4",
141 Self::Sonnet4 => "claude-sonnet-4",
142 Self::Custom(s) => s,
143 }
144 }
145
146 fn validate_custom(s: &str) -> Result<(), ModelError> {
148 if s.starts_with("claude-") {
149 Ok(())
150 } else {
151 Err(ModelError::InvalidPrefix {
152 model: s.to_string(),
153 expected_prefixes: &["claude-"],
154 })
155 }
156 }
157}
158
159impl std::str::FromStr for ClaudeModel {
160 type Err = ModelError;
161
162 fn from_str(s: &str) -> Result<Self, Self::Err> {
163 match s.to_lowercase().as_str() {
164 "opus" | "opus-4.6" | "opus46" | "claude-opus-4.6" | "claude-opus-4-6" => {
166 Ok(Self::Opus46)
167 }
168 "sonnet" | "sonnet-4.6" | "sonnet46" | "claude-sonnet-4.6" | "claude-sonnet-4-6" => {
170 Ok(Self::Sonnet46)
171 }
172 "haiku"
174 | "haiku-4.5"
175 | "haiku45"
176 | "claude-haiku-4.5"
177 | "claude-haiku-4-5-20251001" => Ok(Self::Haiku45),
178 "opus-4.5" | "opus45" | "claude-opus-4.5" | "claude-opus-4-5-20251101" => {
180 Ok(Self::Opus45)
181 }
182 "sonnet-4.5" | "sonnet45" | "claude-sonnet-4.5" | "claude-sonnet-4-5-20250929" => {
184 Ok(Self::Sonnet45)
185 }
186 "opus-4.1" | "opus41" | "claude-opus-4.1" | "claude-opus-4-1-20250805" => {
188 Ok(Self::Opus41)
189 }
190 "opus-4" | "opus4" | "claude-opus-4" | "claude-opus-4-20250514" => Ok(Self::Opus4),
192 "sonnet-4" | "sonnet4" | "claude-sonnet-4" | "claude-sonnet-4-20250514" => {
194 Ok(Self::Sonnet4)
195 }
196 _ => {
198 Self::validate_custom(s)?;
199 Ok(Self::Custom(s.to_string()))
200 }
201 }
202 }
203}
204
205impl fmt::Display for ClaudeModel {
206 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
207 write!(f, "{}", self.as_api_id())
208 }
209}
210
211#[derive(Debug, Clone, PartialEq, Eq, Hash)]
229pub enum GeminiModel {
230 Pro31,
232 Flash3,
234 Pro3,
236 Flash25,
238 Pro25,
240 FlashLite25,
242 Flash20,
244 Custom(String),
246}
247
248impl Default for GeminiModel {
249 fn default() -> Self {
250 Self::Flash25
251 }
252}
253
254impl GeminiModel {
255 pub fn as_api_id(&self) -> &str {
257 match self {
258 Self::Pro31 => "gemini-3.1-pro-preview",
259 Self::Flash3 => "gemini-3-flash-preview",
260 Self::Pro3 => "gemini-3-pro-preview",
261 Self::Flash25 => "gemini-2.5-flash",
262 Self::Pro25 => "gemini-2.5-pro",
263 Self::FlashLite25 => "gemini-2.5-flash-lite",
264 Self::Flash20 => "gemini-2.0-flash",
265 Self::Custom(s) => s,
266 }
267 }
268
269 pub fn as_cli_name(&self) -> &str {
271 match self {
272 Self::Pro31 => "pro-3.1",
273 Self::Flash3 => "flash-3",
274 Self::Pro3 => "pro-3",
275 Self::Flash25 => "flash",
276 Self::Pro25 => "pro",
277 Self::FlashLite25 => "flash-lite",
278 Self::Flash20 => "flash-2.0",
279 Self::Custom(s) => s,
280 }
281 }
282
283 fn validate_custom(s: &str) -> Result<(), ModelError> {
284 if s.starts_with("gemini-") {
285 Ok(())
286 } else {
287 Err(ModelError::InvalidPrefix {
288 model: s.to_string(),
289 expected_prefixes: &["gemini-"],
290 })
291 }
292 }
293}
294
295impl std::str::FromStr for GeminiModel {
296 type Err = ModelError;
297
298 fn from_str(s: &str) -> Result<Self, Self::Err> {
299 match s.to_lowercase().as_str() {
300 "pro-3.1" | "pro31" | "gemini-3.1-pro-preview" => Ok(Self::Pro31),
302 "flash-3" | "flash3" | "gemini-3-flash-preview" | "gemini-3-flash" => Ok(Self::Flash3),
304 "pro-3" | "pro3" | "gemini-3-pro-preview" | "gemini-3-pro" => Ok(Self::Pro3),
306 "flash" | "flash-2.5" | "flash25" | "gemini-2.5-flash" => Ok(Self::Flash25),
308 "pro" | "pro-2.5" | "pro25" | "gemini-2.5-pro" => Ok(Self::Pro25),
310 "flash-lite" | "lite" | "gemini-2.5-flash-lite" => Ok(Self::FlashLite25),
312 "flash-2.0" | "flash20" | "flash-2" | "gemini-2.0-flash" => Ok(Self::Flash20),
314 _ => {
316 Self::validate_custom(s)?;
317 Ok(Self::Custom(s.to_string()))
318 }
319 }
320 }
321}
322
323impl fmt::Display for GeminiModel {
324 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
325 write!(f, "{}", self.as_api_id())
326 }
327}
328
329#[derive(Debug, Clone, PartialEq, Eq, Hash)]
347pub enum OpenAIModel {
348 Gpt52,
351 Gpt52Pro,
353 Gpt51,
355 Gpt5,
357 Gpt5Mini,
359
360 Gpt52Codex,
363 Gpt51Codex,
365 Gpt51CodexMini,
367 Gpt5Codex,
369 Gpt5CodexMini,
371
372 Gpt41,
375 Gpt41Mini,
377 Gpt4o,
379 Gpt4oMini,
381
382 O3Pro,
385 O3,
387 O3Mini,
389 O1,
391 O1Pro,
393
394 Custom(String),
396}
397
398impl Default for OpenAIModel {
399 fn default() -> Self {
400 Self::Gpt5
401 }
402}
403
404impl OpenAIModel {
405 pub fn as_api_id(&self) -> &str {
407 match self {
408 Self::Gpt52 => "gpt-5.2",
410 Self::Gpt52Pro => "gpt-5.2-pro",
411 Self::Gpt51 => "gpt-5.1",
412 Self::Gpt5 => "gpt-5",
413 Self::Gpt5Mini => "gpt-5-mini",
414 Self::Gpt52Codex => "gpt-5.2-codex",
416 Self::Gpt51Codex => "gpt-5.1-codex",
417 Self::Gpt51CodexMini => "gpt-5.1-codex-mini",
418 Self::Gpt5Codex => "gpt-5-codex",
419 Self::Gpt5CodexMini => "gpt-5-codex-mini",
420 Self::Gpt41 => "gpt-4.1",
422 Self::Gpt41Mini => "gpt-4.1-mini",
423 Self::Gpt4o => "gpt-4o",
424 Self::Gpt4oMini => "gpt-4o-mini",
425 Self::O3Pro => "o3-pro",
427 Self::O3 => "o3",
428 Self::O3Mini => "o3-mini",
429 Self::O1 => "o1",
430 Self::O1Pro => "o1-pro",
431 Self::Custom(s) => s,
433 }
434 }
435
436 pub fn as_cli_name(&self) -> &str {
438 self.as_api_id() }
440
441 fn validate_custom(s: &str) -> Result<(), ModelError> {
442 const VALID_PREFIXES: &[&str] = &["gpt-", "o1-", "o3-"];
443 if VALID_PREFIXES.iter().any(|p| s.starts_with(p)) {
444 Ok(())
445 } else {
446 Err(ModelError::InvalidPrefix {
447 model: s.to_string(),
448 expected_prefixes: VALID_PREFIXES,
449 })
450 }
451 }
452}
453
454impl std::str::FromStr for OpenAIModel {
455 type Err = ModelError;
456
457 fn from_str(s: &str) -> Result<Self, Self::Err> {
458 match s.to_lowercase().as_str() {
459 "5.2" | "gpt-5.2" | "gpt52" => Ok(Self::Gpt52),
461 "5.2-pro" | "gpt-5.2-pro" => Ok(Self::Gpt52Pro),
463 "5.1" | "gpt-5.1" | "gpt51" => Ok(Self::Gpt51),
465 "5" | "gpt-5" | "gpt5" => Ok(Self::Gpt5),
467 "5-mini" | "gpt-5-mini" => Ok(Self::Gpt5Mini),
469 "5.2-codex" | "gpt-5.2-codex" | "codex" => Ok(Self::Gpt52Codex),
471 "5.1-codex" | "gpt-5.1-codex" => Ok(Self::Gpt51Codex),
473 "5.1-codex-mini" | "gpt-5.1-codex-mini" | "codex-mini" => Ok(Self::Gpt51CodexMini),
474 "5-codex" | "gpt-5-codex" => Ok(Self::Gpt5Codex),
476 "5-codex-mini" | "gpt-5-codex-mini" => Ok(Self::Gpt5CodexMini),
477 "4.1" | "gpt-4.1" | "gpt41" => Ok(Self::Gpt41),
479 "4.1-mini" | "gpt-4.1-mini" => Ok(Self::Gpt41Mini),
480 "4o" | "gpt-4o" => Ok(Self::Gpt4o),
482 "4o-mini" | "gpt-4o-mini" => Ok(Self::Gpt4oMini),
483 "o3-pro" => Ok(Self::O3Pro),
485 "o3" => Ok(Self::O3),
486 "o3-mini" => Ok(Self::O3Mini),
487 "o1" => Ok(Self::O1),
488 "o1-pro" => Ok(Self::O1Pro),
489 _ => {
491 Self::validate_custom(s)?;
492 Ok(Self::Custom(s.to_string()))
493 }
494 }
495 }
496}
497
498impl fmt::Display for OpenAIModel {
499 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
500 write!(f, "{}", self.as_api_id())
501 }
502}
503
504#[derive(Debug, Clone, PartialEq, Eq, Hash)]
512pub enum Model {
513 Claude(ClaudeModel),
514 Gemini(GeminiModel),
515 OpenAI(OpenAIModel),
516}
517
518impl Model {
519 pub fn as_api_id(&self) -> &str {
521 match self {
522 Self::Claude(m) => m.as_api_id(),
523 Self::Gemini(m) => m.as_api_id(),
524 Self::OpenAI(m) => m.as_api_id(),
525 }
526 }
527
528 pub fn as_cli_name(&self) -> &str {
530 match self {
531 Self::Claude(m) => m.as_cli_name(),
532 Self::Gemini(m) => m.as_cli_name(),
533 Self::OpenAI(m) => m.as_cli_name(),
534 }
535 }
536}
537
538impl From<ClaudeModel> for Model {
539 fn from(m: ClaudeModel) -> Self {
540 Self::Claude(m)
541 }
542}
543
544impl From<GeminiModel> for Model {
545 fn from(m: GeminiModel) -> Self {
546 Self::Gemini(m)
547 }
548}
549
550impl From<OpenAIModel> for Model {
551 fn from(m: OpenAIModel) -> Self {
552 Self::OpenAI(m)
553 }
554}
555
556impl fmt::Display for Model {
557 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
558 write!(f, "{}", self.as_api_id())
559 }
560}
561
562#[cfg(test)]
567mod tests {
568 use super::*;
569
570 mod claude_model {
571 use super::*;
572
573 #[test]
574 fn test_default() {
575 assert_eq!(ClaudeModel::default(), ClaudeModel::Sonnet46);
576 }
577
578 #[test]
579 fn test_api_id() {
580 assert_eq!(ClaudeModel::Opus46.as_api_id(), "claude-opus-4-6");
581 assert_eq!(ClaudeModel::Sonnet46.as_api_id(), "claude-sonnet-4-6");
582 assert_eq!(
583 ClaudeModel::Haiku45.as_api_id(),
584 "claude-haiku-4-5-20251001"
585 );
586 assert_eq!(ClaudeModel::Opus45.as_api_id(), "claude-opus-4-5-20251101");
587 assert_eq!(
588 ClaudeModel::Sonnet45.as_api_id(),
589 "claude-sonnet-4-5-20250929"
590 );
591 }
592
593 #[test]
594 fn test_cli_name() {
595 assert_eq!(ClaudeModel::Opus46.as_cli_name(), "claude-opus-4.6");
596 assert_eq!(ClaudeModel::Sonnet46.as_cli_name(), "claude-sonnet-4.6");
597 assert_eq!(ClaudeModel::Haiku45.as_cli_name(), "claude-haiku-4.5");
598 assert_eq!(ClaudeModel::Opus45.as_cli_name(), "claude-opus-4.5");
599 assert_eq!(ClaudeModel::Sonnet4.as_cli_name(), "claude-sonnet-4");
600 }
601
602 #[test]
603 fn test_parse_shorthand() {
604 assert_eq!("opus".parse::<ClaudeModel>().unwrap(), ClaudeModel::Opus46);
605 assert_eq!(
606 "sonnet".parse::<ClaudeModel>().unwrap(),
607 ClaudeModel::Sonnet46
608 );
609 assert_eq!(
610 "haiku".parse::<ClaudeModel>().unwrap(),
611 ClaudeModel::Haiku45
612 );
613 }
614
615 #[test]
616 fn test_parse_versioned_shorthand() {
617 assert_eq!(
618 "opus-4.6".parse::<ClaudeModel>().unwrap(),
619 ClaudeModel::Opus46
620 );
621 assert_eq!(
622 "opus-4.5".parse::<ClaudeModel>().unwrap(),
623 ClaudeModel::Opus45
624 );
625 assert_eq!(
626 "sonnet-4.6".parse::<ClaudeModel>().unwrap(),
627 ClaudeModel::Sonnet46
628 );
629 assert_eq!(
630 "sonnet-4.5".parse::<ClaudeModel>().unwrap(),
631 ClaudeModel::Sonnet45
632 );
633 assert_eq!(
634 "haiku-4.5".parse::<ClaudeModel>().unwrap(),
635 ClaudeModel::Haiku45
636 );
637 }
638
639 #[test]
640 fn test_parse_full_api_id() {
641 assert_eq!(
642 "claude-opus-4-6".parse::<ClaudeModel>().unwrap(),
643 ClaudeModel::Opus46
644 );
645 assert_eq!(
646 "claude-opus-4-5-20251101".parse::<ClaudeModel>().unwrap(),
647 ClaudeModel::Opus45
648 );
649 assert_eq!(
650 "claude-sonnet-4".parse::<ClaudeModel>().unwrap(),
651 ClaudeModel::Sonnet4
652 );
653 }
654
655 #[test]
656 fn test_parse_custom_valid() {
657 let model: ClaudeModel = "claude-future-model-2027".parse().unwrap();
658 assert_eq!(
659 model,
660 ClaudeModel::Custom("claude-future-model-2027".to_string())
661 );
662 }
663
664 #[test]
665 fn test_parse_custom_invalid() {
666 let result: Result<ClaudeModel, _> = "gpt-4o".parse();
667 assert!(result.is_err());
668 }
669 }
670
671 mod gemini_model {
672 use super::*;
673
674 #[test]
675 fn test_default() {
676 assert_eq!(GeminiModel::default(), GeminiModel::Flash25);
677 }
678
679 #[test]
680 fn test_api_id() {
681 assert_eq!(GeminiModel::Pro31.as_api_id(), "gemini-3.1-pro-preview");
682 assert_eq!(GeminiModel::Flash3.as_api_id(), "gemini-3-flash-preview");
683 assert_eq!(GeminiModel::Pro3.as_api_id(), "gemini-3-pro-preview");
684 assert_eq!(
685 GeminiModel::FlashLite25.as_api_id(),
686 "gemini-2.5-flash-lite"
687 );
688 }
689
690 #[test]
691 fn test_parse() {
692 assert_eq!(
693 "flash".parse::<GeminiModel>().unwrap(),
694 GeminiModel::Flash25
695 );
696 assert_eq!("pro".parse::<GeminiModel>().unwrap(), GeminiModel::Pro25);
697 assert_eq!(
698 "flash-3".parse::<GeminiModel>().unwrap(),
699 GeminiModel::Flash3
700 );
701 assert_eq!(
702 "pro-3.1".parse::<GeminiModel>().unwrap(),
703 GeminiModel::Pro31
704 );
705 assert_eq!(
706 "flash-lite".parse::<GeminiModel>().unwrap(),
707 GeminiModel::FlashLite25
708 );
709 }
710
711 #[test]
712 fn test_parse_legacy_api_id() {
713 assert_eq!(
715 "gemini-3-flash".parse::<GeminiModel>().unwrap(),
716 GeminiModel::Flash3
717 );
718 assert_eq!(
719 "gemini-3-pro".parse::<GeminiModel>().unwrap(),
720 GeminiModel::Pro3
721 );
722 }
723
724 #[test]
725 fn test_custom_invalid() {
726 let result: Result<GeminiModel, _> = "claude-opus".parse();
727 assert!(result.is_err());
728 }
729 }
730
731 mod openai_model {
732 use super::*;
733
734 #[test]
735 fn test_default() {
736 assert_eq!(OpenAIModel::default(), OpenAIModel::Gpt5);
737 }
738
739 #[test]
740 fn test_api_id() {
741 assert_eq!(OpenAIModel::Gpt52Pro.as_api_id(), "gpt-5.2-pro");
742 assert_eq!(OpenAIModel::Gpt52Codex.as_api_id(), "gpt-5.2-codex");
743 assert_eq!(OpenAIModel::Gpt5.as_api_id(), "gpt-5");
744 }
745
746 #[test]
747 fn test_parse() {
748 assert_eq!("5".parse::<OpenAIModel>().unwrap(), OpenAIModel::Gpt5);
749 assert_eq!(
750 "gpt-5.2".parse::<OpenAIModel>().unwrap(),
751 OpenAIModel::Gpt52
752 );
753 assert_eq!(
754 "5.2-pro".parse::<OpenAIModel>().unwrap(),
755 OpenAIModel::Gpt52Pro
756 );
757 assert_eq!("o3".parse::<OpenAIModel>().unwrap(), OpenAIModel::O3);
758 assert_eq!(
759 "codex".parse::<OpenAIModel>().unwrap(),
760 OpenAIModel::Gpt52Codex
761 );
762 }
763
764 #[test]
765 fn test_parse_legacy() {
766 assert_eq!("4o".parse::<OpenAIModel>().unwrap(), OpenAIModel::Gpt4o);
767 assert_eq!(
768 "gpt-4.1".parse::<OpenAIModel>().unwrap(),
769 OpenAIModel::Gpt41
770 );
771 }
772
773 #[test]
774 fn test_custom_valid() {
775 let model: OpenAIModel = "o3-deep-research".parse().unwrap();
776 assert_eq!(model, OpenAIModel::Custom("o3-deep-research".to_string()));
777 }
778
779 #[test]
780 fn test_custom_invalid() {
781 let result: Result<OpenAIModel, _> = "gemini-pro".parse();
782 assert!(result.is_err());
783 }
784 }
785}