961 lines · 32.0 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::{    AuthenticateError, 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)            .width(rems(20.))            .max_height(Some(rems(20.).into()))    } else {        Picker::list(delegate, window, cx).show_scrollbar(true)    }} 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,    _authenticate_all_providers_task: Task<()>,    _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),            _authenticate_all_providers_task: Self::authenticate_all_providers(cx),            _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)    }     /// Authenticates all providers in the [`LanguageModelRegistry`].    ///    /// We do this so that we can populate the language selector with all of the    /// models from the configured providers.    fn authenticate_all_providers(cx: &mut App) -> Task<()> {        let authenticate_all_providers = LanguageModelRegistry::global(cx)            .read(cx)            .visible_providers()            .iter()            .map(|provider| (provider.id(), provider.name(), provider.authenticate(cx)))            .collect::<Vec<_>>();         cx.spawn(async move |_cx| {            for (provider_id, provider_name, authenticate_task) in authenticate_all_providers {                if let Err(err) = authenticate_task.await {                    if matches!(err, AuthenticateError::CredentialsNotFound) {                        // Since we're authenticating these providers in the                        // background for the purposes of populating the                        // language selector, we don't care about providers                        // where the credentials are not found.                    } else {                        // Some providers have noisy failure states that we                        // don't want to spam the logs with every time the                        // language model selector is initialized.                        //                        // Ideally these should have more clear failure modes                        // that we know are safe to ignore here, like what we do                        // with `CredentialsNotFound` above.                        match provider_id.0.as_ref() {                            "lmstudio" | "ollama" => {                                // LM Studio and Ollama both make fetch requests to the local APIs to determine if they are "authenticated".                                //                                // These fail noisily, so we don't log them.                            }                            "copilot_chat" => {                                // Copilot Chat returns an error if Copilot is not enabled, so we don't log those errors.                            }                            _ => {                                log::error!(                                    "Failed to authenticate provider: {}: {err:#}",                                    provider_name.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 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(        &mut 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

Lines 1–500 of 961.