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    Zed(ZedModel),
 16}
 17
 18#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema, EnumIter)]
 19pub enum ZedModel {
 20    #[serde(rename = "qwen2-7b-instruct")]
 21    Qwen2_7bInstruct,
 22}
 23
 24impl ZedModel {
 25    pub fn id(&self) -> &str {
 26        match self {
 27            ZedModel::Qwen2_7bInstruct => "qwen2-7b-instruct",
 28        }
 29    }
 30
 31    pub fn display_name(&self) -> &str {
 32        match self {
 33            ZedModel::Qwen2_7bInstruct => "Qwen2 7B Instruct",
 34        }
 35    }
 36
 37    pub fn max_token_count(&self) -> usize {
 38        match self {
 39            ZedModel::Qwen2_7bInstruct => 28000,
 40        }
 41    }
 42}
 43
 44impl Default for CloudModel {
 45    fn default() -> Self {
 46        Self::Anthropic(anthropic::Model::default())
 47    }
 48}
 49
 50impl CloudModel {
 51    pub fn id(&self) -> &str {
 52        match self {
 53            Self::Anthropic(model) => model.id(),
 54            Self::OpenAi(model) => model.id(),
 55            Self::Google(model) => model.id(),
 56            Self::Zed(model) => model.id(),
 57        }
 58    }
 59
 60    pub fn display_name(&self) -> &str {
 61        match self {
 62            Self::Anthropic(model) => model.display_name(),
 63            Self::OpenAi(model) => model.display_name(),
 64            Self::Google(model) => model.display_name(),
 65            Self::Zed(model) => model.display_name(),
 66        }
 67    }
 68
 69    pub fn icon(&self) -> Option<IconName> {
 70        match self {
 71            Self::Anthropic(_) => Some(IconName::AiAnthropicHosted),
 72            _ => None,
 73        }
 74    }
 75
 76    pub fn max_token_count(&self) -> usize {
 77        match self {
 78            Self::Anthropic(model) => model.max_token_count(),
 79            Self::OpenAi(model) => model.max_token_count(),
 80            Self::Google(model) => model.max_token_count(),
 81            Self::Zed(model) => model.max_token_count(),
 82        }
 83    }
 84
 85    /// Returns the availability of this model.
 86    pub fn availability(&self) -> LanguageModelAvailability {
 87        match self {
 88            Self::Anthropic(model) => match model {
 89                anthropic::Model::Claude3_5Sonnet => {
 90                    LanguageModelAvailability::RequiresPlan(Plan::Free)
 91                }
 92                anthropic::Model::Claude3Opus
 93                | anthropic::Model::Claude3Sonnet
 94                | anthropic::Model::Claude3Haiku
 95                | anthropic::Model::Custom { .. } => {
 96                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
 97                }
 98            },
 99            Self::OpenAi(model) => match model {
100                open_ai::Model::ThreePointFiveTurbo
101                | open_ai::Model::Four
102                | open_ai::Model::FourTurbo
103                | open_ai::Model::FourOmni
104                | open_ai::Model::FourOmniMini
105                | open_ai::Model::Custom { .. } => {
106                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
107                }
108            },
109            Self::Google(model) => match model {
110                google_ai::Model::Gemini15Pro
111                | google_ai::Model::Gemini15Flash
112                | google_ai::Model::Custom { .. } => {
113                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
114                }
115            },
116            Self::Zed(model) => match model {
117                ZedModel::Qwen2_7bInstruct => LanguageModelAvailability::RequiresPlan(Plan::ZedPro),
118            },
119        }
120    }
121}