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                language_model_selector(
 33                    {
 34                        let model_context = model_usage_context.clone();
 35                        move |cx| model_context.configured_model(cx)
 36                    },
 37                    {
 38                        let fs = fs.clone();
 39                        move |model, cx| {
 40                            let provider = model.provider_id().0.to_string();
 41                            let model_id = model.id().0.to_string();
 42                            match &model_usage_context {
 43                                ModelUsageContext::InlineAssistant => {
 44                                    update_settings_file(fs.clone(), cx, move |settings, _cx| {
 45                                        settings
 46                                            .agent
 47                                            .get_or_insert_default()
 48                                            .set_inline_assistant_model(provider.clone(), model_id);
 49                                    });
 50                                }
 51                            }
 52                        }
 53                    },
 54                    {
 55                        let fs = fs.clone();
 56                        move |model, should_be_favorite, cx| {
 57                            crate::favorite_models::toggle_in_settings(
 58                                model,
 59                                should_be_favorite,
 60                                fs.clone(),
 61                                cx,
 62                            );
 63                        }
 64                    },
 65                    true, // Use popover styles for picker
 66                    focus_handle_clone,
 67                    window,
 68                    cx,
 69                )
 70            }),
 71            menu_handle,
 72            focus_handle,
 73        }
 74    }
 75
 76    pub fn toggle(&self, window: &mut Window, cx: &mut Context<Self>) {
 77        self.menu_handle.toggle(window, cx);
 78    }
 79
 80    pub fn active_model(&self, cx: &App) -> Option<language_model::ConfiguredModel> {
 81        self.selector.read(cx).delegate.active_model(cx)
 82    }
 83}
 84
 85impl Render for AgentModelSelector {
 86    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
 87        let model = self.selector.read(cx).delegate.active_model(cx);
 88        let model_name = model
 89            .as_ref()
 90            .map(|model| model.model.name().0)
 91            .unwrap_or_else(|| SharedString::from("Select a Model"));
 92
 93        let provider_icon_path = model.as_ref().and_then(|model| model.provider.icon_path());
 94        let provider_icon_name = model.as_ref().map(|model| model.provider.icon());
 95        let color = if self.menu_handle.is_deployed() {
 96            Color::Accent
 97        } else {
 98            Color::Muted
 99        };
100
101        let focus_handle = self.focus_handle.clone();
102
103        PickerPopoverMenu::new(
104            self.selector.clone(),
105            ButtonLike::new("active-model")
106                .when_some(provider_icon_path.clone(), |this, icon_path| {
107                    this.child(
108                        Icon::from_external_svg(icon_path)
109                            .color(color)
110                            .size(IconSize::XSmall),
111                    )
112                })
113                .when(provider_icon_path.is_none(), |this| {
114                    this.when_some(provider_icon_name, |this, icon| {
115                        this.child(Icon::new(icon).color(color).size(IconSize::XSmall))
116                    })
117                })
118                .selected_style(ButtonStyle::Tinted(TintColor::Accent))
119                .child(
120                    Label::new(model_name)
121                        .color(color)
122                        .size(LabelSize::Small)
123                        .ml_0p5(),
124                )
125                .child(
126                    Icon::new(IconName::ChevronDown)
127                        .color(color)
128                        .size(IconSize::XSmall),
129                ),
130            move |_window, cx| {
131                Tooltip::for_action_in("Change Model", &ToggleModelSelector, &focus_handle, cx)
132            },
133            gpui::Corner::TopRight,
134            cx,
135        )
136        .with_handle(self.menu_handle.clone())
137        .offset(gpui::Point {
138            x: px(0.0),
139            y: px(2.0),
140        })
141        .render(window, cx)
142    }
143}