use super::open_ai::count_open_ai_tokens;use anthropic::AnthropicError;use anyhow::{anyhow, Result};use client::{ zed_urls, Client, PerformCompletionParams, UserStore, EXPIRED_LLM_TOKEN_HEADER_NAME, MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME,};use collections::BTreeMap;use feature_flags::{FeatureFlagAppExt, LlmClosedBeta, ZedPro};use futures::{ future::BoxFuture, stream::BoxStream, AsyncBufReadExt, FutureExt, Stream, StreamExt, TryStreamExt as _,};use gpui::{ AnyElement, AnyView, App, AsyncApp, Context, Entity, EventEmitter, Global, ReadGlobal, Subscription, Task,};use http_client::{AsyncBody, HttpClient, Method, Response, StatusCode};use language_model::{ CloudModel, LanguageModel, LanguageModelCacheConfiguration, LanguageModelId, LanguageModelName, LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, LanguageModelProviderTosView, LanguageModelRequest, RateLimiter, ZED_CLOUD_PROVIDER_ID,};use language_model::{ LanguageModelAvailability, LanguageModelCompletionEvent, LanguageModelProvider,};use proto::TypedEnvelope;use schemars::JsonSchema;use serde::{de::DeserializeOwned, Deserialize, Serialize};use serde_json::value::RawValue;use settings::{Settings, SettingsStore};use smol::{ io::{AsyncReadExt, BufReader}, lock::{RwLock, RwLockUpgradableReadGuard, RwLockWriteGuard},};use std::fmt;use std::{ future, sync::{Arc, LazyLock},};use strum::IntoEnumIterator;use thiserror::Error;use ui::{prelude::*, TintColor}; use crate::provider::anthropic::map_to_language_model_completion_events;use crate::AllLanguageModelSettings; use super::anthropic::count_anthropic_tokens; pub const PROVIDER_NAME: &str = "Zed"; const ZED_CLOUD_PROVIDER_ADDITIONAL_MODELS_JSON: Option<&str> = option_env!("ZED_CLOUD_PROVIDER_ADDITIONAL_MODELS_JSON"); fn zed_cloud_provider_additional_models() -> &'static [AvailableModel] { static ADDITIONAL_MODELS: LazyLock<Vec<AvailableModel>> = LazyLock::new(|| { ZED_CLOUD_PROVIDER_ADDITIONAL_MODELS_JSON .map(|json| serde_json::from_str(json).unwrap()) .unwrap_or_default() }); ADDITIONAL_MODELS.as_slice()} #[derive(Default, Clone, Debug, PartialEq)]pub struct ZedDotDevSettings { pub available_models: Vec<AvailableModel>,} #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]#[serde(rename_all = "lowercase")]pub enum AvailableProvider { Anthropic, OpenAi, Google,} #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]pub struct AvailableModel { /// The provider of the language model. pub provider: AvailableProvider, /// The model's name in the provider's API. e.g. claude-3-5-sonnet-20240620 pub name: String, /// The name displayed in the UI, such as in the assistant panel model dropdown menu. pub display_name: Option<String>, /// The size of the context window, indicating the maximum number of tokens the model can process. pub max_tokens: usize, /// The maximum number of output tokens allowed by the model. pub max_output_tokens: Option<u32>, /// The maximum number of completion tokens allowed by the model (o1-* only) pub max_completion_tokens: Option<u32>, /// Override this model with a different Anthropic model for tool calls. pub tool_override: Option<String>, /// Indicates whether this custom model supports caching. pub cache_configuration: Option<LanguageModelCacheConfiguration>, /// The default temperature to use for this model. pub default_temperature: Option<f32>, /// Any extra beta headers to provide when using the model. #[serde(default)] pub extra_beta_headers: Vec<String>,} 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)) }} pub struct CloudLanguageModelProvider { client: Arc<Client>, state: gpui::Entity<State>, _maintain_client_status: Task<()>,} pub struct State { client: Arc<Client>, llm_api_token: LlmApiToken, user_store: Entity<UserStore>, status: client::Status, accept_terms: Option<Task<Result<()>>>, _settings_subscription: Subscription, _llm_token_subscription: Subscription,} impl State { fn new( client: Arc<Client>, user_store: Entity<UserStore>, status: client::Status, cx: &mut Context<Self>, ) -> Self { let refresh_llm_token_listener = RefreshLlmTokenListener::global(cx); Self { client: client.clone(), llm_api_token: LlmApiToken::default(), user_store, status, accept_terms: None, _settings_subscription: cx.observe_global::<SettingsStore>(|_, cx| { cx.notify(); }), _llm_token_subscription: cx.subscribe( &refresh_llm_token_listener, |this, _listener, _event, cx| { let client = this.client.clone(); let llm_api_token = this.llm_api_token.clone(); cx.spawn(|_this, _cx| async move { llm_api_token.refresh(&client).await?; anyhow::Ok(()) }) .detach_and_log_err(cx); }, ), } } fn is_signed_out(&self) -> bool { self.status.is_signed_out() } fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<()>> { let client = self.client.clone(); cx.spawn(move |this, mut cx| async move { client.authenticate_and_connect(true, &cx).await?; this.update(&mut cx, |_, cx| cx.notify()) }) } fn has_accepted_terms_of_service(&self, cx: &App) -> bool { self.user_store .read(cx) .current_user_has_accepted_terms() .unwrap_or(false) } fn accept_terms_of_service(&mut self, cx: &mut Context<Self>) { let user_store = self.user_store.clone(); self.accept_terms = Some(cx.spawn(move |this, mut cx| async move { let _ = user_store .update(&mut cx, |store, cx| store.accept_terms_of_service(cx))? .await; this.update(&mut cx, |this, cx| { this.accept_terms = None; cx.notify() }) })); }} impl CloudLanguageModelProvider { pub fn new(user_store: Entity<UserStore>, client: Arc<Client>, cx: &mut App) -> Self { let mut status_rx = client.status(); let status = *status_rx.borrow(); let state = cx.new(|cx| State::new(client.clone(), user_store.clone(), status, cx)); let state_ref = state.downgrade(); let maintain_client_status = cx.spawn(|mut cx| async move { while let Some(status) = status_rx.next().await { if let Some(this) = state_ref.upgrade() { _ = this.update(&mut cx, |this, cx| { if this.status != status { this.status = status; cx.notify(); } }); } else { break; } } }); Self { client, state: state.clone(), _maintain_client_status: maintain_client_status, } }} impl LanguageModelProviderState for CloudLanguageModelProvider { type ObservableEntity = State; fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> { Some(self.state.clone()) }} impl LanguageModelProvider for CloudLanguageModelProvider { fn id(&self) -> LanguageModelProviderId { LanguageModelProviderId(ZED_CLOUD_PROVIDER_ID.into()) } fn name(&self) -> LanguageModelProviderName { LanguageModelProviderName(PROVIDER_NAME.into()) } fn icon(&self) -> IconName { IconName::AiZed } fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> { let mut models = BTreeMap::default(); if cx.is_staff() { for model in anthropic::Model::iter() { if !matches!(model, anthropic::Model::Custom { .. }) { models.insert(model.id().to_string(), CloudModel::Anthropic(model)); } } for model in open_ai::Model::iter() { if !matches!(model, open_ai::Model::Custom { .. }) { models.insert(model.id().to_string(), CloudModel::OpenAi(model)); } } for model in google_ai::Model::iter() { if !matches!(model, google_ai::Model::Custom { .. }) { models.insert(model.id().to_string(), CloudModel::Google(model)); } } } else { models.insert( anthropic::Model::Claude3_5Sonnet.id().to_string(), CloudModel::Anthropic(anthropic::Model::Claude3_5Sonnet), ); } let llm_closed_beta_models = if cx.has_flag::<LlmClosedBeta>() { zed_cloud_provider_additional_models() } else { &[] }; // Override with available models from settings for model in AllLanguageModelSettings::get_global(cx) .zed_dot_dev .available_models .iter() .chain(llm_closed_beta_models) .cloned() { let model = match model.provider { AvailableProvider::Anthropic => CloudModel::Anthropic(anthropic::Model::Custom { name: model.name.clone(), display_name: model.display_name.clone(), max_tokens: model.max_tokens, tool_override: model.tool_override.clone(), cache_configuration: model.cache_configuration.as_ref().map(|config| { anthropic::AnthropicModelCacheConfiguration { max_cache_anchors: config.max_cache_anchors, should_speculate: config.should_speculate, min_total_token: config.min_total_token, } }), default_temperature: model.default_temperature, max_output_tokens: model.max_output_tokens, extra_beta_headers: model.extra_beta_headers.clone(), }), AvailableProvider::OpenAi => CloudModel::OpenAi(open_ai::Model::Custom { name: model.name.clone(), display_name: model.display_name.clone(), max_tokens: model.max_tokens, max_output_tokens: model.max_output_tokens, max_completion_tokens: model.max_completion_tokens, }), AvailableProvider::Google => CloudModel::Google(google_ai::Model::Custom { name: model.name.clone(), display_name: model.display_name.clone(), max_tokens: model.max_tokens, }), }; models.insert(model.id().to_string(), model.clone()); } let llm_api_token = self.state.read(cx).llm_api_token.clone(); models .into_values() .map(|model| { Arc::new(CloudLanguageModel { id: LanguageModelId::from(model.id().to_string()), model, llm_api_token: llm_api_token.clone(), client: self.client.clone(), request_limiter: RateLimiter::new(4), }) as Arc<dyn LanguageModel> }) .collect() } fn is_authenticated(&self, cx: &App) -> bool { !self.state.read(cx).is_signed_out() } fn authenticate(&self, _cx: &mut App) -> Task<Result<()>> { Task::ready(Ok(())) } fn configuration_view(&self, _: &mut Window, cx: &mut App) -> AnyView { cx.new(|_| ConfigurationView { state: self.state.clone(), }) .into() } fn must_accept_terms(&self, cx: &App) -> bool { !self.state.read(cx).has_accepted_terms_of_service(cx) } fn render_accept_terms( &self, view: LanguageModelProviderTosView, cx: &mut App, ) -> Option<AnyElement> { render_accept_terms(self.state.clone(), view, cx) } fn reset_credentials(&self, _cx: &mut App) -> Task<Result<()>> { Task::ready(Ok(())) }} fn render_accept_terms( state: Entity<State>, view_kind: LanguageModelProviderTosView, cx: &mut App,) -> Option<AnyElement> { if state.read(cx).has_accepted_terms_of_service(cx) { return None; } let accept_terms_disabled = state.read(cx).accept_terms.is_some(); let terms_button = Button::new("terms_of_service", "Terms of Service") .style(ButtonStyle::Subtle) .icon(IconName::ArrowUpRight) .icon_color(Color::Muted) .icon_size(IconSize::XSmall) .on_click(move |_, _window, cx| cx.open_url("https://zed.dev/terms-of-service")); let text = "To start using Zed AI, please read and accept the"; let form = v_flex() .w_full() .gap_2() .when( view_kind == LanguageModelProviderTosView::ThreadEmptyState, |form| form.items_center(), ) .child( h_flex() .flex_wrap() .when( view_kind == LanguageModelProviderTosView::ThreadEmptyState, |form| form.justify_center(), ) .child(Label::new(text)) .child(terms_button), ) .child({ let button_container = h_flex().w_full().child( Button::new("accept_terms", "I accept the Terms of Service") .style(ButtonStyle::Tinted(TintColor::Accent)) .disabled(accept_terms_disabled) .on_click({ let state = state.downgrade(); move |_, _window, cx| { state .update(cx, |state, cx| state.accept_terms_of_service(cx)) .ok(); } }), ); match view_kind { LanguageModelProviderTosView::ThreadEmptyState => button_container.justify_center(), LanguageModelProviderTosView::PromptEditorPopup => button_container.justify_end(), LanguageModelProviderTosView::Configuration => button_container.justify_start(), } }); Some(form.into_any())} pub struct CloudLanguageModel { id: LanguageModelId, model: CloudModel, llm_api_token: LlmApiToken, client: Arc<Client>, request_limiter: RateLimiter,} #[derive(Clone, Default)]pub struct LlmApiToken(Arc<RwLock<Option<String>>>); #[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." ) }} impl CloudLanguageModel { async fn perform_llm_completion( client: Arc<Client>, llm_api_token: LlmApiToken, body: PerformCompletionParams, ) -> Result<Response<AsyncBody>> { let http_client = &client.http_client(); let mut token = llm_api_token.acquire(&client).await?; let mut did_retry = false;Lines 1–500 of 943. Load more
use super::open_ai::count_open_ai_tokens; use anthropic::AnthropicError; use anyhow::{anyhow, Result}; use client::{ zed_urls, Client, PerformCompletionParams, UserStore, EXPIRED_LLM_TOKEN_HEADER_NAME, MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME, }; use collections::BTreeMap; use feature_flags::{FeatureFlagAppExt, LlmClosedBeta, ZedPro}; use futures::{ future::BoxFuture, stream::BoxStream, AsyncBufReadExt, FutureExt, Stream, StreamExt, TryStreamExt as _, }; use gpui::{ AnyElement, AnyView, App, AsyncApp, Context, Entity, EventEmitter, Global, ReadGlobal, Subscription, Task, }; use http_client::{AsyncBody, HttpClient, Method, Response, StatusCode}; use language_model::{ CloudModel, LanguageModel, LanguageModelCacheConfiguration, LanguageModelId, LanguageModelName, LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, LanguageModelProviderTosView, LanguageModelRequest, RateLimiter, ZED_CLOUD_PROVIDER_ID, }; use language_model::{ LanguageModelAvailability, LanguageModelCompletionEvent, LanguageModelProvider, }; use proto::TypedEnvelope; use schemars::JsonSchema; use serde::{de::DeserializeOwned, Deserialize, Serialize}; use serde_json::value::RawValue; use settings::{Settings, SettingsStore}; use smol::{ io::{AsyncReadExt, BufReader}, lock::{RwLock, RwLockUpgradableReadGuard, RwLockWriteGuard}, }; use std::fmt; use std::{ future, sync::{Arc, LazyLock}, }; use strum::IntoEnumIterator; use thiserror::Error; use ui::{prelude::*, TintColor}; use crate::provider::anthropic::map_to_language_model_completion_events; use crate::AllLanguageModelSettings; use super::anthropic::count_anthropic_tokens; pub const PROVIDER_NAME: &str = "Zed"; const ZED_CLOUD_PROVIDER_ADDITIONAL_MODELS_JSON: Option<&str> = option_env!("ZED_CLOUD_PROVIDER_ADDITIONAL_MODELS_JSON"); fn zed_cloud_provider_additional_models() -> &'static [AvailableModel] { static ADDITIONAL_MODELS: LazyLock<Vec<AvailableModel>> = LazyLock::new(|| { ZED_CLOUD_PROVIDER_ADDITIONAL_MODELS_JSON .map(|json| serde_json::from_str(json).unwrap()) .unwrap_or_default() }); ADDITIONAL_MODELS.as_slice() } #[derive(Default, Clone, Debug, PartialEq)] pub struct ZedDotDevSettings { pub available_models: Vec<AvailableModel>, } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)] #[serde(rename_all = "lowercase")] pub enum AvailableProvider { Anthropic, OpenAi, Google, } #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)] pub struct AvailableModel { /// The provider of the language model. pub provider: AvailableProvider, /// The model's name in the provider's API. e.g. claude-3-5-sonnet-20240620 pub name: String, /// The name displayed in the UI, such as in the assistant panel model dropdown menu. pub display_name: Option<String>, /// The size of the context window, indicating the maximum number of tokens the model can process. pub max_tokens: usize, /// The maximum number of output tokens allowed by the model. pub max_output_tokens: Option<u32>, /// The maximum number of completion tokens allowed by the model (o1-* only) pub max_completion_tokens: Option<u32>, /// Override this model with a different Anthropic model for tool calls. pub tool_override: Option<String>, /// Indicates whether this custom model supports caching. pub cache_configuration: Option<LanguageModelCacheConfiguration>, /// The default temperature to use for this model. pub default_temperature: Option<f32>, /// Any extra beta headers to provide when using the model. #[serde(default)] pub extra_beta_headers: Vec<String>, } 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)) } } pub struct CloudLanguageModelProvider { client: Arc<Client>, state: gpui::Entity<State>, _maintain_client_status: Task<()>, } pub struct State { client: Arc<Client>, llm_api_token: LlmApiToken, user_store: Entity<UserStore>, status: client::Status, accept_terms: Option<Task<Result<()>>>, _settings_subscription: Subscription, _llm_token_subscription: Subscription, } impl State { fn new( client: Arc<Client>, user_store: Entity<UserStore>, status: client::Status, cx: &mut Context<Self>, ) -> Self { let refresh_llm_token_listener = RefreshLlmTokenListener::global(cx); Self { client: client.clone(), llm_api_token: LlmApiToken::default(), user_store, status, accept_terms: None, _settings_subscription: cx.observe_global::<SettingsStore>(|_, cx| { cx.notify(); }), _llm_token_subscription: cx.subscribe( &refresh_llm_token_listener, |this, _listener, _event, cx| { let client = this.client.clone(); let llm_api_token = this.llm_api_token.clone(); cx.spawn(|_this, _cx| async move { llm_api_token.refresh(&client).await?; anyhow::Ok(()) }) .detach_and_log_err(cx); }, ), } } fn is_signed_out(&self) -> bool { self.status.is_signed_out() } fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<()>> { let client = self.client.clone(); cx.spawn(move |this, mut cx| async move { client.authenticate_and_connect(true, &cx).await?; this.update(&mut cx, |_, cx| cx.notify()) }) } fn has_accepted_terms_of_service(&self, cx: &App) -> bool { self.user_store .read(cx) .current_user_has_accepted_terms() .unwrap_or(false) } fn accept_terms_of_service(&mut self, cx: &mut Context<Self>) { let user_store = self.user_store.clone(); self.accept_terms = Some(cx.spawn(move |this, mut cx| async move { let _ = user_store .update(&mut cx, |store, cx| store.accept_terms_of_service(cx))? .await; this.update(&mut cx, |this, cx| { this.accept_terms = None; cx.notify() }) })); } } impl CloudLanguageModelProvider { pub fn new(user_store: Entity<UserStore>, client: Arc<Client>, cx: &mut App) -> Self { let mut status_rx = client.status(); let status = *status_rx.borrow(); let state = cx.new(|cx| State::new(client.clone(), user_store.clone(), status, cx)); let state_ref = state.downgrade(); let maintain_client_status = cx.spawn(|mut cx| async move { while let Some(status) = status_rx.next().await { if let Some(this) = state_ref.upgrade() { _ = this.update(&mut cx, |this, cx| { if this.status != status { this.status = status; cx.notify(); } }); } else { break; } } }); Self { client, state: state.clone(), _maintain_client_status: maintain_client_status, } } } impl LanguageModelProviderState for CloudLanguageModelProvider { type ObservableEntity = State; fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> { Some(self.state.clone()) } } impl LanguageModelProvider for CloudLanguageModelProvider { fn id(&self) -> LanguageModelProviderId { LanguageModelProviderId(ZED_CLOUD_PROVIDER_ID.into()) } fn name(&self) -> LanguageModelProviderName { LanguageModelProviderName(PROVIDER_NAME.into()) } fn icon(&self) -> IconName { IconName::AiZed } fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> { let mut models = BTreeMap::default(); if cx.is_staff() { for model in anthropic::Model::iter() { if !matches!(model, anthropic::Model::Custom { .. }) { models.insert(model.id().to_string(), CloudModel::Anthropic(model)); } } for model in open_ai::Model::iter() { if !matches!(model, open_ai::Model::Custom { .. }) { models.insert(model.id().to_string(), CloudModel::OpenAi(model)); } } for model in google_ai::Model::iter() { if !matches!(model, google_ai::Model::Custom { .. }) { models.insert(model.id().to_string(), CloudModel::Google(model)); } } } else { models.insert( anthropic::Model::Claude3_5Sonnet.id().to_string(), CloudModel::Anthropic(anthropic::Model::Claude3_5Sonnet), ); } let llm_closed_beta_models = if cx.has_flag::<LlmClosedBeta>() { zed_cloud_provider_additional_models() } else { &[] }; // Override with available models from settings for model in AllLanguageModelSettings::get_global(cx) .zed_dot_dev .available_models .iter() .chain(llm_closed_beta_models) .cloned() { let model = match model.provider { AvailableProvider::Anthropic => CloudModel::Anthropic(anthropic::Model::Custom { name: model.name.clone(), display_name: model.display_name.clone(), max_tokens: model.max_tokens, tool_override: model.tool_override.clone(), cache_configuration: model.cache_configuration.as_ref().map(|config| { anthropic::AnthropicModelCacheConfiguration { max_cache_anchors: config.max_cache_anchors, should_speculate: config.should_speculate, min_total_token: config.min_total_token, } }), default_temperature: model.default_temperature, max_output_tokens: model.max_output_tokens, extra_beta_headers: model.extra_beta_headers.clone(), }), AvailableProvider::OpenAi => CloudModel::OpenAi(open_ai::Model::Custom { name: model.name.clone(), display_name: model.display_name.clone(), max_tokens: model.max_tokens, max_output_tokens: model.max_output_tokens, max_completion_tokens: model.max_completion_tokens, }), AvailableProvider::Google => CloudModel::Google(google_ai::Model::Custom { name: model.name.clone(), display_name: model.display_name.clone(), max_tokens: model.max_tokens, }), }; models.insert(model.id().to_string(), model.clone()); } let llm_api_token = self.state.read(cx).llm_api_token.clone(); models .into_values() .map(|model| { Arc::new(CloudLanguageModel { id: LanguageModelId::from(model.id().to_string()), model, llm_api_token: llm_api_token.clone(), client: self.client.clone(), request_limiter: RateLimiter::new(4), }) as Arc<dyn LanguageModel> }) .collect() } fn is_authenticated(&self, cx: &App) -> bool { !self.state.read(cx).is_signed_out() } fn authenticate(&self, _cx: &mut App) -> Task<Result<()>> { Task::ready(Ok(())) } fn configuration_view(&self, _: &mut Window, cx: &mut App) -> AnyView { cx.new(|_| ConfigurationView { state: self.state.clone(), }) .into() } fn must_accept_terms(&self, cx: &App) -> bool { !self.state.read(cx).has_accepted_terms_of_service(cx) } fn render_accept_terms( &self, view: LanguageModelProviderTosView, cx: &mut App, ) -> Option<AnyElement> { render_accept_terms(self.state.clone(), view, cx) } fn reset_credentials(&self, _cx: &mut App) -> Task<Result<()>> { Task::ready(Ok(())) } } fn render_accept_terms( state: Entity<State>, view_kind: LanguageModelProviderTosView, cx: &mut App, ) -> Option<AnyElement> { if state.read(cx).has_accepted_terms_of_service(cx) { return None; } let accept_terms_disabled = state.read(cx).accept_terms.is_some(); let terms_button = Button::new("terms_of_service", "Terms of Service") .style(ButtonStyle::Subtle) .icon(IconName::ArrowUpRight) .icon_color(Color::Muted) .icon_size(IconSize::XSmall) .on_click(move |_, _window, cx| cx.open_url("https://zed.dev/terms-of-service")); let text = "To start using Zed AI, please read and accept the"; let form = v_flex() .w_full() .gap_2() .when( view_kind == LanguageModelProviderTosView::ThreadEmptyState, |form| form.items_center(), ) .child( h_flex() .flex_wrap() .when( view_kind == LanguageModelProviderTosView::ThreadEmptyState, |form| form.justify_center(), ) .child(Label::new(text)) .child(terms_button), ) .child({ let button_container = h_flex().w_full().child( Button::new("accept_terms", "I accept the Terms of Service") .style(ButtonStyle::Tinted(TintColor::Accent)) .disabled(accept_terms_disabled) .on_click({ let state = state.downgrade(); move |_, _window, cx| { state .update(cx, |state, cx| state.accept_terms_of_service(cx)) .ok(); } }), ); match view_kind { LanguageModelProviderTosView::ThreadEmptyState => button_container.justify_center(), LanguageModelProviderTosView::PromptEditorPopup => button_container.justify_end(), LanguageModelProviderTosView::Configuration => button_container.justify_start(), } }); Some(form.into_any()) } pub struct CloudLanguageModel { id: LanguageModelId, model: CloudModel, llm_api_token: LlmApiToken, client: Arc<Client>, request_limiter: RateLimiter, } #[derive(Clone, Default)] pub struct LlmApiToken(Arc<RwLock<Option<String>>>); #[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." ) } } impl CloudLanguageModel { async fn perform_llm_completion( client: Arc<Client>, llm_api_token: LlmApiToken, body: PerformCompletionParams, ) -> Result<Response<AsyncBody>> { let http_client = &client.http_client(); let mut token = llm_api_token.acquire(&client).await?; let mut did_retry = false;