1use crate::LanguageModelRequest;
  2pub use anthropic::Model as AnthropicModel;
  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 id(&self) -> &str {
 43        match self {
 44            Self::Gpt3Point5Turbo => "gpt-3.5-turbo",
 45            Self::Gpt4 => "gpt-4",
 46            Self::Gpt4Turbo => "gpt-4-turbo-preview",
 47            Self::Gpt4Omni => "gpt-4o",
 48            Self::Gpt4OmniMini => "gpt-4o-mini",
 49            Self::Claude3_5Sonnet => "claude-3-5-sonnet",
 50            Self::Claude3Opus => "claude-3-opus",
 51            Self::Claude3Sonnet => "claude-3-sonnet",
 52            Self::Claude3Haiku => "claude-3-haiku",
 53            Self::Gemini15Pro => "gemini-1.5-pro",
 54            Self::Gemini15Flash => "gemini-1.5-flash",
 55            Self::Custom { name, .. } => name,
 56        }
 57    }
 58
 59    pub fn display_name(&self) -> &str {
 60        match self {
 61            Self::Gpt3Point5Turbo => "GPT 3.5 Turbo",
 62            Self::Gpt4 => "GPT 4",
 63            Self::Gpt4Turbo => "GPT 4 Turbo",
 64            Self::Gpt4Omni => "GPT 4 Omni",
 65            Self::Gpt4OmniMini => "GPT 4 Omni 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 max_token_count(&self) -> usize {
 77        match self {
 78            Self::Gpt3Point5Turbo => 2048,
 79            Self::Gpt4 => 4096,
 80            Self::Gpt4Turbo | Self::Gpt4Omni => 128000,
 81            Self::Gpt4OmniMini => 128000,
 82            Self::Claude3_5Sonnet
 83            | Self::Claude3Opus
 84            | Self::Claude3Sonnet
 85            | Self::Claude3Haiku => 200000,
 86            Self::Gemini15Pro => 128000,
 87            Self::Gemini15Flash => 32000,
 88            Self::Custom { max_tokens, .. } => max_tokens.unwrap_or(200_000),
 89        }
 90    }
 91
 92    pub fn preprocess_request(&self, request: &mut LanguageModelRequest) {
 93        match self {
 94            Self::Claude3Opus
 95            | Self::Claude3Sonnet
 96            | Self::Claude3Haiku
 97            | Self::Claude3_5Sonnet => {
 98                request.preprocess_anthropic();
 99            }
100            Self::Custom { name, .. } if name.starts_with("anthropic/") => {
101                request.preprocess_anthropic();
102            }
103            _ => {}
104        }
105    }
106}