assistant_model_selector.rs

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