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