ZIm/crates/language_model/src/model/cloud_model.rs

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 crate::{LanguageModelAvailability, LanguageModelToolSchemaFormat};

#[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 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
                | anthropic::Model::Claude3_7SonnetThinking => {
                    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::FourPointOne
                | open_ai::Model::FourPointOneMini
                | open_ai::Model::FourPointOneNano
                | 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::Gemini25ProExp0325
                | google_ai::Model::Gemini25ProPreview0325
                | google_ai::Model::Custom { .. } => {
                    LanguageModelAvailability::RequiresPlan(Plan::ZedPro)
                }
            },
        }
    }

    pub fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
        match self {
            Self::Anthropic(_) | Self::OpenAi(_) => LanguageModelToolSchemaFormat::JsonSchema,
            Self::Google(_) => LanguageModelToolSchemaFormat::JsonSchemaSubset,
        }
    }
}

#[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))
    }
}