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, AppContext, AsyncAppContext, EventEmitter, FontWeight, Global, Model, ModelContext, ReadGlobal, Subscription, Task,};use http_client::{AsyncBody, HttpClient, Method, Response, StatusCode};use language_model::{ CloudModel, LanguageModel, LanguageModelCacheConfiguration, LanguageModelId, LanguageModelName, LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, 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>,} struct GlobalRefreshLlmTokenListener(Model<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 AppContext) { let listener = cx.new_model(|cx| RefreshLlmTokenListener::new(client, cx)); cx.set_global(GlobalRefreshLlmTokenListener(listener)); } pub fn global(cx: &AppContext) -> Model<Self> { GlobalRefreshLlmTokenListener::global(cx).0.clone() } fn new(client: Arc<Client>, cx: &mut ModelContext<Self>) -> Self { Self { _llm_token_subscription: client .add_message_handler(cx.weak_model(), Self::handle_refresh_llm_token), } } async fn handle_refresh_llm_token( this: Model<Self>, _: TypedEnvelope<proto::RefreshLlmToken>, mut cx: AsyncAppContext, ) -> Result<()> { this.update(&mut cx, |_this, cx| cx.emit(RefreshLlmTokenEvent)) }} pub struct CloudLanguageModelProvider { client: Arc<Client>, state: gpui::Model<State>, _maintain_client_status: Task<()>,} pub struct State { client: Arc<Client>, llm_api_token: LlmApiToken, user_store: Model<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: Model<UserStore>, status: client::Status, cx: &mut ModelContext<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 ModelContext<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: &AppContext) -> bool { self.user_store .read(cx) .current_user_has_accepted_terms() .unwrap_or(false) } fn accept_terms_of_service(&mut self, cx: &mut ModelContext<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: Model<UserStore>, client: Arc<Client>, cx: &mut AppContext) -> Self { let mut status_rx = client.status(); let status = *status_rx.borrow(); let state = cx.new_model(|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::Model<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: &AppContext) -> 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, }), 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: &AppContext) -> bool { !self.state.read(cx).is_signed_out() } fn authenticate(&self, _cx: &mut AppContext) -> Task<Result<()>> { Task::ready(Ok(())) } fn configuration_view(&self, cx: &mut WindowContext) -> AnyView { cx.new_view(|_cx| ConfigurationView { state: self.state.clone(), }) .into() } fn must_accept_terms(&self, cx: &AppContext) -> bool { !self.state.read(cx).has_accepted_terms_of_service(cx) } fn render_accept_terms(&self, cx: &mut WindowContext) -> Option<AnyElement> { let state = self.state.read(cx); let terms = [( "terms_of_service", "Terms of Service", "https://zed.dev/terms-of-service", )] .map(|(id, label, url)| { Button::new(id, label) .style(ButtonStyle::Subtle) .icon(IconName::ExternalLink) .icon_size(IconSize::XSmall) .icon_color(Color::Muted) .on_click(move |_, cx| cx.open_url(url)) }); if state.has_accepted_terms_of_service(cx) { None } else { let disabled = state.accept_terms.is_some(); Some( v_flex() .gap_2() .child( v_flex() .child(Label::new("Terms and Conditions").weight(FontWeight::MEDIUM)) .child( Label::new( "Please read and accept our terms and conditions to continue.", ) .size(LabelSize::Small), ), ) .child(v_flex().gap_1().children(terms)) .child( h_flex().justify_end().child( Button::new("accept_terms", "I've read it and accept it") .disabled(disabled) .on_click({ let state = self.state.downgrade(); move |_, cx| { state .update(cx, |state, cx| { state.accept_terms_of_service(cx) }) .ok(); } }), ), ) .into_any(), ) } } fn reset_credentials(&self, _cx: &mut AppContext) -> Task<Result<()>> { Task::ready(Ok(())) }} 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; let response = loop { let request_builder = http_client::Request::builder(); let request = request_builder .method(Method::POST) .uri(http_client.build_zed_llm_url("/completion", &[])?.as_ref()) .header("Content-Type", "application/json") .header("Authorization", format!("Bearer {token}")) .body(serde_json::to_string(&body)?.into())?; let mut response = http_client.send(request).await?; if response.status().is_success() { break response; } else if !did_retry && response .headers() .get(EXPIRED_LLM_TOKEN_HEADER_NAME) .is_some() {Lines 1–500 of 959. 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, AppContext, AsyncAppContext, EventEmitter, FontWeight, Global, Model, ModelContext, ReadGlobal, Subscription, Task, }; use http_client::{AsyncBody, HttpClient, Method, Response, StatusCode}; use language_model::{ CloudModel, LanguageModel, LanguageModelCacheConfiguration, LanguageModelId, LanguageModelName, LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, 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>, } struct GlobalRefreshLlmTokenListener(Model<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 AppContext) { let listener = cx.new_model(|cx| RefreshLlmTokenListener::new(client, cx)); cx.set_global(GlobalRefreshLlmTokenListener(listener)); } pub fn global(cx: &AppContext) -> Model<Self> { GlobalRefreshLlmTokenListener::global(cx).0.clone() } fn new(client: Arc<Client>, cx: &mut ModelContext<Self>) -> Self { Self { _llm_token_subscription: client .add_message_handler(cx.weak_model(), Self::handle_refresh_llm_token), } } async fn handle_refresh_llm_token( this: Model<Self>, _: TypedEnvelope<proto::RefreshLlmToken>, mut cx: AsyncAppContext, ) -> Result<()> { this.update(&mut cx, |_this, cx| cx.emit(RefreshLlmTokenEvent)) } } pub struct CloudLanguageModelProvider { client: Arc<Client>, state: gpui::Model<State>, _maintain_client_status: Task<()>, } pub struct State { client: Arc<Client>, llm_api_token: LlmApiToken, user_store: Model<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: Model<UserStore>, status: client::Status, cx: &mut ModelContext<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 ModelContext<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: &AppContext) -> bool { self.user_store .read(cx) .current_user_has_accepted_terms() .unwrap_or(false) } fn accept_terms_of_service(&mut self, cx: &mut ModelContext<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: Model<UserStore>, client: Arc<Client>, cx: &mut AppContext) -> Self { let mut status_rx = client.status(); let status = *status_rx.borrow(); let state = cx.new_model(|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::Model<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: &AppContext) -> 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, }), 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: &AppContext) -> bool { !self.state.read(cx).is_signed_out() } fn authenticate(&self, _cx: &mut AppContext) -> Task<Result<()>> { Task::ready(Ok(())) } fn configuration_view(&self, cx: &mut WindowContext) -> AnyView { cx.new_view(|_cx| ConfigurationView { state: self.state.clone(), }) .into() } fn must_accept_terms(&self, cx: &AppContext) -> bool { !self.state.read(cx).has_accepted_terms_of_service(cx) } fn render_accept_terms(&self, cx: &mut WindowContext) -> Option<AnyElement> { let state = self.state.read(cx); let terms = [( "terms_of_service", "Terms of Service", "https://zed.dev/terms-of-service", )] .map(|(id, label, url)| { Button::new(id, label) .style(ButtonStyle::Subtle) .icon(IconName::ExternalLink) .icon_size(IconSize::XSmall) .icon_color(Color::Muted) .on_click(move |_, cx| cx.open_url(url)) }); if state.has_accepted_terms_of_service(cx) { None } else { let disabled = state.accept_terms.is_some(); Some( v_flex() .gap_2() .child( v_flex() .child(Label::new("Terms and Conditions").weight(FontWeight::MEDIUM)) .child( Label::new( "Please read and accept our terms and conditions to continue.", ) .size(LabelSize::Small), ), ) .child(v_flex().gap_1().children(terms)) .child( h_flex().justify_end().child( Button::new("accept_terms", "I've read it and accept it") .disabled(disabled) .on_click({ let state = self.state.downgrade(); move |_, cx| { state .update(cx, |state, cx| { state.accept_terms_of_service(cx) }) .ok(); } }), ), ) .into_any(), ) } } fn reset_credentials(&self, _cx: &mut AppContext) -> Task<Result<()>> { Task::ready(Ok(())) } } 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; let response = loop { let request_builder = http_client::Request::builder(); let request = request_builder .method(Method::POST) .uri(http_client.build_zed_llm_url("/completion", &[])?.as_ref()) .header("Content-Type", "application/json") .header("Authorization", format!("Bearer {token}")) .body(serde_json::to_string(&body)?.into())?; let mut response = http_client.send(request).await?; if response.status().is_success() { break response; } else if !did_retry && response .headers() .get(EXPIRED_LLM_TOKEN_HEADER_NAME) .is_some() {