diff --git a/crates/acp_thread/src/connection.rs b/crates/acp_thread/src/connection.rs index a99ae8e8637..a730e6c39d9 100644 --- a/crates/acp_thread/src/connection.rs +++ b/crates/acp_thread/src/connection.rs @@ -4,7 +4,7 @@ use anyhow::Result; use chrono::{DateTime, Utc}; use collections::{HashMap, HashSet, IndexMap}; use gpui::{Entity, SharedString, Task}; -use language_model::LanguageModelProviderId; +use language_model::{DisabledReason, LanguageModelProviderId}; use project::{AgentId, Project}; use serde::{Deserialize, Serialize}; use std::{any::Any, error::Error, fmt, path::PathBuf, rc::Rc}; @@ -493,6 +493,7 @@ pub struct AgentModelInfo { pub icon: Option, pub is_latest: bool, pub cost: Option, + pub disabled: Option, } #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -1064,6 +1065,7 @@ mod test_support { icon: Some(AgentModelIcon::Named(ui::IconName::ZedAssistant)), is_latest: false, cost: None, + disabled: None, })), } } diff --git a/crates/agent/src/agent.rs b/crates/agent/src/agent.rs index 377328da965..172b8b9a1ec 100644 --- a/crates/agent/src/agent.rs +++ b/crates/agent/src/agent.rs @@ -310,6 +310,7 @@ impl LanguageModels { }), is_latest: model.is_latest(), cost: model.model_cost_info().map(|cost| cost.to_shared_string()), + disabled: model.is_disabled(), } } @@ -5545,6 +5546,7 @@ mod internal_tests { ui::IconName::ZedAssistant )), is_latest: false, + disabled: None, cost: None, }] )]) diff --git a/crates/agent_ui/src/model_selector.rs b/crates/agent_ui/src/model_selector.rs index fa82018b8a3..c0d4441b160 100644 --- a/crates/agent_ui/src/model_selector.rs +++ b/crates/agent_ui/src/model_selector.rs @@ -271,6 +271,7 @@ impl PickerDelegate for ModelPickerDelegate { fn confirm(&mut self, _secondary: bool, window: &mut Window, cx: &mut Context>) { if let Some(ModelPickerEntry::Model(model_info, _)) = self.filtered_entries.get(self.selected_index) + && model_info.disabled.is_none() { self.selector .select_model(model_info.id.clone(), cx) @@ -336,6 +337,7 @@ impl PickerDelegate for ModelPickerDelegate { Some(AgentModelIcon::Named(icon)) => this.icon(*icon), None => this, }) + .disabled(model_info.disabled.clone()) .is_selected(is_selected) .is_focused(selected) .is_latest(model_info.is_latest) @@ -511,6 +513,7 @@ mod tests { description: None, icon: None, is_latest: false, + disabled: None, cost: None, }) .collect::>(), @@ -618,6 +621,7 @@ mod tests { description: None, icon: None, is_latest: false, + disabled: None, cost: None, }, AgentModelInfo { @@ -626,6 +630,7 @@ mod tests { description: None, icon: None, is_latest: false, + disabled: None, cost: None, }, ]; @@ -810,6 +815,7 @@ mod tests { description: None, icon: None, is_latest: false, + disabled: None, cost: None, }, acp_thread::AgentModelInfo { @@ -818,6 +824,7 @@ mod tests { description: None, icon: None, is_latest: false, + disabled: None, cost: None, }, ]); @@ -860,6 +867,7 @@ mod tests { description: None, icon: None, is_latest: false, + disabled: None, cost: None, }, acp_thread::AgentModelInfo { @@ -868,6 +876,7 @@ mod tests { description: None, icon: None, is_latest: false, + disabled: None, cost: None, }, ]); diff --git a/crates/agent_ui/src/ui/model_selector_components.rs b/crates/agent_ui/src/ui/model_selector_components.rs index 37e303ff1e1..e194c582dec 100644 --- a/crates/agent_ui/src/ui/model_selector_components.rs +++ b/crates/agent_ui/src/ui/model_selector_components.rs @@ -1,4 +1,5 @@ use gpui::{Action, ClickEvent, FocusHandle, prelude::*}; +use language_model::DisabledReason; use ui::{Chip, ElevationIndex, KeyBinding, ListItem, ListItemSpacing, Tooltip, prelude::*}; use zed_actions::agent::ToggleModelSelector; @@ -52,6 +53,7 @@ pub struct ModelSelectorListItem { is_focused: bool, is_latest: bool, is_favorite: bool, + disabled: Option, on_toggle_favorite: Option>, cost_info: Option, } @@ -66,6 +68,7 @@ impl ModelSelectorListItem { is_focused: false, is_latest: false, is_favorite: false, + disabled: None, on_toggle_favorite: None, cost_info: None, } @@ -86,6 +89,11 @@ impl ModelSelectorListItem { self } + pub fn disabled(mut self, disabled: Option) -> Self { + self.disabled = disabled; + self + } + pub fn is_focused(mut self, is_focused: bool) -> Self { self.is_focused = is_focused; self @@ -117,8 +125,12 @@ impl ModelSelectorListItem { impl RenderOnce for ModelSelectorListItem { fn render(self, _window: &mut Window, _cx: &mut App) -> impl IntoElement { + let is_disabled = self.disabled.is_some(); + let model_icon_color = if self.is_selected { Color::Accent + } else if is_disabled { + Color::Disabled } else { Color::Muted }; @@ -129,6 +141,15 @@ impl RenderOnce for ModelSelectorListItem { .inset(true) .spacing(ListItemSpacing::Sparse) .toggle_state(self.is_focused) + .when_some(self.disabled, |this, disabled_reason| { + this.disabled(true) + .tooltip(Tooltip::text(disabled_reason.0)) + .end_slot( + div() + .pr_2() + .child(Icon::new(IconName::Info).color(Color::Muted)), + ) + }) .child( h_flex() .w_full() @@ -143,7 +164,11 @@ impl RenderOnce for ModelSelectorListItem { .size(IconSize::Small), ) }) - .child(Label::new(self.title).truncate()) + .child( + Label::new(self.title) + .when(is_disabled, |this| this.color(Color::Disabled)) + .truncate(), + ) .when(self.is_latest, |parent| parent.child(Chip::new("Latest"))) .when_some(self.cost_info, |this, cost_info| { let tooltip_text = if cost_info.ends_with('×') { @@ -157,26 +182,34 @@ impl RenderOnce for ModelSelectorListItem { this.child(Chip::new(cost_info).tooltip(Tooltip::text(tooltip_text))) }), ) - .end_slot(div().pr_2().when(self.is_selected, |this| { - this.child(Icon::new(IconName::Check).color(Color::Accent)) - })) - .end_slot_on_hover(div().pr_1p5().when_some(self.on_toggle_favorite, { - |this, handle_click| { - let (icon, color, tooltip) = if is_favorite { - (IconName::StarFilled, Color::Accent, "Unfavorite Model") - } else { - (IconName::Star, Color::Default, "Favorite Model") - }; - this.child( - IconButton::new(("toggle-favorite", self.index), icon) - .layer(ElevationIndex::ElevatedSurface) - .icon_color(color) - .icon_size(IconSize::Small) - .tooltip(Tooltip::text(tooltip)) - .on_click(move |event, window, cx| (handle_click)(event, window, cx)), - ) - } - })) + .when(self.is_selected, |this| { + this.end_slot( + div() + .pr_2() + .child(Icon::new(IconName::Check).color(Color::Accent)), + ) + }) + .when(!is_disabled, |this| { + this.end_slot_on_hover(div().pr_1p5().when_some(self.on_toggle_favorite, { + |this, handle_click| { + let (icon, color, tooltip) = if is_favorite { + (IconName::StarFilled, Color::Accent, "Unfavorite Model") + } else { + (IconName::Star, Color::Default, "Favorite Model") + }; + this.child( + IconButton::new(("toggle-favorite", self.index), icon) + .layer(ElevationIndex::ElevatedSurface) + .icon_color(color) + .icon_size(IconSize::Small) + .tooltip(Tooltip::text(tooltip)) + .on_click(move |event, window, cx| { + (handle_click)(event, window, cx) + }), + ) + } + })) + }) } } diff --git a/crates/language_model/src/language_model.rs b/crates/language_model/src/language_model.rs index ff56c35ec83..df911d1b77b 100644 --- a/crates/language_model/src/language_model.rs +++ b/crates/language_model/src/language_model.rs @@ -24,6 +24,15 @@ pub fn init(cx: &mut App) { registry::init(cx); } +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct DisabledReason(pub SharedString); + +impl DisabledReason { + pub fn new(reason: impl Into) -> Self { + Self(reason.into()) + } +} + pub struct LanguageModelTextStream { pub message_id: Option, pub stream: BoxStream<'static, Result>, @@ -58,6 +67,11 @@ pub trait LanguageModel: Send + Sync { false } + /// Whether the model is currently disabled and, if so, why this is the case. + fn is_disabled(&self) -> Option { + None + } + /// Whether requests to this model require the user to consent to the /// upstream provider retaining inference logs (i.e. the model cannot be /// offered with Zero Data Retention). diff --git a/crates/language_models/src/provider/ollama.rs b/crates/language_models/src/provider/ollama.rs index ad866cb63fa..43a14c33e87 100644 --- a/crates/language_models/src/provider/ollama.rs +++ b/crates/language_models/src/provider/ollama.rs @@ -7,11 +7,12 @@ use futures::{Stream, TryFutureExt, stream}; use gpui::{AnyView, App, AsyncApp, Context, CursorStyle, Entity, Task, TaskExt}; use http_client::{CustomHeaders, HttpClient}; use language_model::{ - ApiKeyState, AuthenticateError, EnvVar, IconOrSvg, LanguageModel, LanguageModelCompletionError, - LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, LanguageModelProvider, - LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState, - LanguageModelRequest, LanguageModelRequestTool, LanguageModelToolChoice, LanguageModelToolUse, - LanguageModelToolUseId, MessageContent, RateLimiter, Role, StopReason, TokenUsage, env_var, + ApiKeyState, AuthenticateError, DisabledReason, EnvVar, IconOrSvg, LanguageModel, + LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, + LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName, + LanguageModelProviderState, LanguageModelRequest, LanguageModelRequestTool, + LanguageModelToolChoice, LanguageModelToolUse, LanguageModelToolUseId, MessageContent, + RateLimiter, Role, StopReason, TokenUsage, env_var, }; use menu; use ollama::{ @@ -139,23 +140,32 @@ impl State { let extra_headers = extra_headers.clone(); async move { let name = model.name.as_str(); - let model = show_model( + + show_model( http_client.as_ref(), &api_url, api_key.as_deref(), name, &extra_headers, ) - .await?; - let ollama_model = ollama::Model::new( - name, - None, - model.context_length, - Some(model.supports_tools()), - Some(model.supports_vision()), - Some(model.supports_thinking()), - ); - Ok(ollama_model) + .await + .map_or_else( + |error| { + ollama::Model::new_disabled( + name, + format!("Failed to fetch model from API: {error}",), + ) + }, + |model| { + ollama::Model::new( + name, + model.context_length, + Some(model.supports_tools()), + Some(model.supports_vision()), + Some(model.supports_thinking()), + ) + }, + ) } }); @@ -163,10 +173,8 @@ impl State { // since there is an arbitrary number of models available let mut ollama_models: Vec<_> = futures::stream::iter(tasks) .buffer_unordered(5) - .collect::>>() - .await - .into_iter() - .collect::>>()?; + .collect() + .await; ollama_models.sort_by(|a, b| a.name.cmp(&b.name)); @@ -308,6 +316,7 @@ impl LanguageModelProvider for OllamaLanguageModelProvider { .map(|model| { Arc::new(OllamaLanguageModel { id: LanguageModelId::from(model.name.clone()), + disabled: model.disabled.as_ref().map(|d| DisabledReason::new(d)), model, http_client: self.http_client.clone(), request_limiter: RateLimiter::new(4), @@ -350,6 +359,7 @@ pub struct OllamaLanguageModel { http_client: Arc, request_limiter: RateLimiter, state: Entity, + disabled: Option, } impl OllamaLanguageModel { @@ -510,6 +520,10 @@ impl LanguageModel for OllamaLanguageModel { format!("ollama/{}", self.model.id()) } + fn is_disabled(&self) -> Option { + self.disabled.clone() + } + fn max_token_count(&self) -> u64 { self.model.max_token_count() } @@ -1089,6 +1103,7 @@ fn merge_settings_into_models( supports_tools: setting_model.supports_tools, supports_vision: setting_model.supports_images, supports_thinking: setting_model.supports_thinking, + disabled: None, }, ); } @@ -1126,6 +1141,7 @@ mod tests { supports_tools: None, supports_vision: None, supports_thinking: None, + disabled: None, }, ); models.insert( @@ -1138,6 +1154,7 @@ mod tests { supports_tools: None, supports_vision: None, supports_thinking: None, + disabled: None, }, ); diff --git a/crates/ollama/src/ollama.rs b/crates/ollama/src/ollama.rs index 356bce5d74d..0bbb210b126 100644 --- a/crates/ollama/src/ollama.rs +++ b/crates/ollama/src/ollama.rs @@ -21,6 +21,7 @@ pub struct Model { pub supports_tools: Option, pub supports_vision: Option, pub supports_thinking: Option, + pub disabled: Option, } fn get_max_tokens(_name: &str) -> u64 { @@ -31,7 +32,6 @@ fn get_max_tokens(_name: &str) -> u64 { impl Model { pub fn new( name: &str, - display_name: Option<&str>, max_tokens: Option, supports_tools: Option, supports_vision: Option, @@ -39,14 +39,26 @@ impl Model { ) -> Self { Self { name: name.to_owned(), - display_name: display_name - .map(ToString::to_string) - .or_else(|| name.strip_suffix(":latest").map(ToString::to_string)), + display_name: name.strip_suffix(":latest").map(ToString::to_string), max_tokens: max_tokens.unwrap_or_else(|| get_max_tokens(name)), keep_alive: Some(KeepAlive::indefinite()), supports_tools, supports_vision, supports_thinking, + disabled: None, + } + } + + pub fn new_disabled(name: &str, reason: String) -> Self { + Self { + name: name.to_owned(), + display_name: name.strip_suffix(":latest").map(ToString::to_string), + max_tokens: get_max_tokens(name), + keep_alive: Some(KeepAlive::indefinite()), + supports_tools: None, + supports_vision: None, + supports_thinking: None, + disabled: Some(reason), } }