cloud.rs

  1use anthropic::{AnthropicError, AnthropicModelMode};
  2use anyhow::{Result, anyhow};
  3use client::{
  4    Client, EXPIRED_LLM_TOKEN_HEADER_NAME, MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME,
  5    PerformCompletionParams, UserStore, zed_urls,
  6};
  7use collections::BTreeMap;
  8use feature_flags::{FeatureFlagAppExt, LlmClosedBeta, ZedPro};
  9use futures::{
 10    AsyncBufReadExt, FutureExt, Stream, StreamExt, TryStreamExt as _, future::BoxFuture,
 11    stream::BoxStream,
 12};
 13use gpui::{AnyElement, AnyView, App, AsyncApp, Context, Entity, Subscription, Task};
 14use http_client::{AsyncBody, HttpClient, Method, Response, StatusCode};
 15use language_model::{
 16    AuthenticateError, CloudModel, LanguageModel, LanguageModelCacheConfiguration, LanguageModelId,
 17    LanguageModelName, LanguageModelProviderId, LanguageModelProviderName,
 18    LanguageModelProviderState, LanguageModelProviderTosView, LanguageModelRequest,
 19    LanguageModelToolSchemaFormat, RateLimiter, ZED_CLOUD_PROVIDER_ID,
 20};
 21use language_model::{
 22    LanguageModelAvailability, LanguageModelCompletionEvent, LanguageModelProvider, LlmApiToken,
 23    MaxMonthlySpendReachedError, PaymentRequiredError, RefreshLlmTokenListener,
 24};
 25use schemars::JsonSchema;
 26use serde::{Deserialize, Serialize, de::DeserializeOwned};
 27use serde_json::value::RawValue;
 28use settings::{Settings, SettingsStore};
 29use smol::Timer;
 30use smol::io::{AsyncReadExt, BufReader};
 31use std::{
 32    future,
 33    sync::{Arc, LazyLock},
 34    time::Duration,
 35};
 36use strum::IntoEnumIterator;
 37use ui::{TintColor, prelude::*};
 38
 39use crate::AllLanguageModelSettings;
 40use crate::provider::anthropic::{
 41    count_anthropic_tokens, into_anthropic, map_to_language_model_completion_events,
 42};
 43use crate::provider::google::into_google;
 44use crate::provider::open_ai::{count_open_ai_tokens, into_open_ai};
 45
 46pub const PROVIDER_NAME: &str = "Zed";
 47
 48const ZED_CLOUD_PROVIDER_ADDITIONAL_MODELS_JSON: Option<&str> =
 49    option_env!("ZED_CLOUD_PROVIDER_ADDITIONAL_MODELS_JSON");
 50
 51fn zed_cloud_provider_additional_models() -> &'static [AvailableModel] {
 52    static ADDITIONAL_MODELS: LazyLock<Vec<AvailableModel>> = LazyLock::new(|| {
 53        ZED_CLOUD_PROVIDER_ADDITIONAL_MODELS_JSON
 54            .map(|json| serde_json::from_str(json).unwrap())
 55            .unwrap_or_default()
 56    });
 57    ADDITIONAL_MODELS.as_slice()
 58}
 59
 60#[derive(Default, Clone, Debug, PartialEq)]
 61pub struct ZedDotDevSettings {
 62    pub available_models: Vec<AvailableModel>,
 63}
 64
 65#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
 66#[serde(rename_all = "lowercase")]
 67pub enum AvailableProvider {
 68    Anthropic,
 69    OpenAi,
 70    Google,
 71}
 72
 73#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
 74pub struct AvailableModel {
 75    /// The provider of the language model.
 76    pub provider: AvailableProvider,
 77    /// The model's name in the provider's API. e.g. claude-3-5-sonnet-20240620
 78    pub name: String,
 79    /// The name displayed in the UI, such as in the assistant panel model dropdown menu.
 80    pub display_name: Option<String>,
 81    /// The size of the context window, indicating the maximum number of tokens the model can process.
 82    pub max_tokens: usize,
 83    /// The maximum number of output tokens allowed by the model.
 84    pub max_output_tokens: Option<u32>,
 85    /// The maximum number of completion tokens allowed by the model (o1-* only)
 86    pub max_completion_tokens: Option<u32>,
 87    /// Override this model with a different Anthropic model for tool calls.
 88    pub tool_override: Option<String>,
 89    /// Indicates whether this custom model supports caching.
 90    pub cache_configuration: Option<LanguageModelCacheConfiguration>,
 91    /// The default temperature to use for this model.
 92    pub default_temperature: Option<f32>,
 93    /// Any extra beta headers to provide when using the model.
 94    #[serde(default)]
 95    pub extra_beta_headers: Vec<String>,
 96    /// The model's mode (e.g. thinking)
 97    pub mode: Option<ModelMode>,
 98}
 99
100#[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
101#[serde(tag = "type", rename_all = "lowercase")]
102pub enum ModelMode {
103    #[default]
104    Default,
105    Thinking {
106        /// The maximum number of tokens to use for reasoning. Must be lower than the model's `max_output_tokens`.
107        budget_tokens: Option<u32>,
108    },
109}
110
111impl From<ModelMode> for AnthropicModelMode {
112    fn from(value: ModelMode) -> Self {
113        match value {
114            ModelMode::Default => AnthropicModelMode::Default,
115            ModelMode::Thinking { budget_tokens } => AnthropicModelMode::Thinking { budget_tokens },
116        }
117    }
118}
119
120pub struct CloudLanguageModelProvider {
121    client: Arc<Client>,
122    state: gpui::Entity<State>,
123    _maintain_client_status: Task<()>,
124}
125
126pub struct State {
127    client: Arc<Client>,
128    llm_api_token: LlmApiToken,
129    user_store: Entity<UserStore>,
130    status: client::Status,
131    accept_terms: Option<Task<Result<()>>>,
132    _settings_subscription: Subscription,
133    _llm_token_subscription: Subscription,
134}
135
136impl State {
137    fn new(
138        client: Arc<Client>,
139        user_store: Entity<UserStore>,
140        status: client::Status,
141        cx: &mut Context<Self>,
142    ) -> Self {
143        let refresh_llm_token_listener = RefreshLlmTokenListener::global(cx);
144
145        Self {
146            client: client.clone(),
147            llm_api_token: LlmApiToken::default(),
148            user_store,
149            status,
150            accept_terms: None,
151            _settings_subscription: cx.observe_global::<SettingsStore>(|_, cx| {
152                cx.notify();
153            }),
154            _llm_token_subscription: cx.subscribe(
155                &refresh_llm_token_listener,
156                |this, _listener, _event, cx| {
157                    let client = this.client.clone();
158                    let llm_api_token = this.llm_api_token.clone();
159                    cx.spawn(async move |_this, _cx| {
160                        llm_api_token.refresh(&client).await?;
161                        anyhow::Ok(())
162                    })
163                    .detach_and_log_err(cx);
164                },
165            ),
166        }
167    }
168
169    fn is_signed_out(&self) -> bool {
170        self.status.is_signed_out()
171    }
172
173    fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
174        let client = self.client.clone();
175        cx.spawn(async move |this, cx| {
176            client.authenticate_and_connect(true, &cx).await?;
177            this.update(cx, |_, cx| cx.notify())
178        })
179    }
180
181    fn has_accepted_terms_of_service(&self, cx: &App) -> bool {
182        self.user_store
183            .read(cx)
184            .current_user_has_accepted_terms()
185            .unwrap_or(false)
186    }
187
188    fn accept_terms_of_service(&mut self, cx: &mut Context<Self>) {
189        let user_store = self.user_store.clone();
190        self.accept_terms = Some(cx.spawn(async move |this, cx| {
191            let _ = user_store
192                .update(cx, |store, cx| store.accept_terms_of_service(cx))?
193                .await;
194            this.update(cx, |this, cx| {
195                this.accept_terms = None;
196                cx.notify()
197            })
198        }));
199    }
200}
201
202impl CloudLanguageModelProvider {
203    pub fn new(user_store: Entity<UserStore>, client: Arc<Client>, cx: &mut App) -> Self {
204        let mut status_rx = client.status();
205        let status = *status_rx.borrow();
206
207        let state = cx.new(|cx| State::new(client.clone(), user_store.clone(), status, cx));
208
209        let state_ref = state.downgrade();
210        let maintain_client_status = cx.spawn(async move |cx| {
211            while let Some(status) = status_rx.next().await {
212                if let Some(this) = state_ref.upgrade() {
213                    _ = this.update(cx, |this, cx| {
214                        if this.status != status {
215                            this.status = status;
216                            cx.notify();
217                        }
218                    });
219                } else {
220                    break;
221                }
222            }
223        });
224
225        Self {
226            client,
227            state: state.clone(),
228            _maintain_client_status: maintain_client_status,
229        }
230    }
231}
232
233impl LanguageModelProviderState for CloudLanguageModelProvider {
234    type ObservableEntity = State;
235
236    fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> {
237        Some(self.state.clone())
238    }
239}
240
241impl LanguageModelProvider for CloudLanguageModelProvider {
242    fn id(&self) -> LanguageModelProviderId {
243        LanguageModelProviderId(ZED_CLOUD_PROVIDER_ID.into())
244    }
245
246    fn name(&self) -> LanguageModelProviderName {
247        LanguageModelProviderName(PROVIDER_NAME.into())
248    }
249
250    fn icon(&self) -> IconName {
251        IconName::AiZed
252    }
253
254    fn default_model(&self, cx: &App) -> Option<Arc<dyn LanguageModel>> {
255        let llm_api_token = self.state.read(cx).llm_api_token.clone();
256        let model = CloudModel::Anthropic(anthropic::Model::default());
257        Some(Arc::new(CloudLanguageModel {
258            id: LanguageModelId::from(model.id().to_string()),
259            model,
260            llm_api_token: llm_api_token.clone(),
261            client: self.client.clone(),
262            request_limiter: RateLimiter::new(4),
263        }))
264    }
265
266    fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
267        let mut models = BTreeMap::default();
268
269        if cx.is_staff() {
270            for model in anthropic::Model::iter() {
271                if !matches!(model, anthropic::Model::Custom { .. }) {
272                    models.insert(model.id().to_string(), CloudModel::Anthropic(model));
273                }
274            }
275            for model in open_ai::Model::iter() {
276                if !matches!(model, open_ai::Model::Custom { .. }) {
277                    models.insert(model.id().to_string(), CloudModel::OpenAi(model));
278                }
279            }
280            for model in google_ai::Model::iter() {
281                if !matches!(model, google_ai::Model::Custom { .. }) {
282                    models.insert(model.id().to_string(), CloudModel::Google(model));
283                }
284            }
285        } else {
286            models.insert(
287                anthropic::Model::Claude3_5Sonnet.id().to_string(),
288                CloudModel::Anthropic(anthropic::Model::Claude3_5Sonnet),
289            );
290            models.insert(
291                anthropic::Model::Claude3_7Sonnet.id().to_string(),
292                CloudModel::Anthropic(anthropic::Model::Claude3_7Sonnet),
293            );
294            models.insert(
295                anthropic::Model::Claude3_7SonnetThinking.id().to_string(),
296                CloudModel::Anthropic(anthropic::Model::Claude3_7SonnetThinking),
297            );
298        }
299
300        let llm_closed_beta_models = if cx.has_flag::<LlmClosedBeta>() {
301            zed_cloud_provider_additional_models()
302        } else {
303            &[]
304        };
305
306        // Override with available models from settings
307        for model in AllLanguageModelSettings::get_global(cx)
308            .zed_dot_dev
309            .available_models
310            .iter()
311            .chain(llm_closed_beta_models)
312            .cloned()
313        {
314            let model = match model.provider {
315                AvailableProvider::Anthropic => CloudModel::Anthropic(anthropic::Model::Custom {
316                    name: model.name.clone(),
317                    display_name: model.display_name.clone(),
318                    max_tokens: model.max_tokens,
319                    tool_override: model.tool_override.clone(),
320                    cache_configuration: model.cache_configuration.as_ref().map(|config| {
321                        anthropic::AnthropicModelCacheConfiguration {
322                            max_cache_anchors: config.max_cache_anchors,
323                            should_speculate: config.should_speculate,
324                            min_total_token: config.min_total_token,
325                        }
326                    }),
327                    default_temperature: model.default_temperature,
328                    max_output_tokens: model.max_output_tokens,
329                    extra_beta_headers: model.extra_beta_headers.clone(),
330                    mode: model.mode.unwrap_or_default().into(),
331                }),
332                AvailableProvider::OpenAi => CloudModel::OpenAi(open_ai::Model::Custom {
333                    name: model.name.clone(),
334                    display_name: model.display_name.clone(),
335                    max_tokens: model.max_tokens,
336                    max_output_tokens: model.max_output_tokens,
337                    max_completion_tokens: model.max_completion_tokens,
338                }),
339                AvailableProvider::Google => CloudModel::Google(google_ai::Model::Custom {
340                    name: model.name.clone(),
341                    display_name: model.display_name.clone(),
342                    max_tokens: model.max_tokens,
343                }),
344            };
345            models.insert(model.id().to_string(), model.clone());
346        }
347
348        let llm_api_token = self.state.read(cx).llm_api_token.clone();
349        models
350            .into_values()
351            .map(|model| {
352                Arc::new(CloudLanguageModel {
353                    id: LanguageModelId::from(model.id().to_string()),
354                    model,
355                    llm_api_token: llm_api_token.clone(),
356                    client: self.client.clone(),
357                    request_limiter: RateLimiter::new(4),
358                }) as Arc<dyn LanguageModel>
359            })
360            .collect()
361    }
362
363    fn is_authenticated(&self, cx: &App) -> bool {
364        !self.state.read(cx).is_signed_out()
365    }
366
367    fn authenticate(&self, _cx: &mut App) -> Task<Result<(), AuthenticateError>> {
368        Task::ready(Ok(()))
369    }
370
371    fn configuration_view(&self, _: &mut Window, cx: &mut App) -> AnyView {
372        cx.new(|_| ConfigurationView {
373            state: self.state.clone(),
374        })
375        .into()
376    }
377
378    fn must_accept_terms(&self, cx: &App) -> bool {
379        !self.state.read(cx).has_accepted_terms_of_service(cx)
380    }
381
382    fn render_accept_terms(
383        &self,
384        view: LanguageModelProviderTosView,
385        cx: &mut App,
386    ) -> Option<AnyElement> {
387        render_accept_terms(self.state.clone(), view, cx)
388    }
389
390    fn reset_credentials(&self, _cx: &mut App) -> Task<Result<()>> {
391        Task::ready(Ok(()))
392    }
393}
394
395fn render_accept_terms(
396    state: Entity<State>,
397    view_kind: LanguageModelProviderTosView,
398    cx: &mut App,
399) -> Option<AnyElement> {
400    if state.read(cx).has_accepted_terms_of_service(cx) {
401        return None;
402    }
403
404    let accept_terms_disabled = state.read(cx).accept_terms.is_some();
405
406    let terms_button = Button::new("terms_of_service", "Terms of Service")
407        .style(ButtonStyle::Subtle)
408        .icon(IconName::ArrowUpRight)
409        .icon_color(Color::Muted)
410        .icon_size(IconSize::XSmall)
411        .on_click(move |_, _window, cx| cx.open_url("https://zed.dev/terms-of-service"));
412
413    let text = "To start using Zed AI, please read and accept the";
414
415    let form = v_flex()
416        .w_full()
417        .gap_2()
418        .child(
419            h_flex()
420                .flex_wrap()
421                .items_start()
422                .child(Label::new(text))
423                .child(terms_button),
424        )
425        .child({
426            let button_container = h_flex().w_full().child(
427                Button::new("accept_terms", "I accept the Terms of Service")
428                    .style(ButtonStyle::Tinted(TintColor::Accent))
429                    .disabled(accept_terms_disabled)
430                    .on_click({
431                        let state = state.downgrade();
432                        move |_, _window, cx| {
433                            state
434                                .update(cx, |state, cx| state.accept_terms_of_service(cx))
435                                .ok();
436                        }
437                    }),
438            );
439
440            match view_kind {
441                LanguageModelProviderTosView::PromptEditorPopup => button_container.justify_end(),
442                LanguageModelProviderTosView::Configuration
443                | LanguageModelProviderTosView::ThreadEmptyState => {
444                    button_container.justify_start()
445                }
446            }
447        });
448
449    Some(form.into_any())
450}
451
452pub struct CloudLanguageModel {
453    id: LanguageModelId,
454    model: CloudModel,
455    llm_api_token: LlmApiToken,
456    client: Arc<Client>,
457    request_limiter: RateLimiter,
458}
459
460impl CloudLanguageModel {
461    const MAX_RETRIES: usize = 3;
462
463    async fn perform_llm_completion(
464        client: Arc<Client>,
465        llm_api_token: LlmApiToken,
466        body: PerformCompletionParams,
467    ) -> Result<Response<AsyncBody>> {
468        let http_client = &client.http_client();
469
470        let mut token = llm_api_token.acquire(&client).await?;
471        let mut retries_remaining = Self::MAX_RETRIES;
472        let mut retry_delay = Duration::from_secs(1);
473
474        loop {
475            let request_builder = http_client::Request::builder();
476            let request = request_builder
477                .method(Method::POST)
478                .uri(http_client.build_zed_llm_url("/completion", &[])?.as_ref())
479                .header("Content-Type", "application/json")
480                .header("Authorization", format!("Bearer {token}"))
481                .body(serde_json::to_string(&body)?.into())?;
482            let mut response = http_client.send(request).await?;
483            let status = response.status();
484            if status.is_success() {
485                return Ok(response);
486            } else if response
487                .headers()
488                .get(EXPIRED_LLM_TOKEN_HEADER_NAME)
489                .is_some()
490            {
491                retries_remaining -= 1;
492                token = llm_api_token.refresh(&client).await?;
493            } else if status == StatusCode::FORBIDDEN
494                && response
495                    .headers()
496                    .get(MAX_LLM_MONTHLY_SPEND_REACHED_HEADER_NAME)
497                    .is_some()
498            {
499                return Err(anyhow!(MaxMonthlySpendReachedError));
500            } else if status.as_u16() >= 500 && status.as_u16() < 600 {
501                // If we encounter an error in the 500 range, retry after a delay.
502                // We've seen at least these in the wild from API providers:
503                // * 500 Internal Server Error
504                // * 502 Bad Gateway
505                // * 529 Service Overloaded
506
507                if retries_remaining == 0 {
508                    let mut body = String::new();
509                    response.body_mut().read_to_string(&mut body).await?;
510                    return Err(anyhow!(
511                        "cloud language model completion failed after {} retries with status {status}: {body}",
512                        Self::MAX_RETRIES
513                    ));
514                }
515
516                Timer::after(retry_delay).await;
517
518                retries_remaining -= 1;
519                retry_delay *= 2; // If it fails again, wait longer.
520            } else if status == StatusCode::PAYMENT_REQUIRED {
521                return Err(anyhow!(PaymentRequiredError));
522            } else {
523                let mut body = String::new();
524                response.body_mut().read_to_string(&mut body).await?;
525                return Err(anyhow!(
526                    "cloud language model completion failed with status {status}: {body}",
527                ));
528            }
529        }
530    }
531}
532
533impl LanguageModel for CloudLanguageModel {
534    fn id(&self) -> LanguageModelId {
535        self.id.clone()
536    }
537
538    fn name(&self) -> LanguageModelName {
539        LanguageModelName::from(self.model.display_name().to_string())
540    }
541
542    fn icon(&self) -> Option<IconName> {
543        self.model.icon()
544    }
545
546    fn provider_id(&self) -> LanguageModelProviderId {
547        LanguageModelProviderId(ZED_CLOUD_PROVIDER_ID.into())
548    }
549
550    fn provider_name(&self) -> LanguageModelProviderName {
551        LanguageModelProviderName(PROVIDER_NAME.into())
552    }
553
554    fn telemetry_id(&self) -> String {
555        format!("zed.dev/{}", self.model.id())
556    }
557
558    fn availability(&self) -> LanguageModelAvailability {
559        self.model.availability()
560    }
561
562    fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
563        self.model.tool_input_format()
564    }
565
566    fn max_token_count(&self) -> usize {
567        self.model.max_token_count()
568    }
569
570    fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
571        match &self.model {
572            CloudModel::Anthropic(model) => {
573                model
574                    .cache_configuration()
575                    .map(|cache| LanguageModelCacheConfiguration {
576                        max_cache_anchors: cache.max_cache_anchors,
577                        should_speculate: cache.should_speculate,
578                        min_total_token: cache.min_total_token,
579                    })
580            }
581            CloudModel::OpenAi(_) | CloudModel::Google(_) => None,
582        }
583    }
584
585    fn count_tokens(
586        &self,
587        request: LanguageModelRequest,
588        cx: &App,
589    ) -> BoxFuture<'static, Result<usize>> {
590        match self.model.clone() {
591            CloudModel::Anthropic(_) => count_anthropic_tokens(request, cx),
592            CloudModel::OpenAi(model) => count_open_ai_tokens(request, model, cx),
593            CloudModel::Google(model) => {
594                let client = self.client.clone();
595                let request = into_google(request, model.id().into());
596                let request = google_ai::CountTokensRequest {
597                    contents: request.contents,
598                };
599                async move {
600                    let request = serde_json::to_string(&request)?;
601                    let response = client
602                        .request(proto::CountLanguageModelTokens {
603                            provider: proto::LanguageModelProvider::Google as i32,
604                            request,
605                        })
606                        .await?;
607                    Ok(response.token_count as usize)
608                }
609                .boxed()
610            }
611        }
612    }
613
614    fn stream_completion(
615        &self,
616        request: LanguageModelRequest,
617        _cx: &AsyncApp,
618    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<LanguageModelCompletionEvent>>>> {
619        match &self.model {
620            CloudModel::Anthropic(model) => {
621                let request = into_anthropic(
622                    request,
623                    model.request_id().into(),
624                    model.default_temperature(),
625                    model.max_output_tokens(),
626                    model.mode(),
627                );
628                let client = self.client.clone();
629                let llm_api_token = self.llm_api_token.clone();
630                let future = self.request_limiter.stream(async move {
631                    let response = Self::perform_llm_completion(
632                        client.clone(),
633                        llm_api_token,
634                        PerformCompletionParams {
635                            provider: client::LanguageModelProvider::Anthropic,
636                            model: request.model.clone(),
637                            provider_request: RawValue::from_string(serde_json::to_string(
638                                &request,
639                            )?)?,
640                        },
641                    )
642                    .await?;
643                    Ok(map_to_language_model_completion_events(Box::pin(
644                        response_lines(response).map_err(AnthropicError::Other),
645                    )))
646                });
647                async move { Ok(future.await?.boxed()) }.boxed()
648            }
649            CloudModel::OpenAi(model) => {
650                let client = self.client.clone();
651                let request = into_open_ai(request, model.id().into(), model.max_output_tokens());
652                let llm_api_token = self.llm_api_token.clone();
653                let future = self.request_limiter.stream(async move {
654                    let response = Self::perform_llm_completion(
655                        client.clone(),
656                        llm_api_token,
657                        PerformCompletionParams {
658                            provider: client::LanguageModelProvider::OpenAi,
659                            model: request.model.clone(),
660                            provider_request: RawValue::from_string(serde_json::to_string(
661                                &request,
662                            )?)?,
663                        },
664                    )
665                    .await?;
666                    Ok(open_ai::extract_text_from_events(response_lines(response)))
667                });
668                async move {
669                    Ok(future
670                        .await?
671                        .map(|result| result.map(LanguageModelCompletionEvent::Text))
672                        .boxed())
673                }
674                .boxed()
675            }
676            CloudModel::Google(model) => {
677                let client = self.client.clone();
678                let request = into_google(request, model.id().into());
679                let llm_api_token = self.llm_api_token.clone();
680                let future = self.request_limiter.stream(async move {
681                    let response = Self::perform_llm_completion(
682                        client.clone(),
683                        llm_api_token,
684                        PerformCompletionParams {
685                            provider: client::LanguageModelProvider::Google,
686                            model: request.model.clone(),
687                            provider_request: RawValue::from_string(serde_json::to_string(
688                                &request,
689                            )?)?,
690                        },
691                    )
692                    .await?;
693                    Ok(google_ai::extract_text_from_events(response_lines(
694                        response,
695                    )))
696                });
697                async move {
698                    Ok(future
699                        .await?
700                        .map(|result| result.map(LanguageModelCompletionEvent::Text))
701                        .boxed())
702                }
703                .boxed()
704            }
705        }
706    }
707
708    fn use_any_tool(
709        &self,
710        request: LanguageModelRequest,
711        tool_name: String,
712        tool_description: String,
713        input_schema: serde_json::Value,
714        _cx: &AsyncApp,
715    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
716        let client = self.client.clone();
717        let llm_api_token = self.llm_api_token.clone();
718
719        match &self.model {
720            CloudModel::Anthropic(model) => {
721                let mut request = into_anthropic(
722                    request,
723                    model.tool_model_id().into(),
724                    model.default_temperature(),
725                    model.max_output_tokens(),
726                    model.mode(),
727                );
728                request.tool_choice = Some(anthropic::ToolChoice::Tool {
729                    name: tool_name.clone(),
730                });
731                request.tools = vec![anthropic::Tool {
732                    name: tool_name.clone(),
733                    description: tool_description,
734                    input_schema,
735                }];
736
737                self.request_limiter
738                    .run(async move {
739                        let response = Self::perform_llm_completion(
740                            client.clone(),
741                            llm_api_token,
742                            PerformCompletionParams {
743                                provider: client::LanguageModelProvider::Anthropic,
744                                model: request.model.clone(),
745                                provider_request: RawValue::from_string(serde_json::to_string(
746                                    &request,
747                                )?)?,
748                            },
749                        )
750                        .await?;
751
752                        Ok(anthropic::extract_tool_args_from_events(
753                            tool_name,
754                            Box::pin(response_lines(response)),
755                        )
756                        .await?
757                        .boxed())
758                    })
759                    .boxed()
760            }
761            CloudModel::OpenAi(model) => {
762                let mut request =
763                    into_open_ai(request, model.id().into(), model.max_output_tokens());
764                request.tool_choice = Some(open_ai::ToolChoice::Other(
765                    open_ai::ToolDefinition::Function {
766                        function: open_ai::FunctionDefinition {
767                            name: tool_name.clone(),
768                            description: None,
769                            parameters: None,
770                        },
771                    },
772                ));
773                request.tools = vec![open_ai::ToolDefinition::Function {
774                    function: open_ai::FunctionDefinition {
775                        name: tool_name.clone(),
776                        description: Some(tool_description),
777                        parameters: Some(input_schema),
778                    },
779                }];
780
781                self.request_limiter
782                    .run(async move {
783                        let response = Self::perform_llm_completion(
784                            client.clone(),
785                            llm_api_token,
786                            PerformCompletionParams {
787                                provider: client::LanguageModelProvider::OpenAi,
788                                model: request.model.clone(),
789                                provider_request: RawValue::from_string(serde_json::to_string(
790                                    &request,
791                                )?)?,
792                            },
793                        )
794                        .await?;
795
796                        Ok(open_ai::extract_tool_args_from_events(
797                            tool_name,
798                            Box::pin(response_lines(response)),
799                        )
800                        .await?
801                        .boxed())
802                    })
803                    .boxed()
804            }
805            CloudModel::Google(_) => {
806                future::ready(Err(anyhow!("tool use not implemented for Google AI"))).boxed()
807            }
808        }
809    }
810}
811
812fn response_lines<T: DeserializeOwned>(
813    response: Response<AsyncBody>,
814) -> impl Stream<Item = Result<T>> {
815    futures::stream::try_unfold(
816        (String::new(), BufReader::new(response.into_body())),
817        move |(mut line, mut body)| async {
818            match body.read_line(&mut line).await {
819                Ok(0) => Ok(None),
820                Ok(_) => {
821                    let event: T = serde_json::from_str(&line)?;
822                    line.clear();
823                    Ok(Some((event, (line, body))))
824                }
825                Err(e) => Err(e.into()),
826            }
827        },
828    )
829}
830
831struct ConfigurationView {
832    state: gpui::Entity<State>,
833}
834
835impl ConfigurationView {
836    fn authenticate(&mut self, cx: &mut Context<Self>) {
837        self.state.update(cx, |state, cx| {
838            state.authenticate(cx).detach_and_log_err(cx);
839        });
840        cx.notify();
841    }
842}
843
844impl Render for ConfigurationView {
845    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
846        const ZED_AI_URL: &str = "https://zed.dev/ai";
847
848        let is_connected = !self.state.read(cx).is_signed_out();
849        let plan = self.state.read(cx).user_store.read(cx).current_plan();
850        let has_accepted_terms = self.state.read(cx).has_accepted_terms_of_service(cx);
851
852        let is_pro = plan == Some(proto::Plan::ZedPro);
853        let subscription_text = Label::new(if is_pro {
854            "You have full access to Zed's hosted LLMs, which include models from Anthropic, OpenAI, and Google. They come with faster speeds and higher limits through Zed Pro."
855        } else {
856            "You have basic access to models from Anthropic through the Zed AI Free plan."
857        });
858        let manage_subscription_button = if is_pro {
859            Some(
860                h_flex().child(
861                    Button::new("manage_settings", "Manage Subscription")
862                        .style(ButtonStyle::Tinted(TintColor::Accent))
863                        .on_click(
864                            cx.listener(|_, _, _, cx| cx.open_url(&zed_urls::account_url(cx))),
865                        ),
866                ),
867            )
868        } else if cx.has_flag::<ZedPro>() {
869            Some(
870                h_flex()
871                    .gap_2()
872                    .child(
873                        Button::new("learn_more", "Learn more")
874                            .style(ButtonStyle::Subtle)
875                            .on_click(cx.listener(|_, _, _, cx| cx.open_url(ZED_AI_URL))),
876                    )
877                    .child(
878                        Button::new("upgrade", "Upgrade")
879                            .style(ButtonStyle::Subtle)
880                            .color(Color::Accent)
881                            .on_click(
882                                cx.listener(|_, _, _, cx| cx.open_url(&zed_urls::account_url(cx))),
883                            ),
884                    ),
885            )
886        } else {
887            None
888        };
889
890        if is_connected {
891            v_flex()
892                .gap_3()
893                .w_full()
894                .children(render_accept_terms(
895                    self.state.clone(),
896                    LanguageModelProviderTosView::Configuration,
897                    cx,
898                ))
899                .when(has_accepted_terms, |this| {
900                    this.child(subscription_text)
901                        .children(manage_subscription_button)
902                })
903        } else {
904            v_flex()
905                .gap_2()
906                .child(Label::new("Use Zed AI to access hosted language models."))
907                .child(
908                    Button::new("sign_in", "Sign In")
909                        .icon_color(Color::Muted)
910                        .icon(IconName::Github)
911                        .icon_position(IconPosition::Start)
912                        .on_click(cx.listener(move |this, _, _, cx| this.authenticate(cx))),
913                )
914        }
915    }
916}