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 Opus45,
87 Sonnet45,
89 Opus41,
91 Opus4,
93 Sonnet4,
95 Haiku35,
97 Custom(String),
99}
100
101impl Default for ClaudeModel {
102 fn default() -> Self {
103 Self::Sonnet45
104 }
105}
106
107impl ClaudeModel {
108 pub fn as_api_id(&self) -> &str {
112 match self {
113 Self::Opus45 => "claude-opus-4-5-20251124",
114 Self::Sonnet45 => "claude-sonnet-4-5-20250929",
115 Self::Opus41 => "claude-opus-4-1-20250805",
116 Self::Opus4 => "claude-opus-4-20250514",
117 Self::Sonnet4 => "claude-sonnet-4-20250514",
118 Self::Haiku35 => "claude-3-5-haiku-20241022",
119 Self::Custom(s) => s,
120 }
121 }
122
123 pub fn as_cli_name(&self) -> &str {
127 match self {
128 Self::Opus45 => "claude-opus-4.5",
129 Self::Sonnet45 => "claude-sonnet-4.5",
130 Self::Opus41 => "claude-opus-4.1",
131 Self::Opus4 => "claude-opus-4",
132 Self::Sonnet4 => "claude-sonnet-4",
133 Self::Haiku35 => "claude-haiku-3.5",
134 Self::Custom(s) => s,
135 }
136 }
137
138 fn validate_custom(s: &str) -> Result<(), ModelError> {
140 if s.starts_with("claude-") {
141 Ok(())
142 } else {
143 Err(ModelError::InvalidPrefix {
144 model: s.to_string(),
145 expected_prefixes: &["claude-"],
146 })
147 }
148 }
149}
150
151impl std::str::FromStr for ClaudeModel {
152 type Err = ModelError;
153
154 fn from_str(s: &str) -> Result<Self, Self::Err> {
155 match s.to_lowercase().as_str() {
156 "opus" | "opus-4.5" | "opus45" | "claude-opus-4.5" | "claude-opus-4-5-20251124" => {
158 Ok(Self::Opus45)
159 }
160 "sonnet"
162 | "sonnet-4.5"
163 | "sonnet45"
164 | "claude-sonnet-4.5"
165 | "claude-sonnet-4-5-20250929" => Ok(Self::Sonnet45),
166 "opus-4.1" | "opus41" | "claude-opus-4.1" | "claude-opus-4-1-20250805" => {
168 Ok(Self::Opus41)
169 }
170 "opus-4" | "opus4" | "claude-opus-4" | "claude-opus-4-20250514" => Ok(Self::Opus4),
172 "sonnet-4" | "sonnet4" | "claude-sonnet-4" | "claude-sonnet-4-20250514" => {
174 Ok(Self::Sonnet4)
175 }
176 "haiku"
178 | "haiku-3.5"
179 | "haiku35"
180 | "claude-haiku-3.5"
181 | "claude-3-5-haiku-20241022" => Ok(Self::Haiku35),
182 _ => {
184 Self::validate_custom(s)?;
185 Ok(Self::Custom(s.to_string()))
186 }
187 }
188 }
189}
190
191impl fmt::Display for ClaudeModel {
192 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
193 write!(f, "{}", self.as_api_id())
194 }
195}
196
197#[derive(Debug, Clone, PartialEq, Eq, Hash)]
215pub enum GeminiModel {
216 Flash3,
218 Pro3,
220 Flash25,
222 Pro25,
224 Flash20,
226 Custom(String),
228}
229
230impl Default for GeminiModel {
231 fn default() -> Self {
232 Self::Flash25
233 }
234}
235
236impl GeminiModel {
237 pub fn as_api_id(&self) -> &str {
239 match self {
240 Self::Flash3 => "gemini-3-flash",
241 Self::Pro3 => "gemini-3-pro",
242 Self::Flash25 => "gemini-2.5-flash",
243 Self::Pro25 => "gemini-2.5-pro",
244 Self::Flash20 => "gemini-2.0-flash",
245 Self::Custom(s) => s,
246 }
247 }
248
249 pub fn as_cli_name(&self) -> &str {
251 match self {
252 Self::Flash3 => "flash-3",
253 Self::Pro3 => "pro-3",
254 Self::Flash25 => "flash",
255 Self::Pro25 => "pro",
256 Self::Flash20 => "flash-2.0",
257 Self::Custom(s) => s,
258 }
259 }
260
261 fn validate_custom(s: &str) -> Result<(), ModelError> {
262 if s.starts_with("gemini-") {
263 Ok(())
264 } else {
265 Err(ModelError::InvalidPrefix {
266 model: s.to_string(),
267 expected_prefixes: &["gemini-"],
268 })
269 }
270 }
271}
272
273impl std::str::FromStr for GeminiModel {
274 type Err = ModelError;
275
276 fn from_str(s: &str) -> Result<Self, Self::Err> {
277 match s.to_lowercase().as_str() {
278 "flash-3" | "flash3" | "gemini-3-flash" => Ok(Self::Flash3),
280 "pro-3" | "pro3" | "gemini-3-pro" => Ok(Self::Pro3),
282 "flash" | "flash-2.5" | "flash25" | "gemini-2.5-flash" => Ok(Self::Flash25),
284 "pro" | "pro-2.5" | "pro25" | "gemini-2.5-pro" => Ok(Self::Pro25),
286 "flash-2.0" | "flash20" | "flash-2" | "gemini-2.0-flash" => Ok(Self::Flash20),
288 _ => {
290 Self::validate_custom(s)?;
291 Ok(Self::Custom(s.to_string()))
292 }
293 }
294 }
295}
296
297impl fmt::Display for GeminiModel {
298 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
299 write!(f, "{}", self.as_api_id())
300 }
301}
302
303#[derive(Debug, Clone, PartialEq, Eq, Hash)]
321pub enum OpenAIModel {
322 Gpt52,
325 Gpt51,
327 Gpt5,
329 Gpt5Mini,
331 Gpt5Nano,
333
334 Gpt51Codex,
337 Gpt51CodexMini,
339 Gpt5Codex,
341 Gpt5CodexMini,
343
344 Gpt41,
347 Gpt41Mini,
349 Gpt4o,
351 Gpt4oMini,
353
354 O3Pro,
357 O3,
359 O3Mini,
361 O1,
363 O1Pro,
365
366 Custom(String),
368}
369
370impl Default for OpenAIModel {
371 fn default() -> Self {
372 Self::Gpt4o }
374}
375
376impl OpenAIModel {
377 pub fn as_api_id(&self) -> &str {
379 match self {
380 Self::Gpt52 => "gpt-5.2",
382 Self::Gpt51 => "gpt-5.1",
383 Self::Gpt5 => "gpt-5",
384 Self::Gpt5Mini => "gpt-5-mini",
385 Self::Gpt5Nano => "gpt-5-nano",
386 Self::Gpt51Codex => "gpt-5.1-codex",
388 Self::Gpt51CodexMini => "gpt-5.1-codex-mini",
389 Self::Gpt5Codex => "gpt-5-codex",
390 Self::Gpt5CodexMini => "gpt-5-codex-mini",
391 Self::Gpt41 => "gpt-4.1",
393 Self::Gpt41Mini => "gpt-4.1-mini",
394 Self::Gpt4o => "gpt-4o",
395 Self::Gpt4oMini => "gpt-4o-mini",
396 Self::O3Pro => "o3-pro",
398 Self::O3 => "o3",
399 Self::O3Mini => "o3-mini",
400 Self::O1 => "o1",
401 Self::O1Pro => "o1-pro",
402 Self::Custom(s) => s,
404 }
405 }
406
407 pub fn as_cli_name(&self) -> &str {
409 self.as_api_id() }
411
412 fn validate_custom(s: &str) -> Result<(), ModelError> {
413 const VALID_PREFIXES: &[&str] = &["gpt-", "o1-", "o3-", "o4-"];
414 if VALID_PREFIXES.iter().any(|p| s.starts_with(p)) {
415 Ok(())
416 } else {
417 Err(ModelError::InvalidPrefix {
418 model: s.to_string(),
419 expected_prefixes: VALID_PREFIXES,
420 })
421 }
422 }
423}
424
425impl std::str::FromStr for OpenAIModel {
426 type Err = ModelError;
427
428 fn from_str(s: &str) -> Result<Self, Self::Err> {
429 match s.to_lowercase().as_str() {
430 "5.2" | "gpt-5.2" | "gpt52" => Ok(Self::Gpt52),
432 "5.1" | "gpt-5.1" | "gpt51" => Ok(Self::Gpt51),
434 "5" | "gpt-5" | "gpt5" => Ok(Self::Gpt5),
436 "5-mini" | "gpt-5-mini" => Ok(Self::Gpt5Mini),
438 "5-nano" | "gpt-5-nano" => Ok(Self::Gpt5Nano),
439 "5.1-codex" | "gpt-5.1-codex" | "codex" => Ok(Self::Gpt51Codex),
441 "5.1-codex-mini" | "gpt-5.1-codex-mini" | "codex-mini" => Ok(Self::Gpt51CodexMini),
442 "5-codex" | "gpt-5-codex" => Ok(Self::Gpt5Codex),
444 "5-codex-mini" | "gpt-5-codex-mini" => Ok(Self::Gpt5CodexMini),
445 "4.1" | "gpt-4.1" | "gpt41" => Ok(Self::Gpt41),
447 "4.1-mini" | "gpt-4.1-mini" => Ok(Self::Gpt41Mini),
448 "4o" | "gpt-4o" => Ok(Self::Gpt4o),
450 "4o-mini" | "gpt-4o-mini" => Ok(Self::Gpt4oMini),
451 "o3-pro" => Ok(Self::O3Pro),
453 "o3" => Ok(Self::O3),
454 "o3-mini" => Ok(Self::O3Mini),
455 "o1" => Ok(Self::O1),
456 "o1-pro" => Ok(Self::O1Pro),
457 _ => {
459 Self::validate_custom(s)?;
460 Ok(Self::Custom(s.to_string()))
461 }
462 }
463 }
464}
465
466impl fmt::Display for OpenAIModel {
467 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
468 write!(f, "{}", self.as_api_id())
469 }
470}
471
472#[derive(Debug, Clone, PartialEq, Eq, Hash)]
480pub enum Model {
481 Claude(ClaudeModel),
482 Gemini(GeminiModel),
483 OpenAI(OpenAIModel),
484}
485
486impl Model {
487 pub fn as_api_id(&self) -> &str {
489 match self {
490 Self::Claude(m) => m.as_api_id(),
491 Self::Gemini(m) => m.as_api_id(),
492 Self::OpenAI(m) => m.as_api_id(),
493 }
494 }
495
496 pub fn as_cli_name(&self) -> &str {
498 match self {
499 Self::Claude(m) => m.as_cli_name(),
500 Self::Gemini(m) => m.as_cli_name(),
501 Self::OpenAI(m) => m.as_cli_name(),
502 }
503 }
504}
505
506impl From<ClaudeModel> for Model {
507 fn from(m: ClaudeModel) -> Self {
508 Self::Claude(m)
509 }
510}
511
512impl From<GeminiModel> for Model {
513 fn from(m: GeminiModel) -> Self {
514 Self::Gemini(m)
515 }
516}
517
518impl From<OpenAIModel> for Model {
519 fn from(m: OpenAIModel) -> Self {
520 Self::OpenAI(m)
521 }
522}
523
524impl fmt::Display for Model {
525 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
526 write!(f, "{}", self.as_api_id())
527 }
528}
529
530#[cfg(test)]
535mod tests {
536 use super::*;
537
538 mod claude_model {
539 use super::*;
540
541 #[test]
542 fn test_default() {
543 assert_eq!(ClaudeModel::default(), ClaudeModel::Sonnet45);
544 }
545
546 #[test]
547 fn test_api_id() {
548 assert_eq!(ClaudeModel::Opus45.as_api_id(), "claude-opus-4-5-20251124");
549 assert_eq!(
550 ClaudeModel::Sonnet45.as_api_id(),
551 "claude-sonnet-4-5-20250929"
552 );
553 assert_eq!(
554 ClaudeModel::Haiku35.as_api_id(),
555 "claude-3-5-haiku-20241022"
556 );
557 }
558
559 #[test]
560 fn test_cli_name() {
561 assert_eq!(ClaudeModel::Opus45.as_cli_name(), "claude-opus-4.5");
562 assert_eq!(ClaudeModel::Sonnet4.as_cli_name(), "claude-sonnet-4");
563 }
564
565 #[test]
566 fn test_parse_shorthand() {
567 assert_eq!("opus".parse::<ClaudeModel>().unwrap(), ClaudeModel::Opus45);
568 assert_eq!(
569 "sonnet".parse::<ClaudeModel>().unwrap(),
570 ClaudeModel::Sonnet45
571 );
572 assert_eq!(
573 "haiku".parse::<ClaudeModel>().unwrap(),
574 ClaudeModel::Haiku35
575 );
576 }
577
578 #[test]
579 fn test_parse_full_name() {
580 assert_eq!(
581 "claude-opus-4-5-20251124".parse::<ClaudeModel>().unwrap(),
582 ClaudeModel::Opus45
583 );
584 assert_eq!(
585 "claude-sonnet-4".parse::<ClaudeModel>().unwrap(),
586 ClaudeModel::Sonnet4
587 );
588 }
589
590 #[test]
591 fn test_parse_custom_valid() {
592 let model: ClaudeModel = "claude-future-model-2026".parse().unwrap();
593 assert_eq!(
594 model,
595 ClaudeModel::Custom("claude-future-model-2026".to_string())
596 );
597 }
598
599 #[test]
600 fn test_parse_custom_invalid() {
601 let result: Result<ClaudeModel, _> = "gpt-4o".parse();
602 assert!(result.is_err());
603 }
604 }
605
606 mod gemini_model {
607 use super::*;
608
609 #[test]
610 fn test_default() {
611 assert_eq!(GeminiModel::default(), GeminiModel::Flash25);
612 }
613
614 #[test]
615 fn test_parse() {
616 assert_eq!(
617 "flash".parse::<GeminiModel>().unwrap(),
618 GeminiModel::Flash25
619 );
620 assert_eq!("pro".parse::<GeminiModel>().unwrap(), GeminiModel::Pro25);
621 assert_eq!(
622 "flash-3".parse::<GeminiModel>().unwrap(),
623 GeminiModel::Flash3
624 );
625 }
626
627 #[test]
628 fn test_custom_invalid() {
629 let result: Result<GeminiModel, _> = "claude-opus".parse();
630 assert!(result.is_err());
631 }
632 }
633
634 mod openai_model {
635 use super::*;
636
637 #[test]
638 fn test_default() {
639 assert_eq!(OpenAIModel::default(), OpenAIModel::Gpt4o);
640 }
641
642 #[test]
643 fn test_parse() {
644 assert_eq!("4o".parse::<OpenAIModel>().unwrap(), OpenAIModel::Gpt4o);
645 assert_eq!(
646 "gpt-5.2".parse::<OpenAIModel>().unwrap(),
647 OpenAIModel::Gpt52
648 );
649 assert_eq!("o3".parse::<OpenAIModel>().unwrap(), OpenAIModel::O3);
650 assert_eq!(
651 "codex".parse::<OpenAIModel>().unwrap(),
652 OpenAIModel::Gpt51Codex
653 );
654 }
655
656 #[test]
657 fn test_o_series_validation() {
658 let model: OpenAIModel = "o4-mini".parse().unwrap();
660 assert_eq!(model, OpenAIModel::Custom("o4-mini".to_string()));
661 }
662
663 #[test]
664 fn test_custom_invalid() {
665 let result: Result<OpenAIModel, _> = "gemini-pro".parse();
666 assert!(result.is_err());
667 }
668 }
669}