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        Self {
 29            selector: cx.new(move |cx| {
 30                let fs = fs.clone();
 31                language_model_selector(
 32                    {
 33                        let model_context = model_usage_context.clone();
 34                        move |cx| model_context.configured_model(cx)
 35                    },
 36                    move |model, cx| {
 37                        let provider = model.provider_id().0.to_string();
 38                        let model_id = model.id().0.to_string();
 39                        match &model_usage_context {
 40                            ModelUsageContext::InlineAssistant => {
 41                                update_settings_file(fs.clone(), cx, move |settings, _cx| {
 42                                    settings
 43                                        .agent
 44                                        .get_or_insert_default()
 45                                        .set_inline_assistant_model(provider.clone(), model_id);
 46                                });
 47                            }
 48                        }
 49                    },
 50                    window,
 51                    cx,
 52                )
 53            }),
 54            menu_handle,
 55            focus_handle,
 56        }
 57    }
 58
 59    pub fn toggle(&self, window: &mut Window, cx: &mut Context<Self>) {
 60        self.menu_handle.toggle(window, cx);
 61    }
 62}
 63
 64impl Render for AgentModelSelector {
 65    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
 66        let model = self.selector.read(cx).delegate.active_model(cx);
 67        let model_name = model
 68            .as_ref()
 69            .map(|model| model.model.name().0)
 70            .unwrap_or_else(|| SharedString::from("Select a Model"));
 71
 72        let provider_icon = model.as_ref().map(|model| model.provider.icon());
 73        let color = if self.menu_handle.is_deployed() {
 74            Color::Accent
 75        } else {
 76            Color::Muted
 77        };
 78
 79        let focus_handle = self.focus_handle.clone();
 80
 81        PickerPopoverMenu::new(
 82            self.selector.clone(),
 83            ButtonLike::new("active-model")
 84                .when_some(provider_icon, |this, icon| {
 85                    this.child(Icon::new(icon).color(color).size(IconSize::XSmall))
 86                })
 87                .selected_style(ButtonStyle::Tinted(TintColor::Accent))
 88                .child(
 89                    Label::new(model_name)
 90                        .color(color)
 91                        .size(LabelSize::Small)
 92                        .ml_0p5(),
 93                )
 94                .child(
 95                    Icon::new(IconName::ChevronDown)
 96                        .color(color)
 97                        .size(IconSize::XSmall),
 98                ),
 99            move |window, cx| {
100                Tooltip::for_action_in(
101                    "Change Model",
102                    &ToggleModelSelector,
103                    &focus_handle,
104                    window,
105                    cx,
106                )
107            },
108            gpui::Corner::TopRight,
109            cx,
110        )
111        .with_handle(self.menu_handle.clone())
112        .offset(gpui::Point {
113            x: px(0.0),
114            y: px(2.0),
115        })
116        .render(window, cx)
117    }
118}