1use proto::Plan;
  2use schemars::JsonSchema;
  3use serde::{Deserialize, Serialize};
  4use strum::EnumIter;
  5use ui::IconName;
  6
  7use crate::LanguageModelAvailability;
  8
  9#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
 10#[serde(tag = "provider", rename_all = "lowercase")]
 11pub enum CloudModel {
 12    Anthropic(anthropic::Model),
 13    OpenAi(open_ai::Model),
 14    Google(google_ai::Model),
 15}
 16
 17#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema, EnumIter)]
 18pub enum ZedModel {
 19    #[serde(rename = "Qwen/Qwen2-7B-Instruct")]
 20    Qwen2_7bInstruct,
 21}
 22
 23impl Default for CloudModel {
 24    fn default() -> Self {
 25        Self::Anthropic(anthropic::Model::default())
 26    }
 27}
 28
 29impl CloudModel {
 30    pub fn id(&self) -> &str {
 31        match self {
 32            Self::Anthropic(model) => model.id(),
 33            Self::OpenAi(model) => model.id(),
 34            Self::Google(model) => model.id(),
 35        }
 36    }
 37
 38    pub fn display_name(&self) -> &str {
 39        match self {
 40            Self::Anthropic(model) => model.display_name(),
 41            Self::OpenAi(model) => model.display_name(),
 42            Self::Google(model) => model.display_name(),
 43        }
 44    }
 45
 46    pub fn icon(&self) -> Option<IconName> {
 47        match self {
 48            Self::Anthropic(_) => Some(IconName::AiAnthropicHosted),
 49            _ => None,
 50        }
 51    }
 52
 53    pub fn max_token_count(&self) -> usize {
 54        match self {
 55            Self::Anthropic(model) => model.max_token_count(),
 56            Self::OpenAi(model) => model.max_token_count(),
 57            Self::Google(model) => model.max_token_count(),
 58        }
 59    }
 60
 61    /// Returns the availability of this model.
 62    pub fn availability(&self) -> LanguageModelAvailability {
 63        match self {
 64            Self::Anthropic(model) => match model {
 65                anthropic::Model::Claude3_5Sonnet => {
 66                    LanguageModelAvailability::RequiresPlan(Plan::Free)
 67                }
 68                anthropic::Model::Claude3Opus
 69                | anthropic::Model::Claude3Sonnet
 70                | anthropic::Model::Claude3Haiku
 71                | anthropic::Model::Claude3_5Haiku
 72                | anthropic::Model::Custom { .. } => {
 73                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
 74                }
 75            },
 76            Self::OpenAi(model) => match model {
 77                open_ai::Model::ThreePointFiveTurbo
 78                | open_ai::Model::Four
 79                | open_ai::Model::FourTurbo
 80                | open_ai::Model::FourOmni
 81                | open_ai::Model::FourOmniMini
 82                | open_ai::Model::O1Mini
 83                | open_ai::Model::O1Preview
 84                | open_ai::Model::O1
 85                | open_ai::Model::O3Mini
 86                | open_ai::Model::Custom { .. } => {
 87                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
 88                }
 89            },
 90            Self::Google(model) => match model {
 91                google_ai::Model::Gemini15Pro
 92                | google_ai::Model::Gemini15Flash
 93                | google_ai::Model::Gemini20Pro
 94                | google_ai::Model::Gemini20Flash
 95                | google_ai::Model::Gemini20FlashThinking
 96                | google_ai::Model::Gemini20FlashLite
 97                | google_ai::Model::Custom { .. } => {
 98                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
 99                }
100            },
101        }
102    }
103}