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::Custom { .. } => {
72                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
73                }
74            },
75            Self::OpenAi(model) => match model {
76                open_ai::Model::ThreePointFiveTurbo
77                | open_ai::Model::Four
78                | open_ai::Model::FourTurbo
79                | open_ai::Model::FourOmni
80                | open_ai::Model::FourOmniMini
81                | open_ai::Model::O1Mini
82                | open_ai::Model::O1Preview
83                | open_ai::Model::Custom { .. } => {
84                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
85                }
86            },
87            Self::Google(model) => match model {
88                google_ai::Model::Gemini15Pro
89                | google_ai::Model::Gemini15Flash
90                | google_ai::Model::Custom { .. } => {
91                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
92                }
93            },
94        }
95    }
96}