tenant.openagents/omega

1use std::sync::Arc;
2
3use ::settings::{Settings, SettingsStore};
4use client::{Client, UserStore};
5use collections::{HashMap, HashSet};
6use credentials_provider::CredentialsProvider;
7use gpui::{App, Context, Entity};
8use language_model::{LanguageModelProviderId, LanguageModelRegistry};
9use provider::deepseek::DeepSeekLanguageModelProvider;
10
11pub mod extension;
12pub mod provider;
13mod settings;
14
15pub use crate::extension::init_proxy as init_extension_proxy;
16
17use crate::provider::anthropic::AnthropicLanguageModelProvider;
18use crate::provider::anthropic_compatible::AnthropicCompatibleLanguageModelProvider;
19use crate::provider::bedrock::BedrockLanguageModelProvider;
20use crate::provider::cloud::CloudLanguageModelProvider;
21use crate::provider::copilot_chat::CopilotChatLanguageModelProvider;
22use crate::provider::google::GoogleLanguageModelProvider;
23use crate::provider::llama_cpp::LlamaCppLanguageModelProvider;
24use crate::provider::lmstudio::LmStudioLanguageModelProvider;
25pub use crate::provider::mistral::MistralLanguageModelProvider;
26use crate::provider::ollama::OllamaLanguageModelProvider;
27use crate::provider::open_ai::OpenAiLanguageModelProvider;
28use crate::provider::open_ai_compatible::OpenAiCompatibleLanguageModelProvider;
29use crate::provider::open_router::OpenRouterLanguageModelProvider;
30use crate::provider::openai_subscribed::OpenAiSubscribedProvider;
31use crate::provider::opencode::OpenCodeLanguageModelProvider;
32use crate::provider::vercel_ai_gateway::VercelAiGatewayLanguageModelProvider;
33use crate::provider::x_ai::XAiLanguageModelProvider;
34pub use crate::settings::*;
35
36pub fn init(user_store: Entity<UserStore>, client: Arc<Client>, cx: &mut App) {
37    let credentials_provider = client.credentials_provider();
38    let registry = LanguageModelRegistry::global(cx);
39    registry.update(cx, |registry, cx| {
40        register_language_model_providers(
41            registry,
42            user_store,
43            client.clone(),
44            credentials_provider.clone(),
45            cx,
46        );
47    });
48
49    // Subscribe to extension store events to track LLM extension installations
50    if let Some(extension_store) = extension_host::ExtensionStore::try_global(cx) {
51        cx.subscribe(&extension_store, {
52            let registry = registry.downgrade();
53            move |extension_store, event, cx| {
54                let Some(registry) = registry.upgrade() else {
55                    return;
56                };
57                match event {
58                    extension_host::Event::ExtensionInstalled(extension_id) => {
59                        if let Some(manifest) = extension_store
60                            .read(cx)
61                            .extension_manifest_for_id(extension_id)
62                        {
63                            if !manifest.language_model_providers.is_empty() {
64                                registry.update(cx, |registry, cx| {
65                                    registry.extension_installed(extension_id.clone(), cx);
66                                });
67                            }
68                        }
69                    }
70                    extension_host::Event::ExtensionUninstalled(extension_id) => {
71                        registry.update(cx, |registry, cx| {
72                            registry.extension_uninstalled(extension_id, cx);
73                        });
74                    }
75                    extension_host::Event::ExtensionsUpdated => {
76                        let mut new_ids = HashSet::default();
77                        for (extension_id, entry) in extension_store.read(cx).installed_extensions()
78                        {
79                            if !entry.manifest.language_model_providers.is_empty() {
80                                new_ids.insert(extension_id.clone());
81                            }
82                        }
83                        registry.update(cx, |registry, cx| {
84                            registry.sync_installed_llm_extensions(new_ids, cx);
85                        });
86                    }
87                    _ => {}
88                }
89            }
90        })
91        .detach();
92
93        // Initialize with currently installed extensions
94        registry.update(cx, |registry, cx| {
95            let mut initial_ids = HashSet::default();
96            for (extension_id, entry) in extension_store.read(cx).installed_extensions() {
97                if !entry.manifest.language_model_providers.is_empty() {
98                    initial_ids.insert(extension_id.clone());
99                }
100            }
101            registry.sync_installed_llm_extensions(initial_ids, cx);
102        });
103    }
104
105    let mut compatible_providers = CompatibleProviders::from_settings(cx);
106
107    registry.update(cx, |registry, cx| {
108        register_compatible_providers(
109            registry,
110            &CompatibleProviders::default(),
111            &compatible_providers,
112            &client,
113            &credentials_provider,
114            cx,
115        );
116    });
117
118    let registry = registry.downgrade();
119    cx.observe_global::<SettingsStore>(move |cx| {
120        let Some(registry) = registry.upgrade() else {
121            return;
122        };
123        let compatible_providers_new = CompatibleProviders::from_settings(cx);
124        if compatible_providers_new != compatible_providers {
125            registry.update(cx, |registry, cx| {
126                register_compatible_providers(
127                    registry,
128                    &compatible_providers,
129                    &compatible_providers_new,
130                    &client,
131                    &credentials_provider,
132                    cx,
133                );
134            });
135            compatible_providers = compatible_providers_new;
136        }
137    })
138    .detach();
139}
140
141#[derive(Default, PartialEq, Eq)]
142struct CompatibleProviders(HashMap<Arc<str>, CompatibleProviderKind>);
143
144#[derive(Copy, Clone, Debug, PartialEq, Eq)]
145enum CompatibleProviderKind {
146    OpenAi,
147    Anthropic,
148}
149
150impl CompatibleProviders {
151    fn from_settings(cx: &App) -> Self {
152        let settings = AllLanguageModelSettings::get_global(cx);
153        let mut providers: HashMap<Arc<str>, CompatibleProviderKind> = settings
154            .openai_compatible
155            .keys()
156            .map(|id| (id.clone(), CompatibleProviderKind::OpenAi))
157            .collect();
158        for id in settings.anthropic_compatible.keys() {
159            // The registry has a single provider ID namespace, so a name can
160            // only refer to one provider. OpenAI-compatible entries win
161            // collisions because they predate Anthropic-compatible ones, so
162            // existing configurations keep working.
163            if providers.contains_key(id) {
164                log::warn!(
165                    "ignoring `anthropic_compatible` provider `{id}`: \
166                     an `openai_compatible` provider with the same name exists"
167                );
168            } else {
169                providers.insert(id.clone(), CompatibleProviderKind::Anthropic);
170            }
171        }
172        Self(providers)
173    }
174}
175
176fn register_compatible_providers(
177    registry: &mut LanguageModelRegistry,
178    old: &CompatibleProviders,
179    new: &CompatibleProviders,
180    client: &Arc<Client>,
181    credentials_provider: &Arc<dyn CredentialsProvider>,
182    cx: &mut Context<LanguageModelRegistry>,
183) {
184    for (provider_id, old_kind) in &old.0 {
185        if new.0.get(provider_id) != Some(old_kind) {
186            registry.unregister_provider(LanguageModelProviderId::from(provider_id.clone()), cx);
187        }
188    }
189
190    for (provider_id, kind) in &new.0 {
191        if old.0.get(provider_id) != Some(kind) {
192            match kind {
193                CompatibleProviderKind::OpenAi => registry.register_provider(
194                    Arc::new(OpenAiCompatibleLanguageModelProvider::new(
195                        provider_id.clone(),
196                        client.http_client(),
197                        credentials_provider.clone(),
198                        cx,
199                    )),
200                    cx,
201                ),
202                CompatibleProviderKind::Anthropic => registry.register_provider(
203                    Arc::new(AnthropicCompatibleLanguageModelProvider::new(
204                        provider_id.clone(),
205                        client.http_client(),
206                        credentials_provider.clone(),
207                        cx,
208                    )),
209                    cx,
210                ),
211            }
212        }
213    }
214}
215
216fn register_language_model_providers(
217    registry: &mut LanguageModelRegistry,
218    user_store: Entity<UserStore>,
219    client: Arc<Client>,
220    credentials_provider: Arc<dyn CredentialsProvider>,
221    cx: &mut Context<LanguageModelRegistry>,
222) {
223    if app_identity::zed_production_services_enabled() {
224        registry.register_provider(
225            Arc::new(CloudLanguageModelProvider::new(
226                user_store,
227                client.clone(),
228                cx,
229            )),
230            cx,
231        );
232    } else {
233        log::info!(
234            "Skipping Zed cloud language model provider; Omega production services isolation is enabled"
235        );
236    }
237    registry.register_provider(
238        Arc::new(AnthropicLanguageModelProvider::new(
239            client.http_client(),
240            credentials_provider.clone(),
241            cx,
242        )),
243        cx,
244    );
245    registry.register_provider(
246        Arc::new(OpenAiLanguageModelProvider::new(
247            client.http_client(),
248            credentials_provider.clone(),
249            cx,
250        )),
251        cx,
252    );
253    registry.register_provider(
254        Arc::new(OllamaLanguageModelProvider::new(
255            client.http_client(),
256            credentials_provider.clone(),
257            cx,
258        )),
259        cx,
260    );
261    registry.register_provider(
262        Arc::new(LmStudioLanguageModelProvider::new(
263            client.http_client(),
264            credentials_provider.clone(),
265            cx,
266        )),
267        cx,
268    );
269    registry.register_provider(
270        Arc::new(LlamaCppLanguageModelProvider::new(
271            client.http_client(),
272            credentials_provider.clone(),
273            cx,
274        )),
275        cx,
276    );
277    registry.register_provider(
278        Arc::new(DeepSeekLanguageModelProvider::new(
279            client.http_client(),
280            credentials_provider.clone(),
281            cx,
282        )),
283        cx,
284    );
285    registry.register_provider(
286        Arc::new(GoogleLanguageModelProvider::new(
287            client.http_client(),
288            credentials_provider.clone(),
289            cx,
290        )),
291        cx,
292    );
293    registry.register_provider(
294        MistralLanguageModelProvider::global(
295            client.http_client(),
296            credentials_provider.clone(),
297            cx,
298        ),
299        cx,
300    );
301    registry.register_provider(
302        Arc::new(BedrockLanguageModelProvider::new(
303            client.http_client(),
304            credentials_provider.clone(),
305            cx,
306        )),
307        cx,
308    );
309    registry.register_provider(
310        Arc::new(OpenRouterLanguageModelProvider::new(
311            client.http_client(),
312            credentials_provider.clone(),
313            cx,
314        )),
315        cx,
316    );
317    registry.register_provider(
318        Arc::new(VercelAiGatewayLanguageModelProvider::new(
319            client.http_client(),
320            credentials_provider.clone(),
321            cx,
322        )),
323        cx,
324    );
325    registry.register_provider(
326        Arc::new(XAiLanguageModelProvider::new(
327            client.http_client(),
328            credentials_provider.clone(),
329            cx,
330        )),
331        cx,
332    );
333    registry.register_provider(
334        Arc::new(OpenCodeLanguageModelProvider::new(
335            client.http_client(),
336            credentials_provider.clone(),
337            cx,
338        )),
339        cx,
340    );
341    registry.register_provider(Arc::new(CopilotChatLanguageModelProvider::new(cx)), cx);
342    registry.register_provider(
343        Arc::new(OpenAiSubscribedProvider::new(
344            client.http_client(),
345            credentials_provider,
346            cx,
347        )),
348        cx,
349    );
350}
351
352#[cfg(test)]
353mod tests {
354    use super::*;
355    use anyhow::Result;
356    use clock::FakeSystemClock;
357    use feature_flags::FeatureFlagAppExt as _;
358    use gpui::{AppContext as _, AsyncApp, BorrowAppContext as _};
359    use http_client::FakeHttpClient;
360    use language_model::IconOrSvg;
361    use release_channel::AppVersion;
362    use std::future::Future;
363    use std::pin::Pin;
364    use ui::IconName;
365
366    struct FakeCredentialsProvider;
367
368    impl CredentialsProvider for FakeCredentialsProvider {
369        fn read_credentials<'a>(
370            &'a self,
371            _url: &'a str,
372            _cx: &'a AsyncApp,
373        ) -> Pin<Box<dyn Future<Output = Result<Option<(String, Vec<u8>)>>> + 'a>> {
374            Box::pin(async { Ok(None) })
375        }
376
377        fn write_credentials<'a>(
378            &'a self,
379            _url: &'a str,
380            _username: &'a str,
381            _password: &'a [u8],
382            _cx: &'a AsyncApp,
383        ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
384            Box::pin(async { Ok(()) })
385        }
386
387        fn delete_credentials<'a>(
388            &'a self,
389            _url: &'a str,
390            _cx: &'a AsyncApp,
391        ) -> Pin<Box<dyn Future<Output = Result<()>> + 'a>> {
392            Box::pin(async { Ok(()) })
393        }
394    }
395
396    fn init_test(cx: &mut App) -> (Arc<Client>, Arc<dyn CredentialsProvider>) {
397        let settings_store = SettingsStore::test(cx);
398        cx.set_global(settings_store);
399        cx.set_global(db::AppDatabase::test_new());
400        let app_version = AppVersion::global(cx);
401        release_channel::init_test(app_version, release_channel::ReleaseChannel::Dev, cx);
402        gpui_tokio::init(cx);
403        cx.update_flags(false, Vec::new());
404
405        let client = Client::new(
406            Arc::new(FakeSystemClock::new()),
407            FakeHttpClient::with_404_response(),
408            cx,
409        );
410        (client, Arc::new(FakeCredentialsProvider))
411    }
412
413    fn update_compatible_provider_settings(
414        openai: &[&str],
415        anthropic: &[&str],
416        cx: &mut App,
417    ) -> CompatibleProviders {
418        fn section(ids: &[&str]) -> serde_json::Value {
419            ids.iter()
420                .map(|id| {
421                    (
422                        id.to_string(),
423                        serde_json::json!({
424                            "api_url": "https://example.com",
425                            "available_models": [],
426                        }),
427                    )
428                })
429                .collect::<serde_json::Map<String, serde_json::Value>>()
430                .into()
431        }
432
433        let content = serde_json::json!({
434            "language_models": {
435                "openai_compatible": section(openai),
436                "anthropic_compatible": section(anthropic),
437            }
438        })
439        .to_string();
440        cx.update_global::<SettingsStore, _>(|store, cx| {
441            store
442                .set_user_settings(&content, cx)
443                .expect("failed to parse test settings");
444        });
445        CompatibleProviders::from_settings(cx)
446    }
447
448    fn provider_icons(registry: &LanguageModelRegistry, id: &str) -> Vec<IconOrSvg> {
449        registry
450            .providers()
451            .into_iter()
452            .filter(|provider| provider.id().0.as_ref() == id)
453            .map(|provider| provider.icon())
454            .collect()
455    }
456
457    #[gpui::test]
458    fn test_compatible_provider_id_collision_resolves_when_one_entry_is_removed(cx: &mut App) {
459        let (client, credentials_provider) = init_test(cx);
460        let registry = cx.new(|_| LanguageModelRegistry::default());
461
462        // The same provider name is configured in both `openai_compatible`
463        // and `anthropic_compatible` settings sections; the OpenAI-compatible
464        // entry wins the collision.
465        let both = update_compatible_provider_settings(&["acme"], &["acme"], cx);
466        registry.update(cx, |registry, cx| {
467            register_compatible_providers(
468                registry,
469                &CompatibleProviders::default(),
470                &both,
471                &client,
472                &credentials_provider,
473                cx,
474            );
475        });
476        assert_eq!(
477            registry.read_with(cx, |registry, _| provider_icons(registry, "acme")),
478            vec![IconOrSvg::Icon(IconName::AiOpenAiCompat)],
479            "the OpenAI-compatible provider should win the name collision"
480        );
481
482        // The user removes the `anthropic_compatible` entry; the remaining
483        // `openai_compatible` entry must stay registered.
484        let openai_only = update_compatible_provider_settings(&["acme"], &[], cx);
485        registry.update(cx, |registry, cx| {
486            register_compatible_providers(
487                registry,
488                &both,
489                &openai_only,
490                &client,
491                &credentials_provider,
492                cx,
493            );
494        });
495        assert_eq!(
496            registry.read_with(cx, |registry, _| provider_icons(registry, "acme")),
497            vec![IconOrSvg::Icon(IconName::AiOpenAiCompat)],
498            "the provider registered for `acme` should be the OpenAI-compatible one"
499        );
500    }
501
502    #[gpui::test]
503    fn test_compatible_provider_changes_kind_and_unregisters(cx: &mut App) {
504        let (client, credentials_provider) = init_test(cx);
505        let registry = cx.new(|_| LanguageModelRegistry::default());
506
507        let both = update_compatible_provider_settings(&["acme"], &["acme"], cx);
508        registry.update(cx, |registry, cx| {
509            register_compatible_providers(
510                registry,
511                &CompatibleProviders::default(),
512                &both,
513                &client,
514                &credentials_provider,
515                cx,
516            );
517        });
518
519        // Removing the `openai_compatible` entry hands the name over to the
520        // remaining `anthropic_compatible` entry.
521        let anthropic_only = update_compatible_provider_settings(&[], &["acme"], cx);
522        registry.update(cx, |registry, cx| {
523            register_compatible_providers(
524                registry,
525                &both,
526                &anthropic_only,
527                &client,
528                &credentials_provider,
529                cx,
530            );
531        });
532        assert_eq!(
533            registry.read_with(cx, |registry, _| provider_icons(registry, "acme")),
534            vec![IconOrSvg::Icon(IconName::AiAnthropicCompat)],
535            "after removing the openai_compatible entry, the anthropic_compatible provider should be registered"
536        );
537
538        // Removing the last entry unregisters the provider entirely.
539        let none = update_compatible_provider_settings(&[], &[], cx);
540        registry.update(cx, |registry, cx| {
541            register_compatible_providers(
542                registry,
543                &anthropic_only,
544                &none,
545                &client,
546                &credentials_provider,
547                cx,
548            );
549        });
550        assert_eq!(
551            registry.read_with(cx, |registry, _| provider_icons(registry, "acme")),
552            Vec::new(),
553            "removing all entries should unregister the provider"
554        );
555    }
556}
557