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