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::Custom { .. } => {
86                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
87                }
88            },
89            Self::Google(model) => match model {
90                google_ai::Model::Gemini15Pro
91                | google_ai::Model::Gemini15Flash
92                | google_ai::Model::Gemini20Flash
93                | google_ai::Model::Custom { .. } => {
94                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
95                }
96            },
97        }
98    }
99}