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