902 lines · 29.2 KB · Rust
RawDownload
use std::{cmp::Reverse, sync::Arc}; use agent_settings::AgentSettings;use collections::{HashMap, HashSet, IndexMap};use fuzzy::{StringMatch, StringMatchCandidate, match_strings};use gpui::{    Action, AnyElement, App, BackgroundExecutor, DismissEvent, FocusHandle, ForegroundExecutor,    Subscription, Task,};use language_model::{    ConfiguredModel, IconOrSvg, LanguageModel, LanguageModelId, LanguageModelProvider,    LanguageModelProviderId, LanguageModelRegistry,};use ordered_float::OrderedFloat;use picker::{Picker, PickerDelegate};use settings::Settings;use ui::prelude::*;use zed_actions::agent::OpenSettings; use crate::ui::{ModelSelectorFooter, ModelSelectorHeader, ModelSelectorListItem}; type OnModelChanged = Arc<dyn Fn(Arc<dyn LanguageModel>, &mut App) + 'static>;type GetActiveModel = Arc<dyn Fn(&App) -> Option<ConfiguredModel> + 'static>;type OnToggleFavorite = Arc<dyn Fn(Arc<dyn LanguageModel>, bool, &mut App) + 'static>; pub type LanguageModelSelector = Picker<LanguageModelPickerDelegate>; pub fn language_model_selector(    get_active_model: impl Fn(&App) -> Option<ConfiguredModel> + 'static,    on_model_changed: impl Fn(Arc<dyn LanguageModel>, &mut App) + 'static,    on_toggle_favorite: impl Fn(Arc<dyn LanguageModel>, bool, &mut App) + 'static,    popover_styles: bool,    focus_handle: FocusHandle,    window: &mut Window,    cx: &mut Context<LanguageModelSelector>,) -> LanguageModelSelector {    let delegate = LanguageModelPickerDelegate::new(        get_active_model,        on_model_changed,        on_toggle_favorite,        popover_styles,        focus_handle,        window,        cx,    );     if popover_styles {        Picker::list(delegate, window, cx)            .show_scrollbar(true)            .initial_width(rems(20.))            .popover()    } else {        Picker::list(delegate, window, cx)            .show_scrollbar(true)            .embedded()    }} fn all_models(cx: &App) -> GroupedModels {    let lm_registry = LanguageModelRegistry::global(cx).read(cx);    let providers = lm_registry.visible_providers();     let mut favorites_index = FavoritesIndex::default();     for sel in &AgentSettings::get_global(cx).favorite_models {        favorites_index            .entry(sel.provider.0.clone().into())            .or_default()            .insert(sel.model.clone().into());    }     let recommended = providers        .iter()        .flat_map(|provider| {            provider                .recommended_models(cx)                .into_iter()                .map(|model| ModelInfo::new(&**provider, model, &favorites_index))        })        .collect();     let all = providers        .iter()        .flat_map(|provider| {            provider                .provided_models(cx)                .into_iter()                .map(|model| ModelInfo::new(&**provider, model, &favorites_index))        })        .collect();     GroupedModels::new(all, recommended)} type FavoritesIndex = HashMap<LanguageModelProviderId, HashSet<LanguageModelId>>; #[derive(Clone)]struct ModelInfo {    model: Arc<dyn LanguageModel>,    icon: IconOrSvg,    is_favorite: bool,} impl ModelInfo {    fn new(        provider: &dyn LanguageModelProvider,        model: Arc<dyn LanguageModel>,        favorites_index: &FavoritesIndex,    ) -> Self {        let is_favorite = favorites_index            .get(&provider.id())            .map_or(false, |set| set.contains(&model.id()));         Self {            model,            icon: provider.icon(),            is_favorite,        }    }} pub struct LanguageModelPickerDelegate {    on_model_changed: OnModelChanged,    get_active_model: GetActiveModel,    on_toggle_favorite: OnToggleFavorite,    all_models: Arc<GroupedModels>,    filtered_entries: Vec<LanguageModelPickerEntry>,    selected_index: usize,    _subscriptions: Vec<Subscription>,    popover_styles: bool,    focus_handle: FocusHandle,} impl LanguageModelPickerDelegate {    fn new(        get_active_model: impl Fn(&App) -> Option<ConfiguredModel> + 'static,        on_model_changed: impl Fn(Arc<dyn LanguageModel>, &mut App) + 'static,        on_toggle_favorite: impl Fn(Arc<dyn LanguageModel>, bool, &mut App) + 'static,        popover_styles: bool,        focus_handle: FocusHandle,        window: &mut Window,        cx: &mut Context<Picker<Self>>,    ) -> Self {        let on_model_changed = Arc::new(on_model_changed);        let models = all_models(cx);        let entries = models.entries();         Self {            on_model_changed,            all_models: Arc::new(models),            selected_index: Self::get_active_model_index(&entries, get_active_model(cx)),            filtered_entries: entries,            get_active_model: Arc::new(get_active_model),            on_toggle_favorite: Arc::new(on_toggle_favorite),            _subscriptions: vec![cx.subscribe_in(                &LanguageModelRegistry::global(cx),                window,                |picker, _, event, window, cx| {                    match event {                        language_model::Event::ProviderStateChanged(_)                        | language_model::Event::AddedProvider(_)                        | language_model::Event::RemovedProvider(_) => {                            let query = picker.query(cx);                            picker.delegate.all_models = Arc::new(all_models(cx));                            // Update matches will automatically drop the previous task                            // if we get a provider event again                            picker.update_matches(query, window, cx)                        }                        _ => {}                    }                },            )],            popover_styles,            focus_handle,        }    }     fn get_active_model_index(        entries: &[LanguageModelPickerEntry],        active_model: Option<ConfiguredModel>,    ) -> usize {        entries            .iter()            .position(|entry| {                if let LanguageModelPickerEntry::Model(model) = entry {                    active_model                        .as_ref()                        .map(|active_model| {                            active_model.model.id() == model.model.id()                                && active_model.provider.id() == model.model.provider_id()                        })                        .unwrap_or_default()                } else {                    false                }            })            .unwrap_or(0)    }     pub fn active_model(&self, cx: &App) -> Option<ConfiguredModel> {        (self.get_active_model)(cx)    }     pub fn favorites_count(&self) -> usize {        self.all_models.favorites.len()    }     pub fn cycle_favorite_models(&mut self, window: &mut Window, cx: &mut Context<Picker<Self>>) {        if self.all_models.favorites.is_empty() {            return;        }         let active_model = (self.get_active_model)(cx);        let active_provider_id = active_model.as_ref().map(|m| m.provider.id());        let active_model_id = active_model.as_ref().map(|m| m.model.id());         let current_index = self            .all_models            .favorites            .iter()            .position(|info| {                Some(info.model.provider_id()) == active_provider_id                    && Some(info.model.id()) == active_model_id            })            .unwrap_or(usize::MAX);         let next_index = if current_index == usize::MAX {            0        } else {            (current_index + 1) % self.all_models.favorites.len()        };         let next_model = self.all_models.favorites[next_index].model.clone();         (self.on_model_changed)(next_model, cx);         // Align the picker selection with the newly-active model        let new_index =            Self::get_active_model_index(&self.filtered_entries, (self.get_active_model)(cx));        self.set_selected_index(new_index, window, cx);    }} struct GroupedModels {    favorites: Vec<ModelInfo>,    recommended: Vec<ModelInfo>,    all: IndexMap<LanguageModelProviderId, Vec<ModelInfo>>,} impl GroupedModels {    pub fn new(all: Vec<ModelInfo>, recommended: Vec<ModelInfo>) -> Self {        let favorites = all            .iter()            .filter(|info| info.is_favorite)            .cloned()            .collect();         let mut all_by_provider: IndexMap<_, Vec<ModelInfo>> = IndexMap::default();        for model in all {            let provider = model.model.provider_id();            if let Some(models) = all_by_provider.get_mut(&provider) {                models.push(model);            } else {                all_by_provider.insert(provider, vec![model]);            }        }         Self {            favorites,            recommended,            all: all_by_provider,        }    }     fn entries(&self) -> Vec<LanguageModelPickerEntry> {        let mut entries = Vec::new();         if !self.favorites.is_empty() {            entries.push(LanguageModelPickerEntry::Separator("Favorite".into()));            for info in &self.favorites {                entries.push(LanguageModelPickerEntry::Model(info.clone()));            }        }         if !self.recommended.is_empty() {            entries.push(LanguageModelPickerEntry::Separator("Recommended".into()));            for info in &self.recommended {                entries.push(LanguageModelPickerEntry::Model(info.clone()));            }        }         for models in self.all.values() {            if models.is_empty() {                continue;            }            entries.push(LanguageModelPickerEntry::Separator(                models[0].model.provider_name().0,            ));            for info in models {                entries.push(LanguageModelPickerEntry::Model(info.clone()));            }        }         entries    }} enum LanguageModelPickerEntry {    Model(ModelInfo),    Separator(SharedString),} struct ModelMatcher {    models: Vec<ModelInfo>,    fg_executor: ForegroundExecutor,    bg_executor: BackgroundExecutor,    candidates: Vec<StringMatchCandidate>,} impl ModelMatcher {    fn new(        models: Vec<ModelInfo>,        fg_executor: ForegroundExecutor,        bg_executor: BackgroundExecutor,    ) -> ModelMatcher {        let candidates = Self::make_match_candidates(&models);        Self {            models,            fg_executor,            bg_executor,            candidates,        }    }     pub fn fuzzy_search(&self, query: &str) -> Vec<ModelInfo> {        let mut matches = self.fg_executor.block_on(match_strings(            &self.candidates,            query,            false,            true,            100,            &Default::default(),            self.bg_executor.clone(),        ));         let sorting_key = |mat: &StringMatch| {            let candidate = &self.candidates[mat.candidate_id];            (Reverse(OrderedFloat(mat.score)), candidate.id)        };        matches.sort_unstable_by_key(sorting_key);         let matched_models: Vec<_> = matches            .into_iter()            .map(|mat| self.models[mat.candidate_id].clone())            .collect();         matched_models    }     pub fn exact_search(&self, query: &str) -> Vec<ModelInfo> {        self.models            .iter()            .filter(|m| {                m.model                    .name()                    .0                    .to_lowercase()                    .contains(&query.to_lowercase())            })            .cloned()            .collect::<Vec<_>>()    }     fn make_match_candidates(model_infos: &Vec<ModelInfo>) -> Vec<StringMatchCandidate> {        model_infos            .iter()            .enumerate()            .map(|(index, model)| {                StringMatchCandidate::new(                    index,                    &format!(                        "{}/{}",                        &model.model.provider_name().0,                        &model.model.name().0                    ),                )            })            .collect::<Vec<_>>()    }} impl PickerDelegate for LanguageModelPickerDelegate {    type ListItem = AnyElement;     fn name() -> &'static str {        "language model selector"    }     fn match_count(&self) -> usize {        self.filtered_entries.len()    }     fn selected_index(&self) -> usize {        self.selected_index    }     fn set_selected_index(&mut self, ix: usize, _: &mut Window, cx: &mut Context<Picker<Self>>) {        self.selected_index = ix.min(self.filtered_entries.len().saturating_sub(1));        cx.notify();    }     fn can_select(&self, ix: usize, _window: &mut Window, _cx: &mut Context<Picker<Self>>) -> bool {        match self.filtered_entries.get(ix) {            Some(LanguageModelPickerEntry::Model(_)) => true,            Some(LanguageModelPickerEntry::Separator(_)) | None => false,        }    }     fn placeholder_text(&self, _window: &mut Window, _cx: &mut App) -> Arc<str> {        "Select a model…".into()    }     fn update_matches(        &mut self,        query: String,        window: &mut Window,        cx: &mut Context<Picker<Self>>,    ) -> Task<()> {        let all_models = self.all_models.clone();        let active_model = (self.get_active_model)(cx);        let fg_executor = cx.foreground_executor();        let bg_executor = cx.background_executor();         let language_model_registry = LanguageModelRegistry::global(cx);         let configured_providers = language_model_registry            .read(cx)            .visible_providers()            .into_iter()            .filter(|provider| provider.is_authenticated(cx))            .collect::<Vec<_>>();         let configured_provider_ids = configured_providers            .iter()            .map(|provider| provider.id())            .collect::<Vec<_>>();         let recommended_models = all_models            .recommended            .iter()            .filter(|m| configured_provider_ids.contains(&m.model.provider_id()))            .cloned()            .collect::<Vec<_>>();         let available_models = all_models            .all            .values()            .flat_map(|models| models.iter())            .filter(|m| configured_provider_ids.contains(&m.model.provider_id()))            .cloned()            .collect::<Vec<_>>();         let matcher_rec =            ModelMatcher::new(recommended_models, fg_executor.clone(), bg_executor.clone());        let matcher_all =            ModelMatcher::new(available_models, fg_executor.clone(), bg_executor.clone());         let recommended = matcher_rec.exact_search(&query);        let all = matcher_all.fuzzy_search(&query);         let filtered_models = GroupedModels::new(all, recommended);         cx.spawn_in(window, async move |this, cx| {            this.update_in(cx, |this, window, cx| {                this.delegate.filtered_entries = filtered_models.entries();                // Finds the currently selected model in the list                let new_index =                    Self::get_active_model_index(&this.delegate.filtered_entries, active_model);                this.set_selected_index(new_index, Some(picker::Direction::Down), true, window, cx);                cx.notify();            })            .ok();        })    }     fn confirm(&mut self, _secondary: bool, window: &mut Window, cx: &mut Context<Picker<Self>>) {        if let Some(LanguageModelPickerEntry::Model(model_info)) =            self.filtered_entries.get(self.selected_index)        {            let model = model_info.model.clone();            (self.on_model_changed)(model.clone(), cx);             let current_index = self.selected_index;            self.set_selected_index(current_index, window, cx);             cx.emit(DismissEvent);        }    }     fn dismissed(&mut self, _: &mut Window, cx: &mut Context<Picker<Self>>) {

Lines 1–500 of 902.