1use crate::{
  2    ModelUsageContext,
  3    language_model_selector::{
  4        LanguageModelSelector, ToggleModelSelector, language_model_selector,
  5    },
  6};
  7use agent_settings::AgentSettings;
  8use fs::Fs;
  9use gpui::{Entity, FocusHandle, SharedString};
 10use language_model::{ConfiguredModel, LanguageModelRegistry};
 11use picker::popover_menu::PickerPopoverMenu;
 12use settings::update_settings_file;
 13use std::sync::Arc;
 14use ui::{PopoverMenuHandle, Tooltip, prelude::*};
 15
 16pub struct AgentModelSelector {
 17    selector: Entity<LanguageModelSelector>,
 18    menu_handle: PopoverMenuHandle<LanguageModelSelector>,
 19    focus_handle: FocusHandle,
 20}
 21
 22impl AgentModelSelector {
 23    pub(crate) fn new(
 24        fs: Arc<dyn Fs>,
 25        menu_handle: PopoverMenuHandle<LanguageModelSelector>,
 26        focus_handle: FocusHandle,
 27        model_usage_context: ModelUsageContext,
 28        window: &mut Window,
 29        cx: &mut Context<Self>,
 30    ) -> Self {
 31        Self {
 32            selector: cx.new(move |cx| {
 33                let fs = fs.clone();
 34                language_model_selector(
 35                    {
 36                        let model_context = model_usage_context.clone();
 37                        move |cx| model_context.configured_model(cx)
 38                    },
 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::Thread(thread) => {
 44                                thread.update(cx, |thread, cx| {
 45                                    let registry = LanguageModelRegistry::read_global(cx);
 46                                    if let Some(provider) = registry.provider(&model.provider_id())
 47                                    {
 48                                        thread.set_configured_model(
 49                                            Some(ConfiguredModel {
 50                                                provider,
 51                                                model: model.clone(),
 52                                            }),
 53                                            cx,
 54                                        );
 55                                    }
 56                                });
 57                                update_settings_file::<AgentSettings>(
 58                                    fs.clone(),
 59                                    cx,
 60                                    move |settings, _cx| {
 61                                        settings.set_model(model.clone());
 62                                    },
 63                                );
 64                            }
 65                            ModelUsageContext::InlineAssistant => {
 66                                update_settings_file::<AgentSettings>(
 67                                    fs.clone(),
 68                                    cx,
 69                                    move |settings, _cx| {
 70                                        settings.set_inline_assistant_model(
 71                                            provider.clone(),
 72                                            model_id.clone(),
 73                                        );
 74                                    },
 75                                );
 76                            }
 77                        }
 78                    },
 79                    window,
 80                    cx,
 81                )
 82            }),
 83            menu_handle,
 84            focus_handle,
 85        }
 86    }
 87
 88    pub fn toggle(&self, window: &mut Window, cx: &mut Context<Self>) {
 89        self.menu_handle.toggle(window, cx);
 90    }
 91}
 92
 93impl Render for AgentModelSelector {
 94    fn render(&mut self, window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
 95        let model = self.selector.read(cx).delegate.active_model(cx);
 96        let model_name = model
 97            .map(|model| model.model.name().0)
 98            .unwrap_or_else(|| SharedString::from("No model selected"));
 99
100        let focus_handle = self.focus_handle.clone();
101
102        PickerPopoverMenu::new(
103            self.selector.clone(),
104            Button::new("active-model", model_name)
105                .label_size(LabelSize::Small)
106                .color(Color::Muted)
107                .icon(IconName::ChevronDown)
108                .icon_size(IconSize::XSmall)
109                .icon_position(IconPosition::End)
110                .icon_color(Color::Muted),
111            move |window, cx| {
112                Tooltip::for_action_in(
113                    "Change Model",
114                    &ToggleModelSelector,
115                    &focus_handle,
116                    window,
117                    cx,
118                )
119            },
120            gpui::Corner::BottomRight,
121            cx,
122        )
123        .with_handle(self.menu_handle.clone())
124        .render(window, cx)
125    }
126}