1use collections::HashMap;
2use extension::{
3 ExtensionHostProxy, ExtensionLanguageModelProviderProxy, LanguageModelProviderRegistration,
4};
5use gpui::{App, Entity};
6use language_model::{LanguageModelProviderId, LanguageModelRegistry};
7use std::sync::{Arc, LazyLock};
8
9/// Maps built-in provider IDs to their corresponding extension IDs.
10/// When an extension with this ID is installed, the built-in provider should be hidden.
11static BUILTIN_TO_EXTENSION_MAP: LazyLock<HashMap<&'static str, &'static str>> =
12 LazyLock::new(|| {
13 let mut map = HashMap::default();
14 map.insert("anthropic", "anthropic");
15 map.insert("openai", "openai");
16 map.insert("google", "google-ai");
17 map.insert("openrouter", "openrouter");
18 map.insert("copilot_chat", "copilot-chat");
19 map
20 });
21
22/// Returns the extension ID that should hide the given built-in provider.
23pub fn extension_for_builtin_provider(provider_id: &str) -> Option<&'static str> {
24 BUILTIN_TO_EXTENSION_MAP.get(provider_id).copied()
25}
26
27/// Proxy that registers extension language model providers with the LanguageModelRegistry.
28pub struct LanguageModelProviderRegistryProxy {
29 registry: Entity<LanguageModelRegistry>,
30}
31
32impl LanguageModelProviderRegistryProxy {
33 pub fn new(registry: Entity<LanguageModelRegistry>) -> Self {
34 Self { registry }
35 }
36}
37
38impl ExtensionLanguageModelProviderProxy for LanguageModelProviderRegistryProxy {
39 fn register_language_model_provider(
40 &self,
41 _provider_id: Arc<str>,
42 register_fn: LanguageModelProviderRegistration,
43 cx: &mut App,
44 ) {
45 register_fn(cx);
46 }
47
48 fn unregister_language_model_provider(&self, provider_id: Arc<str>, cx: &mut App) {
49 self.registry.update(cx, |registry, cx| {
50 registry.unregister_provider(LanguageModelProviderId::from(provider_id), cx);
51 });
52 }
53}
54
55/// Initialize the extension language model provider proxy.
56/// This must be called BEFORE extension_host::init to ensure the proxy is available
57/// when extensions try to register their language model providers.
58pub fn init_proxy(cx: &mut App) {
59 let proxy = ExtensionHostProxy::default_global(cx);
60 let registry = LanguageModelRegistry::global(cx);
61
62 registry.update(cx, |registry, _cx| {
63 registry.set_builtin_provider_hiding_fn(Box::new(extension_for_builtin_provider));
64 });
65
66 proxy.register_language_model_provider_proxy(LanguageModelProviderRegistryProxy::new(registry));
67}