1use crate::{
2 settings::AllLanguageModelSettings, LanguageModel, LanguageModelId, LanguageModelName,
3 LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
4 LanguageModelProviderState, LanguageModelRequest, RateLimiter, Role,
5};
6use anyhow::{anyhow, Context as _, Result};
7use collections::BTreeMap;
8use editor::{Editor, EditorElement, EditorStyle};
9use futures::{future::BoxFuture, stream::BoxStream, FutureExt, StreamExt};
10use gpui::{
11 AnyView, AppContext, AsyncAppContext, FontStyle, Subscription, Task, TextStyle, View,
12 WhiteSpace,
13};
14use http_client::HttpClient;
15use schemars::JsonSchema;
16use serde::{Deserialize, Serialize};
17use settings::{Settings, SettingsStore};
18use std::{sync::Arc, time::Duration};
19use strum::IntoEnumIterator;
20use theme::ThemeSettings;
21use ui::prelude::*;
22use util::ResultExt;
23
24const PROVIDER_ID: &str = "anthropic";
25const PROVIDER_NAME: &str = "Anthropic";
26
27#[derive(Default, Clone, Debug, PartialEq)]
28pub struct AnthropicSettings {
29 pub api_url: String,
30 pub low_speed_timeout: Option<Duration>,
31 pub available_models: Vec<AvailableModel>,
32 pub needs_setting_migration: bool,
33}
34
35#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
36pub struct AvailableModel {
37 pub name: String,
38 pub max_tokens: usize,
39 pub tool_override: Option<String>,
40}
41
42pub struct AnthropicLanguageModelProvider {
43 http_client: Arc<dyn HttpClient>,
44 state: gpui::Model<State>,
45}
46
47pub struct State {
48 api_key: Option<String>,
49 _subscription: Subscription,
50}
51
52impl AnthropicLanguageModelProvider {
53 pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut AppContext) -> Self {
54 let state = cx.new_model(|cx| State {
55 api_key: None,
56 _subscription: cx.observe_global::<SettingsStore>(|_, cx| {
57 cx.notify();
58 }),
59 });
60
61 Self { http_client, state }
62 }
63}
64
65impl LanguageModelProviderState for AnthropicLanguageModelProvider {
66 type ObservableEntity = State;
67
68 fn observable_entity(&self) -> Option<gpui::Model<Self::ObservableEntity>> {
69 Some(self.state.clone())
70 }
71}
72
73impl LanguageModelProvider for AnthropicLanguageModelProvider {
74 fn id(&self) -> LanguageModelProviderId {
75 LanguageModelProviderId(PROVIDER_ID.into())
76 }
77
78 fn name(&self) -> LanguageModelProviderName {
79 LanguageModelProviderName(PROVIDER_NAME.into())
80 }
81
82 fn provided_models(&self, cx: &AppContext) -> Vec<Arc<dyn LanguageModel>> {
83 let mut models = BTreeMap::default();
84
85 // Add base models from anthropic::Model::iter()
86 for model in anthropic::Model::iter() {
87 if !matches!(model, anthropic::Model::Custom { .. }) {
88 models.insert(model.id().to_string(), model);
89 }
90 }
91
92 // Override with available models from settings
93 for model in AllLanguageModelSettings::get_global(cx)
94 .anthropic
95 .available_models
96 .iter()
97 {
98 models.insert(
99 model.name.clone(),
100 anthropic::Model::Custom {
101 name: model.name.clone(),
102 max_tokens: model.max_tokens,
103 tool_override: model.tool_override.clone(),
104 },
105 );
106 }
107
108 models
109 .into_values()
110 .map(|model| {
111 Arc::new(AnthropicModel {
112 id: LanguageModelId::from(model.id().to_string()),
113 model,
114 state: self.state.clone(),
115 http_client: self.http_client.clone(),
116 request_limiter: RateLimiter::new(4),
117 }) as Arc<dyn LanguageModel>
118 })
119 .collect()
120 }
121
122 fn is_authenticated(&self, cx: &AppContext) -> bool {
123 self.state.read(cx).api_key.is_some()
124 }
125
126 fn authenticate(&self, cx: &mut AppContext) -> Task<Result<()>> {
127 if self.is_authenticated(cx) {
128 Task::ready(Ok(()))
129 } else {
130 let api_url = AllLanguageModelSettings::get_global(cx)
131 .anthropic
132 .api_url
133 .clone();
134 let state = self.state.clone();
135 cx.spawn(|mut cx| async move {
136 let api_key = if let Ok(api_key) = std::env::var("ANTHROPIC_API_KEY") {
137 api_key
138 } else {
139 let (_, api_key) = cx
140 .update(|cx| cx.read_credentials(&api_url))?
141 .await?
142 .ok_or_else(|| anyhow!("credentials not found"))?;
143 String::from_utf8(api_key)?
144 };
145
146 state.update(&mut cx, |this, cx| {
147 this.api_key = Some(api_key);
148 cx.notify();
149 })
150 })
151 }
152 }
153
154 fn authentication_prompt(&self, cx: &mut WindowContext) -> AnyView {
155 cx.new_view(|cx| AuthenticationPrompt::new(self.state.clone(), cx))
156 .into()
157 }
158
159 fn reset_credentials(&self, cx: &mut AppContext) -> Task<Result<()>> {
160 let state = self.state.clone();
161 let delete_credentials =
162 cx.delete_credentials(&AllLanguageModelSettings::get_global(cx).anthropic.api_url);
163 cx.spawn(|mut cx| async move {
164 delete_credentials.await.log_err();
165 state.update(&mut cx, |this, cx| {
166 this.api_key = None;
167 cx.notify();
168 })
169 })
170 }
171}
172
173pub struct AnthropicModel {
174 id: LanguageModelId,
175 model: anthropic::Model,
176 state: gpui::Model<State>,
177 http_client: Arc<dyn HttpClient>,
178 request_limiter: RateLimiter,
179}
180
181pub fn count_anthropic_tokens(
182 request: LanguageModelRequest,
183 cx: &AppContext,
184) -> BoxFuture<'static, Result<usize>> {
185 cx.background_executor()
186 .spawn(async move {
187 let messages = request
188 .messages
189 .into_iter()
190 .map(|message| tiktoken_rs::ChatCompletionRequestMessage {
191 role: match message.role {
192 Role::User => "user".into(),
193 Role::Assistant => "assistant".into(),
194 Role::System => "system".into(),
195 },
196 content: Some(message.content),
197 name: None,
198 function_call: None,
199 })
200 .collect::<Vec<_>>();
201
202 // Tiktoken doesn't yet support these models, so we manually use the
203 // same tokenizer as GPT-4.
204 tiktoken_rs::num_tokens_from_messages("gpt-4", &messages)
205 })
206 .boxed()
207}
208
209impl AnthropicModel {
210 fn request_completion(
211 &self,
212 request: anthropic::Request,
213 cx: &AsyncAppContext,
214 ) -> BoxFuture<'static, Result<anthropic::Response>> {
215 let http_client = self.http_client.clone();
216
217 let Ok((api_key, api_url)) = cx.read_model(&self.state, |state, cx| {
218 let settings = &AllLanguageModelSettings::get_global(cx).anthropic;
219 (state.api_key.clone(), settings.api_url.clone())
220 }) else {
221 return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
222 };
223
224 async move {
225 let api_key = api_key.ok_or_else(|| anyhow!("missing api key"))?;
226 anthropic::complete(http_client.as_ref(), &api_url, &api_key, request).await
227 }
228 .boxed()
229 }
230
231 fn stream_completion(
232 &self,
233 request: anthropic::Request,
234 cx: &AsyncAppContext,
235 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<anthropic::Event>>>> {
236 let http_client = self.http_client.clone();
237
238 let Ok((api_key, api_url, low_speed_timeout)) = cx.read_model(&self.state, |state, cx| {
239 let settings = &AllLanguageModelSettings::get_global(cx).anthropic;
240 (
241 state.api_key.clone(),
242 settings.api_url.clone(),
243 settings.low_speed_timeout,
244 )
245 }) else {
246 return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
247 };
248
249 async move {
250 let api_key = api_key.ok_or_else(|| anyhow!("missing api key"))?;
251 let request = anthropic::stream_completion(
252 http_client.as_ref(),
253 &api_url,
254 &api_key,
255 request,
256 low_speed_timeout,
257 );
258 request.await
259 }
260 .boxed()
261 }
262}
263
264impl LanguageModel for AnthropicModel {
265 fn id(&self) -> LanguageModelId {
266 self.id.clone()
267 }
268
269 fn name(&self) -> LanguageModelName {
270 LanguageModelName::from(self.model.display_name().to_string())
271 }
272
273 fn provider_id(&self) -> LanguageModelProviderId {
274 LanguageModelProviderId(PROVIDER_ID.into())
275 }
276
277 fn provider_name(&self) -> LanguageModelProviderName {
278 LanguageModelProviderName(PROVIDER_NAME.into())
279 }
280
281 fn telemetry_id(&self) -> String {
282 format!("anthropic/{}", self.model.id())
283 }
284
285 fn max_token_count(&self) -> usize {
286 self.model.max_token_count()
287 }
288
289 fn count_tokens(
290 &self,
291 request: LanguageModelRequest,
292 cx: &AppContext,
293 ) -> BoxFuture<'static, Result<usize>> {
294 count_anthropic_tokens(request, cx)
295 }
296
297 fn stream_completion(
298 &self,
299 request: LanguageModelRequest,
300 cx: &AsyncAppContext,
301 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
302 let request = request.into_anthropic(self.model.id().into());
303 let request = self.stream_completion(request, cx);
304 let future = self.request_limiter.stream(async move {
305 let response = request.await?;
306 Ok(anthropic::extract_text_from_events(response))
307 });
308 async move { Ok(future.await?.boxed()) }.boxed()
309 }
310
311 fn use_any_tool(
312 &self,
313 request: LanguageModelRequest,
314 tool_name: String,
315 tool_description: String,
316 input_schema: serde_json::Value,
317 cx: &AsyncAppContext,
318 ) -> BoxFuture<'static, Result<serde_json::Value>> {
319 let mut request = request.into_anthropic(self.model.tool_model_id().into());
320 request.tool_choice = Some(anthropic::ToolChoice::Tool {
321 name: tool_name.clone(),
322 });
323 request.tools = vec![anthropic::Tool {
324 name: tool_name.clone(),
325 description: tool_description,
326 input_schema,
327 }];
328
329 let response = self.request_completion(request, cx);
330 self.request_limiter
331 .run(async move {
332 let response = response.await?;
333 response
334 .content
335 .into_iter()
336 .find_map(|content| {
337 if let anthropic::Content::ToolUse { name, input, .. } = content {
338 if name == tool_name {
339 Some(input)
340 } else {
341 None
342 }
343 } else {
344 None
345 }
346 })
347 .context("tool not used")
348 })
349 .boxed()
350 }
351}
352
353struct AuthenticationPrompt {
354 api_key: View<Editor>,
355 state: gpui::Model<State>,
356}
357
358impl AuthenticationPrompt {
359 fn new(state: gpui::Model<State>, cx: &mut WindowContext) -> Self {
360 Self {
361 api_key: cx.new_view(|cx| {
362 let mut editor = Editor::single_line(cx);
363 editor.set_placeholder_text(
364 "sk-000000000000000000000000000000000000000000000000",
365 cx,
366 );
367 editor
368 }),
369 state,
370 }
371 }
372
373 fn save_api_key(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
374 let api_key = self.api_key.read(cx).text(cx);
375 if api_key.is_empty() {
376 return;
377 }
378
379 let write_credentials = cx.write_credentials(
380 AllLanguageModelSettings::get_global(cx)
381 .anthropic
382 .api_url
383 .as_str(),
384 "Bearer",
385 api_key.as_bytes(),
386 );
387 let state = self.state.clone();
388 cx.spawn(|_, mut cx| async move {
389 write_credentials.await?;
390
391 state.update(&mut cx, |this, cx| {
392 this.api_key = Some(api_key);
393 cx.notify();
394 })
395 })
396 .detach_and_log_err(cx);
397 }
398
399 fn render_api_key_editor(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
400 let settings = ThemeSettings::get_global(cx);
401 let text_style = TextStyle {
402 color: cx.theme().colors().text,
403 font_family: settings.ui_font.family.clone(),
404 font_features: settings.ui_font.features.clone(),
405 font_fallbacks: settings.ui_font.fallbacks.clone(),
406 font_size: rems(0.875).into(),
407 font_weight: settings.ui_font.weight,
408 font_style: FontStyle::Normal,
409 line_height: relative(1.3),
410 background_color: None,
411 underline: None,
412 strikethrough: None,
413 white_space: WhiteSpace::Normal,
414 };
415 EditorElement::new(
416 &self.api_key,
417 EditorStyle {
418 background: cx.theme().colors().editor_background,
419 local_player: cx.theme().players().local(),
420 text: text_style,
421 ..Default::default()
422 },
423 )
424 }
425}
426
427impl Render for AuthenticationPrompt {
428 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
429 const INSTRUCTIONS: [&str; 4] = [
430 "To use the assistant panel or inline assistant, you need to add your Anthropic API key.",
431 "You can create an API key at: https://console.anthropic.com/settings/keys",
432 "",
433 "Paste your Anthropic API key below and hit enter to use the assistant:",
434 ];
435
436 v_flex()
437 .p_4()
438 .size_full()
439 .on_action(cx.listener(Self::save_api_key))
440 .children(
441 INSTRUCTIONS.map(|instruction| Label::new(instruction).size(LabelSize::Small)),
442 )
443 .child(
444 h_flex()
445 .w_full()
446 .my_2()
447 .px_2()
448 .py_1()
449 .bg(cx.theme().colors().editor_background)
450 .rounded_md()
451 .child(self.render_api_key_editor(cx)),
452 )
453 .child(
454 Label::new(
455 "You can also assign the ANTHROPIC_API_KEY environment variable and restart Zed.",
456 )
457 .size(LabelSize::Small),
458 )
459 .child(
460 h_flex()
461 .gap_2()
462 .child(Label::new("Click on").size(LabelSize::Small))
463 .child(Icon::new(IconName::ZedAssistant).size(IconSize::XSmall))
464 .child(
465 Label::new("in the status bar to close this panel.").size(LabelSize::Small),
466 ),
467 )
468 .into_any()
469 }
470}