ZIm/crates/assistant/src/model_selector.rs
2024-07-30 19:38:52 +02:00

149 lines
5.8 KiB
Rust

use std::sync::Arc;
use crate::assistant_settings::AssistantSettings;
use fs::Fs;
use gpui::SharedString;
use language_model::LanguageModelRegistry;
use settings::update_settings_file;
use ui::{prelude::*, ContextMenu, PopoverMenu, PopoverMenuHandle, PopoverTrigger};
#[derive(IntoElement)]
pub struct ModelSelector<T: PopoverTrigger> {
handle: Option<PopoverMenuHandle<ContextMenu>>,
fs: Arc<dyn Fs>,
trigger: T,
info_text: Option<SharedString>,
}
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<ContextMenu>) -> 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<T: PopoverTrigger> RenderOnce for ModelSelector<T> {
fn render(self, _: &mut WindowContext) -> impl IntoElement {
let mut menu = PopoverMenu::new("model-switcher");
if let Some(handle) = self.handle {
menu = menu.with_handle(handle);
}
let info_text = self.info_text.clone();
menu.menu(move |cx| {
ContextMenu::build(cx, |mut menu, cx| {
if let Some(info_text) = info_text.clone() {
menu = menu
.custom_row(move |_cx| {
Label::new(info_text.clone())
.color(Color::Muted)
.into_any_element()
})
.separator();
}
for (index, provider) in LanguageModelRegistry::global(cx)
.read(cx)
.providers()
.into_iter()
.enumerate()
{
if index > 0 {
menu = menu.separator();
}
menu = menu.header(provider.name().0);
let available_models = provider.provided_models(cx);
if available_models.is_empty() {
menu = menu.custom_entry(
{
move |_| {
h_flex()
.w_full()
.gap_1()
.child(Icon::new(IconName::Settings))
.child(Label::new("Configure"))
.into_any()
}
},
{
let provider = provider.clone();
move |cx| {
LanguageModelRegistry::global(cx).update(
cx,
|completion_provider, cx| {
completion_provider
.set_active_provider(Some(provider.clone()), cx);
},
);
}
},
);
}
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());
for available_model in available_models {
menu = menu.custom_entry(
{
let id = available_model.id();
let provider_id = available_model.provider_id();
let model_name = available_model.name().0.clone();
let selected_model = selected_model.clone();
let selected_provider = selected_provider.clone();
move |_| {
h_flex()
.w_full()
.justify_between()
.child(Label::new(model_name.clone()))
.when(
selected_model.as_ref() == Some(&id)
&& selected_provider.as_ref() == Some(&provider_id),
|this| this.child(Icon::new(IconName::Check)),
)
.into_any()
}
},
{
let fs = self.fs.clone();
let model = available_model.clone();
move |cx| {
let model = model.clone();
update_settings_file::<AssistantSettings>(
fs.clone(),
cx,
move |settings, _| settings.set_model(model),
);
}
},
);
}
}
menu
})
.into()
})
.trigger(self.trigger)
.attach(gpui::AnchorCorner::BottomLeft)
}
}