use anyhow::{Result, anyhow}; use convert_case::{Case, Casing}; use futures::{FutureExt, StreamExt, future, future::BoxFuture}; use gpui::{AnyView, App, AsyncApp, Context, Entity, SharedString, Task, Window}; use http_client::HttpClient; use language_model::{ AuthenticateError, LanguageModel, LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice, LanguageModelToolSchemaFormat, RateLimiter, }; use menu; use open_ai::{ResponseStreamEvent, stream_completion}; use settings::{Settings, SettingsStore}; use std::sync::Arc; use ui::{ElevationIndex, Tooltip, prelude::*}; use ui_input::SingleLineInput; use util::{ResultExt, truncate_and_trailoff}; use zed_env_vars::EnvVar; use crate::api_key::ApiKeyState; use crate::provider::open_ai::{OpenAiEventMapper, into_open_ai}; pub use settings::OpenAiCompatibleAvailableModel as AvailableModel; pub use settings::OpenAiCompatibleModelCapabilities as ModelCapabilities; #[derive(Default, Clone, Debug, PartialEq)] pub struct OpenAiCompatibleSettings { pub api_url: String, pub available_models: Vec, } pub struct OpenAiCompatibleLanguageModelProvider { id: LanguageModelProviderId, name: LanguageModelProviderName, http_client: Arc, state: Entity, } pub struct State { id: Arc, api_key_env_var: EnvVar, api_key_state: ApiKeyState, settings: OpenAiCompatibleSettings, } impl State { fn is_authenticated(&self) -> bool { self.api_key_state.has_key() } fn set_api_key(&mut self, api_key: Option, cx: &mut Context) -> Task> { let api_url = SharedString::new(self.settings.api_url.as_str()); self.api_key_state .store(api_url, api_key, |this| &mut this.api_key_state, cx) } fn authenticate(&mut self, cx: &mut Context) -> Task> { let api_url = SharedString::new(self.settings.api_url.clone()); self.api_key_state.load_if_needed( api_url, &self.api_key_env_var, |this| &mut this.api_key_state, cx, ) } } impl OpenAiCompatibleLanguageModelProvider { pub fn new(id: Arc, http_client: Arc, cx: &mut App) -> Self { fn resolve_settings<'a>(id: &'a str, cx: &'a App) -> Option<&'a OpenAiCompatibleSettings> { crate::AllLanguageModelSettings::get_global(cx) .openai_compatible .get(id) } let api_key_env_var_name = format!("{}_API_KEY", id).to_case(Case::UpperSnake).into(); let state = cx.new(|cx| { cx.observe_global::(|this: &mut State, cx| { let Some(settings) = resolve_settings(&this.id, cx).cloned() else { return; }; if &this.settings != &settings { let api_url = SharedString::new(settings.api_url.as_str()); this.api_key_state.handle_url_change( api_url, &this.api_key_env_var, |this| &mut this.api_key_state, cx, ); this.settings = settings; cx.notify(); } }) .detach(); let settings = resolve_settings(&id, cx).cloned().unwrap_or_default(); State { id: id.clone(), api_key_env_var: EnvVar::new(api_key_env_var_name), api_key_state: ApiKeyState::new(SharedString::new(settings.api_url.as_str())), settings, } }); Self { id: id.clone().into(), name: id.into(), http_client, state, } } fn create_language_model(&self, model: AvailableModel) -> Arc { Arc::new(OpenAiCompatibleLanguageModel { id: LanguageModelId::from(model.name.clone()), provider_id: self.id.clone(), provider_name: self.name.clone(), model, state: self.state.clone(), http_client: self.http_client.clone(), request_limiter: RateLimiter::new(4), }) } } impl LanguageModelProviderState for OpenAiCompatibleLanguageModelProvider { type ObservableEntity = State; fn observable_entity(&self) -> Option> { Some(self.state.clone()) } } impl LanguageModelProvider for OpenAiCompatibleLanguageModelProvider { fn id(&self) -> LanguageModelProviderId { self.id.clone() } fn name(&self) -> LanguageModelProviderName { self.name.clone() } fn icon(&self) -> IconName { IconName::AiOpenAiCompat } fn default_model(&self, cx: &App) -> Option> { self.state .read(cx) .settings .available_models .first() .map(|model| self.create_language_model(model.clone())) } fn default_fast_model(&self, _cx: &App) -> Option> { None } fn provided_models(&self, cx: &App) -> Vec> { self.state .read(cx) .settings .available_models .iter() .map(|model| self.create_language_model(model.clone())) .collect() } fn is_authenticated(&self, cx: &App) -> bool { self.state.read(cx).is_authenticated() } fn authenticate(&self, cx: &mut App) -> Task> { self.state.update(cx, |state, cx| state.authenticate(cx)) } fn configuration_view( &self, _target_agent: language_model::ConfigurationViewTargetAgent, window: &mut Window, cx: &mut App, ) -> AnyView { cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx)) .into() } fn reset_credentials(&self, cx: &mut App) -> Task> { self.state .update(cx, |state, cx| state.set_api_key(None, cx)) } } pub struct OpenAiCompatibleLanguageModel { id: LanguageModelId, provider_id: LanguageModelProviderId, provider_name: LanguageModelProviderName, model: AvailableModel, state: Entity, http_client: Arc, request_limiter: RateLimiter, } impl OpenAiCompatibleLanguageModel { fn stream_completion( &self, request: open_ai::Request, cx: &AsyncApp, ) -> BoxFuture<'static, Result>>> { let http_client = self.http_client.clone(); let Ok((api_key, api_url)) = self.state.read_with(cx, |state, _cx| { let api_url = &state.settings.api_url; ( state.api_key_state.key(api_url), state.settings.api_url.clone(), ) }) else { return future::ready(Err(anyhow!("App state dropped"))).boxed(); }; let provider = self.provider_name.clone(); let future = self.request_limiter.stream(async move { let Some(api_key) = api_key else { return Err(LanguageModelCompletionError::NoApiKey { provider }); }; let request = stream_completion(http_client.as_ref(), &api_url, &api_key, request); let response = request.await?; Ok(response) }); async move { Ok(future.await?.boxed()) }.boxed() } } impl LanguageModel for OpenAiCompatibleLanguageModel { fn id(&self) -> LanguageModelId { self.id.clone() } fn name(&self) -> LanguageModelName { LanguageModelName::from( self.model .display_name .clone() .unwrap_or_else(|| self.model.name.clone()), ) } fn provider_id(&self) -> LanguageModelProviderId { self.provider_id.clone() } fn provider_name(&self) -> LanguageModelProviderName { self.provider_name.clone() } fn supports_tools(&self) -> bool { self.model.capabilities.tools } fn tool_input_format(&self) -> LanguageModelToolSchemaFormat { LanguageModelToolSchemaFormat::JsonSchemaSubset } fn supports_images(&self) -> bool { self.model.capabilities.images } fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool { match choice { LanguageModelToolChoice::Auto => self.model.capabilities.tools, LanguageModelToolChoice::Any => self.model.capabilities.tools, LanguageModelToolChoice::None => true, } } fn telemetry_id(&self) -> String { format!("openai/{}", self.model.name) } fn max_token_count(&self) -> u64 { self.model.max_tokens } fn max_output_tokens(&self) -> Option { self.model.max_output_tokens } fn count_tokens( &self, request: LanguageModelRequest, cx: &App, ) -> BoxFuture<'static, Result> { let max_token_count = self.max_token_count(); cx.background_spawn(async move { let messages = super::open_ai::collect_tiktoken_messages(request); let model = if max_token_count >= 100_000 { // If the max tokens is 100k or more, it is likely the o200k_base tokenizer from gpt4o "gpt-4o" } else { // Otherwise fallback to gpt-4, since only cl100k_base and o200k_base are // supported with this tiktoken method "gpt-4" }; tiktoken_rs::num_tokens_from_messages(model, &messages).map(|tokens| tokens as u64) }) .boxed() } fn stream_completion( &self, request: LanguageModelRequest, cx: &AsyncApp, ) -> BoxFuture< 'static, Result< futures::stream::BoxStream< 'static, Result, >, LanguageModelCompletionError, >, > { let request = into_open_ai( request, &self.model.name, self.model.capabilities.parallel_tool_calls, self.model.capabilities.prompt_cache_key, self.max_output_tokens(), None, ); let completions = self.stream_completion(request, cx); async move { let mapper = OpenAiEventMapper::new(); Ok(mapper.map_stream(completions.await?).boxed()) } .boxed() } } struct ConfigurationView { api_key_editor: Entity, state: Entity, load_credentials_task: Option>, } impl ConfigurationView { fn new(state: Entity, window: &mut Window, cx: &mut Context) -> Self { let api_key_editor = cx.new(|cx| { SingleLineInput::new( window, cx, "000000000000000000000000000000000000000000000000000", ) }); cx.observe(&state, |_, _, cx| { cx.notify(); }) .detach(); let load_credentials_task = Some(cx.spawn_in(window, { let state = state.clone(); async move |this, cx| { if let Some(task) = state .update(cx, |state, cx| state.authenticate(cx)) .log_err() { // We don't log an error, because "not signed in" is also an error. let _ = task.await; } this.update(cx, |this, cx| { this.load_credentials_task = None; cx.notify(); }) .log_err(); } })); Self { api_key_editor, state, load_credentials_task, } } fn save_api_key(&mut self, _: &menu::Confirm, window: &mut Window, cx: &mut Context) { let api_key = self.api_key_editor.read(cx).text(cx).trim().to_string(); if api_key.is_empty() { return; } // url changes can cause the editor to be displayed again self.api_key_editor .update(cx, |input, cx| input.set_text("", window, cx)); let state = self.state.clone(); cx.spawn_in(window, async move |_, cx| { state .update(cx, |state, cx| state.set_api_key(Some(api_key), cx))? .await }) .detach_and_log_err(cx); } fn reset_api_key(&mut self, window: &mut Window, cx: &mut Context) { self.api_key_editor .update(cx, |input, cx| input.set_text("", window, cx)); let state = self.state.clone(); cx.spawn_in(window, async move |_, cx| { state .update(cx, |state, cx| state.set_api_key(None, cx))? .await }) .detach_and_log_err(cx); } fn should_render_editor(&self, cx: &Context) -> bool { !self.state.read(cx).is_authenticated() } } impl Render for ConfigurationView { fn render(&mut self, _: &mut Window, cx: &mut Context) -> impl IntoElement { let state = self.state.read(cx); let env_var_set = state.api_key_state.is_from_env_var(); let env_var_name = &state.api_key_env_var.name; let api_key_section = if self.should_render_editor(cx) { v_flex() .on_action(cx.listener(Self::save_api_key)) .child(Label::new("To use Zed's agent with an OpenAI-compatible provider, you need to add an API key.")) .child( div() .pt(DynamicSpacing::Base04.rems(cx)) .child(self.api_key_editor.clone()) ) .child( Label::new( format!("You can also assign the {env_var_name} environment variable and restart Zed."), ) .size(LabelSize::Small).color(Color::Muted), ) .into_any() } else { h_flex() .mt_1() .p_1() .justify_between() .rounded_md() .border_1() .border_color(cx.theme().colors().border) .bg(cx.theme().colors().background) .child( h_flex() .gap_1() .child(Icon::new(IconName::Check).color(Color::Success)) .child(Label::new(if env_var_set { format!("API key set in {env_var_name} environment variable") } else { format!("API key configured for {}", truncate_and_trailoff(&state.settings.api_url, 32)) })), ) .child( Button::new("reset-api-key", "Reset API Key") .label_size(LabelSize::Small) .icon(IconName::Undo) .icon_size(IconSize::Small) .icon_position(IconPosition::Start) .layer(ElevationIndex::ModalSurface) .when(env_var_set, |this| { this.tooltip(Tooltip::text(format!("To reset your API key, unset the {env_var_name} environment variable."))) }) .on_click(cx.listener(|this, _, window, cx| this.reset_api_key(window, cx))), ) .into_any() }; if self.load_credentials_task.is_some() { div().child(Label::new("Loading credentials…")).into_any() } else { v_flex().size_full().child(api_key_section).into_any() } } }