1use crate::{
2 provider::{
3 anthropic::AnthropicLanguageModelProvider, cloud::CloudLanguageModelProvider,
4 copilot_chat::CopilotChatLanguageModelProvider, google::GoogleLanguageModelProvider,
5 ollama::OllamaLanguageModelProvider, open_ai::OpenAiLanguageModelProvider,
6 },
7 LanguageModel, LanguageModelId, LanguageModelProvider, LanguageModelProviderId,
8 LanguageModelProviderState,
9};
10use client::Client;
11use collections::BTreeMap;
12use gpui::{AppContext, EventEmitter, Global, Model, ModelContext};
13use std::sync::Arc;
14use ui::Context;
15
16pub fn init(client: Arc<Client>, cx: &mut AppContext) {
17 let registry = cx.new_model(|cx| {
18 let mut registry = LanguageModelRegistry::default();
19 register_language_model_providers(&mut registry, client, cx);
20 registry
21 });
22 cx.set_global(GlobalLanguageModelRegistry(registry));
23}
24
25fn register_language_model_providers(
26 registry: &mut LanguageModelRegistry,
27 client: Arc<Client>,
28 cx: &mut ModelContext<LanguageModelRegistry>,
29) {
30 use feature_flags::FeatureFlagAppExt;
31
32 registry.register_provider(
33 AnthropicLanguageModelProvider::new(client.http_client(), cx),
34 cx,
35 );
36 registry.register_provider(
37 OpenAiLanguageModelProvider::new(client.http_client(), cx),
38 cx,
39 );
40 registry.register_provider(
41 OllamaLanguageModelProvider::new(client.http_client(), cx),
42 cx,
43 );
44 registry.register_provider(
45 GoogleLanguageModelProvider::new(client.http_client(), cx),
46 cx,
47 );
48 registry.register_provider(CopilotChatLanguageModelProvider::new(cx), cx);
49
50 cx.observe_flag::<feature_flags::LanguageModels, _>(move |enabled, cx| {
51 let client = client.clone();
52 LanguageModelRegistry::global(cx).update(cx, move |registry, cx| {
53 if enabled {
54 registry.register_provider(CloudLanguageModelProvider::new(client.clone(), cx), cx);
55 } else {
56 registry.unregister_provider(
57 &LanguageModelProviderId::from(
58 crate::provider::cloud::PROVIDER_NAME.to_string(),
59 ),
60 cx,
61 );
62 }
63 });
64 })
65 .detach();
66}
67
68struct GlobalLanguageModelRegistry(Model<LanguageModelRegistry>);
69
70impl Global for GlobalLanguageModelRegistry {}
71
72#[derive(Default)]
73pub struct LanguageModelRegistry {
74 active_model: Option<ActiveModel>,
75 providers: BTreeMap<LanguageModelProviderId, Arc<dyn LanguageModelProvider>>,
76}
77
78pub struct ActiveModel {
79 provider: Arc<dyn LanguageModelProvider>,
80 model: Option<Arc<dyn LanguageModel>>,
81}
82
83pub struct ActiveModelChanged;
84
85impl EventEmitter<ActiveModelChanged> for LanguageModelRegistry {}
86
87impl LanguageModelRegistry {
88 pub fn global(cx: &AppContext) -> Model<Self> {
89 cx.global::<GlobalLanguageModelRegistry>().0.clone()
90 }
91
92 pub fn read_global(cx: &AppContext) -> &Self {
93 cx.global::<GlobalLanguageModelRegistry>().0.read(cx)
94 }
95
96 #[cfg(any(test, feature = "test-support"))]
97 pub fn test(cx: &mut AppContext) -> crate::provider::fake::FakeLanguageModelProvider {
98 let fake_provider = crate::provider::fake::FakeLanguageModelProvider::default();
99 let registry = cx.new_model(|cx| {
100 let mut registry = Self::default();
101 registry.register_provider(fake_provider.clone(), cx);
102 let model = fake_provider.provided_models(cx)[0].clone();
103 registry.set_active_model(Some(model), cx);
104 registry
105 });
106 cx.set_global(GlobalLanguageModelRegistry(registry));
107 fake_provider
108 }
109
110 pub fn register_provider<T: LanguageModelProvider + LanguageModelProviderState>(
111 &mut self,
112 provider: T,
113 cx: &mut ModelContext<Self>,
114 ) {
115 let name = provider.id();
116
117 if let Some(subscription) = provider.subscribe(cx) {
118 subscription.detach();
119 }
120
121 self.providers.insert(name, Arc::new(provider));
122 cx.notify();
123 }
124
125 pub fn unregister_provider(
126 &mut self,
127 name: &LanguageModelProviderId,
128 cx: &mut ModelContext<Self>,
129 ) {
130 if self.providers.remove(name).is_some() {
131 cx.notify();
132 }
133 }
134
135 pub fn providers(&self) -> impl Iterator<Item = &Arc<dyn LanguageModelProvider>> {
136 self.providers.values()
137 }
138
139 pub fn available_models(&self, cx: &AppContext) -> Vec<Arc<dyn LanguageModel>> {
140 self.providers
141 .values()
142 .flat_map(|provider| provider.provided_models(cx))
143 .collect()
144 }
145
146 pub fn provider(
147 &self,
148 name: &LanguageModelProviderId,
149 ) -> Option<Arc<dyn LanguageModelProvider>> {
150 self.providers.get(name).cloned()
151 }
152
153 pub fn select_active_model(
154 &mut self,
155 provider: &LanguageModelProviderId,
156 model_id: &LanguageModelId,
157 cx: &mut ModelContext<Self>,
158 ) {
159 let Some(provider) = self.provider(&provider) else {
160 return;
161 };
162
163 let models = provider.provided_models(cx);
164 if let Some(model) = models.iter().find(|model| &model.id() == model_id).cloned() {
165 self.set_active_model(Some(model), cx);
166 }
167 }
168
169 pub fn set_active_provider(
170 &mut self,
171 provider: Option<Arc<dyn LanguageModelProvider>>,
172 cx: &mut ModelContext<Self>,
173 ) {
174 self.active_model = provider.map(|provider| ActiveModel {
175 provider,
176 model: None,
177 });
178 cx.emit(ActiveModelChanged);
179 }
180
181 pub fn set_active_model(
182 &mut self,
183 model: Option<Arc<dyn LanguageModel>>,
184 cx: &mut ModelContext<Self>,
185 ) {
186 if let Some(model) = model {
187 let provider_id = model.provider_id();
188 if let Some(provider) = self.providers.get(&provider_id).cloned() {
189 self.active_model = Some(ActiveModel {
190 provider,
191 model: Some(model),
192 });
193 cx.emit(ActiveModelChanged);
194 } else {
195 log::warn!("Active model's provider not found in registry");
196 }
197 } else {
198 self.active_model = None;
199 cx.emit(ActiveModelChanged);
200 }
201 }
202
203 pub fn active_provider(&self) -> Option<Arc<dyn LanguageModelProvider>> {
204 Some(self.active_model.as_ref()?.provider.clone())
205 }
206
207 pub fn active_model(&self) -> Option<Arc<dyn LanguageModel>> {
208 self.active_model.as_ref()?.model.clone()
209 }
210}
211
212#[cfg(test)]
213mod tests {
214 use super::*;
215 use crate::provider::fake::FakeLanguageModelProvider;
216
217 #[gpui::test]
218 fn test_register_providers(cx: &mut AppContext) {
219 let registry = cx.new_model(|_| LanguageModelRegistry::default());
220
221 registry.update(cx, |registry, cx| {
222 registry.register_provider(FakeLanguageModelProvider::default(), cx);
223 });
224
225 let providers = registry.read(cx).providers().collect::<Vec<_>>();
226 assert_eq!(providers.len(), 1);
227 assert_eq!(providers[0].id(), crate::provider::fake::provider_id());
228
229 registry.update(cx, |registry, cx| {
230 registry.unregister_provider(&crate::provider::fake::provider_id(), cx);
231 });
232
233 let providers = registry.read(cx).providers().collect::<Vec<_>>();
234 assert!(providers.is_empty());
235 }
236}