1pub use anthropic::Model as AnthropicModel;
  2use anyhow::{anyhow, Result};
  3pub use ollama::Model as OllamaModel;
  4pub use open_ai::Model as OpenAiModel;
  5use schemars::JsonSchema;
  6use serde::{Deserialize, Serialize};
  7use strum::EnumIter;
  8
  9#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize, JsonSchema, EnumIter)]
 10pub enum CloudModel {
 11    #[serde(rename = "gpt-3.5-turbo")]
 12    Gpt3Point5Turbo,
 13    #[serde(rename = "gpt-4")]
 14    Gpt4,
 15    #[serde(rename = "gpt-4-turbo-preview")]
 16    Gpt4Turbo,
 17    #[serde(rename = "gpt-4o")]
 18    #[default]
 19    Gpt4Omni,
 20    #[serde(rename = "gpt-4o-mini")]
 21    Gpt4OmniMini,
 22    #[serde(rename = "claude-3-5-sonnet")]
 23    Claude3_5Sonnet,
 24    #[serde(rename = "claude-3-opus")]
 25    Claude3Opus,
 26    #[serde(rename = "claude-3-sonnet")]
 27    Claude3Sonnet,
 28    #[serde(rename = "claude-3-haiku")]
 29    Claude3Haiku,
 30    #[serde(rename = "gemini-1.5-pro")]
 31    Gemini15Pro,
 32    #[serde(rename = "gemini-1.5-flash")]
 33    Gemini15Flash,
 34    #[serde(rename = "custom")]
 35    Custom {
 36        name: String,
 37        max_tokens: Option<usize>,
 38    },
 39}
 40
 41impl CloudModel {
 42    pub fn from_id(value: &str) -> Result<Self> {
 43        match value {
 44            "gpt-3.5-turbo" => Ok(Self::Gpt3Point5Turbo),
 45            "gpt-4" => Ok(Self::Gpt4),
 46            "gpt-4-turbo-preview" => Ok(Self::Gpt4Turbo),
 47            "gpt-4o" => Ok(Self::Gpt4Omni),
 48            "gpt-4o-mini" => Ok(Self::Gpt4OmniMini),
 49            "claude-3-5-sonnet" => Ok(Self::Claude3_5Sonnet),
 50            "claude-3-opus" => Ok(Self::Claude3Opus),
 51            "claude-3-sonnet" => Ok(Self::Claude3Sonnet),
 52            "claude-3-haiku" => Ok(Self::Claude3Haiku),
 53            "gemini-1.5-pro" => Ok(Self::Gemini15Pro),
 54            "gemini-1.5-flash" => Ok(Self::Gemini15Flash),
 55            _ => Err(anyhow!("invalid model id")),
 56        }
 57    }
 58
 59    pub fn id(&self) -> &str {
 60        match self {
 61            Self::Gpt3Point5Turbo => "gpt-3.5-turbo",
 62            Self::Gpt4 => "gpt-4",
 63            Self::Gpt4Turbo => "gpt-4-turbo-preview",
 64            Self::Gpt4Omni => "gpt-4o",
 65            Self::Gpt4OmniMini => "gpt-4o-mini",
 66            Self::Claude3_5Sonnet => "claude-3-5-sonnet",
 67            Self::Claude3Opus => "claude-3-opus",
 68            Self::Claude3Sonnet => "claude-3-sonnet",
 69            Self::Claude3Haiku => "claude-3-haiku",
 70            Self::Gemini15Pro => "gemini-1.5-pro",
 71            Self::Gemini15Flash => "gemini-1.5-flash",
 72            Self::Custom { name, .. } => name,
 73        }
 74    }
 75
 76    pub fn display_name(&self) -> &str {
 77        match self {
 78            Self::Gpt3Point5Turbo => "GPT 3.5 Turbo",
 79            Self::Gpt4 => "GPT 4",
 80            Self::Gpt4Turbo => "GPT 4 Turbo",
 81            Self::Gpt4Omni => "GPT 4 Omni",
 82            Self::Gpt4OmniMini => "GPT 4 Omni Mini",
 83            Self::Claude3_5Sonnet => "Claude 3.5 Sonnet",
 84            Self::Claude3Opus => "Claude 3 Opus",
 85            Self::Claude3Sonnet => "Claude 3 Sonnet",
 86            Self::Claude3Haiku => "Claude 3 Haiku",
 87            Self::Gemini15Pro => "Gemini 1.5 Pro",
 88            Self::Gemini15Flash => "Gemini 1.5 Flash",
 89            Self::Custom { name, .. } => name,
 90        }
 91    }
 92
 93    pub fn max_token_count(&self) -> usize {
 94        match self {
 95            Self::Gpt3Point5Turbo => 2048,
 96            Self::Gpt4 => 4096,
 97            Self::Gpt4Turbo | Self::Gpt4Omni => 128000,
 98            Self::Gpt4OmniMini => 128000,
 99            Self::Claude3_5Sonnet
100            | Self::Claude3Opus
101            | Self::Claude3Sonnet
102            | Self::Claude3Haiku => 200000,
103            Self::Gemini15Pro => 128000,
104            Self::Gemini15Flash => 32000,
105            Self::Custom { max_tokens, .. } => max_tokens.unwrap_or(200_000),
106        }
107    }
108}