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 focus_handle = self.focus_handle.clone();
84
85 let model_registry = LanguageModelRegistry::read_global(cx);
86 let model = match self.model_type {
87 ModelType::Default => model_registry.default_model(),
88 ModelType::InlineAssistant => model_registry.inline_assistant_model(),
89 };
90 let (model_name, model_icon) = match model {
91 Some(model) => (model.model.name().0, Some(model.provider.icon())),
92 _ => (SharedString::from("No model selected"), None),
93 };
94
95 LanguageModelSelectorPopoverMenu::new(
96 self.selector.clone(),
97 ButtonLike::new("active-model")
98 .style(ButtonStyle::Subtle)
99 .child(
100 h_flex()
101 .gap_0p5()
102 .children(
103 model_icon.map(|icon| {
104 Icon::new(icon).color(Color::Muted).size(IconSize::Small)
105 }),
106 )
107 .child(
108 Label::new(model_name)
109 .size(LabelSize::Small)
110 .color(Color::Muted)
111 .ml_1(),
112 )
113 .child(
114 Icon::new(IconName::ChevronDown)
115 .color(Color::Muted)
116 .size(IconSize::XSmall),
117 ),
118 ),
119 move |window, cx| {
120 Tooltip::for_action_in(
121 "Change Model",
122 &ToggleModelSelector,
123 &focus_handle,
124 window,
125 cx,
126 )
127 },
128 gpui::Corner::BottomRight,
129 )
130 .with_handle(self.menu_handle.clone())
131 }
132}