1use std::sync::Arc;
2
3use fuzzy::StringMatch;
4use gpui::{AnyElement, App, Context, DismissEvent, ReadGlobal, SharedString, Task, Window, px};
5use picker::{Picker, PickerDelegate};
6use settings::SettingsStore;
7use ui::{ListItem, ListItemSpacing, PopoverMenu, prelude::*};
8use util::ResultExt;
9
10use crate::{
11 SettingField, SettingsFieldMetadata, SettingsUiFile, render_picker_trigger_button,
12 update_settings_file,
13};
14
15type OllamaModelPicker = Picker<OllamaModelPickerDelegate>;
16
17struct OllamaModelPickerDelegate {
18 models: Vec<SharedString>,
19 filtered_models: Vec<StringMatch>,
20 selected_index: usize,
21 on_model_changed: Arc<dyn Fn(SharedString, &mut Window, &mut App) + 'static>,
22}
23
24impl OllamaModelPickerDelegate {
25 fn new(
26 current_model: SharedString,
27 on_model_changed: impl Fn(SharedString, &mut Window, &mut App) + 'static,
28 cx: &mut Context<OllamaModelPicker>,
29 ) -> Self {
30 let mut models = edit_prediction::ollama::fetch_models(cx);
31
32 let current_in_list = models.contains(¤t_model);
33 if !current_model.is_empty() && !current_in_list {
34 models.insert(0, current_model.clone());
35 }
36
37 let selected_index = models
38 .iter()
39 .position(|model| *model == current_model)
40 .unwrap_or(0);
41
42 let filtered_models = models
43 .iter()
44 .enumerate()
45 .map(|(index, model)| StringMatch {
46 candidate_id: index,
47 string: model.to_string(),
48 positions: Vec::new(),
49 score: 0.0,
50 })
51 .collect();
52
53 Self {
54 models,
55 filtered_models,
56 selected_index,
57 on_model_changed: Arc::new(on_model_changed),
58 }
59 }
60}
61
62impl PickerDelegate for OllamaModelPickerDelegate {
63 type ListItem = AnyElement;
64
65 fn match_count(&self) -> usize {
66 self.filtered_models.len()
67 }
68
69 fn selected_index(&self) -> usize {
70 self.selected_index
71 }
72
73 fn set_selected_index(
74 &mut self,
75 ix: usize,
76 _: &mut Window,
77 cx: &mut Context<OllamaModelPicker>,
78 ) {
79 self.selected_index = ix.min(self.filtered_models.len().saturating_sub(1));
80 cx.notify();
81 }
82
83 fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
84 "Search models…".into()
85 }
86
87 fn update_matches(
88 &mut self,
89 query: String,
90 _window: &mut Window,
91 cx: &mut Context<OllamaModelPicker>,
92 ) -> Task<()> {
93 let query_lower = query.to_lowercase();
94
95 self.filtered_models = self
96 .models
97 .iter()
98 .enumerate()
99 .filter(|(_, model)| query.is_empty() || model.to_lowercase().contains(&query_lower))
100 .map(|(index, model)| StringMatch {
101 candidate_id: index,
102 string: model.to_string(),
103 positions: Vec::new(),
104 score: 0.0,
105 })
106 .collect();
107
108 self.selected_index = 0;
109 cx.notify();
110
111 Task::ready(())
112 }
113
114 fn confirm(
115 &mut self,
116 _secondary: bool,
117 window: &mut Window,
118 cx: &mut Context<OllamaModelPicker>,
119 ) {
120 let Some(model_match) = self.filtered_models.get(self.selected_index) else {
121 return;
122 };
123
124 (self.on_model_changed)(model_match.string.clone().into(), window, cx);
125 cx.emit(DismissEvent);
126 }
127
128 fn dismissed(&mut self, window: &mut Window, cx: &mut Context<OllamaModelPicker>) {
129 cx.defer_in(window, |picker, window, cx| {
130 picker.set_query("", window, cx);
131 });
132 cx.emit(DismissEvent);
133 }
134
135 fn render_match(
136 &self,
137 ix: usize,
138 selected: bool,
139 _window: &mut Window,
140 _cx: &mut Context<OllamaModelPicker>,
141 ) -> Option<Self::ListItem> {
142 let model_match = self.filtered_models.get(ix)?;
143
144 Some(
145 ListItem::new(ix)
146 .inset(true)
147 .spacing(ListItemSpacing::Sparse)
148 .toggle_state(selected)
149 .child(Label::new(model_match.string.clone()))
150 .into_any_element(),
151 )
152 }
153}
154
155pub fn render_ollama_model_picker(
156 field: SettingField<settings::OllamaModelName>,
157 file: SettingsUiFile,
158 _metadata: Option<&SettingsFieldMetadata>,
159 _window: &mut Window,
160 cx: &mut App,
161) -> AnyElement {
162 let (_, value) = SettingsStore::global(cx).get_value_from_file(file.to_settings(), field.pick);
163 let current_value: SharedString = value
164 .map(|m| m.0.clone().into())
165 .unwrap_or_else(|| "".into());
166
167 PopoverMenu::new("ollama-model-picker")
168 .trigger(render_picker_trigger_button(
169 "ollama_model_picker_trigger".into(),
170 if current_value.is_empty() {
171 "Select a model…".into()
172 } else {
173 current_value.clone()
174 },
175 ))
176 .menu(move |window, cx| {
177 Some(cx.new(|cx| {
178 let file = file.clone();
179 let current_value = current_value.clone();
180 let delegate = OllamaModelPickerDelegate::new(
181 current_value,
182 move |model_name, window, cx| {
183 update_settings_file(
184 file.clone(),
185 field.json_path,
186 window,
187 cx,
188 move |settings, _cx| {
189 (field.write)(
190 settings,
191 Some(settings::OllamaModelName(model_name.to_string())),
192 );
193 },
194 )
195 .log_err();
196 },
197 cx,
198 );
199
200 Picker::uniform_list(delegate, window, cx)
201 .show_scrollbar(true)
202 .width(rems_from_px(210.))
203 .max_height(Some(rems(18.).into()))
204 }))
205 })
206 .anchor(gpui::Corner::TopLeft)
207 .offset(gpui::Point {
208 x: px(0.0),
209 y: px(2.0),
210 })
211 .with_handle(ui::PopoverMenuHandle::default())
212 .into_any_element()
213}