1use super::open_ai::count_open_ai_tokens;
2use crate::{
3 settings::AllLanguageModelSettings, CloudModel, LanguageModel, LanguageModelId,
4 LanguageModelName, LanguageModelProviderId, LanguageModelProviderName,
5 LanguageModelProviderState, LanguageModelRequest, RateLimiter,
6};
7use anyhow::{anyhow, Context as _, Result};
8use client::Client;
9use collections::BTreeMap;
10use futures::{future::BoxFuture, stream::BoxStream, FutureExt, StreamExt};
11use gpui::{AnyView, AppContext, AsyncAppContext, ModelContext, Subscription, Task};
12use schemars::JsonSchema;
13use serde::{Deserialize, Serialize};
14use settings::{Settings, SettingsStore};
15use std::{future, sync::Arc};
16use strum::IntoEnumIterator;
17use ui::prelude::*;
18
19use crate::LanguageModelProvider;
20
21use super::anthropic::count_anthropic_tokens;
22
23pub const PROVIDER_ID: &str = "zed.dev";
24pub const PROVIDER_NAME: &str = "zed.dev";
25
26#[derive(Default, Clone, Debug, PartialEq)]
27pub struct ZedDotDevSettings {
28 pub available_models: Vec<AvailableModel>,
29}
30
31#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
32#[serde(rename_all = "lowercase")]
33pub enum AvailableProvider {
34 Anthropic,
35 OpenAi,
36 Google,
37}
38
39#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
40pub struct AvailableModel {
41 provider: AvailableProvider,
42 name: String,
43 max_tokens: usize,
44 tool_override: Option<String>,
45}
46
47pub struct CloudLanguageModelProvider {
48 client: Arc<Client>,
49 state: gpui::Model<State>,
50 _maintain_client_status: Task<()>,
51}
52
53pub struct State {
54 client: Arc<Client>,
55 status: client::Status,
56 _subscription: Subscription,
57}
58
59impl State {
60 fn authenticate(&self, cx: &mut ModelContext<Self>) -> Task<Result<()>> {
61 let client = self.client.clone();
62 cx.spawn(move |this, mut cx| async move {
63 client.authenticate_and_connect(true, &cx).await?;
64 this.update(&mut cx, |_, cx| cx.notify())
65 })
66 }
67}
68
69impl CloudLanguageModelProvider {
70 pub fn new(client: Arc<Client>, cx: &mut AppContext) -> Self {
71 let mut status_rx = client.status();
72 let status = *status_rx.borrow();
73
74 let state = cx.new_model(|cx| State {
75 client: client.clone(),
76 status,
77 _subscription: cx.observe_global::<SettingsStore>(|_, cx| {
78 cx.notify();
79 }),
80 });
81
82 let state_ref = state.downgrade();
83 let maintain_client_status = cx.spawn(|mut cx| async move {
84 while let Some(status) = status_rx.next().await {
85 if let Some(this) = state_ref.upgrade() {
86 _ = this.update(&mut cx, |this, cx| {
87 this.status = status;
88 cx.notify();
89 });
90 } else {
91 break;
92 }
93 }
94 });
95
96 Self {
97 client,
98 state,
99 _maintain_client_status: maintain_client_status,
100 }
101 }
102}
103
104impl LanguageModelProviderState for CloudLanguageModelProvider {
105 type ObservableEntity = State;
106
107 fn observable_entity(&self) -> Option<gpui::Model<Self::ObservableEntity>> {
108 Some(self.state.clone())
109 }
110}
111
112impl LanguageModelProvider for CloudLanguageModelProvider {
113 fn id(&self) -> LanguageModelProviderId {
114 LanguageModelProviderId(PROVIDER_ID.into())
115 }
116
117 fn name(&self) -> LanguageModelProviderName {
118 LanguageModelProviderName(PROVIDER_NAME.into())
119 }
120
121 fn provided_models(&self, cx: &AppContext) -> Vec<Arc<dyn LanguageModel>> {
122 let mut models = BTreeMap::default();
123
124 for model in anthropic::Model::iter() {
125 if !matches!(model, anthropic::Model::Custom { .. }) {
126 models.insert(model.id().to_string(), CloudModel::Anthropic(model));
127 }
128 }
129 for model in open_ai::Model::iter() {
130 if !matches!(model, open_ai::Model::Custom { .. }) {
131 models.insert(model.id().to_string(), CloudModel::OpenAi(model));
132 }
133 }
134 for model in google_ai::Model::iter() {
135 if !matches!(model, google_ai::Model::Custom { .. }) {
136 models.insert(model.id().to_string(), CloudModel::Google(model));
137 }
138 }
139
140 // Override with available models from settings
141 for model in &AllLanguageModelSettings::get_global(cx)
142 .zed_dot_dev
143 .available_models
144 {
145 let model = match model.provider {
146 AvailableProvider::Anthropic => CloudModel::Anthropic(anthropic::Model::Custom {
147 name: model.name.clone(),
148 max_tokens: model.max_tokens,
149 tool_override: model.tool_override.clone(),
150 }),
151 AvailableProvider::OpenAi => CloudModel::OpenAi(open_ai::Model::Custom {
152 name: model.name.clone(),
153 max_tokens: model.max_tokens,
154 }),
155 AvailableProvider::Google => CloudModel::Google(google_ai::Model::Custom {
156 name: model.name.clone(),
157 max_tokens: model.max_tokens,
158 }),
159 };
160 models.insert(model.id().to_string(), model.clone());
161 }
162
163 models
164 .into_values()
165 .map(|model| {
166 Arc::new(CloudLanguageModel {
167 id: LanguageModelId::from(model.id().to_string()),
168 model,
169 client: self.client.clone(),
170 request_limiter: RateLimiter::new(4),
171 }) as Arc<dyn LanguageModel>
172 })
173 .collect()
174 }
175
176 fn is_authenticated(&self, cx: &AppContext) -> bool {
177 self.state.read(cx).status.is_connected()
178 }
179
180 fn authenticate(&self, cx: &mut AppContext) -> Task<Result<()>> {
181 self.state.update(cx, |state, cx| state.authenticate(cx))
182 }
183
184 fn authentication_prompt(&self, cx: &mut WindowContext) -> AnyView {
185 cx.new_view(|_cx| AuthenticationPrompt {
186 state: self.state.clone(),
187 })
188 .into()
189 }
190
191 fn reset_credentials(&self, _cx: &mut AppContext) -> Task<Result<()>> {
192 Task::ready(Ok(()))
193 }
194}
195
196pub struct CloudLanguageModel {
197 id: LanguageModelId,
198 model: CloudModel,
199 client: Arc<Client>,
200 request_limiter: RateLimiter,
201}
202
203impl LanguageModel for CloudLanguageModel {
204 fn id(&self) -> LanguageModelId {
205 self.id.clone()
206 }
207
208 fn name(&self) -> LanguageModelName {
209 LanguageModelName::from(self.model.display_name().to_string())
210 }
211
212 fn provider_id(&self) -> LanguageModelProviderId {
213 LanguageModelProviderId(PROVIDER_ID.into())
214 }
215
216 fn provider_name(&self) -> LanguageModelProviderName {
217 LanguageModelProviderName(PROVIDER_NAME.into())
218 }
219
220 fn telemetry_id(&self) -> String {
221 format!("zed.dev/{}", self.model.id())
222 }
223
224 fn max_token_count(&self) -> usize {
225 self.model.max_token_count()
226 }
227
228 fn count_tokens(
229 &self,
230 request: LanguageModelRequest,
231 cx: &AppContext,
232 ) -> BoxFuture<'static, Result<usize>> {
233 match self.model.clone() {
234 CloudModel::Anthropic(_) => count_anthropic_tokens(request, cx),
235 CloudModel::OpenAi(model) => count_open_ai_tokens(request, model, cx),
236 CloudModel::Google(model) => {
237 let client = self.client.clone();
238 let request = request.into_google(model.id().into());
239 let request = google_ai::CountTokensRequest {
240 contents: request.contents,
241 };
242 async move {
243 let request = serde_json::to_string(&request)?;
244 let response = client
245 .request(proto::CountLanguageModelTokens {
246 provider: proto::LanguageModelProvider::Google as i32,
247 request,
248 })
249 .await?;
250 Ok(response.token_count as usize)
251 }
252 .boxed()
253 }
254 }
255 }
256
257 fn stream_completion(
258 &self,
259 request: LanguageModelRequest,
260 _: &AsyncAppContext,
261 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
262 match &self.model {
263 CloudModel::Anthropic(model) => {
264 let client = self.client.clone();
265 let request = request.into_anthropic(model.id().into());
266 let future = self.request_limiter.stream(async move {
267 let request = serde_json::to_string(&request)?;
268 let stream = client
269 .request_stream(proto::StreamCompleteWithLanguageModel {
270 provider: proto::LanguageModelProvider::Anthropic as i32,
271 request,
272 })
273 .await?;
274 Ok(anthropic::extract_text_from_events(
275 stream.map(|item| Ok(serde_json::from_str(&item?.event)?)),
276 ))
277 });
278 async move { Ok(future.await?.boxed()) }.boxed()
279 }
280 CloudModel::OpenAi(model) => {
281 let client = self.client.clone();
282 let request = request.into_open_ai(model.id().into());
283 let future = self.request_limiter.stream(async move {
284 let request = serde_json::to_string(&request)?;
285 let stream = client
286 .request_stream(proto::StreamCompleteWithLanguageModel {
287 provider: proto::LanguageModelProvider::OpenAi as i32,
288 request,
289 })
290 .await?;
291 Ok(open_ai::extract_text_from_events(
292 stream.map(|item| Ok(serde_json::from_str(&item?.event)?)),
293 ))
294 });
295 async move { Ok(future.await?.boxed()) }.boxed()
296 }
297 CloudModel::Google(model) => {
298 let client = self.client.clone();
299 let request = request.into_google(model.id().into());
300 let future = self.request_limiter.stream(async move {
301 let request = serde_json::to_string(&request)?;
302 let stream = client
303 .request_stream(proto::StreamCompleteWithLanguageModel {
304 provider: proto::LanguageModelProvider::Google as i32,
305 request,
306 })
307 .await?;
308 Ok(google_ai::extract_text_from_events(
309 stream.map(|item| Ok(serde_json::from_str(&item?.event)?)),
310 ))
311 });
312 async move { Ok(future.await?.boxed()) }.boxed()
313 }
314 }
315 }
316
317 fn use_any_tool(
318 &self,
319 request: LanguageModelRequest,
320 tool_name: String,
321 tool_description: String,
322 input_schema: serde_json::Value,
323 _cx: &AsyncAppContext,
324 ) -> BoxFuture<'static, Result<serde_json::Value>> {
325 match &self.model {
326 CloudModel::Anthropic(model) => {
327 let client = self.client.clone();
328 let mut request = request.into_anthropic(model.tool_model_id().into());
329 request.tool_choice = Some(anthropic::ToolChoice::Tool {
330 name: tool_name.clone(),
331 });
332 request.tools = vec![anthropic::Tool {
333 name: tool_name.clone(),
334 description: tool_description,
335 input_schema,
336 }];
337
338 self.request_limiter
339 .run(async move {
340 let request = serde_json::to_string(&request)?;
341 let response = client
342 .request(proto::CompleteWithLanguageModel {
343 provider: proto::LanguageModelProvider::Anthropic as i32,
344 request,
345 })
346 .await?;
347 let response: anthropic::Response =
348 serde_json::from_str(&response.completion)?;
349 response
350 .content
351 .into_iter()
352 .find_map(|content| {
353 if let anthropic::Content::ToolUse { name, input, .. } = content {
354 if name == tool_name {
355 Some(input)
356 } else {
357 None
358 }
359 } else {
360 None
361 }
362 })
363 .context("tool not used")
364 })
365 .boxed()
366 }
367 CloudModel::OpenAi(_) => {
368 future::ready(Err(anyhow!("tool use not implemented for OpenAI"))).boxed()
369 }
370 CloudModel::Google(_) => {
371 future::ready(Err(anyhow!("tool use not implemented for Google AI"))).boxed()
372 }
373 }
374 }
375}
376
377struct AuthenticationPrompt {
378 state: gpui::Model<State>,
379}
380
381impl Render for AuthenticationPrompt {
382 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
383 const LABEL: &str = "Generate and analyze code with language models. You can dialog with the assistant in this panel or transform code inline.";
384
385 v_flex().gap_6().p_4().child(Label::new(LABEL)).child(
386 v_flex()
387 .gap_2()
388 .child(
389 Button::new("sign_in", "Sign in")
390 .icon_color(Color::Muted)
391 .icon(IconName::Github)
392 .icon_position(IconPosition::Start)
393 .style(ButtonStyle::Filled)
394 .full_width()
395 .on_click(cx.listener(move |this, _, cx| {
396 this.state.update(cx, |provider, cx| {
397 provider.authenticate(cx).detach_and_log_err(cx);
398 cx.notify();
399 });
400 })),
401 )
402 .child(
403 div().flex().w_full().items_center().child(
404 Label::new("Sign in to enable collaboration.")
405 .color(Color::Muted)
406 .size(LabelSize::Small),
407 ),
408 ),
409 )
410 }
411}