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