agent_model_selector.rs

  1use crate::{
  2    ModelUsageContext,
  3    language_model_selector::{LanguageModelSelector, language_model_selector},
  4};
  5use fs::Fs;
  6use gpui::{Entity, FocusHandle, SharedString};
  7use picker::popover_menu::PickerPopoverMenu;
  8use settings::update_settings_file;
  9use std::sync::Arc;
 10use ui::{ButtonLike, PopoverMenuHandle, TintColor, Tooltip, prelude::*};
 11use zed_actions::agent::ToggleModelSelector;
 12
 13pub struct AgentModelSelector {
 14    selector: Entity<LanguageModelSelector>,
 15    menu_handle: PopoverMenuHandle<LanguageModelSelector>,
 16    focus_handle: FocusHandle,
 17}
 18
 19impl AgentModelSelector {
 20    pub(crate) fn new(
 21        fs: Arc<dyn Fs>,
 22        menu_handle: PopoverMenuHandle<LanguageModelSelector>,
 23        focus_handle: FocusHandle,
 24        model_usage_context: ModelUsageContext,
 25        window: &mut Window,
 26        cx: &mut Context<Self>,
 27    ) -> Self {
 28        let focus_handle_clone = focus_handle.clone();
 29
 30        Self {
 31            selector: cx.new(move |cx| {
 32                let fs = fs.clone();
 33                language_model_selector(
 34                    {
 35                        let model_context = model_usage_context.clone();
 36                        move |cx| model_context.configured_model(cx)
 37                    },
 38                    move |model, cx| {
 39                        let provider = model.provider_id().0.to_string();
 40                        let model_id = model.id().0.to_string();
 41                        match &model_usage_context {
 42                            ModelUsageContext::InlineAssistant => {
 43                                update_settings_file(fs.clone(), cx, move |settings, _cx| {
 44                                    settings
 45                                        .agent
 46                                        .get_or_insert_default()
 47                                        .set_inline_assistant_model(provider.clone(), model_id);
 48                                });
 49                            }
 50                        }
 51                    },
 52                    true, // Use popover styles for picker
 53                    focus_handle_clone,
 54                    window,
 55                    cx,
 56                )
 57            }),
 58            menu_handle,
 59            focus_handle,
 60        }
 61    }
 62
 63    pub fn toggle(&self, window: &mut Window, cx: &mut Context<Self>) {
 64        self.menu_handle.toggle(window, cx);
 65    }
 66}
 67
 68impl Render for AgentModelSelector {
 69    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
 70        let model = self.selector.read(cx).delegate.active_model(cx);
 71        let model_name = model
 72            .as_ref()
 73            .map(|model| model.model.name().0)
 74            .unwrap_or_else(|| SharedString::from("Select a Model"));
 75
 76        let provider_icon = model.as_ref().map(|model| model.provider.icon());
 77        let color = if self.menu_handle.is_deployed() {
 78            Color::Accent
 79        } else {
 80            Color::Muted
 81        };
 82
 83        let focus_handle = self.focus_handle.clone();
 84
 85        PickerPopoverMenu::new(
 86            self.selector.clone(),
 87            ButtonLike::new("active-model")
 88                .when_some(provider_icon, |this, icon| {
 89                    this.child(Icon::new(icon).color(color).size(IconSize::XSmall))
 90                })
 91                .selected_style(ButtonStyle::Tinted(TintColor::Accent))
 92                .child(
 93                    Label::new(model_name)
 94                        .color(color)
 95                        .size(LabelSize::Small)
 96                        .ml_0p5(),
 97                )
 98                .child(
 99                    Icon::new(IconName::ChevronDown)
100                        .color(color)
101                        .size(IconSize::XSmall),
102                ),
103            move |_window, cx| {
104                Tooltip::for_action_in("Change Model", &ToggleModelSelector, &focus_handle, cx)
105            },
106            gpui::Corner::TopRight,
107            cx,
108        )
109        .with_handle(self.menu_handle.clone())
110        .offset(gpui::Point {
111            x: px(0.0),
112            y: px(2.0),
113        })
114        .render(window, cx)
115    }
116}