1use assistant_settings::AssistantSettings;
  2use fs::Fs;
  3use gpui::{Entity, FocusHandle, SharedString};
  4use language_model::LanguageModelRegistry;
  5use language_model_selector::{
  6    LanguageModelSelector, LanguageModelSelectorPopoverMenu, ToggleModelSelector,
  7};
  8use settings::update_settings_file;
  9use std::sync::Arc;
 10use ui::{ButtonLike, PopoverMenuHandle, Tooltip, prelude::*};
 11
 12#[derive(Clone, Copy)]
 13pub enum ModelType {
 14    Default,
 15    InlineAssistant,
 16}
 17
 18pub struct AssistantModelSelector {
 19    selector: Entity<LanguageModelSelector>,
 20    menu_handle: PopoverMenuHandle<LanguageModelSelector>,
 21    focus_handle: FocusHandle,
 22    model_type: ModelType,
 23}
 24
 25impl AssistantModelSelector {
 26    pub(crate) fn new(
 27        fs: Arc<dyn Fs>,
 28        menu_handle: PopoverMenuHandle<LanguageModelSelector>,
 29        focus_handle: FocusHandle,
 30        model_type: ModelType,
 31        window: &mut Window,
 32        cx: &mut App,
 33    ) -> Self {
 34        Self {
 35            selector: cx.new(|cx| {
 36                let fs = fs.clone();
 37                LanguageModelSelector::new(
 38                    move |model, cx| {
 39                        let provider = model.provider_id().0.to_string();
 40                        let model_id = model.id().0.to_string();
 41
 42                        match model_type {
 43                            ModelType::Default => {
 44                                update_settings_file::<AssistantSettings>(
 45                                    fs.clone(),
 46                                    cx,
 47                                    move |settings, _cx| {
 48                                        settings.set_model(model.clone());
 49                                    },
 50                                );
 51                            }
 52                            ModelType::InlineAssistant => {
 53                                update_settings_file::<AssistantSettings>(
 54                                    fs.clone(),
 55                                    cx,
 56                                    move |settings, _cx| {
 57                                        settings.set_inline_assistant_model(
 58                                            provider.clone(),
 59                                            model_id.clone(),
 60                                        );
 61                                    },
 62                                );
 63                            }
 64                        }
 65                    },
 66                    window,
 67                    cx,
 68                )
 69            }),
 70            menu_handle,
 71            focus_handle,
 72            model_type,
 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
 81impl Render for AssistantModelSelector {
 82    fn render(&mut self, _window: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
 83        let model_registry = LanguageModelRegistry::read_global(cx);
 84
 85        let model = match self.model_type {
 86            ModelType::Default => model_registry.default_model(),
 87            ModelType::InlineAssistant => model_registry.inline_assistant_model(),
 88        };
 89
 90        let focus_handle = self.focus_handle.clone();
 91        let model_name = match model {
 92            Some(model) => model.model.name().0,
 93            _ => SharedString::from("No model selected"),
 94        };
 95
 96        LanguageModelSelectorPopoverMenu::new(
 97            self.selector.clone(),
 98            ButtonLike::new("active-model")
 99                .style(ButtonStyle::Subtle)
100                .child(
101                    h_flex()
102                        .gap_0p5()
103                        .child(
104                            Label::new(model_name)
105                                .size(LabelSize::Small)
106                                .color(Color::Muted),
107                        )
108                        .child(
109                            Icon::new(IconName::ChevronDown)
110                                .color(Color::Muted)
111                                .size(IconSize::XSmall),
112                        ),
113                ),
114            move |window, cx| {
115                Tooltip::for_action_in(
116                    "Change Model",
117                    &ToggleModelSelector,
118                    &focus_handle,
119                    window,
120                    cx,
121                )
122            },
123            gpui::Corner::BottomRight,
124        )
125        .with_handle(self.menu_handle.clone())
126    }
127}