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