tenant.openagents/omega

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(&current_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 name() -> &'static str {
66        "ollama model picker"
67    }
68
69    fn match_count(&self) -> usize {
70        self.filtered_models.len()
71    }
72
73    fn selected_index(&self) -> usize {
74        self.selected_index
75    }
76
77    fn set_selected_index(
78        &mut self,
79        ix: usize,
80        _: &mut Window,
81        cx: &mut Context<OllamaModelPicker>,
82    ) {
83        self.selected_index = ix.min(self.filtered_models.len().saturating_sub(1));
84        cx.notify();
85    }
86
87    fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {
88        "Search models…".into()
89    }
90
91    fn update_matches(
92        &mut self,
93        query: String,
94        _window: &mut Window,
95        cx: &mut Context<OllamaModelPicker>,
96    ) -> Task<()> {
97        let query_lower = query.to_lowercase();
98
99        self.filtered_models = self
100            .models
101            .iter()
102            .enumerate()
103            .filter(|(_, model)| query.is_empty() || model.to_lowercase().contains(&query_lower))
104            .map(|(index, model)| StringMatch {
105                candidate_id: index,
106                string: model.to_string(),
107                positions: Vec::new(),
108                score: 0.0,
109            })
110            .collect();
111
112        self.selected_index = 0;
113        cx.notify();
114
115        Task::ready(())
116    }
117
118    fn confirm(
119        &mut self,
120        _secondary: bool,
121        window: &mut Window,
122        cx: &mut Context<OllamaModelPicker>,
123    ) {
124        let Some(model_match) = self.filtered_models.get(self.selected_index) else {
125            return;
126        };
127
128        (self.on_model_changed)(model_match.string.clone().into(), window, cx);
129        cx.emit(DismissEvent);
130    }
131
132    fn dismissed(&mut self, window: &mut Window, cx: &mut Context<OllamaModelPicker>) {
133        cx.defer_in(window, |picker, window, cx| {
134            picker.set_query("", window, cx);
135        });
136        cx.emit(DismissEvent);
137    }
138
139    fn render_match(
140        &self,
141        ix: usize,
142        selected: bool,
143        _window: &mut Window,
144        _cx: &mut Context<OllamaModelPicker>,
145    ) -> Option<Self::ListItem> {
146        let model_match = self.filtered_models.get(ix)?;
147
148        Some(
149            ListItem::new(ix)
150                .inset(true)
151                .spacing(ListItemSpacing::Sparse)
152                .toggle_state(selected)
153                .child(Label::new(model_match.string.clone()))
154                .into_any_element(),
155        )
156    }
157}
158
159pub fn render_ollama_model_picker(
160    field: SettingField<settings::OllamaModelName>,
161    file: SettingsUiFile,
162    _metadata: Option<&SettingsFieldMetadata>,
163    title: &'static str,
164    description: &'static str,
165    _window: &mut Window,
166    cx: &mut App,
167) -> AnyElement {
168    let (_, value) = SettingsStore::global(cx).get_value_from_file(file.to_settings(), field.pick);
169    let current_value: SharedString = value
170        .map(|m| m.0.clone().into())
171        .unwrap_or_else(|| "".into());
172
173    let trigger_value: SharedString = if current_value.is_empty() {
174        "Select a model…".into()
175    } else {
176        current_value.clone()
177    };
178
179    PopoverMenu::new("ollama-model-picker")
180        .trigger(
181            render_picker_trigger_button("ollama_model_picker_trigger".into(), trigger_value)
182                .aria_label(title)
183                .when(!description.is_empty(), |this| {
184                    this.aria_description(description)
185                }),
186        )
187        .menu(move |window, cx| {
188            Some(cx.new(|cx| {
189                let file = file.clone();
190                let current_value = current_value.clone();
191                let delegate = OllamaModelPickerDelegate::new(
192                    current_value,
193                    move |model_name, window, cx| {
194                        update_settings_file(
195                            file.clone(),
196                            field.json_path,
197                            window,
198                            cx,
199                            move |settings, app| {
200                                (field.write)(
201                                    settings,
202                                    Some(settings::OllamaModelName(model_name.to_string())),
203                                    app,
204                                );
205                            },
206                        )
207                        .log_err();
208                    },
209                    cx,
210                );
211
212                Picker::uniform_list(delegate, window, cx)
213                    .show_scrollbar(true)
214                    .initial_width(rems_from_px(210.))
215                    .max_height(rems(18.))
216                    .popover()
217            }))
218        })
219        .anchor(gpui::Anchor::TopLeft)
220        .offset(gpui::Point {
221            x: px(0.0),
222            y: px(2.0),
223        })
224        .with_handle(ui::PopoverMenuHandle::default())
225        .into_any_element()
226}
227