Yehowshua/ZIm
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Factor out language model selector into its own crate (#21113)

Browse source 
This PR factors the language model selector out into its own
`language_model_selector` crate so that it can be reused in
`assistant2`.

Also renamed it from `ModelSelector` to `LanguageModelSelector` to be a
bit more specific.

Release Notes:

- N/A
This commit is contained in:
Marshall Bowers 2024-11-23 12:11:31 -05:00 • committed by GitHub
parent 9adc3b4e82
commit 3a0408953d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 119 additions and 48 deletions
Download patch file Download diff file Expand all files Collapse all files

1
crates/assistant/Cargo.toml
View file

@ -50,6 +50,7 @@ indexed_docs.workspace = true
indoc.workspace = true
language.workspace = true
language_model.workspace = true
language_model_selector.workspace = true
language_models.workspace = true
log.workspace = true
lsp.workspace = true

2
crates/assistant/src/assistant.rs
View file

@ -5,7 +5,6 @@ pub mod assistant_settings;
mod context;
pub mod context_store;
mod inline_assistant;
mod model_selector;
mod patch;
mod prompt_library;
mod prompts;
@ -37,7 +36,6 @@ pub(crate) use inline_assistant::*;
use language_model::{
LanguageModelId, LanguageModelProviderId, LanguageModelRegistry, LanguageModelResponseMessage,
};
pub(crate) use model_selector::*;
pub use patch::*;
pub use prompts::PromptBuilder;
use prompts::PromptLoadingParams;

26
crates/assistant/src/assistant_panel.rs
View file

@ -17,9 +17,9 @@ use crate::{
ContextEvent, ContextId, ContextStore, ContextStoreEvent, CopyCode, CycleMessageRole,
DeployHistory, DeployPromptLibrary, Edit, InlineAssistant, InsertDraggedFiles,
InsertIntoEditor, InvokedSlashCommandId, InvokedSlashCommandStatus, Message, MessageId,
MessageMetadata, MessageStatus, ModelPickerDelegate, ModelSelector, NewContext,
ParsedSlashCommand, PendingSlashCommandStatus, QuoteSelection, RemoteContextMetadata,
RequestType, SavedContextMetadata, Split, ToggleFocus, ToggleModelSelector,
MessageMetadata, MessageStatus, NewContext, ParsedSlashCommand, PendingSlashCommandStatus,
QuoteSelection, RemoteContextMetadata, RequestType, SavedContextMetadata, Split, ToggleFocus,
ToggleModelSelector,
};
use anyhow::Result;
use assistant_slash_command::{SlashCommand, SlashCommandOutputSection};
@ -55,6 +55,7 @@ use language_model::{
LanguageModelProvider, LanguageModelProviderId, LanguageModelRegistry, Role,
ZED_CLOUD_PROVIDER_ID,
};
use language_model_selector::{LanguageModelPickerDelegate, LanguageModelSelector};
use multi_buffer::MultiBufferRow;
use picker::{Picker, PickerDelegate};
use project::lsp_store::LocalLspAdapterDelegate;
@ -142,7 +143,7 @@ pub struct AssistantPanel {
languages: Arc<LanguageRegistry>,
fs: Arc<dyn Fs>,
subscriptions: Vec<Subscription>,
model_selector_menu_handle: PopoverMenuHandle<Picker<ModelPickerDelegate>>,
model_selector_menu_handle: PopoverMenuHandle<Picker<LanguageModelPickerDelegate>>,
model_summary_editor: View<Editor>,
authenticate_provider_task: Option<(LanguageModelProviderId, Task<()>)>,
configuration_subscription: Option<Subscription>,
@ -4457,13 +4458,13 @@ pub struct ContextEditorToolbarItem {
fs: Arc<dyn Fs>,
active_context_editor: Option<WeakView<ContextEditor>>,
model_summary_editor: View<Editor>,
model_selector_menu_handle: PopoverMenuHandle<Picker<ModelPickerDelegate>>,
model_selector_menu_handle: PopoverMenuHandle<Picker<LanguageModelPickerDelegate>>,
}
impl ContextEditorToolbarItem {
pub fn new(
workspace: &Workspace,
model_selector_menu_handle: PopoverMenuHandle<Picker<ModelPickerDelegate>>,
model_selector_menu_handle: PopoverMenuHandle<Picker<LanguageModelPickerDelegate>>,
model_summary_editor: View<Editor>,
) -> Self {
Self {
@ -4559,8 +4560,17 @@ impl Render for ContextEditorToolbarItem {
// .map(|remaining_items| format!("Files to scan: {}", remaining_items))
// })
.child(
ModelSelector::new(
self.fs.clone(),
LanguageModelSelector::new(
{
let fs = self.fs.clone();
move |model, cx| {
update_settings_file::<AssistantSettings>(
fs.clone(),
cx,
move |settings, _| settings.set_model(model.clone()),
);
}
},
ButtonLike::new("active-model")
.style(ButtonStyle::Subtle)
.child(

20
crates/assistant/src/inline_assistant.rs
View file

@ -1,7 +1,7 @@
use crate::{
assistant_settings::AssistantSettings, humanize_token_count, prompts::PromptBuilder,
AssistantPanel, AssistantPanelEvent, CharOperation, CycleNextInlineAssist,
CyclePreviousInlineAssist, LineDiff, LineOperation, ModelSelector, RequestType, StreamingDiff,
CyclePreviousInlineAssist, LineDiff, LineOperation, RequestType, StreamingDiff,
};
use anyhow::{anyhow, Context as _, Result};
use client::{telemetry::Telemetry, ErrorExt};
@ -33,12 +33,13 @@ use language_model::{
LanguageModel, LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage,
LanguageModelTextStream, Role,
};
use language_model_selector::LanguageModelSelector;
use language_models::report_assistant_event;
use multi_buffer::MultiBufferRow;
use parking_lot::Mutex;
use project::{CodeAction, ProjectTransaction};
use rope::Rope;
use settings::{Settings, SettingsStore};
use settings::{update_settings_file, Settings, SettingsStore};
use smol::future::FutureExt;
use std::{
cmp,
@ -1500,8 +1501,17 @@ impl Render for PromptEditor {
.justify_center()
.gap_2()
.child(
ModelSelector::new(
self.fs.clone(),
LanguageModelSelector::new(
{
let fs = self.fs.clone();
move |model, cx| {
update_settings_file::<AssistantSettings>(
fs.clone(),
cx,
move |settings, _| settings.set_model(model.clone()),
);
}
},
IconButton::new("context", IconName::SettingsAlt)
.shape(IconButtonShape::Square)
.icon_size(IconSize::Small)
@ -1521,7 +1531,7 @@ impl Render for PromptEditor {
)
}),
)
.with_info_text(
.info_text(
"Inline edits use context\n\
from the currently selected\n\
assistant panel tab.",

352
crates/assistant/src/model_selector.rs
View file

@ -1,352 +0,0 @@
use feature_flags::ZedPro;
use language_model::{LanguageModel, LanguageModelAvailability, LanguageModelRegistry};
use proto::Plan;
use workspace::ShowConfiguration;
use std::sync::Arc;
use crate::assistant_settings::AssistantSettings;
use fs::Fs;
use gpui::{Action, AnyElement, DismissEvent, SharedString, Task};
use picker::{Picker, PickerDelegate};
use settings::update_settings_file;
use ui::{prelude::*, ListItem, ListItemSpacing, PopoverMenu, PopoverMenuHandle, PopoverTrigger};
const TRY_ZED_PRO_URL: &str = "https://zed.dev/pro";
#[derive(IntoElement)]
pub struct ModelSelector<T: PopoverTrigger> {
handle: Option<PopoverMenuHandle<Picker<ModelPickerDelegate>>>,
fs: Arc<dyn Fs>,
trigger: T,
info_text: Option<SharedString>,
}
pub struct ModelPickerDelegate {
fs: Arc<dyn Fs>,
all_models: Vec<ModelInfo>,
filtered_models: Vec<ModelInfo>,
selected_index: usize,
}
#[derive(Clone)]
struct ModelInfo {
model: Arc<dyn LanguageModel>,
icon: IconName,
availability: LanguageModelAvailability,
is_selected: bool,
}
impl<T: PopoverTrigger> ModelSelector<T> {
pub fn new(fs: Arc<dyn Fs>, trigger: T) -> Self {
ModelSelector {
handle: None,
fs,
trigger,
info_text: None,
}
}
pub fn with_handle(mut self, handle: PopoverMenuHandle<Picker<ModelPickerDelegate>>) -> Self {
self.handle = Some(handle);
self
}
pub fn with_info_text(mut self, text: impl Into<SharedString>) -> Self {
self.info_text = Some(text.into());
self
}
}
impl PickerDelegate for ModelPickerDelegate {
type ListItem = ListItem;
fn match_count(&self) -> usize {
self.filtered_models.len()
}
fn selected_index(&self) -> usize {
self.selected_index
}
fn set_selected_index(&mut self, ix: usize, cx: &mut ViewContext<Picker<Self>>) {
self.selected_index = ix.min(self.filtered_models.len().saturating_sub(1));
cx.notify();
}
fn placeholder_text(&self, _cx: &mut WindowContext) -> Arc<str> {
"Select a model...".into()
}
fn update_matches(&mut self, query: String, cx: &mut ViewContext<Picker<Self>>) -> Task<()> {
let all_models = self.all_models.clone();
let llm_registry = LanguageModelRegistry::global(cx);
let configured_models: Vec<_> = llm_registry
.read(cx)
.providers()
.iter()
.filter(|provider| provider.is_authenticated(cx))
.map(|provider| provider.id())
.collect();
cx.spawn(|this, mut cx| async move {
let filtered_models = cx
.background_executor()
.spawn(async move {
let displayed_models = if configured_models.is_empty() {
all_models
} else {
all_models
.into_iter()
.filter(|model_info| {
configured_models.contains(&model_info.model.provider_id())
})
.collect::<Vec<_>>()
};
if query.is_empty() {
displayed_models
} else {
displayed_models
.into_iter()
.filter(|model_info| {
model_info
.model
.name()
.0
.to_lowercase()
.contains(&query.to_lowercase())
})
.collect()
}
})
.await;
this.update(&mut cx, |this, cx| {
this.delegate.filtered_models = filtered_models;
this.delegate.set_selected_index(0, cx);
cx.notify();
})
.ok();
})
}
fn confirm(&mut self, _secondary: bool, cx: &mut ViewContext<Picker<Self>>) {
if let Some(model_info) = self.filtered_models.get(self.selected_index) {
let model = model_info.model.clone();
update_settings_file::<AssistantSettings>(self.fs.clone(), cx, move |settings, _| {
settings.set_model(model.clone())
});
// Update the selection status
let selected_model_id = model_info.model.id();
let selected_provider_id = model_info.model.provider_id();
for model in &mut self.all_models {
model.is_selected = model.model.id() == selected_model_id
&& model.model.provider_id() == selected_provider_id;
}
for model in &mut self.filtered_models {
model.is_selected = model.model.id() == selected_model_id
&& model.model.provider_id() == selected_provider_id;
}
cx.emit(DismissEvent);
}
}
fn dismissed(&mut self, _cx: &mut ViewContext<Picker<Self>>) {}
fn render_header(&self, cx: &mut ViewContext<Picker<Self>>) -> Option<AnyElement> {
let configured_models_count = LanguageModelRegistry::global(cx)
.read(cx)
.providers()
.iter()
.filter(|provider| provider.is_authenticated(cx))
.count();
if configured_models_count > 0 {
Some(
Label::new("Configured Models")
.size(LabelSize::Small)
.color(Color::Muted)
.mt_1()
.mb_0p5()
.ml_3()
.into_any_element(),
)
} else {
None
}
}
fn render_match(
&self,
ix: usize,
selected: bool,
cx: &mut ViewContext<Picker<Self>>,
) -> Option<Self::ListItem> {
use feature_flags::FeatureFlagAppExt;
let show_badges = cx.has_flag::<ZedPro>();
let model_info = self.filtered_models.get(ix)?;
let provider_name: String = model_info.model.provider_name().0.clone().into();
Some(
ListItem::new(ix)
.inset(true)
.spacing(ListItemSpacing::Sparse)
.selected(selected)
.start_slot(
div().pr_0p5().child(
Icon::new(model_info.icon)
.color(Color::Muted)
.size(IconSize::Medium),
),
)
.child(
h_flex()
.w_full()
.items_center()
.gap_1p5()
.min_w(px(200.))
.child(Label::new(model_info.model.name().0.clone()))
.child(
h_flex()
.gap_0p5()
.child(
Label::new(provider_name)
.size(LabelSize::XSmall)
.color(Color::Muted),
)
.children(match model_info.availability {
LanguageModelAvailability::Public => None,
LanguageModelAvailability::RequiresPlan(Plan::Free) => None,
LanguageModelAvailability::RequiresPlan(Plan::ZedPro) => {
show_badges.then(|| {
Label::new("Pro")
.size(LabelSize::XSmall)
.color(Color::Muted)
})
}
}),
),
)
.end_slot(div().when(model_info.is_selected, |this| {
this.child(
Icon::new(IconName::Check)
.color(Color::Accent)
.size(IconSize::Small),
)
})),
)
}
fn render_footer(&self, cx: &mut ViewContext<Picker<Self>>) -> Option<gpui::AnyElement> {
use feature_flags::FeatureFlagAppExt;
let plan = proto::Plan::ZedPro;
let is_trial = false;
Some(
h_flex()
.w_full()
.border_t_1()
.border_color(cx.theme().colors().border_variant)
.p_1()
.gap_4()
.justify_between()
.when(cx.has_flag::<ZedPro>(), |this| {
this.child(match plan {
// Already a Zed Pro subscriber
Plan::ZedPro => Button::new("zed-pro", "Zed Pro")
.icon(IconName::ZedAssistant)
.icon_size(IconSize::Small)
.icon_color(Color::Muted)
.icon_position(IconPosition::Start)
.on_click(|_, cx| {
cx.dispatch_action(Box::new(zed_actions::OpenAccountSettings))
}),
// Free user
Plan::Free => Button::new(
"try-pro",
if is_trial {
"Upgrade to Pro"
} else {
"Try Pro"
},
)
.on_click(|_, cx| cx.open_url(TRY_ZED_PRO_URL)),
})
})
.child(
Button::new("configure", "Configure")
.icon(IconName::Settings)
.icon_size(IconSize::Small)
.icon_color(Color::Muted)
.icon_position(IconPosition::Start)
.on_click(|_, cx| {
cx.dispatch_action(ShowConfiguration.boxed_clone());
}),
)
.into_any(),
)
}
}
impl<T: PopoverTrigger> RenderOnce for ModelSelector<T> {
fn render(self, cx: &mut WindowContext) -> impl IntoElement {
let selected_provider = LanguageModelRegistry::read_global(cx)
.active_provider()
.map(|m| m.id());
let selected_model = LanguageModelRegistry::read_global(cx)
.active_model()
.map(|m| m.id());
let all_models = LanguageModelRegistry::global(cx)
.read(cx)
.providers()
.iter()
.flat_map(|provider| {
let provider_id = provider.id();
let icon = provider.icon();
let selected_model = selected_model.clone();
let selected_provider = selected_provider.clone();
provider.provided_models(cx).into_iter().map(move |model| {
let model = model.clone();
let icon = model.icon().unwrap_or(icon);
ModelInfo {
model: model.clone(),
icon,
availability: model.availability(),
is_selected: selected_model.as_ref() == Some(&model.id())
&& selected_provider.as_ref() == Some(&provider_id),
}
})
})
.collect::<Vec<_>>();
let delegate = ModelPickerDelegate {
fs: self.fs.clone(),
all_models: all_models.clone(),
filtered_models: all_models,
selected_index: 0,
};
let picker_view = cx.new_view(|cx| {
let picker = Picker::uniform_list(delegate, cx).max_height(Some(rems(20.).into()));
picker
});
PopoverMenu::new("model-switcher")
.menu(move |_cx| Some(picker_view.clone()))
.trigger(self.trigger)
.attach(gpui::AnchorCorner::BottomLeft)
.when_some(self.handle, |menu, handle| menu.with_handle(handle))
}
}

21
crates/assistant/src/terminal_inline_assistant.rs
View file

@ -1,6 +1,7 @@
use crate::assistant_settings::AssistantSettings;
use crate::{
humanize_token_count, prompts::PromptBuilder, AssistantPanel, AssistantPanelEvent,
ModelSelector, RequestType, DEFAULT_CONTEXT_LINES,
humanize_token_count, prompts::PromptBuilder, AssistantPanel, AssistantPanelEvent, RequestType,
DEFAULT_CONTEXT_LINES,
};
use anyhow::{Context as _, Result};
use client::telemetry::Telemetry;
@ -19,8 +20,9 @@ use language::Buffer;
use language_model::{
LanguageModelRegistry, LanguageModelRequest, LanguageModelRequestMessage, Role,
};
use language_model_selector::LanguageModelSelector;
use language_models::report_assistant_event;
use settings::Settings;
use settings::{update_settings_file, Settings};
use std::{
cmp,
sync::Arc,
@ -612,8 +614,17 @@ impl Render for PromptEditor {
.w_12()
.justify_center()
.gap_2()
.child(ModelSelector::new(
self.fs.clone(),
.child(LanguageModelSelector::new(
{
let fs = self.fs.clone();
move |model, cx| {
update_settings_file::<AssistantSettings>(
fs.clone(),
cx,
move |settings, _| settings.set_model(model.clone()),
);
}
},
IconButton::new("context", IconName::SettingsAlt)
.shape(IconButtonShape::Square)
.icon_size(IconSize::Small)