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            "anthropic_terms_of_service",
260            "Anthropic Terms of Service",
261            "https://www.anthropic.com/legal/consumer-terms",
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 & Conditions").weight(FontWeight::MEDIUM))
282                            .child(Label::new("Please read and accept the terms and conditions of Zed AI and our provider partners to continue.").size(LabelSize::Small))
283                    )
284                    .child(v_flex().gap_1().children(terms))
285                    .child(
286                        h_flex().justify_end().child(
287                            Button::new("accept_terms", "I've read it and accept it")
288                                .disabled(disabled)
289                                .on_click({
290                                    let state = self.state.downgrade();
291                                    move |_, cx| {
292                                        state
293                                            .update(cx, |state, cx| {
294                                                state.accept_terms_of_service(cx)
295                                            })
296                                            .ok();
297                                    }
298                                }),
299                        ),
300                    )
301                    .into_any(),
302            )
303        }
304    }
305
306    fn reset_credentials(&self, _cx: &mut AppContext) -> Task<Result<()>> {
307        Task::ready(Ok(()))
308    }
309}
310
311pub struct CloudLanguageModel {
312    id: LanguageModelId,
313    model: CloudModel,
314    llm_api_token: LlmApiToken,
315    client: Arc<Client>,
316    request_limiter: RateLimiter,
317}
318
319#[derive(Clone, Default)]
320struct LlmApiToken(Arc<RwLock<Option<String>>>);
321
322impl CloudLanguageModel {
323    async fn perform_llm_completion(
324        client: Arc<Client>,
325        llm_api_token: LlmApiToken,
326        body: PerformCompletionParams,
327    ) -> Result<Response<AsyncBody>> {
328        let http_client = &client.http_client();
329
330        let mut token = llm_api_token.acquire(&client).await?;
331        let mut did_retry = false;
332
333        let response = loop {
334            let request = http_client::Request::builder()
335                .method(Method::POST)
336                .uri(http_client.build_zed_llm_url("/completion", &[])?.as_ref())
337                .header("Content-Type", "application/json")
338                .header("Authorization", format!("Bearer {token}"))
339                .body(serde_json::to_string(&body)?.into())?;
340            let mut response = http_client.send(request).await?;
341            if response.status().is_success() {
342                break response;
343            } else if !did_retry
344                && response
345                    .headers()
346                    .get(EXPIRED_LLM_TOKEN_HEADER_NAME)
347                    .is_some()
348            {
349                did_retry = true;
350                token = llm_api_token.refresh(&client).await?;
351            } else {
352                let mut body = String::new();
353                response.body_mut().read_to_string(&mut body).await?;
354                break Err(anyhow!(
355                    "cloud language model completion failed with status {}: {body}",
356                    response.status()
357                ))?;
358            }
359        };
360
361        Ok(response)
362    }
363}
364
365impl LanguageModel for CloudLanguageModel {
366    fn id(&self) -> LanguageModelId {
367        self.id.clone()
368    }
369
370    fn name(&self) -> LanguageModelName {
371        LanguageModelName::from(self.model.display_name().to_string())
372    }
373
374    fn provider_id(&self) -> LanguageModelProviderId {
375        LanguageModelProviderId(PROVIDER_ID.into())
376    }
377
378    fn provider_name(&self) -> LanguageModelProviderName {
379        LanguageModelProviderName(PROVIDER_NAME.into())
380    }
381
382    fn telemetry_id(&self) -> String {
383        format!("zed.dev/{}", self.model.id())
384    }
385
386    fn availability(&self) -> LanguageModelAvailability {
387        self.model.availability()
388    }
389
390    fn max_token_count(&self) -> usize {
391        self.model.max_token_count()
392    }
393
394    fn count_tokens(
395        &self,
396        request: LanguageModelRequest,
397        cx: &AppContext,
398    ) -> BoxFuture<'static, Result<usize>> {
399        match self.model.clone() {
400            CloudModel::Anthropic(_) => count_anthropic_tokens(request, cx),
401            CloudModel::OpenAi(model) => count_open_ai_tokens(request, model, cx),
402            CloudModel::Google(model) => {
403                let client = self.client.clone();
404                let request = request.into_google(model.id().into());
405                let request = google_ai::CountTokensRequest {
406                    contents: request.contents,
407                };
408                async move {
409                    let request = serde_json::to_string(&request)?;
410                    let response = client
411                        .request(proto::CountLanguageModelTokens {
412                            provider: proto::LanguageModelProvider::Google as i32,
413                            request,
414                        })
415                        .await?;
416                    Ok(response.token_count as usize)
417                }
418                .boxed()
419            }
420            CloudModel::Zed(_) => {
421                count_open_ai_tokens(request, open_ai::Model::ThreePointFiveTurbo, cx)
422            }
423        }
424    }
425
426    fn stream_completion(
427        &self,
428        request: LanguageModelRequest,
429        _cx: &AsyncAppContext,
430    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
431        match &self.model {
432            CloudModel::Anthropic(model) => {
433                let request = request.into_anthropic(model.id().into());
434                let client = self.client.clone();
435                let llm_api_token = self.llm_api_token.clone();
436                let future = self.request_limiter.stream(async move {
437                    let response = Self::perform_llm_completion(
438                        client.clone(),
439                        llm_api_token,
440                        PerformCompletionParams {
441                            provider: client::LanguageModelProvider::Anthropic,
442                            model: request.model.clone(),
443                            provider_request: RawValue::from_string(serde_json::to_string(
444                                &request,
445                            )?)?,
446                        },
447                    )
448                    .await?;
449                    let body = BufReader::new(response.into_body());
450                    let stream = futures::stream::try_unfold(body, move |mut body| async move {
451                        let mut buffer = String::new();
452                        match body.read_line(&mut buffer).await {
453                            Ok(0) => Ok(None),
454                            Ok(_) => {
455                                let event: anthropic::Event = serde_json::from_str(&buffer)
456                                    .context("failed to parse Anthropic event")?;
457                                Ok(Some((event, body)))
458                            }
459                            Err(err) => Err(AnthropicError::Other(err.into())),
460                        }
461                    });
462
463                    Ok(anthropic::extract_text_from_events(stream))
464                });
465                async move {
466                    Ok(future
467                        .await?
468                        .map(|result| result.map_err(|err| anyhow!(err)))
469                        .boxed())
470                }
471                .boxed()
472            }
473            CloudModel::OpenAi(model) => {
474                let client = self.client.clone();
475                let request = request.into_open_ai(model.id().into());
476                let llm_api_token = self.llm_api_token.clone();
477                let future = self.request_limiter.stream(async move {
478                    let response = Self::perform_llm_completion(
479                        client.clone(),
480                        llm_api_token,
481                        PerformCompletionParams {
482                            provider: client::LanguageModelProvider::OpenAi,
483                            model: request.model.clone(),
484                            provider_request: RawValue::from_string(serde_json::to_string(
485                                &request,
486                            )?)?,
487                        },
488                    )
489                    .await?;
490                    let body = BufReader::new(response.into_body());
491                    let stream = futures::stream::try_unfold(body, move |mut body| async move {
492                        let mut buffer = String::new();
493                        match body.read_line(&mut buffer).await {
494                            Ok(0) => Ok(None),
495                            Ok(_) => {
496                                let event: open_ai::ResponseStreamEvent =
497                                    serde_json::from_str(&buffer)?;
498                                Ok(Some((event, body)))
499                            }
500                            Err(e) => Err(e.into()),
501                        }
502                    });
503
504                    Ok(open_ai::extract_text_from_events(stream))
505                });
506                async move { Ok(future.await?.boxed()) }.boxed()
507            }
508            CloudModel::Google(model) => {
509                let client = self.client.clone();
510                let request = request.into_google(model.id().into());
511                let llm_api_token = self.llm_api_token.clone();
512                let future = self.request_limiter.stream(async move {
513                    let response = Self::perform_llm_completion(
514                        client.clone(),
515                        llm_api_token,
516                        PerformCompletionParams {
517                            provider: client::LanguageModelProvider::Google,
518                            model: request.model.clone(),
519                            provider_request: RawValue::from_string(serde_json::to_string(
520                                &request,
521                            )?)?,
522                        },
523                    )
524                    .await?;
525                    let body = BufReader::new(response.into_body());
526                    let stream = futures::stream::try_unfold(body, move |mut body| async move {
527                        let mut buffer = String::new();
528                        match body.read_line(&mut buffer).await {
529                            Ok(0) => Ok(None),
530                            Ok(_) => {
531                                let event: google_ai::GenerateContentResponse =
532                                    serde_json::from_str(&buffer)?;
533                                Ok(Some((event, body)))
534                            }
535                            Err(e) => Err(e.into()),
536                        }
537                    });
538
539                    Ok(google_ai::extract_text_from_events(stream))
540                });
541                async move { Ok(future.await?.boxed()) }.boxed()
542            }
543            CloudModel::Zed(model) => {
544                let client = self.client.clone();
545                let mut request = request.into_open_ai(model.id().into());
546                request.max_tokens = Some(4000);
547                let llm_api_token = self.llm_api_token.clone();
548                let future = self.request_limiter.stream(async move {
549                    let response = Self::perform_llm_completion(
550                        client.clone(),
551                        llm_api_token,
552                        PerformCompletionParams {
553                            provider: client::LanguageModelProvider::Zed,
554                            model: request.model.clone(),
555                            provider_request: RawValue::from_string(serde_json::to_string(
556                                &request,
557                            )?)?,
558                        },
559                    )
560                    .await?;
561                    let body = BufReader::new(response.into_body());
562                    let stream = futures::stream::try_unfold(body, move |mut body| async move {
563                        let mut buffer = String::new();
564                        match body.read_line(&mut buffer).await {
565                            Ok(0) => Ok(None),
566                            Ok(_) => {
567                                let event: open_ai::ResponseStreamEvent =
568                                    serde_json::from_str(&buffer)?;
569                                Ok(Some((event, body)))
570                            }
571                            Err(e) => Err(e.into()),
572                        }
573                    });
574
575                    Ok(open_ai::extract_text_from_events(stream))
576                });
577                async move { Ok(future.await?.boxed()) }.boxed()
578            }
579        }
580    }
581
582    fn use_any_tool(
583        &self,
584        request: LanguageModelRequest,
585        tool_name: String,
586        tool_description: String,
587        input_schema: serde_json::Value,
588        _cx: &AsyncAppContext,
589    ) -> BoxFuture<'static, Result<serde_json::Value>> {
590        match &self.model {
591            CloudModel::Anthropic(model) => {
592                let client = self.client.clone();
593                let mut request = request.into_anthropic(model.tool_model_id().into());
594                request.tool_choice = Some(anthropic::ToolChoice::Tool {
595                    name: tool_name.clone(),
596                });
597                request.tools = vec![anthropic::Tool {
598                    name: tool_name.clone(),
599                    description: tool_description,
600                    input_schema,
601                }];
602
603                let llm_api_token = self.llm_api_token.clone();
604                self.request_limiter
605                    .run(async move {
606                        let response = Self::perform_llm_completion(
607                            client.clone(),
608                            llm_api_token,
609                            PerformCompletionParams {
610                                provider: client::LanguageModelProvider::Anthropic,
611                                model: request.model.clone(),
612                                provider_request: RawValue::from_string(serde_json::to_string(
613                                    &request,
614                                )?)?,
615                            },
616                        )
617                        .await?;
618
619                        let mut tool_use_index = None;
620                        let mut tool_input = String::new();
621                        let mut body = BufReader::new(response.into_body());
622                        let mut line = String::new();
623                        while body.read_line(&mut line).await? > 0 {
624                            let event: anthropic::Event = serde_json::from_str(&line)?;
625                            line.clear();
626
627                            match event {
628                                anthropic::Event::ContentBlockStart {
629                                    content_block,
630                                    index,
631                                } => {
632                                    if let anthropic::Content::ToolUse { name, .. } = content_block
633                                    {
634                                        if name == tool_name {
635                                            tool_use_index = Some(index);
636                                        }
637                                    }
638                                }
639                                anthropic::Event::ContentBlockDelta { index, delta } => match delta
640                                {
641                                    anthropic::ContentDelta::TextDelta { .. } => {}
642                                    anthropic::ContentDelta::InputJsonDelta { partial_json } => {
643                                        if Some(index) == tool_use_index {
644                                            tool_input.push_str(&partial_json);
645                                        }
646                                    }
647                                },
648                                anthropic::Event::ContentBlockStop { index } => {
649                                    if Some(index) == tool_use_index {
650                                        return Ok(serde_json::from_str(&tool_input)?);
651                                    }
652                                }
653                                _ => {}
654                            }
655                        }
656
657                        if tool_use_index.is_some() {
658                            Err(anyhow!("tool content incomplete"))
659                        } else {
660                            Err(anyhow!("tool not used"))
661                        }
662                    })
663                    .boxed()
664            }
665            CloudModel::OpenAi(model) => {
666                let mut request = request.into_open_ai(model.id().into());
667                let client = self.client.clone();
668                let mut function = open_ai::FunctionDefinition {
669                    name: tool_name.clone(),
670                    description: None,
671                    parameters: None,
672                };
673                let func = open_ai::ToolDefinition::Function {
674                    function: function.clone(),
675                };
676                request.tool_choice = Some(open_ai::ToolChoice::Other(func.clone()));
677                // Fill in description and params separately, as they're not needed for tool_choice field.
678                function.description = Some(tool_description);
679                function.parameters = Some(input_schema);
680                request.tools = vec![open_ai::ToolDefinition::Function { function }];
681
682                let llm_api_token = self.llm_api_token.clone();
683                self.request_limiter
684                    .run(async move {
685                        let response = Self::perform_llm_completion(
686                            client.clone(),
687                            llm_api_token,
688                            PerformCompletionParams {
689                                provider: client::LanguageModelProvider::OpenAi,
690                                model: request.model.clone(),
691                                provider_request: RawValue::from_string(serde_json::to_string(
692                                    &request,
693                                )?)?,
694                            },
695                        )
696                        .await?;
697
698                        let mut body = BufReader::new(response.into_body());
699                        let mut line = String::new();
700                        let mut load_state = None;
701
702                        while body.read_line(&mut line).await? > 0 {
703                            let part: open_ai::ResponseStreamEvent = serde_json::from_str(&line)?;
704                            line.clear();
705
706                            for choice in part.choices {
707                                let Some(tool_calls) = choice.delta.tool_calls else {
708                                    continue;
709                                };
710
711                                for call in tool_calls {
712                                    if let Some(func) = call.function {
713                                        if func.name.as_deref() == Some(tool_name.as_str()) {
714                                            load_state = Some((String::default(), call.index));
715                                        }
716                                        if let Some((arguments, (output, index))) =
717                                            func.arguments.zip(load_state.as_mut())
718                                        {
719                                            if call.index == *index {
720                                                output.push_str(&arguments);
721                                            }
722                                        }
723                                    }
724                                }
725                            }
726                        }
727
728                        if let Some((arguments, _)) = load_state {
729                            return Ok(serde_json::from_str(&arguments)?);
730                        } else {
731                            bail!("tool not used");
732                        }
733                    })
734                    .boxed()
735            }
736            CloudModel::Google(_) => {
737                future::ready(Err(anyhow!("tool use not implemented for Google AI"))).boxed()
738            }
739            CloudModel::Zed(model) => {
740                // All Zed models are OpenAI-based at the time of writing.
741                let mut request = request.into_open_ai(model.id().into());
742                let client = self.client.clone();
743                let mut function = open_ai::FunctionDefinition {
744                    name: tool_name.clone(),
745                    description: None,
746                    parameters: None,
747                };
748                let func = open_ai::ToolDefinition::Function {
749                    function: function.clone(),
750                };
751                request.tool_choice = Some(open_ai::ToolChoice::Other(func.clone()));
752                // Fill in description and params separately, as they're not needed for tool_choice field.
753                function.description = Some(tool_description);
754                function.parameters = Some(input_schema);
755                request.tools = vec![open_ai::ToolDefinition::Function { function }];
756
757                let llm_api_token = self.llm_api_token.clone();
758                self.request_limiter
759                    .run(async move {
760                        let response = Self::perform_llm_completion(
761                            client.clone(),
762                            llm_api_token,
763                            PerformCompletionParams {
764                                provider: client::LanguageModelProvider::Zed,
765                                model: request.model.clone(),
766                                provider_request: RawValue::from_string(serde_json::to_string(
767                                    &request,
768                                )?)?,
769                            },
770                        )
771                        .await?;
772
773                        let mut body = BufReader::new(response.into_body());
774                        let mut line = String::new();
775                        let mut load_state = None;
776
777                        while body.read_line(&mut line).await? > 0 {
778                            let part: open_ai::ResponseStreamEvent = serde_json::from_str(&line)?;
779                            line.clear();
780
781                            for choice in part.choices {
782                                let Some(tool_calls) = choice.delta.tool_calls else {
783                                    continue;
784                                };
785
786                                for call in tool_calls {
787                                    if let Some(func) = call.function {
788                                        if func.name.as_deref() == Some(tool_name.as_str()) {
789                                            load_state = Some((String::default(), call.index));
790                                        }
791                                        if let Some((arguments, (output, index))) =
792                                            func.arguments.zip(load_state.as_mut())
793                                        {
794                                            if call.index == *index {
795                                                output.push_str(&arguments);
796                                            }
797                                        }
798                                    }
799                                }
800                            }
801                        }
802                        if let Some((arguments, _)) = load_state {
803                            return Ok(serde_json::from_str(&arguments)?);
804                        } else {
805                            bail!("tool not used");
806                        }
807                    })
808                    .boxed()
809            }
810        }
811    }
812}
813
814impl LlmApiToken {
815    async fn acquire(&self, client: &Arc<Client>) -> Result<String> {
816        let lock = self.0.upgradable_read().await;
817        if let Some(token) = lock.as_ref() {
818            Ok(token.to_string())
819        } else {
820            Self::fetch(RwLockUpgradableReadGuard::upgrade(lock).await, &client).await
821        }
822    }
823
824    async fn refresh(&self, client: &Arc<Client>) -> Result<String> {
825        Self::fetch(self.0.write().await, &client).await
826    }
827
828    async fn fetch<'a>(
829        mut lock: RwLockWriteGuard<'a, Option<String>>,
830        client: &Arc<Client>,
831    ) -> Result<String> {
832        let response = client.request(proto::GetLlmToken {}).await?;
833        *lock = Some(response.token.clone());
834        Ok(response.token.clone())
835    }
836}
837
838struct ConfigurationView {
839    state: gpui::Model<State>,
840}
841
842impl ConfigurationView {
843    fn authenticate(&mut self, cx: &mut ViewContext<Self>) {
844        self.state.update(cx, |state, cx| {
845            state.authenticate(cx).detach_and_log_err(cx);
846        });
847        cx.notify();
848    }
849}
850
851impl Render for ConfigurationView {
852    fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
853        const ZED_AI_URL: &str = "https://zed.dev/ai";
854        const ACCOUNT_SETTINGS_URL: &str = "https://zed.dev/account";
855
856        let is_connected = !self.state.read(cx).is_signed_out();
857        let plan = self.state.read(cx).user_store.read(cx).current_plan();
858        let must_accept_terms = !self.state.read(cx).has_accepted_terms_of_service(cx);
859
860        let is_pro = plan == Some(proto::Plan::ZedPro);
861
862        if is_connected {
863            v_flex()
864                .gap_3()
865                .max_w_4_5()
866                .when(must_accept_terms, |this| {
867                    this.child(Label::new(
868                        "You must accept the terms of service to use this provider.",
869                    ))
870                })
871                .child(Label::new(
872                    if is_pro {
873                        "You have full access to Zed's hosted models from Anthropic, OpenAI, Google with faster speeds and higher limits through Zed Pro."
874                    } else {
875                        "You have basic access to models from Anthropic through the Zed AI Free plan."
876                    }))
877                .children(if is_pro {
878                    Some(
879                        h_flex().child(
880                            Button::new("manage_settings", "Manage Subscription")
881                                .style(ButtonStyle::Filled)
882                                .on_click(
883                                    cx.listener(|_, _, cx| cx.open_url(ACCOUNT_SETTINGS_URL)),
884                                ),
885                        ),
886                    )
887                } else if cx.has_flag::<ZedPro>() {
888                    Some(
889                        h_flex()
890                            .gap_2()
891                            .child(
892                                Button::new("learn_more", "Learn more")
893                                    .style(ButtonStyle::Subtle)
894                                    .on_click(cx.listener(|_, _, cx| cx.open_url(ZED_AI_URL))),
895                            )
896                            .child(
897                                Button::new("upgrade", "Upgrade")
898                                    .style(ButtonStyle::Subtle)
899                                    .color(Color::Accent)
900                                    .on_click(
901                                        cx.listener(|_, _, cx| cx.open_url(ACCOUNT_SETTINGS_URL)),
902                                    ),
903                            ),
904                    )
905                } else {
906                    None
907                })
908        } else {
909            v_flex()
910                .gap_6()
911                .child(Label::new("Use the zed.dev to access language models."))
912                .child(
913                    v_flex()
914                        .gap_2()
915                        .child(
916                            Button::new("sign_in", "Sign in")
917                                .icon_color(Color::Muted)
918                                .icon(IconName::Github)
919                                .icon_position(IconPosition::Start)
920                                .style(ButtonStyle::Filled)
921                                .full_width()
922                                .on_click(cx.listener(move |this, _, cx| this.authenticate(cx))),
923                        )
924                        .child(
925                            div().flex().w_full().items_center().child(
926                                Label::new("Sign in to enable collaboration.")
927                                    .color(Color::Muted)
928                                    .size(LabelSize::Small),
929                            ),
930                        ),
931                )
932        }
933    }
934}