Piotr OsiewiczRust 1.85 (#25272)e4e758d1 year agoHistory
202 lines · 6.2 KB · Rust
RawDownload
use std::fmt;use std::sync::Arc; use anyhow::Result;use client::Client;use gpui::{    App, AppContext as _, AsyncApp, Context, Entity, EventEmitter, Global, ReadGlobal as _,};use proto::{Plan, TypedEnvelope};use schemars::JsonSchema;use serde::{Deserialize, Serialize};use smol::lock::{RwLock, RwLockUpgradableReadGuard, RwLockWriteGuard};use strum::EnumIter;use thiserror::Error;use ui::IconName; use crate::LanguageModelAvailability; #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]#[serde(tag = "provider", rename_all = "lowercase")]pub enum CloudModel {    Anthropic(anthropic::Model),    OpenAi(open_ai::Model),    Google(google_ai::Model),} #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema, EnumIter)]pub enum ZedModel {    #[serde(rename = "Qwen/Qwen2-7B-Instruct")]    Qwen2_7bInstruct,} impl Default for CloudModel {    fn default() -> Self {        Self::Anthropic(anthropic::Model::default())    }} impl CloudModel {    pub fn id(&self) -> &str {        match self {            Self::Anthropic(model) => model.id(),            Self::OpenAi(model) => model.id(),            Self::Google(model) => model.id(),        }    }     pub fn display_name(&self) -> &str {        match self {            Self::Anthropic(model) => model.display_name(),            Self::OpenAi(model) => model.display_name(),            Self::Google(model) => model.display_name(),        }    }     pub fn icon(&self) -> Option<IconName> {        match self {            Self::Anthropic(_) => Some(IconName::AiAnthropicHosted),            _ => None,        }    }     pub fn max_token_count(&self) -> usize {        match self {            Self::Anthropic(model) => model.max_token_count(),            Self::OpenAi(model) => model.max_token_count(),            Self::Google(model) => model.max_token_count(),        }    }     /// Returns the availability of this model.    pub fn availability(&self) -> LanguageModelAvailability {        match self {            Self::Anthropic(model) => match model {                anthropic::Model::Claude3_5Sonnet | anthropic::Model::Claude3_7Sonnet => {                    LanguageModelAvailability::RequiresPlan(Plan::Free)                }                anthropic::Model::Claude3Opus                | anthropic::Model::Claude3Sonnet                | anthropic::Model::Claude3Haiku                | anthropic::Model::Claude3_5Haiku                | anthropic::Model::Custom { .. } => {                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)                }            },            Self::OpenAi(model) => match model {                open_ai::Model::ThreePointFiveTurbo                | open_ai::Model::Four                | open_ai::Model::FourTurbo                | open_ai::Model::FourOmni                | open_ai::Model::FourOmniMini                | open_ai::Model::O1Mini                | open_ai::Model::O1Preview                | open_ai::Model::O1                | open_ai::Model::O3Mini                | open_ai::Model::Custom { .. } => {                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)                }            },            Self::Google(model) => match model {                google_ai::Model::Gemini15Pro                | google_ai::Model::Gemini15Flash                | google_ai::Model::Gemini20Pro                | google_ai::Model::Gemini20Flash                | google_ai::Model::Gemini20FlashThinking                | google_ai::Model::Gemini20FlashLite                | google_ai::Model::Custom { .. } => {                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)                }            },        }    }} #[derive(Error, Debug)]pub struct PaymentRequiredError; impl fmt::Display for PaymentRequiredError {    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {        write!(            f,            "Payment required to use this language model. Please upgrade your account."        )    }} #[derive(Error, Debug)]pub struct MaxMonthlySpendReachedError; impl fmt::Display for MaxMonthlySpendReachedError {    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {        write!(            f,            "Maximum spending limit reached for this month. For more usage, increase your spending limit."        )    }} #[derive(Clone, Default)]pub struct LlmApiToken(Arc<RwLock<Option<String>>>); impl LlmApiToken {    pub async fn acquire(&self, client: &Arc<Client>) -> Result<String> {        let lock = self.0.upgradable_read().await;        if let Some(token) = lock.as_ref() {            Ok(token.to_string())        } else {            Self::fetch(RwLockUpgradableReadGuard::upgrade(lock).await, client).await        }    }     pub async fn refresh(&self, client: &Arc<Client>) -> Result<String> {        Self::fetch(self.0.write().await, client).await    }     async fn fetch(        mut lock: RwLockWriteGuard<'_, Option<String>>,        client: &Arc<Client>,    ) -> Result<String> {        let response = client.request(proto::GetLlmToken {}).await?;        *lock = Some(response.token.clone());        Ok(response.token.clone())    }} struct GlobalRefreshLlmTokenListener(Entity<RefreshLlmTokenListener>); impl Global for GlobalRefreshLlmTokenListener {} pub struct RefreshLlmTokenEvent; pub struct RefreshLlmTokenListener {    _llm_token_subscription: client::Subscription,} impl EventEmitter<RefreshLlmTokenEvent> for RefreshLlmTokenListener {} impl RefreshLlmTokenListener {    pub fn register(client: Arc<Client>, cx: &mut App) {        let listener = cx.new(|cx| RefreshLlmTokenListener::new(client, cx));        cx.set_global(GlobalRefreshLlmTokenListener(listener));    }     pub fn global(cx: &App) -> Entity<Self> {        GlobalRefreshLlmTokenListener::global(cx).0.clone()    }     fn new(client: Arc<Client>, cx: &mut Context<Self>) -> Self {        Self {            _llm_token_subscription: client                .add_message_handler(cx.weak_entity(), Self::handle_refresh_llm_token),        }    }     async fn handle_refresh_llm_token(        this: Entity<Self>,        _: TypedEnvelope<proto::RefreshLlmToken>,        mut cx: AsyncApp,    ) -> Result<()> {        this.update(&mut cx, |_this, cx| cx.emit(RefreshLlmTokenEvent))    }}