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