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