cloud.rs

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