1use crate::AllLanguageModelSettings;
2use anthropic::{AnthropicError, ContentDelta, Event, ResponseContent};
3use anyhow::{anyhow, Context as _, Result};
4use collections::{BTreeMap, HashMap};
5use editor::{Editor, EditorElement, EditorStyle};
6use futures::Stream;
7use futures::{future::BoxFuture, stream::BoxStream, FutureExt, StreamExt, TryStreamExt as _};
8use gpui::{
9 AnyView, App, AsyncApp, Context, Entity, FontStyle, Subscription, Task, TextStyle, WhiteSpace,
10};
11use http_client::HttpClient;
12use language_model::{
13 LanguageModel, LanguageModelCacheConfiguration, LanguageModelId, LanguageModelName,
14 LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
15 LanguageModelProviderState, LanguageModelRequest, RateLimiter, Role,
16};
17use language_model::{LanguageModelCompletionEvent, LanguageModelToolUse, StopReason};
18use schemars::JsonSchema;
19use serde::{Deserialize, Serialize};
20use settings::{Settings, SettingsStore};
21use std::pin::Pin;
22use std::str::FromStr;
23use std::sync::Arc;
24use strum::IntoEnumIterator;
25use theme::ThemeSettings;
26use ui::{prelude::*, Icon, IconName, Tooltip};
27use util::{maybe, ResultExt};
28
29pub const PROVIDER_ID: &str = "anthropic";
30const PROVIDER_NAME: &str = "Anthropic";
31
32#[derive(Default, Clone, Debug, PartialEq)]
33pub struct AnthropicSettings {
34 pub api_url: String,
35 /// Extend Zed's list of Anthropic models.
36 pub available_models: Vec<AvailableModel>,
37 pub needs_setting_migration: bool,
38}
39
40#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
41pub struct AvailableModel {
42 /// The model's name in the Anthropic API. e.g. claude-3-5-sonnet-latest, claude-3-opus-20240229, etc
43 pub name: String,
44 /// The model's name in Zed's UI, such as in the model selector dropdown menu in the assistant panel.
45 pub display_name: Option<String>,
46 /// The model's context window size.
47 pub max_tokens: usize,
48 /// A model `name` to substitute when calling tools, in case the primary model doesn't support tool calling.
49 pub tool_override: Option<String>,
50 /// Configuration of Anthropic's caching API.
51 pub cache_configuration: Option<LanguageModelCacheConfiguration>,
52 pub max_output_tokens: Option<u32>,
53 pub default_temperature: Option<f32>,
54 #[serde(default)]
55 pub extra_beta_headers: Vec<String>,
56}
57
58pub struct AnthropicLanguageModelProvider {
59 http_client: Arc<dyn HttpClient>,
60 state: gpui::Entity<State>,
61}
62
63const ANTHROPIC_API_KEY_VAR: &str = "ANTHROPIC_API_KEY";
64
65pub struct State {
66 api_key: Option<String>,
67 api_key_from_env: bool,
68 _subscription: Subscription,
69}
70
71impl State {
72 fn reset_api_key(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
73 let delete_credentials =
74 cx.delete_credentials(&AllLanguageModelSettings::get_global(cx).anthropic.api_url);
75 cx.spawn(|this, mut cx| async move {
76 delete_credentials.await.ok();
77 this.update(&mut cx, |this, cx| {
78 this.api_key = None;
79 this.api_key_from_env = false;
80 cx.notify();
81 })
82 })
83 }
84
85 fn set_api_key(&mut self, api_key: String, cx: &mut Context<Self>) -> Task<Result<()>> {
86 let write_credentials = cx.write_credentials(
87 AllLanguageModelSettings::get_global(cx)
88 .anthropic
89 .api_url
90 .as_str(),
91 "Bearer",
92 api_key.as_bytes(),
93 );
94 cx.spawn(|this, mut cx| async move {
95 write_credentials.await?;
96
97 this.update(&mut cx, |this, cx| {
98 this.api_key = Some(api_key);
99 cx.notify();
100 })
101 })
102 }
103
104 fn is_authenticated(&self) -> bool {
105 self.api_key.is_some()
106 }
107
108 fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
109 if self.is_authenticated() {
110 Task::ready(Ok(()))
111 } else {
112 let api_url = AllLanguageModelSettings::get_global(cx)
113 .anthropic
114 .api_url
115 .clone();
116
117 cx.spawn(|this, mut cx| async move {
118 let (api_key, from_env) = if let Ok(api_key) = std::env::var(ANTHROPIC_API_KEY_VAR)
119 {
120 (api_key, true)
121 } else {
122 let (_, api_key) = cx
123 .update(|cx| cx.read_credentials(&api_url))?
124 .await?
125 .ok_or_else(|| anyhow!("credentials not found"))?;
126 (String::from_utf8(api_key)?, false)
127 };
128
129 this.update(&mut cx, |this, cx| {
130 this.api_key = Some(api_key);
131 this.api_key_from_env = from_env;
132 cx.notify();
133 })
134 })
135 }
136 }
137}
138
139impl AnthropicLanguageModelProvider {
140 pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self {
141 let state = cx.new(|cx| State {
142 api_key: None,
143 api_key_from_env: false,
144 _subscription: cx.observe_global::<SettingsStore>(|_, cx| {
145 cx.notify();
146 }),
147 });
148
149 Self { http_client, state }
150 }
151}
152
153impl LanguageModelProviderState for AnthropicLanguageModelProvider {
154 type ObservableEntity = State;
155
156 fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> {
157 Some(self.state.clone())
158 }
159}
160
161impl LanguageModelProvider for AnthropicLanguageModelProvider {
162 fn id(&self) -> LanguageModelProviderId {
163 LanguageModelProviderId(PROVIDER_ID.into())
164 }
165
166 fn name(&self) -> LanguageModelProviderName {
167 LanguageModelProviderName(PROVIDER_NAME.into())
168 }
169
170 fn icon(&self) -> IconName {
171 IconName::AiAnthropic
172 }
173
174 fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
175 let mut models = BTreeMap::default();
176
177 // Add base models from anthropic::Model::iter()
178 for model in anthropic::Model::iter() {
179 if !matches!(model, anthropic::Model::Custom { .. }) {
180 models.insert(model.id().to_string(), model);
181 }
182 }
183
184 // Override with available models from settings
185 for model in AllLanguageModelSettings::get_global(cx)
186 .anthropic
187 .available_models
188 .iter()
189 {
190 models.insert(
191 model.name.clone(),
192 anthropic::Model::Custom {
193 name: model.name.clone(),
194 display_name: model.display_name.clone(),
195 max_tokens: model.max_tokens,
196 tool_override: model.tool_override.clone(),
197 cache_configuration: model.cache_configuration.as_ref().map(|config| {
198 anthropic::AnthropicModelCacheConfiguration {
199 max_cache_anchors: config.max_cache_anchors,
200 should_speculate: config.should_speculate,
201 min_total_token: config.min_total_token,
202 }
203 }),
204 max_output_tokens: model.max_output_tokens,
205 default_temperature: model.default_temperature,
206 extra_beta_headers: model.extra_beta_headers.clone(),
207 },
208 );
209 }
210
211 models
212 .into_values()
213 .map(|model| {
214 Arc::new(AnthropicModel {
215 id: LanguageModelId::from(model.id().to_string()),
216 model,
217 state: self.state.clone(),
218 http_client: self.http_client.clone(),
219 request_limiter: RateLimiter::new(4),
220 }) as Arc<dyn LanguageModel>
221 })
222 .collect()
223 }
224
225 fn is_authenticated(&self, cx: &App) -> bool {
226 self.state.read(cx).is_authenticated()
227 }
228
229 fn authenticate(&self, cx: &mut App) -> Task<Result<()>> {
230 self.state.update(cx, |state, cx| state.authenticate(cx))
231 }
232
233 fn configuration_view(&self, window: &mut Window, cx: &mut App) -> AnyView {
234 cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx))
235 .into()
236 }
237
238 fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> {
239 self.state.update(cx, |state, cx| state.reset_api_key(cx))
240 }
241}
242
243pub struct AnthropicModel {
244 id: LanguageModelId,
245 model: anthropic::Model,
246 state: gpui::Entity<State>,
247 http_client: Arc<dyn HttpClient>,
248 request_limiter: RateLimiter,
249}
250
251pub fn count_anthropic_tokens(
252 request: LanguageModelRequest,
253 cx: &App,
254) -> BoxFuture<'static, Result<usize>> {
255 cx.background_executor()
256 .spawn(async move {
257 let messages = request.messages;
258 let mut tokens_from_images = 0;
259 let mut string_messages = Vec::with_capacity(messages.len());
260
261 for message in messages {
262 use language_model::MessageContent;
263
264 let mut string_contents = String::new();
265
266 for content in message.content {
267 match content {
268 MessageContent::Text(text) => {
269 string_contents.push_str(&text);
270 }
271 MessageContent::Image(image) => {
272 tokens_from_images += image.estimate_tokens();
273 }
274 MessageContent::ToolUse(_tool_use) => {
275 // TODO: Estimate token usage from tool uses.
276 }
277 MessageContent::ToolResult(tool_result) => {
278 string_contents.push_str(&tool_result.content);
279 }
280 }
281 }
282
283 if !string_contents.is_empty() {
284 string_messages.push(tiktoken_rs::ChatCompletionRequestMessage {
285 role: match message.role {
286 Role::User => "user".into(),
287 Role::Assistant => "assistant".into(),
288 Role::System => "system".into(),
289 },
290 content: Some(string_contents),
291 name: None,
292 function_call: None,
293 });
294 }
295 }
296
297 // Tiktoken doesn't yet support these models, so we manually use the
298 // same tokenizer as GPT-4.
299 tiktoken_rs::num_tokens_from_messages("gpt-4", &string_messages)
300 .map(|tokens| tokens + tokens_from_images)
301 })
302 .boxed()
303}
304
305impl AnthropicModel {
306 fn stream_completion(
307 &self,
308 request: anthropic::Request,
309 cx: &AsyncApp,
310 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<anthropic::Event, AnthropicError>>>>
311 {
312 let http_client = self.http_client.clone();
313
314 let Ok((api_key, api_url)) = cx.read_entity(&self.state, |state, cx| {
315 let settings = &AllLanguageModelSettings::get_global(cx).anthropic;
316 (state.api_key.clone(), settings.api_url.clone())
317 }) else {
318 return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
319 };
320
321 async move {
322 let api_key = api_key.ok_or_else(|| anyhow!("Missing Anthropic API Key"))?;
323 let request =
324 anthropic::stream_completion(http_client.as_ref(), &api_url, &api_key, request);
325 request.await.context("failed to stream completion")
326 }
327 .boxed()
328 }
329}
330
331impl LanguageModel for AnthropicModel {
332 fn id(&self) -> LanguageModelId {
333 self.id.clone()
334 }
335
336 fn name(&self) -> LanguageModelName {
337 LanguageModelName::from(self.model.display_name().to_string())
338 }
339
340 fn provider_id(&self) -> LanguageModelProviderId {
341 LanguageModelProviderId(PROVIDER_ID.into())
342 }
343
344 fn provider_name(&self) -> LanguageModelProviderName {
345 LanguageModelProviderName(PROVIDER_NAME.into())
346 }
347
348 fn telemetry_id(&self) -> String {
349 format!("anthropic/{}", self.model.id())
350 }
351
352 fn api_key(&self, cx: &App) -> Option<String> {
353 self.state.read(cx).api_key.clone()
354 }
355
356 fn max_token_count(&self) -> usize {
357 self.model.max_token_count()
358 }
359
360 fn max_output_tokens(&self) -> Option<u32> {
361 Some(self.model.max_output_tokens())
362 }
363
364 fn count_tokens(
365 &self,
366 request: LanguageModelRequest,
367 cx: &App,
368 ) -> BoxFuture<'static, Result<usize>> {
369 count_anthropic_tokens(request, cx)
370 }
371
372 fn stream_completion(
373 &self,
374 request: LanguageModelRequest,
375 cx: &AsyncApp,
376 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<LanguageModelCompletionEvent>>>> {
377 let request = request.into_anthropic(
378 self.model.id().into(),
379 self.model.default_temperature(),
380 self.model.max_output_tokens(),
381 );
382 let request = self.stream_completion(request, cx);
383 let future = self.request_limiter.stream(async move {
384 let response = request.await.map_err(|err| anyhow!(err))?;
385 Ok(map_to_language_model_completion_events(response))
386 });
387 async move { Ok(future.await?.boxed()) }.boxed()
388 }
389
390 fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
391 self.model
392 .cache_configuration()
393 .map(|config| LanguageModelCacheConfiguration {
394 max_cache_anchors: config.max_cache_anchors,
395 should_speculate: config.should_speculate,
396 min_total_token: config.min_total_token,
397 })
398 }
399
400 fn use_any_tool(
401 &self,
402 request: LanguageModelRequest,
403 tool_name: String,
404 tool_description: String,
405 input_schema: serde_json::Value,
406 cx: &AsyncApp,
407 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
408 let mut request = request.into_anthropic(
409 self.model.tool_model_id().into(),
410 self.model.default_temperature(),
411 self.model.max_output_tokens(),
412 );
413 request.tool_choice = Some(anthropic::ToolChoice::Tool {
414 name: tool_name.clone(),
415 });
416 request.tools = vec![anthropic::Tool {
417 name: tool_name.clone(),
418 description: tool_description,
419 input_schema,
420 }];
421
422 let response = self.stream_completion(request, cx);
423 self.request_limiter
424 .run(async move {
425 let response = response.await?;
426 Ok(anthropic::extract_tool_args_from_events(
427 tool_name,
428 Box::pin(response.map_err(|e| anyhow!(e))),
429 )
430 .await?
431 .boxed())
432 })
433 .boxed()
434 }
435}
436
437pub fn map_to_language_model_completion_events(
438 events: Pin<Box<dyn Send + Stream<Item = Result<Event, AnthropicError>>>>,
439) -> impl Stream<Item = Result<LanguageModelCompletionEvent>> {
440 struct RawToolUse {
441 id: String,
442 name: String,
443 input_json: String,
444 }
445
446 struct State {
447 events: Pin<Box<dyn Send + Stream<Item = Result<Event, AnthropicError>>>>,
448 tool_uses_by_index: HashMap<usize, RawToolUse>,
449 }
450
451 futures::stream::unfold(
452 State {
453 events,
454 tool_uses_by_index: HashMap::default(),
455 },
456 |mut state| async move {
457 while let Some(event) = state.events.next().await {
458 match event {
459 Ok(event) => match event {
460 Event::ContentBlockStart {
461 index,
462 content_block,
463 } => match content_block {
464 ResponseContent::Text { text } => {
465 return Some((
466 Some(Ok(LanguageModelCompletionEvent::Text(text))),
467 state,
468 ));
469 }
470 ResponseContent::ToolUse { id, name, .. } => {
471 state.tool_uses_by_index.insert(
472 index,
473 RawToolUse {
474 id,
475 name,
476 input_json: String::new(),
477 },
478 );
479
480 return Some((None, state));
481 }
482 },
483 Event::ContentBlockDelta { index, delta } => match delta {
484 ContentDelta::TextDelta { text } => {
485 return Some((
486 Some(Ok(LanguageModelCompletionEvent::Text(text))),
487 state,
488 ));
489 }
490 ContentDelta::InputJsonDelta { partial_json } => {
491 if let Some(tool_use) = state.tool_uses_by_index.get_mut(&index) {
492 tool_use.input_json.push_str(&partial_json);
493 return Some((None, state));
494 }
495 }
496 },
497 Event::ContentBlockStop { index } => {
498 if let Some(tool_use) = state.tool_uses_by_index.remove(&index) {
499 return Some((
500 Some(maybe!({
501 Ok(LanguageModelCompletionEvent::ToolUse(
502 LanguageModelToolUse {
503 id: tool_use.id.into(),
504 name: tool_use.name,
505 input: if tool_use.input_json.is_empty() {
506 serde_json::Value::Null
507 } else {
508 serde_json::Value::from_str(
509 &tool_use.input_json,
510 )
511 .map_err(|err| anyhow!(err))?
512 },
513 },
514 ))
515 })),
516 state,
517 ));
518 }
519 }
520 Event::MessageStart { message } => {
521 return Some((
522 Some(Ok(LanguageModelCompletionEvent::StartMessage {
523 message_id: message.id,
524 })),
525 state,
526 ))
527 }
528 Event::MessageDelta { delta, .. } => {
529 if let Some(stop_reason) = delta.stop_reason.as_deref() {
530 let stop_reason = match stop_reason {
531 "end_turn" => StopReason::EndTurn,
532 "max_tokens" => StopReason::MaxTokens,
533 "tool_use" => StopReason::ToolUse,
534 _ => StopReason::EndTurn,
535 };
536
537 return Some((
538 Some(Ok(LanguageModelCompletionEvent::Stop(stop_reason))),
539 state,
540 ));
541 }
542 }
543 Event::Error { error } => {
544 return Some((
545 Some(Err(anyhow!(AnthropicError::ApiError(error)))),
546 state,
547 ));
548 }
549 _ => {}
550 },
551 Err(err) => {
552 return Some((Some(Err(anyhow!(err))), state));
553 }
554 }
555 }
556
557 None
558 },
559 )
560 .filter_map(|event| async move { event })
561}
562
563struct ConfigurationView {
564 api_key_editor: Entity<Editor>,
565 state: gpui::Entity<State>,
566 load_credentials_task: Option<Task<()>>,
567}
568
569impl ConfigurationView {
570 const PLACEHOLDER_TEXT: &'static str = "sk-ant-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx";
571
572 fn new(state: gpui::Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
573 cx.observe(&state, |_, _, cx| {
574 cx.notify();
575 })
576 .detach();
577
578 let load_credentials_task = Some(cx.spawn({
579 let state = state.clone();
580 |this, mut cx| async move {
581 if let Some(task) = state
582 .update(&mut cx, |state, cx| state.authenticate(cx))
583 .log_err()
584 {
585 // We don't log an error, because "not signed in" is also an error.
586 let _ = task.await;
587 }
588 this.update(&mut cx, |this, cx| {
589 this.load_credentials_task = None;
590 cx.notify();
591 })
592 .log_err();
593 }
594 }));
595
596 Self {
597 api_key_editor: cx.new(|cx| {
598 let mut editor = Editor::single_line(window, cx);
599 editor.set_placeholder_text(Self::PLACEHOLDER_TEXT, cx);
600 editor
601 }),
602 state,
603 load_credentials_task,
604 }
605 }
606
607 fn save_api_key(&mut self, _: &menu::Confirm, window: &mut Window, cx: &mut Context<Self>) {
608 let api_key = self.api_key_editor.read(cx).text(cx);
609 if api_key.is_empty() {
610 return;
611 }
612
613 let state = self.state.clone();
614 cx.spawn_in(window, |_, mut cx| async move {
615 state
616 .update(&mut cx, |state, cx| state.set_api_key(api_key, cx))?
617 .await
618 })
619 .detach_and_log_err(cx);
620
621 cx.notify();
622 }
623
624 fn reset_api_key(&mut self, window: &mut Window, cx: &mut Context<Self>) {
625 self.api_key_editor
626 .update(cx, |editor, cx| editor.set_text("", window, cx));
627
628 let state = self.state.clone();
629 cx.spawn_in(window, |_, mut cx| async move {
630 state
631 .update(&mut cx, |state, cx| state.reset_api_key(cx))?
632 .await
633 })
634 .detach_and_log_err(cx);
635
636 cx.notify();
637 }
638
639 fn render_api_key_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
640 let settings = ThemeSettings::get_global(cx);
641 let text_style = TextStyle {
642 color: cx.theme().colors().text,
643 font_family: settings.ui_font.family.clone(),
644 font_features: settings.ui_font.features.clone(),
645 font_fallbacks: settings.ui_font.fallbacks.clone(),
646 font_size: rems(0.875).into(),
647 font_weight: settings.ui_font.weight,
648 font_style: FontStyle::Normal,
649 line_height: relative(1.3),
650 white_space: WhiteSpace::Normal,
651 ..Default::default()
652 };
653 EditorElement::new(
654 &self.api_key_editor,
655 EditorStyle {
656 background: cx.theme().colors().editor_background,
657 local_player: cx.theme().players().local(),
658 text: text_style,
659 ..Default::default()
660 },
661 )
662 }
663
664 fn should_render_editor(&self, cx: &mut Context<Self>) -> bool {
665 !self.state.read(cx).is_authenticated()
666 }
667}
668
669impl Render for ConfigurationView {
670 fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
671 const ANTHROPIC_CONSOLE_URL: &str = "https://console.anthropic.com/settings/keys";
672 const INSTRUCTIONS: [&str; 3] = [
673 "To use Zed's assistant with Anthropic, you need to add an API key. Follow these steps:",
674 "- Create one at:",
675 "- Paste your API key below and hit enter to use the assistant:",
676 ];
677 let env_var_set = self.state.read(cx).api_key_from_env;
678
679 if self.load_credentials_task.is_some() {
680 div().child(Label::new("Loading credentials...")).into_any()
681 } else if self.should_render_editor(cx) {
682 v_flex()
683 .size_full()
684 .on_action(cx.listener(Self::save_api_key))
685 .child(Label::new(INSTRUCTIONS[0]))
686 .child(h_flex().child(Label::new(INSTRUCTIONS[1])).child(
687 Button::new("anthropic_console", ANTHROPIC_CONSOLE_URL)
688 .style(ButtonStyle::Subtle)
689 .icon(IconName::ArrowUpRight)
690 .icon_size(IconSize::XSmall)
691 .icon_color(Color::Muted)
692 .on_click(move |_, _, cx| cx.open_url(ANTHROPIC_CONSOLE_URL))
693 )
694 )
695 .child(Label::new(INSTRUCTIONS[2]))
696 .child(
697 h_flex()
698 .w_full()
699 .my_2()
700 .px_2()
701 .py_1()
702 .bg(cx.theme().colors().editor_background)
703 .border_1()
704 .border_color(cx.theme().colors().border_variant)
705 .rounded_md()
706 .child(self.render_api_key_editor(cx)),
707 )
708 .child(
709 Label::new(
710 format!("You can also assign the {ANTHROPIC_API_KEY_VAR} environment variable and restart Zed."),
711 )
712 .size(LabelSize::Small),
713 )
714 .into_any()
715 } else {
716 h_flex()
717 .size_full()
718 .justify_between()
719 .child(
720 h_flex()
721 .gap_1()
722 .child(Icon::new(IconName::Check).color(Color::Success))
723 .child(Label::new(if env_var_set {
724 format!("API key set in {ANTHROPIC_API_KEY_VAR} environment variable.")
725 } else {
726 "API key configured.".to_string()
727 })),
728 )
729 .child(
730 Button::new("reset-key", "Reset key")
731 .icon(Some(IconName::Trash))
732 .icon_size(IconSize::Small)
733 .icon_position(IconPosition::Start)
734 .disabled(env_var_set)
735 .when(env_var_set, |this| {
736 this.tooltip(Tooltip::text(format!("To reset your API key, unset the {ANTHROPIC_API_KEY_VAR} environment variable.")))
737 })
738 .on_click(cx.listener(|this, _, window, cx| this.reset_api_key(window, cx))),
739 )
740 .into_any()
741 }
742 }
743}