google.rs

  1use anyhow::{Context as _, Result, anyhow};
  2use collections::BTreeMap;
  3use credentials_provider::CredentialsProvider;
  4use editor::{Editor, EditorElement, EditorStyle};
  5use futures::{FutureExt, Stream, StreamExt, future::BoxFuture};
  6use google_ai::{
  7    FunctionDeclaration, GenerateContentResponse, GoogleModelMode, Part, SystemInstruction,
  8    ThinkingConfig, UsageMetadata,
  9};
 10use gpui::{
 11    AnyView, App, AsyncApp, Context, Entity, FontStyle, Subscription, Task, TextStyle, WhiteSpace,
 12};
 13use http_client::HttpClient;
 14use language_model::{
 15    AuthenticateError, LanguageModelCompletionError, LanguageModelCompletionEvent,
 16    LanguageModelToolChoice, LanguageModelToolSchemaFormat, LanguageModelToolUse,
 17    LanguageModelToolUseId, MessageContent, StopReason,
 18};
 19use language_model::{
 20    LanguageModel, LanguageModelId, LanguageModelName, LanguageModelProvider,
 21    LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState,
 22    LanguageModelRequest, RateLimiter, Role,
 23};
 24use schemars::JsonSchema;
 25use serde::{Deserialize, Serialize};
 26use settings::{Settings, SettingsStore};
 27use std::pin::Pin;
 28use std::sync::{
 29    Arc,
 30    atomic::{self, AtomicU64},
 31};
 32use strum::IntoEnumIterator;
 33use theme::ThemeSettings;
 34use ui::{Icon, IconName, List, Tooltip, prelude::*};
 35use util::ResultExt;
 36
 37use crate::AllLanguageModelSettings;
 38use crate::ui::InstructionListItem;
 39
 40const PROVIDER_ID: &str = "google";
 41const PROVIDER_NAME: &str = "Google AI";
 42
 43#[derive(Default, Clone, Debug, PartialEq)]
 44pub struct GoogleSettings {
 45    pub api_url: String,
 46    pub available_models: Vec<AvailableModel>,
 47}
 48
 49#[derive(Clone, Copy, Debug, Default, PartialEq, Serialize, Deserialize, JsonSchema)]
 50#[serde(tag = "type", rename_all = "lowercase")]
 51pub enum ModelMode {
 52    #[default]
 53    Default,
 54    Thinking {
 55        /// The maximum number of tokens to use for reasoning. Must be lower than the model's `max_output_tokens`.
 56        budget_tokens: Option<u32>,
 57    },
 58}
 59
 60impl From<ModelMode> for GoogleModelMode {
 61    fn from(value: ModelMode) -> Self {
 62        match value {
 63            ModelMode::Default => GoogleModelMode::Default,
 64            ModelMode::Thinking { budget_tokens } => GoogleModelMode::Thinking { budget_tokens },
 65        }
 66    }
 67}
 68
 69impl From<GoogleModelMode> for ModelMode {
 70    fn from(value: GoogleModelMode) -> Self {
 71        match value {
 72            GoogleModelMode::Default => ModelMode::Default,
 73            GoogleModelMode::Thinking { budget_tokens } => ModelMode::Thinking { budget_tokens },
 74        }
 75    }
 76}
 77
 78#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
 79pub struct AvailableModel {
 80    name: String,
 81    display_name: Option<String>,
 82    max_tokens: usize,
 83    mode: Option<ModelMode>,
 84}
 85
 86pub struct GoogleLanguageModelProvider {
 87    http_client: Arc<dyn HttpClient>,
 88    state: gpui::Entity<State>,
 89}
 90
 91pub struct State {
 92    api_key: Option<String>,
 93    api_key_from_env: bool,
 94    _subscription: Subscription,
 95}
 96
 97const GOOGLE_AI_API_KEY_VAR: &str = "GOOGLE_AI_API_KEY";
 98
 99impl State {
100    fn is_authenticated(&self) -> bool {
101        self.api_key.is_some()
102    }
103
104    fn reset_api_key(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
105        let credentials_provider = <dyn CredentialsProvider>::global(cx);
106        let api_url = AllLanguageModelSettings::get_global(cx)
107            .google
108            .api_url
109            .clone();
110        cx.spawn(async move |this, cx| {
111            credentials_provider
112                .delete_credentials(&api_url, &cx)
113                .await
114                .log_err();
115            this.update(cx, |this, cx| {
116                this.api_key = None;
117                this.api_key_from_env = false;
118                cx.notify();
119            })
120        })
121    }
122
123    fn set_api_key(&mut self, api_key: String, cx: &mut Context<Self>) -> Task<Result<()>> {
124        let credentials_provider = <dyn CredentialsProvider>::global(cx);
125        let api_url = AllLanguageModelSettings::get_global(cx)
126            .google
127            .api_url
128            .clone();
129        cx.spawn(async move |this, cx| {
130            credentials_provider
131                .write_credentials(&api_url, "Bearer", api_key.as_bytes(), &cx)
132                .await?;
133            this.update(cx, |this, cx| {
134                this.api_key = Some(api_key);
135                cx.notify();
136            })
137        })
138    }
139
140    fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
141        if self.is_authenticated() {
142            return Task::ready(Ok(()));
143        }
144
145        let credentials_provider = <dyn CredentialsProvider>::global(cx);
146        let api_url = AllLanguageModelSettings::get_global(cx)
147            .google
148            .api_url
149            .clone();
150
151        cx.spawn(async move |this, cx| {
152            let (api_key, from_env) = if let Ok(api_key) = std::env::var(GOOGLE_AI_API_KEY_VAR) {
153                (api_key, true)
154            } else {
155                let (_, api_key) = credentials_provider
156                    .read_credentials(&api_url, &cx)
157                    .await?
158                    .ok_or(AuthenticateError::CredentialsNotFound)?;
159                (
160                    String::from_utf8(api_key).context("invalid {PROVIDER_NAME} API key")?,
161                    false,
162                )
163            };
164
165            this.update(cx, |this, cx| {
166                this.api_key = Some(api_key);
167                this.api_key_from_env = from_env;
168                cx.notify();
169            })?;
170
171            Ok(())
172        })
173    }
174}
175
176impl GoogleLanguageModelProvider {
177    pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self {
178        let state = cx.new(|cx| State {
179            api_key: None,
180            api_key_from_env: false,
181            _subscription: cx.observe_global::<SettingsStore>(|_, cx| {
182                cx.notify();
183            }),
184        });
185
186        Self { http_client, state }
187    }
188
189    fn create_language_model(&self, model: google_ai::Model) -> Arc<dyn LanguageModel> {
190        Arc::new(GoogleLanguageModel {
191            id: LanguageModelId::from(model.id().to_string()),
192            model,
193            state: self.state.clone(),
194            http_client: self.http_client.clone(),
195            request_limiter: RateLimiter::new(4),
196        })
197    }
198}
199
200impl LanguageModelProviderState for GoogleLanguageModelProvider {
201    type ObservableEntity = State;
202
203    fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> {
204        Some(self.state.clone())
205    }
206}
207
208impl LanguageModelProvider for GoogleLanguageModelProvider {
209    fn id(&self) -> LanguageModelProviderId {
210        LanguageModelProviderId(PROVIDER_ID.into())
211    }
212
213    fn name(&self) -> LanguageModelProviderName {
214        LanguageModelProviderName(PROVIDER_NAME.into())
215    }
216
217    fn icon(&self) -> IconName {
218        IconName::AiGoogle
219    }
220
221    fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
222        Some(self.create_language_model(google_ai::Model::default()))
223    }
224
225    fn default_fast_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
226        Some(self.create_language_model(google_ai::Model::default_fast()))
227    }
228
229    fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
230        let mut models = BTreeMap::default();
231
232        // Add base models from google_ai::Model::iter()
233        for model in google_ai::Model::iter() {
234            if !matches!(model, google_ai::Model::Custom { .. }) {
235                models.insert(model.id().to_string(), model);
236            }
237        }
238
239        // Override with available models from settings
240        for model in &AllLanguageModelSettings::get_global(cx)
241            .google
242            .available_models
243        {
244            models.insert(
245                model.name.clone(),
246                google_ai::Model::Custom {
247                    name: model.name.clone(),
248                    display_name: model.display_name.clone(),
249                    max_tokens: model.max_tokens,
250                    mode: model.mode.unwrap_or_default().into(),
251                },
252            );
253        }
254
255        models
256            .into_values()
257            .map(|model| {
258                Arc::new(GoogleLanguageModel {
259                    id: LanguageModelId::from(model.id().to_string()),
260                    model,
261                    state: self.state.clone(),
262                    http_client: self.http_client.clone(),
263                    request_limiter: RateLimiter::new(4),
264                }) as Arc<dyn LanguageModel>
265            })
266            .collect()
267    }
268
269    fn is_authenticated(&self, cx: &App) -> bool {
270        self.state.read(cx).is_authenticated()
271    }
272
273    fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
274        self.state.update(cx, |state, cx| state.authenticate(cx))
275    }
276
277    fn configuration_view(&self, window: &mut Window, cx: &mut App) -> AnyView {
278        cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx))
279            .into()
280    }
281
282    fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> {
283        self.state.update(cx, |state, cx| state.reset_api_key(cx))
284    }
285}
286
287pub struct GoogleLanguageModel {
288    id: LanguageModelId,
289    model: google_ai::Model,
290    state: gpui::Entity<State>,
291    http_client: Arc<dyn HttpClient>,
292    request_limiter: RateLimiter,
293}
294
295impl GoogleLanguageModel {
296    fn stream_completion(
297        &self,
298        request: google_ai::GenerateContentRequest,
299        cx: &AsyncApp,
300    ) -> BoxFuture<
301        'static,
302        Result<futures::stream::BoxStream<'static, Result<GenerateContentResponse>>>,
303    > {
304        let http_client = self.http_client.clone();
305
306        let Ok((api_key, api_url)) = cx.read_entity(&self.state, |state, cx| {
307            let settings = &AllLanguageModelSettings::get_global(cx).google;
308            (state.api_key.clone(), settings.api_url.clone())
309        }) else {
310            return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
311        };
312
313        async move {
314            let api_key = api_key.context("Missing Google API key")?;
315            let request = google_ai::stream_generate_content(
316                http_client.as_ref(),
317                &api_url,
318                &api_key,
319                request,
320            );
321            request.await.context("failed to stream completion")
322        }
323        .boxed()
324    }
325}
326
327impl LanguageModel for GoogleLanguageModel {
328    fn id(&self) -> LanguageModelId {
329        self.id.clone()
330    }
331
332    fn name(&self) -> LanguageModelName {
333        LanguageModelName::from(self.model.display_name().to_string())
334    }
335
336    fn provider_id(&self) -> LanguageModelProviderId {
337        LanguageModelProviderId(PROVIDER_ID.into())
338    }
339
340    fn provider_name(&self) -> LanguageModelProviderName {
341        LanguageModelProviderName(PROVIDER_NAME.into())
342    }
343
344    fn supports_tools(&self) -> bool {
345        true
346    }
347
348    fn supports_images(&self) -> bool {
349        true
350    }
351
352    fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
353        match choice {
354            LanguageModelToolChoice::Auto
355            | LanguageModelToolChoice::Any
356            | LanguageModelToolChoice::None => true,
357        }
358    }
359
360    fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
361        LanguageModelToolSchemaFormat::JsonSchemaSubset
362    }
363
364    fn telemetry_id(&self) -> String {
365        format!("google/{}", self.model.id())
366    }
367
368    fn max_token_count(&self) -> usize {
369        self.model.max_token_count()
370    }
371
372    fn count_tokens(
373        &self,
374        request: LanguageModelRequest,
375        cx: &App,
376    ) -> BoxFuture<'static, Result<usize>> {
377        let model_id = self.model.id().to_string();
378        let request = into_google(request, model_id.clone(), self.model.mode());
379        let http_client = self.http_client.clone();
380        let api_key = self.state.read(cx).api_key.clone();
381
382        let settings = &AllLanguageModelSettings::get_global(cx).google;
383        let api_url = settings.api_url.clone();
384
385        async move {
386            let api_key = api_key.context("Missing Google API key")?;
387            let response = google_ai::count_tokens(
388                http_client.as_ref(),
389                &api_url,
390                &api_key,
391                google_ai::CountTokensRequest {
392                    generate_content_request: request,
393                },
394            )
395            .await?;
396            Ok(response.total_tokens)
397        }
398        .boxed()
399    }
400
401    fn stream_completion(
402        &self,
403        request: LanguageModelRequest,
404        cx: &AsyncApp,
405    ) -> BoxFuture<
406        'static,
407        Result<
408            futures::stream::BoxStream<
409                'static,
410                Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
411            >,
412        >,
413    > {
414        let request = into_google(request, self.model.id().to_string(), self.model.mode());
415        let request = self.stream_completion(request, cx);
416        let future = self.request_limiter.stream(async move {
417            let response = request
418                .await
419                .map_err(|err| LanguageModelCompletionError::Other(anyhow!(err)))?;
420            Ok(GoogleEventMapper::new().map_stream(response))
421        });
422        async move { Ok(future.await?.boxed()) }.boxed()
423    }
424}
425
426pub fn into_google(
427    mut request: LanguageModelRequest,
428    model_id: String,
429    mode: GoogleModelMode,
430) -> google_ai::GenerateContentRequest {
431    fn map_content(content: Vec<MessageContent>) -> Vec<Part> {
432        content
433            .into_iter()
434            .flat_map(|content| match content {
435                language_model::MessageContent::Text(text)
436                | language_model::MessageContent::Thinking { text, .. } => {
437                    if !text.is_empty() {
438                        vec![Part::TextPart(google_ai::TextPart { text })]
439                    } else {
440                        vec![]
441                    }
442                }
443                language_model::MessageContent::RedactedThinking(_) => vec![],
444                language_model::MessageContent::Image(image) => {
445                    vec![Part::InlineDataPart(google_ai::InlineDataPart {
446                        inline_data: google_ai::GenerativeContentBlob {
447                            mime_type: "image/png".to_string(),
448                            data: image.source.to_string(),
449                        },
450                    })]
451                }
452                language_model::MessageContent::ToolUse(tool_use) => {
453                    vec![Part::FunctionCallPart(google_ai::FunctionCallPart {
454                        function_call: google_ai::FunctionCall {
455                            name: tool_use.name.to_string(),
456                            args: tool_use.input,
457                        },
458                    })]
459                }
460                language_model::MessageContent::ToolResult(tool_result) => {
461                    match tool_result.content {
462                        language_model::LanguageModelToolResultContent::Text(text) => {
463                            vec![Part::FunctionResponsePart(
464                                google_ai::FunctionResponsePart {
465                                    function_response: google_ai::FunctionResponse {
466                                        name: tool_result.tool_name.to_string(),
467                                        // The API expects a valid JSON object
468                                        response: serde_json::json!({
469                                            "output": text
470                                        }),
471                                    },
472                                },
473                            )]
474                        }
475                        language_model::LanguageModelToolResultContent::Image(image) => {
476                            vec![
477                                Part::FunctionResponsePart(google_ai::FunctionResponsePart {
478                                    function_response: google_ai::FunctionResponse {
479                                        name: tool_result.tool_name.to_string(),
480                                        // The API expects a valid JSON object
481                                        response: serde_json::json!({
482                                            "output": "Tool responded with an image"
483                                        }),
484                                    },
485                                }),
486                                Part::InlineDataPart(google_ai::InlineDataPart {
487                                    inline_data: google_ai::GenerativeContentBlob {
488                                        mime_type: "image/png".to_string(),
489                                        data: image.source.to_string(),
490                                    },
491                                }),
492                            ]
493                        }
494                    }
495                }
496            })
497            .collect()
498    }
499
500    let system_instructions = if request
501        .messages
502        .first()
503        .map_or(false, |msg| matches!(msg.role, Role::System))
504    {
505        let message = request.messages.remove(0);
506        Some(SystemInstruction {
507            parts: map_content(message.content),
508        })
509    } else {
510        None
511    };
512
513    google_ai::GenerateContentRequest {
514        model: google_ai::ModelName { model_id },
515        system_instruction: system_instructions,
516        contents: request
517            .messages
518            .into_iter()
519            .filter_map(|message| {
520                let parts = map_content(message.content);
521                if parts.is_empty() {
522                    None
523                } else {
524                    Some(google_ai::Content {
525                        parts,
526                        role: match message.role {
527                            Role::User => google_ai::Role::User,
528                            Role::Assistant => google_ai::Role::Model,
529                            Role::System => google_ai::Role::User, // Google AI doesn't have a system role
530                        },
531                    })
532                }
533            })
534            .collect(),
535        generation_config: Some(google_ai::GenerationConfig {
536            candidate_count: Some(1),
537            stop_sequences: Some(request.stop),
538            max_output_tokens: None,
539            temperature: request.temperature.map(|t| t as f64).or(Some(1.0)),
540            thinking_config: match mode {
541                GoogleModelMode::Thinking { budget_tokens } => {
542                    budget_tokens.map(|thinking_budget| ThinkingConfig { thinking_budget })
543                }
544                GoogleModelMode::Default => None,
545            },
546            top_p: None,
547            top_k: None,
548        }),
549        safety_settings: None,
550        tools: (request.tools.len() > 0).then(|| {
551            vec![google_ai::Tool {
552                function_declarations: request
553                    .tools
554                    .into_iter()
555                    .map(|tool| FunctionDeclaration {
556                        name: tool.name,
557                        description: tool.description,
558                        parameters: tool.input_schema,
559                    })
560                    .collect(),
561            }]
562        }),
563        tool_config: request.tool_choice.map(|choice| google_ai::ToolConfig {
564            function_calling_config: google_ai::FunctionCallingConfig {
565                mode: match choice {
566                    LanguageModelToolChoice::Auto => google_ai::FunctionCallingMode::Auto,
567                    LanguageModelToolChoice::Any => google_ai::FunctionCallingMode::Any,
568                    LanguageModelToolChoice::None => google_ai::FunctionCallingMode::None,
569                },
570                allowed_function_names: None,
571            },
572        }),
573    }
574}
575
576pub struct GoogleEventMapper {
577    usage: UsageMetadata,
578    stop_reason: StopReason,
579}
580
581impl GoogleEventMapper {
582    pub fn new() -> Self {
583        Self {
584            usage: UsageMetadata::default(),
585            stop_reason: StopReason::EndTurn,
586        }
587    }
588
589    pub fn map_stream(
590        mut self,
591        events: Pin<Box<dyn Send + Stream<Item = Result<GenerateContentResponse>>>>,
592    ) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
593    {
594        events
595            .map(Some)
596            .chain(futures::stream::once(async { None }))
597            .flat_map(move |event| {
598                futures::stream::iter(match event {
599                    Some(Ok(event)) => self.map_event(event),
600                    Some(Err(error)) => {
601                        vec![Err(LanguageModelCompletionError::Other(anyhow!(error)))]
602                    }
603                    None => vec![Ok(LanguageModelCompletionEvent::Stop(self.stop_reason))],
604                })
605            })
606    }
607
608    pub fn map_event(
609        &mut self,
610        event: GenerateContentResponse,
611    ) -> Vec<Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
612        static TOOL_CALL_COUNTER: AtomicU64 = AtomicU64::new(0);
613
614        let mut events: Vec<_> = Vec::new();
615        let mut wants_to_use_tool = false;
616        if let Some(usage_metadata) = event.usage_metadata {
617            update_usage(&mut self.usage, &usage_metadata);
618            events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(
619                convert_usage(&self.usage),
620            )))
621        }
622        if let Some(candidates) = event.candidates {
623            for candidate in candidates {
624                if let Some(finish_reason) = candidate.finish_reason.as_deref() {
625                    self.stop_reason = match finish_reason {
626                        "STOP" => StopReason::EndTurn,
627                        "MAX_TOKENS" => StopReason::MaxTokens,
628                        _ => {
629                            log::error!("Unexpected google finish_reason: {finish_reason}");
630                            StopReason::EndTurn
631                        }
632                    };
633                }
634                candidate
635                    .content
636                    .parts
637                    .into_iter()
638                    .for_each(|part| match part {
639                        Part::TextPart(text_part) => {
640                            events.push(Ok(LanguageModelCompletionEvent::Text(text_part.text)))
641                        }
642                        Part::InlineDataPart(_) => {}
643                        Part::FunctionCallPart(function_call_part) => {
644                            wants_to_use_tool = true;
645                            let name: Arc<str> = function_call_part.function_call.name.into();
646                            let next_tool_id =
647                                TOOL_CALL_COUNTER.fetch_add(1, atomic::Ordering::SeqCst);
648                            let id: LanguageModelToolUseId =
649                                format!("{}-{}", name, next_tool_id).into();
650
651                            events.push(Ok(LanguageModelCompletionEvent::ToolUse(
652                                LanguageModelToolUse {
653                                    id,
654                                    name,
655                                    is_input_complete: true,
656                                    raw_input: function_call_part.function_call.args.to_string(),
657                                    input: function_call_part.function_call.args,
658                                },
659                            )));
660                        }
661                        Part::FunctionResponsePart(_) => {}
662                        Part::ThoughtPart(_) => {}
663                    });
664            }
665        }
666
667        // Even when Gemini wants to use a Tool, the API
668        // responds with `finish_reason: STOP`
669        if wants_to_use_tool {
670            self.stop_reason = StopReason::ToolUse;
671            events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse)));
672        }
673        events
674    }
675}
676
677pub fn count_google_tokens(
678    request: LanguageModelRequest,
679    cx: &App,
680) -> BoxFuture<'static, Result<usize>> {
681    // We couldn't use the GoogleLanguageModelProvider to count tokens because the github copilot doesn't have the access to google_ai directly.
682    // So we have to use tokenizer from tiktoken_rs to count tokens.
683    cx.background_spawn(async move {
684        let messages = request
685            .messages
686            .into_iter()
687            .map(|message| tiktoken_rs::ChatCompletionRequestMessage {
688                role: match message.role {
689                    Role::User => "user".into(),
690                    Role::Assistant => "assistant".into(),
691                    Role::System => "system".into(),
692                },
693                content: Some(message.string_contents()),
694                name: None,
695                function_call: None,
696            })
697            .collect::<Vec<_>>();
698
699        // Tiktoken doesn't yet support these models, so we manually use the
700        // same tokenizer as GPT-4.
701        tiktoken_rs::num_tokens_from_messages("gpt-4", &messages)
702    })
703    .boxed()
704}
705
706fn update_usage(usage: &mut UsageMetadata, new: &UsageMetadata) {
707    if let Some(prompt_token_count) = new.prompt_token_count {
708        usage.prompt_token_count = Some(prompt_token_count);
709    }
710    if let Some(cached_content_token_count) = new.cached_content_token_count {
711        usage.cached_content_token_count = Some(cached_content_token_count);
712    }
713    if let Some(candidates_token_count) = new.candidates_token_count {
714        usage.candidates_token_count = Some(candidates_token_count);
715    }
716    if let Some(tool_use_prompt_token_count) = new.tool_use_prompt_token_count {
717        usage.tool_use_prompt_token_count = Some(tool_use_prompt_token_count);
718    }
719    if let Some(thoughts_token_count) = new.thoughts_token_count {
720        usage.thoughts_token_count = Some(thoughts_token_count);
721    }
722    if let Some(total_token_count) = new.total_token_count {
723        usage.total_token_count = Some(total_token_count);
724    }
725}
726
727fn convert_usage(usage: &UsageMetadata) -> language_model::TokenUsage {
728    let prompt_tokens = usage.prompt_token_count.unwrap_or(0) as u32;
729    let cached_tokens = usage.cached_content_token_count.unwrap_or(0) as u32;
730    let input_tokens = prompt_tokens - cached_tokens;
731    let output_tokens = usage.candidates_token_count.unwrap_or(0) as u32;
732
733    language_model::TokenUsage {
734        input_tokens,
735        output_tokens,
736        cache_read_input_tokens: cached_tokens,
737        cache_creation_input_tokens: 0,
738    }
739}
740
741struct ConfigurationView {
742    api_key_editor: Entity<Editor>,
743    state: gpui::Entity<State>,
744    load_credentials_task: Option<Task<()>>,
745}
746
747impl ConfigurationView {
748    fn new(state: gpui::Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
749        cx.observe(&state, |_, _, cx| {
750            cx.notify();
751        })
752        .detach();
753
754        let load_credentials_task = Some(cx.spawn_in(window, {
755            let state = state.clone();
756            async move |this, cx| {
757                if let Some(task) = state
758                    .update(cx, |state, cx| state.authenticate(cx))
759                    .log_err()
760                {
761                    // We don't log an error, because "not signed in" is also an error.
762                    let _ = task.await;
763                }
764                this.update(cx, |this, cx| {
765                    this.load_credentials_task = None;
766                    cx.notify();
767                })
768                .log_err();
769            }
770        }));
771
772        Self {
773            api_key_editor: cx.new(|cx| {
774                let mut editor = Editor::single_line(window, cx);
775                editor.set_placeholder_text("AIzaSy...", cx);
776                editor
777            }),
778            state,
779            load_credentials_task,
780        }
781    }
782
783    fn save_api_key(&mut self, _: &menu::Confirm, window: &mut Window, cx: &mut Context<Self>) {
784        let api_key = self.api_key_editor.read(cx).text(cx);
785        if api_key.is_empty() {
786            return;
787        }
788
789        let state = self.state.clone();
790        cx.spawn_in(window, async move |_, cx| {
791            state
792                .update(cx, |state, cx| state.set_api_key(api_key, cx))?
793                .await
794        })
795        .detach_and_log_err(cx);
796
797        cx.notify();
798    }
799
800    fn reset_api_key(&mut self, window: &mut Window, cx: &mut Context<Self>) {
801        self.api_key_editor
802            .update(cx, |editor, cx| editor.set_text("", window, cx));
803
804        let state = self.state.clone();
805        cx.spawn_in(window, async move |_, cx| {
806            state.update(cx, |state, cx| state.reset_api_key(cx))?.await
807        })
808        .detach_and_log_err(cx);
809
810        cx.notify();
811    }
812
813    fn render_api_key_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
814        let settings = ThemeSettings::get_global(cx);
815        let text_style = TextStyle {
816            color: cx.theme().colors().text,
817            font_family: settings.ui_font.family.clone(),
818            font_features: settings.ui_font.features.clone(),
819            font_fallbacks: settings.ui_font.fallbacks.clone(),
820            font_size: rems(0.875).into(),
821            font_weight: settings.ui_font.weight,
822            font_style: FontStyle::Normal,
823            line_height: relative(1.3),
824            white_space: WhiteSpace::Normal,
825            ..Default::default()
826        };
827        EditorElement::new(
828            &self.api_key_editor,
829            EditorStyle {
830                background: cx.theme().colors().editor_background,
831                local_player: cx.theme().players().local(),
832                text: text_style,
833                ..Default::default()
834            },
835        )
836    }
837
838    fn should_render_editor(&self, cx: &mut Context<Self>) -> bool {
839        !self.state.read(cx).is_authenticated()
840    }
841}
842
843impl Render for ConfigurationView {
844    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
845        let env_var_set = self.state.read(cx).api_key_from_env;
846
847        if self.load_credentials_task.is_some() {
848            div().child(Label::new("Loading credentials...")).into_any()
849        } else if self.should_render_editor(cx) {
850            v_flex()
851                .size_full()
852                .on_action(cx.listener(Self::save_api_key))
853                .child(Label::new("To use Zed's assistant with Google AI, you need to add an API key. Follow these steps:"))
854                .child(
855                    List::new()
856                        .child(InstructionListItem::new(
857                            "Create one by visiting",
858                            Some("Google AI's console"),
859                            Some("https://aistudio.google.com/app/apikey"),
860                        ))
861                        .child(InstructionListItem::text_only(
862                            "Paste your API key below and hit enter to start using the assistant",
863                        )),
864                )
865                .child(
866                    h_flex()
867                        .w_full()
868                        .my_2()
869                        .px_2()
870                        .py_1()
871                        .bg(cx.theme().colors().editor_background)
872                        .border_1()
873                        .border_color(cx.theme().colors().border)
874                        .rounded_sm()
875                        .child(self.render_api_key_editor(cx)),
876                )
877                .child(
878                    Label::new(
879                        format!("You can also assign the {GOOGLE_AI_API_KEY_VAR} environment variable and restart Zed."),
880                    )
881                    .size(LabelSize::Small).color(Color::Muted),
882                )
883                .into_any()
884        } else {
885            h_flex()
886                .mt_1()
887                .p_1()
888                .justify_between()
889                .rounded_md()
890                .border_1()
891                .border_color(cx.theme().colors().border)
892                .bg(cx.theme().colors().background)
893                .child(
894                    h_flex()
895                        .gap_1()
896                        .child(Icon::new(IconName::Check).color(Color::Success))
897                        .child(Label::new(if env_var_set {
898                            format!("API key set in {GOOGLE_AI_API_KEY_VAR} environment variable.")
899                        } else {
900                            "API key configured.".to_string()
901                        })),
902                )
903                .child(
904                    Button::new("reset-key", "Reset Key")
905                        .label_size(LabelSize::Small)
906                        .icon(Some(IconName::Trash))
907                        .icon_size(IconSize::Small)
908                        .icon_position(IconPosition::Start)
909                        .disabled(env_var_set)
910                        .when(env_var_set, |this| {
911                            this.tooltip(Tooltip::text(format!("To reset your API key, unset the {GOOGLE_AI_API_KEY_VAR} environment variable.")))
912                        })
913                        .on_click(cx.listener(|this, _, window, cx| this.reset_api_key(window, cx))),
914                )
915                .into_any()
916        }
917    }
918}