use serde::{Deserialize, Serialize, Serializer};
use std::fmt;
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum Model {
Known(KnownModel),
Custom(String),
}
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
pub enum KnownModel {
Claude37SonnetLatest,
Claude37Sonnet20250219,
Claude35HaikuLatest,
Claude35Haiku20241022,
ClaudeSonnet420250514,
ClaudeSonnet40,
Claude4Sonnet20250514,
Claude35SonnetLatest,
Claude35Sonnet20241022,
Claude35Sonnet20240620,
ClaudeOpus40,
ClaudeOpus420250514,
Claude4Opus20250514,
Claude3OpusLatest,
Claude3Opus20240229,
Claude3Sonnet20240229,
Claude3Haiku20240307,
Claude21,
Claude20,
}
impl fmt::Display for Model {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Model::Known(known_model) => write!(f, "{}", known_model),
Model::Custom(custom) => write!(f, "{}", custom),
}
}
}
impl fmt::Display for KnownModel {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
KnownModel::Claude37SonnetLatest => write!(f, "claude-3-7-sonnet-latest"),
KnownModel::Claude37Sonnet20250219 => write!(f, "claude-3-7-sonnet-20250219"),
KnownModel::Claude35HaikuLatest => write!(f, "claude-3-5-haiku-latest"),
KnownModel::Claude35Haiku20241022 => write!(f, "claude-3-5-haiku-20241022"),
KnownModel::ClaudeSonnet420250514 => write!(f, "claude-sonnet-4-20250514"),
KnownModel::ClaudeSonnet40 => write!(f, "claude-sonnet-4-0"),
KnownModel::Claude4Sonnet20250514 => write!(f, "claude-4-sonnet-20250514"),
KnownModel::Claude35SonnetLatest => write!(f, "claude-3-5-sonnet-latest"),
KnownModel::Claude35Sonnet20241022 => write!(f, "claude-3-5-sonnet-20241022"),
KnownModel::Claude35Sonnet20240620 => write!(f, "claude-3-5-sonnet-20240620"),
KnownModel::ClaudeOpus40 => write!(f, "claude-opus-4-0"),
KnownModel::ClaudeOpus420250514 => write!(f, "claude-opus-4-20250514"),
KnownModel::Claude4Opus20250514 => write!(f, "claude-4-opus-20250514"),
KnownModel::Claude3OpusLatest => write!(f, "claude-3-opus-latest"),
KnownModel::Claude3Opus20240229 => write!(f, "claude-3-opus-20240229"),
KnownModel::Claude3Sonnet20240229 => write!(f, "claude-3-sonnet-20240229"),
KnownModel::Claude3Haiku20240307 => write!(f, "claude-3-haiku-20240307"),
KnownModel::Claude21 => write!(f, "claude-2.1"),
KnownModel::Claude20 => write!(f, "claude-2.0"),
}
}
}
impl Serialize for Model {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(&self.to_string())
}
}
impl<'de> Deserialize<'de> for Model {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
let s = String::deserialize(deserializer)?;
match s.as_str() {
"claude-3-7-sonnet-latest" => Ok(Model::Known(KnownModel::Claude37SonnetLatest)),
"claude-3-7-sonnet-20250219" => Ok(Model::Known(KnownModel::Claude37Sonnet20250219)),
"claude-3-5-haiku-latest" => Ok(Model::Known(KnownModel::Claude35HaikuLatest)),
"claude-3-5-haiku-20241022" => Ok(Model::Known(KnownModel::Claude35Haiku20241022)),
"claude-sonnet-4-20250514" => Ok(Model::Known(KnownModel::ClaudeSonnet420250514)),
"claude-sonnet-4-0" => Ok(Model::Known(KnownModel::ClaudeSonnet40)),
"claude-4-sonnet-20250514" => Ok(Model::Known(KnownModel::Claude4Sonnet20250514)),
"claude-3-5-sonnet-latest" => Ok(Model::Known(KnownModel::Claude35SonnetLatest)),
"claude-3-5-sonnet-20241022" => Ok(Model::Known(KnownModel::Claude35Sonnet20241022)),
"claude-3-5-sonnet-20240620" => Ok(Model::Known(KnownModel::Claude35Sonnet20240620)),
"claude-opus-4-0" => Ok(Model::Known(KnownModel::ClaudeOpus40)),
"claude-opus-4-20250514" => Ok(Model::Known(KnownModel::ClaudeOpus420250514)),
"claude-4-opus-20250514" => Ok(Model::Known(KnownModel::Claude4Opus20250514)),
"claude-3-opus-latest" => Ok(Model::Known(KnownModel::Claude3OpusLatest)),
"claude-3-opus-20240229" => Ok(Model::Known(KnownModel::Claude3Opus20240229)),
"claude-3-sonnet-20240229" => Ok(Model::Known(KnownModel::Claude3Sonnet20240229)),
"claude-3-haiku-20240307" => Ok(Model::Known(KnownModel::Claude3Haiku20240307)),
"claude-2.1" => Ok(Model::Known(KnownModel::Claude21)),
"claude-2.0" => Ok(Model::Known(KnownModel::Claude20)),
_ => Ok(Model::Custom(s)),
}
}
}
impl From<KnownModel> for Model {
fn from(model: KnownModel) -> Self {
Model::Known(model)
}
}
impl From<String> for Model {
fn from(model: String) -> Self {
Model::Custom(model)
}
}
impl From<&str> for Model {
fn from(model: &str) -> Self {
Model::Custom(model.to_string())
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn known_model_serialization() {
let model = Model::Known(KnownModel::Claude37SonnetLatest);
let json = serde_json::to_string(&model).unwrap();
assert_eq!(json, r#""claude-3-7-sonnet-latest""#);
let model = Model::Known(KnownModel::Claude35Sonnet20240620);
let json = serde_json::to_string(&model).unwrap();
assert_eq!(json, r#""claude-3-5-sonnet-20240620""#);
}
#[test]
fn custom_model_serialization() {
let model = Model::Custom("claude-4-custom".to_string());
let json = serde_json::to_string(&model).unwrap();
assert_eq!(json, r#""claude-4-custom""#);
}
#[test]
fn model_deserialization() {
let json = r#""claude-3-7-sonnet-latest""#;
let model: Model = serde_json::from_str(json).unwrap();
assert_eq!(model, Model::Known(KnownModel::Claude37SonnetLatest));
let json = r#""claude-4-custom""#;
let model: Model = serde_json::from_str(json).unwrap();
assert_eq!(model, Model::Custom("claude-4-custom".to_string()));
}
#[test]
fn claude_4_models() {
let model = Model::Known(KnownModel::ClaudeSonnet420250514);
assert_eq!(model.to_string(), "claude-sonnet-4-20250514");
let json = serde_json::to_string(&model).unwrap();
assert_eq!(json, r#""claude-sonnet-4-20250514""#);
let model = Model::Known(KnownModel::Claude4Sonnet20250514);
assert_eq!(model.to_string(), "claude-4-sonnet-20250514");
let model = Model::Known(KnownModel::ClaudeOpus40);
assert_eq!(model.to_string(), "claude-opus-4-0");
let model = Model::Known(KnownModel::Claude4Opus20250514);
assert_eq!(model.to_string(), "claude-4-opus-20250514");
let json = r#""claude-sonnet-4-20250514""#;
let model: Model = serde_json::from_str(json).unwrap();
assert_eq!(model, Model::Known(KnownModel::ClaudeSonnet420250514));
let json = r#""claude-4-opus-20250514""#;
let model: Model = serde_json::from_str(json).unwrap();
assert_eq!(model, Model::Known(KnownModel::Claude4Opus20250514));
}
#[test]
fn display() {
let model = Model::Known(KnownModel::Claude37SonnetLatest);
assert_eq!(model.to_string(), "claude-3-7-sonnet-latest");
let model = Model::Custom("claude-4-custom".to_string());
assert_eq!(model.to_string(), "claude-4-custom");
}
}