1use anyhow::{anyhow, Context as _, Result};
2use collections::BTreeMap;
3use credentials_provider::CredentialsProvider;
4use editor::{Editor, EditorElement, EditorStyle};
5use futures::{future::BoxFuture, FutureExt, StreamExt};
6use google_ai::stream_generate_content;
7use gpui::{
8 AnyView, App, AsyncApp, Context, Entity, FontStyle, Subscription, Task, TextStyle, WhiteSpace,
9};
10use http_client::HttpClient;
11use language_model::{AuthenticateError, LanguageModelCompletionEvent};
12use language_model::{
13 LanguageModel, LanguageModelId, LanguageModelName, LanguageModelProvider,
14 LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState,
15 LanguageModelRequest, RateLimiter, Role,
16};
17use schemars::JsonSchema;
18use serde::{Deserialize, Serialize};
19use settings::{Settings, SettingsStore};
20use std::{future, sync::Arc};
21use strum::IntoEnumIterator;
22use theme::ThemeSettings;
23use ui::{prelude::*, Icon, IconName, List, Tooltip};
24use util::ResultExt;
25
26use crate::ui::InstructionListItem;
27use crate::AllLanguageModelSettings;
28
29const PROVIDER_ID: &str = "google";
30const PROVIDER_NAME: &str = "Google AI";
31
32#[derive(Default, Clone, Debug, PartialEq)]
33pub struct GoogleSettings {
34 pub api_url: String,
35 pub available_models: Vec<AvailableModel>,
36}
37
38#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
39pub struct AvailableModel {
40 name: String,
41 display_name: Option<String>,
42 max_tokens: usize,
43}
44
45pub struct GoogleLanguageModelProvider {
46 http_client: Arc<dyn HttpClient>,
47 state: gpui::Entity<State>,
48}
49
50pub struct State {
51 api_key: Option<String>,
52 api_key_from_env: bool,
53 _subscription: Subscription,
54}
55
56const GOOGLE_AI_API_KEY_VAR: &str = "GOOGLE_AI_API_KEY";
57
58impl State {
59 fn is_authenticated(&self) -> bool {
60 self.api_key.is_some()
61 }
62
63 fn reset_api_key(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
64 let credentials_provider = <dyn CredentialsProvider>::global(cx);
65 let api_url = AllLanguageModelSettings::get_global(cx)
66 .google
67 .api_url
68 .clone();
69 cx.spawn(async move |this, cx| {
70 credentials_provider
71 .delete_credentials(&api_url, &cx)
72 .await
73 .log_err();
74 this.update(cx, |this, cx| {
75 this.api_key = None;
76 this.api_key_from_env = false;
77 cx.notify();
78 })
79 })
80 }
81
82 fn set_api_key(&mut self, api_key: String, cx: &mut Context<Self>) -> Task<Result<()>> {
83 let credentials_provider = <dyn CredentialsProvider>::global(cx);
84 let api_url = AllLanguageModelSettings::get_global(cx)
85 .google
86 .api_url
87 .clone();
88 cx.spawn(async move |this, cx| {
89 credentials_provider
90 .write_credentials(&api_url, "Bearer", api_key.as_bytes(), &cx)
91 .await?;
92 this.update(cx, |this, cx| {
93 this.api_key = Some(api_key);
94 cx.notify();
95 })
96 })
97 }
98
99 fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
100 if self.is_authenticated() {
101 return Task::ready(Ok(()));
102 }
103
104 let credentials_provider = <dyn CredentialsProvider>::global(cx);
105 let api_url = AllLanguageModelSettings::get_global(cx)
106 .google
107 .api_url
108 .clone();
109
110 cx.spawn(async move |this, cx| {
111 let (api_key, from_env) = if let Ok(api_key) = std::env::var(GOOGLE_AI_API_KEY_VAR) {
112 (api_key, true)
113 } else {
114 let (_, api_key) = credentials_provider
115 .read_credentials(&api_url, &cx)
116 .await?
117 .ok_or(AuthenticateError::CredentialsNotFound)?;
118 (
119 String::from_utf8(api_key).context("invalid {PROVIDER_NAME} API key")?,
120 false,
121 )
122 };
123
124 this.update(cx, |this, cx| {
125 this.api_key = Some(api_key);
126 this.api_key_from_env = from_env;
127 cx.notify();
128 })?;
129
130 Ok(())
131 })
132 }
133}
134
135impl GoogleLanguageModelProvider {
136 pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self {
137 let state = cx.new(|cx| State {
138 api_key: None,
139 api_key_from_env: false,
140 _subscription: cx.observe_global::<SettingsStore>(|_, cx| {
141 cx.notify();
142 }),
143 });
144
145 Self { http_client, state }
146 }
147}
148
149impl LanguageModelProviderState for GoogleLanguageModelProvider {
150 type ObservableEntity = State;
151
152 fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> {
153 Some(self.state.clone())
154 }
155}
156
157impl LanguageModelProvider for GoogleLanguageModelProvider {
158 fn id(&self) -> LanguageModelProviderId {
159 LanguageModelProviderId(PROVIDER_ID.into())
160 }
161
162 fn name(&self) -> LanguageModelProviderName {
163 LanguageModelProviderName(PROVIDER_NAME.into())
164 }
165
166 fn icon(&self) -> IconName {
167 IconName::AiGoogle
168 }
169
170 fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
171 let model = google_ai::Model::default();
172 Some(Arc::new(GoogleLanguageModel {
173 id: LanguageModelId::from(model.id().to_string()),
174 model,
175 state: self.state.clone(),
176 http_client: self.http_client.clone(),
177 rate_limiter: RateLimiter::new(4),
178 }))
179 }
180
181 fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
182 let mut models = BTreeMap::default();
183
184 // Add base models from google_ai::Model::iter()
185 for model in google_ai::Model::iter() {
186 if !matches!(model, google_ai::Model::Custom { .. }) {
187 models.insert(model.id().to_string(), model);
188 }
189 }
190
191 // Override with available models from settings
192 for model in &AllLanguageModelSettings::get_global(cx)
193 .google
194 .available_models
195 {
196 models.insert(
197 model.name.clone(),
198 google_ai::Model::Custom {
199 name: model.name.clone(),
200 display_name: model.display_name.clone(),
201 max_tokens: model.max_tokens,
202 },
203 );
204 }
205
206 models
207 .into_values()
208 .map(|model| {
209 Arc::new(GoogleLanguageModel {
210 id: LanguageModelId::from(model.id().to_string()),
211 model,
212 state: self.state.clone(),
213 http_client: self.http_client.clone(),
214 rate_limiter: RateLimiter::new(4),
215 }) as Arc<dyn LanguageModel>
216 })
217 .collect()
218 }
219
220 fn is_authenticated(&self, cx: &App) -> bool {
221 self.state.read(cx).is_authenticated()
222 }
223
224 fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
225 self.state.update(cx, |state, cx| state.authenticate(cx))
226 }
227
228 fn configuration_view(&self, window: &mut Window, cx: &mut App) -> AnyView {
229 cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx))
230 .into()
231 }
232
233 fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> {
234 self.state.update(cx, |state, cx| state.reset_api_key(cx))
235 }
236}
237
238pub struct GoogleLanguageModel {
239 id: LanguageModelId,
240 model: google_ai::Model,
241 state: gpui::Entity<State>,
242 http_client: Arc<dyn HttpClient>,
243 rate_limiter: RateLimiter,
244}
245
246impl LanguageModel for GoogleLanguageModel {
247 fn id(&self) -> LanguageModelId {
248 self.id.clone()
249 }
250
251 fn name(&self) -> LanguageModelName {
252 LanguageModelName::from(self.model.display_name().to_string())
253 }
254
255 fn provider_id(&self) -> LanguageModelProviderId {
256 LanguageModelProviderId(PROVIDER_ID.into())
257 }
258
259 fn provider_name(&self) -> LanguageModelProviderName {
260 LanguageModelProviderName(PROVIDER_NAME.into())
261 }
262
263 fn telemetry_id(&self) -> String {
264 format!("google/{}", self.model.id())
265 }
266
267 fn max_token_count(&self) -> usize {
268 self.model.max_token_count()
269 }
270
271 fn count_tokens(
272 &self,
273 request: LanguageModelRequest,
274 cx: &App,
275 ) -> BoxFuture<'static, Result<usize>> {
276 let request = into_google(request, self.model.id().to_string());
277 let http_client = self.http_client.clone();
278 let api_key = self.state.read(cx).api_key.clone();
279
280 let settings = &AllLanguageModelSettings::get_global(cx).google;
281 let api_url = settings.api_url.clone();
282
283 async move {
284 let api_key = api_key.ok_or_else(|| anyhow!("Missing Google API key"))?;
285 let response = google_ai::count_tokens(
286 http_client.as_ref(),
287 &api_url,
288 &api_key,
289 google_ai::CountTokensRequest {
290 contents: request.contents,
291 },
292 )
293 .await?;
294 Ok(response.total_tokens)
295 }
296 .boxed()
297 }
298
299 fn stream_completion(
300 &self,
301 request: LanguageModelRequest,
302 cx: &AsyncApp,
303 ) -> BoxFuture<
304 'static,
305 Result<futures::stream::BoxStream<'static, Result<LanguageModelCompletionEvent>>>,
306 > {
307 let request = into_google(request, self.model.id().to_string());
308
309 let http_client = self.http_client.clone();
310 let Ok((api_key, api_url)) = cx.read_entity(&self.state, |state, cx| {
311 let settings = &AllLanguageModelSettings::get_global(cx).google;
312 (state.api_key.clone(), settings.api_url.clone())
313 }) else {
314 return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
315 };
316
317 let future = self.rate_limiter.stream(async move {
318 let api_key = api_key.ok_or_else(|| anyhow!("Missing Google API Key"))?;
319 let response =
320 stream_generate_content(http_client.as_ref(), &api_url, &api_key, request);
321 let events = response.await?;
322 Ok(google_ai::extract_text_from_events(events).boxed())
323 });
324 async move {
325 Ok(future
326 .await?
327 .map(|result| result.map(LanguageModelCompletionEvent::Text))
328 .boxed())
329 }
330 .boxed()
331 }
332
333 fn use_any_tool(
334 &self,
335 _request: LanguageModelRequest,
336 _name: String,
337 _description: String,
338 _schema: serde_json::Value,
339 _cx: &AsyncApp,
340 ) -> BoxFuture<'static, Result<futures::stream::BoxStream<'static, Result<String>>>> {
341 future::ready(Err(anyhow!("not implemented"))).boxed()
342 }
343}
344
345pub fn into_google(
346 request: LanguageModelRequest,
347 model: String,
348) -> google_ai::GenerateContentRequest {
349 google_ai::GenerateContentRequest {
350 model,
351 contents: request
352 .messages
353 .into_iter()
354 .map(|msg| google_ai::Content {
355 parts: vec![google_ai::Part::TextPart(google_ai::TextPart {
356 text: msg.string_contents(),
357 })],
358 role: match msg.role {
359 Role::User => google_ai::Role::User,
360 Role::Assistant => google_ai::Role::Model,
361 Role::System => google_ai::Role::User, // Google AI doesn't have a system role
362 },
363 })
364 .collect(),
365 generation_config: Some(google_ai::GenerationConfig {
366 candidate_count: Some(1),
367 stop_sequences: Some(request.stop),
368 max_output_tokens: None,
369 temperature: request.temperature.map(|t| t as f64).or(Some(1.0)),
370 top_p: None,
371 top_k: None,
372 }),
373 safety_settings: None,
374 }
375}
376
377pub fn count_google_tokens(
378 request: LanguageModelRequest,
379 cx: &App,
380) -> BoxFuture<'static, Result<usize>> {
381 // We couldn't use the GoogleLanguageModelProvider to count tokens because the github copilot doesn't have the access to google_ai directly.
382 // So we have to use tokenizer from tiktoken_rs to count tokens.
383 cx.background_spawn(async move {
384 let messages = request
385 .messages
386 .into_iter()
387 .map(|message| tiktoken_rs::ChatCompletionRequestMessage {
388 role: match message.role {
389 Role::User => "user".into(),
390 Role::Assistant => "assistant".into(),
391 Role::System => "system".into(),
392 },
393 content: Some(message.string_contents()),
394 name: None,
395 function_call: None,
396 })
397 .collect::<Vec<_>>();
398
399 // Tiktoken doesn't yet support these models, so we manually use the
400 // same tokenizer as GPT-4.
401 tiktoken_rs::num_tokens_from_messages("gpt-4", &messages)
402 })
403 .boxed()
404}
405
406struct ConfigurationView {
407 api_key_editor: Entity<Editor>,
408 state: gpui::Entity<State>,
409 load_credentials_task: Option<Task<()>>,
410}
411
412impl ConfigurationView {
413 fn new(state: gpui::Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
414 cx.observe(&state, |_, _, cx| {
415 cx.notify();
416 })
417 .detach();
418
419 let load_credentials_task = Some(cx.spawn_in(window, {
420 let state = state.clone();
421 async move |this, cx| {
422 if let Some(task) = state
423 .update(cx, |state, cx| state.authenticate(cx))
424 .log_err()
425 {
426 // We don't log an error, because "not signed in" is also an error.
427 let _ = task.await;
428 }
429 this.update(cx, |this, cx| {
430 this.load_credentials_task = None;
431 cx.notify();
432 })
433 .log_err();
434 }
435 }));
436
437 Self {
438 api_key_editor: cx.new(|cx| {
439 let mut editor = Editor::single_line(window, cx);
440 editor.set_placeholder_text("AIzaSy...", cx);
441 editor
442 }),
443 state,
444 load_credentials_task,
445 }
446 }
447
448 fn save_api_key(&mut self, _: &menu::Confirm, window: &mut Window, cx: &mut Context<Self>) {
449 let api_key = self.api_key_editor.read(cx).text(cx);
450 if api_key.is_empty() {
451 return;
452 }
453
454 let state = self.state.clone();
455 cx.spawn_in(window, async move |_, cx| {
456 state
457 .update(cx, |state, cx| state.set_api_key(api_key, cx))?
458 .await
459 })
460 .detach_and_log_err(cx);
461
462 cx.notify();
463 }
464
465 fn reset_api_key(&mut self, window: &mut Window, cx: &mut Context<Self>) {
466 self.api_key_editor
467 .update(cx, |editor, cx| editor.set_text("", window, cx));
468
469 let state = self.state.clone();
470 cx.spawn_in(window, async move |_, cx| {
471 state.update(cx, |state, cx| state.reset_api_key(cx))?.await
472 })
473 .detach_and_log_err(cx);
474
475 cx.notify();
476 }
477
478 fn render_api_key_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
479 let settings = ThemeSettings::get_global(cx);
480 let text_style = TextStyle {
481 color: cx.theme().colors().text,
482 font_family: settings.ui_font.family.clone(),
483 font_features: settings.ui_font.features.clone(),
484 font_fallbacks: settings.ui_font.fallbacks.clone(),
485 font_size: rems(0.875).into(),
486 font_weight: settings.ui_font.weight,
487 font_style: FontStyle::Normal,
488 line_height: relative(1.3),
489 white_space: WhiteSpace::Normal,
490 ..Default::default()
491 };
492 EditorElement::new(
493 &self.api_key_editor,
494 EditorStyle {
495 background: cx.theme().colors().editor_background,
496 local_player: cx.theme().players().local(),
497 text: text_style,
498 ..Default::default()
499 },
500 )
501 }
502
503 fn should_render_editor(&self, cx: &mut Context<Self>) -> bool {
504 !self.state.read(cx).is_authenticated()
505 }
506}
507
508impl Render for ConfigurationView {
509 fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
510 let env_var_set = self.state.read(cx).api_key_from_env;
511
512 if self.load_credentials_task.is_some() {
513 div().child(Label::new("Loading credentials...")).into_any()
514 } else if self.should_render_editor(cx) {
515 v_flex()
516 .size_full()
517 .on_action(cx.listener(Self::save_api_key))
518 .child(Label::new("To use Zed's assistant with Google AI, you need to add an API key. Follow these steps:"))
519 .child(
520 List::new()
521 .child(InstructionListItem::new(
522 "Create one by visiting",
523 Some("Google AI's console"),
524 Some("https://aistudio.google.com/app/apikey"),
525 ))
526 .child(InstructionListItem::text_only(
527 "Paste your API key below and hit enter to start using the assistant",
528 )),
529 )
530 .child(
531 h_flex()
532 .w_full()
533 .my_2()
534 .px_2()
535 .py_1()
536 .bg(cx.theme().colors().editor_background)
537 .border_1()
538 .border_color(cx.theme().colors().border_variant)
539 .rounded_sm()
540 .child(self.render_api_key_editor(cx)),
541 )
542 .child(
543 Label::new(
544 format!("You can also assign the {GOOGLE_AI_API_KEY_VAR} environment variable and restart Zed."),
545 )
546 .size(LabelSize::Small).color(Color::Muted),
547 )
548 .into_any()
549 } else {
550 h_flex()
551 .size_full()
552 .justify_between()
553 .child(
554 h_flex()
555 .gap_1()
556 .child(Icon::new(IconName::Check).color(Color::Success))
557 .child(Label::new(if env_var_set {
558 format!("API key set in {GOOGLE_AI_API_KEY_VAR} environment variable.")
559 } else {
560 "API key configured.".to_string()
561 })),
562 )
563 .child(
564 Button::new("reset-key", "Reset key")
565 .icon(Some(IconName::Trash))
566 .icon_size(IconSize::Small)
567 .icon_position(IconPosition::Start)
568 .disabled(env_var_set)
569 .when(env_var_set, |this| {
570 this.tooltip(Tooltip::text(format!("To reset your API key, unset the {GOOGLE_AI_API_KEY_VAR} environment variable.")))
571 })
572 .on_click(cx.listener(|this, _, window, cx| this.reset_api_key(window, cx))),
573 )
574 .into_any()
575 }
576 }
577}