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_spawn(async move {
256 let messages = request.messages;
257 let mut tokens_from_images = 0;
258 let mut string_messages = Vec::with_capacity(messages.len());
259
260 for message in messages {
261 use language_model::MessageContent;
262
263 let mut string_contents = String::new();
264
265 for content in message.content {
266 match content {
267 MessageContent::Text(text) => {
268 string_contents.push_str(&text);
269 }
270 MessageContent::Image(image) => {
271 tokens_from_images += image.estimate_tokens();
272 }
273 MessageContent::ToolUse(_tool_use) => {
274 // TODO: Estimate token usage from tool uses.
275 }
276 MessageContent::ToolResult(tool_result) => {
277 string_contents.push_str(&tool_result.content);
278 }
279 }
280 }
281
282 if !string_contents.is_empty() {
283 string_messages.push(tiktoken_rs::ChatCompletionRequestMessage {
284 role: match message.role {
285 Role::User => "user".into(),
286 Role::Assistant => "assistant".into(),
287 Role::System => "system".into(),
288 },
289 content: Some(string_contents),
290 name: None,
291 function_call: None,
292 });
293 }
294 }
295
296 // Tiktoken doesn't yet support these models, so we manually use the
297 // same tokenizer as GPT-4.
298 tiktoken_rs::num_tokens_from_messages("gpt-4", &string_messages)
299 .map(|tokens| tokens + tokens_from_images)
300 })
301 .boxed()
302}
303
304impl AnthropicModel {
305 fn stream_completion(
306 &self,
307 request: anthropic::Request,
308 cx: &AsyncApp,
309 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<anthropic::Event, AnthropicError>>>>
310 {
311 let http_client = self.http_client.clone();
312
313 let Ok((api_key, api_url)) = cx.read_entity(&self.state, |state, cx| {
314 let settings = &AllLanguageModelSettings::get_global(cx).anthropic;
315 (state.api_key.clone(), settings.api_url.clone())
316 }) else {
317 return futures::future::ready(Err(anyhow!("App state dropped"))).boxed();
318 };
319
320 async move {
321 let api_key = api_key.ok_or_else(|| anyhow!("Missing Anthropic API Key"))?;
322 let request =
323 anthropic::stream_completion(http_client.as_ref(), &api_url, &api_key, request);
324 request.await.context("failed to stream completion")
325 }
326 .boxed()
327 }
328}
329
330impl LanguageModel for AnthropicModel {
331 fn id(&self) -> LanguageModelId {
332 self.id.clone()
333 }
334
335 fn name(&self) -> LanguageModelName {
336 LanguageModelName::from(self.model.display_name().to_string())
337 }
338
339 fn provider_id(&self) -> LanguageModelProviderId {
340 LanguageModelProviderId(PROVIDER_ID.into())
341 }
342
343 fn provider_name(&self) -> LanguageModelProviderName {
344 LanguageModelProviderName(PROVIDER_NAME.into())
345 }
346
347 fn telemetry_id(&self) -> String {
348 format!("anthropic/{}", self.model.id())
349 }
350
351 fn api_key(&self, cx: &App) -> Option<String> {
352 self.state.read(cx).api_key.clone()
353 }
354
355 fn max_token_count(&self) -> usize {
356 self.model.max_token_count()
357 }
358
359 fn max_output_tokens(&self) -> Option<u32> {
360 Some(self.model.max_output_tokens())
361 }
362
363 fn count_tokens(
364 &self,
365 request: LanguageModelRequest,
366 cx: &App,
367 ) -> BoxFuture<'static, Result<usize>> {
368 count_anthropic_tokens(request, cx)
369 }
370
371 fn stream_completion(
372 &self,
373 request: LanguageModelRequest,
374 cx: &AsyncApp,
375 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<LanguageModelCompletionEvent>>>> {
376 let request = request.into_anthropic(
377 self.model.id().into(),
378 self.model.default_temperature(),
379 self.model.max_output_tokens(),
380 );
381 let request = self.stream_completion(request, cx);
382 let future = self.request_limiter.stream(async move {
383 let response = request.await.map_err(|err| anyhow!(err))?;
384 Ok(map_to_language_model_completion_events(response))
385 });
386 async move { Ok(future.await?.boxed()) }.boxed()
387 }
388
389 fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
390 self.model
391 .cache_configuration()
392 .map(|config| LanguageModelCacheConfiguration {
393 max_cache_anchors: config.max_cache_anchors,
394 should_speculate: config.should_speculate,
395 min_total_token: config.min_total_token,
396 })
397 }
398
399 fn use_any_tool(
400 &self,
401 request: LanguageModelRequest,
402 tool_name: String,
403 tool_description: String,
404 input_schema: serde_json::Value,
405 cx: &AsyncApp,
406 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
407 let mut request = request.into_anthropic(
408 self.model.tool_model_id().into(),
409 self.model.default_temperature(),
410 self.model.max_output_tokens(),
411 );
412 request.tool_choice = Some(anthropic::ToolChoice::Tool {
413 name: tool_name.clone(),
414 });
415 request.tools = vec![anthropic::Tool {
416 name: tool_name.clone(),
417 description: tool_description,
418 input_schema,
419 }];
420
421 let response = self.stream_completion(request, cx);
422 self.request_limiter
423 .run(async move {
424 let response = response.await?;
425 Ok(anthropic::extract_tool_args_from_events(
426 tool_name,
427 Box::pin(response.map_err(|e| anyhow!(e))),
428 )
429 .await?
430 .boxed())
431 })
432 .boxed()
433 }
434}
435
436pub fn map_to_language_model_completion_events(
437 events: Pin<Box<dyn Send + Stream<Item = Result<Event, AnthropicError>>>>,
438) -> impl Stream<Item = Result<LanguageModelCompletionEvent>> {
439 struct RawToolUse {
440 id: String,
441 name: String,
442 input_json: String,
443 }
444
445 struct State {
446 events: Pin<Box<dyn Send + Stream<Item = Result<Event, AnthropicError>>>>,
447 tool_uses_by_index: HashMap<usize, RawToolUse>,
448 }
449
450 futures::stream::unfold(
451 State {
452 events,
453 tool_uses_by_index: HashMap::default(),
454 },
455 |mut state| async move {
456 while let Some(event) = state.events.next().await {
457 match event {
458 Ok(event) => match event {
459 Event::ContentBlockStart {
460 index,
461 content_block,
462 } => match content_block {
463 ResponseContent::Text { text } => {
464 return Some((
465 Some(Ok(LanguageModelCompletionEvent::Text(text))),
466 state,
467 ));
468 }
469 ResponseContent::ToolUse { id, name, .. } => {
470 state.tool_uses_by_index.insert(
471 index,
472 RawToolUse {
473 id,
474 name,
475 input_json: String::new(),
476 },
477 );
478
479 return Some((None, state));
480 }
481 },
482 Event::ContentBlockDelta { index, delta } => match delta {
483 ContentDelta::TextDelta { text } => {
484 return Some((
485 Some(Ok(LanguageModelCompletionEvent::Text(text))),
486 state,
487 ));
488 }
489 ContentDelta::InputJsonDelta { partial_json } => {
490 if let Some(tool_use) = state.tool_uses_by_index.get_mut(&index) {
491 tool_use.input_json.push_str(&partial_json);
492 return Some((None, state));
493 }
494 }
495 },
496 Event::ContentBlockStop { index } => {
497 if let Some(tool_use) = state.tool_uses_by_index.remove(&index) {
498 return Some((
499 Some(maybe!({
500 Ok(LanguageModelCompletionEvent::ToolUse(
501 LanguageModelToolUse {
502 id: tool_use.id.into(),
503 name: tool_use.name,
504 input: if tool_use.input_json.is_empty() {
505 serde_json::Value::Null
506 } else {
507 serde_json::Value::from_str(
508 &tool_use.input_json,
509 )
510 .map_err(|err| anyhow!(err))?
511 },
512 },
513 ))
514 })),
515 state,
516 ));
517 }
518 }
519 Event::MessageStart { message } => {
520 return Some((
521 Some(Ok(LanguageModelCompletionEvent::StartMessage {
522 message_id: message.id,
523 })),
524 state,
525 ))
526 }
527 Event::MessageDelta { delta, .. } => {
528 if let Some(stop_reason) = delta.stop_reason.as_deref() {
529 let stop_reason = match stop_reason {
530 "end_turn" => StopReason::EndTurn,
531 "max_tokens" => StopReason::MaxTokens,
532 "tool_use" => StopReason::ToolUse,
533 _ => StopReason::EndTurn,
534 };
535
536 return Some((
537 Some(Ok(LanguageModelCompletionEvent::Stop(stop_reason))),
538 state,
539 ));
540 }
541 }
542 Event::Error { error } => {
543 return Some((
544 Some(Err(anyhow!(AnthropicError::ApiError(error)))),
545 state,
546 ));
547 }
548 _ => {}
549 },
550 Err(err) => {
551 return Some((Some(Err(anyhow!(err))), state));
552 }
553 }
554 }
555
556 None
557 },
558 )
559 .filter_map(|event| async move { event })
560}
561
562struct ConfigurationView {
563 api_key_editor: Entity<Editor>,
564 state: gpui::Entity<State>,
565 load_credentials_task: Option<Task<()>>,
566}
567
568impl ConfigurationView {
569 const PLACEHOLDER_TEXT: &'static str = "sk-ant-xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx";
570
571 fn new(state: gpui::Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
572 cx.observe(&state, |_, _, cx| {
573 cx.notify();
574 })
575 .detach();
576
577 let load_credentials_task = Some(cx.spawn({
578 let state = state.clone();
579 |this, mut cx| async move {
580 if let Some(task) = state
581 .update(&mut cx, |state, cx| state.authenticate(cx))
582 .log_err()
583 {
584 // We don't log an error, because "not signed in" is also an error.
585 let _ = task.await;
586 }
587 this.update(&mut cx, |this, cx| {
588 this.load_credentials_task = None;
589 cx.notify();
590 })
591 .log_err();
592 }
593 }));
594
595 Self {
596 api_key_editor: cx.new(|cx| {
597 let mut editor = Editor::single_line(window, cx);
598 editor.set_placeholder_text(Self::PLACEHOLDER_TEXT, cx);
599 editor
600 }),
601 state,
602 load_credentials_task,
603 }
604 }
605
606 fn save_api_key(&mut self, _: &menu::Confirm, window: &mut Window, cx: &mut Context<Self>) {
607 let api_key = self.api_key_editor.read(cx).text(cx);
608 if api_key.is_empty() {
609 return;
610 }
611
612 let state = self.state.clone();
613 cx.spawn_in(window, |_, mut cx| async move {
614 state
615 .update(&mut cx, |state, cx| state.set_api_key(api_key, cx))?
616 .await
617 })
618 .detach_and_log_err(cx);
619
620 cx.notify();
621 }
622
623 fn reset_api_key(&mut self, window: &mut Window, cx: &mut Context<Self>) {
624 self.api_key_editor
625 .update(cx, |editor, cx| editor.set_text("", window, cx));
626
627 let state = self.state.clone();
628 cx.spawn_in(window, |_, mut cx| async move {
629 state
630 .update(&mut cx, |state, cx| state.reset_api_key(cx))?
631 .await
632 })
633 .detach_and_log_err(cx);
634
635 cx.notify();
636 }
637
638 fn render_api_key_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
639 let settings = ThemeSettings::get_global(cx);
640 let text_style = TextStyle {
641 color: cx.theme().colors().text,
642 font_family: settings.ui_font.family.clone(),
643 font_features: settings.ui_font.features.clone(),
644 font_fallbacks: settings.ui_font.fallbacks.clone(),
645 font_size: rems(0.875).into(),
646 font_weight: settings.ui_font.weight,
647 font_style: FontStyle::Normal,
648 line_height: relative(1.3),
649 white_space: WhiteSpace::Normal,
650 ..Default::default()
651 };
652 EditorElement::new(
653 &self.api_key_editor,
654 EditorStyle {
655 background: cx.theme().colors().editor_background,
656 local_player: cx.theme().players().local(),
657 text: text_style,
658 ..Default::default()
659 },
660 )
661 }
662
663 fn should_render_editor(&self, cx: &mut Context<Self>) -> bool {
664 !self.state.read(cx).is_authenticated()
665 }
666}
667
668impl Render for ConfigurationView {
669 fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
670 const ANTHROPIC_CONSOLE_URL: &str = "https://console.anthropic.com/settings/keys";
671 const INSTRUCTIONS: [&str; 3] = [
672 "To use Zed's assistant with Anthropic, you need to add an API key. Follow these steps:",
673 "- Create one at:",
674 "- Paste your API key below and hit enter to use the assistant:",
675 ];
676 let env_var_set = self.state.read(cx).api_key_from_env;
677
678 if self.load_credentials_task.is_some() {
679 div().child(Label::new("Loading credentials...")).into_any()
680 } else if self.should_render_editor(cx) {
681 v_flex()
682 .size_full()
683 .on_action(cx.listener(Self::save_api_key))
684 .child(Label::new(INSTRUCTIONS[0]))
685 .child(h_flex().child(Label::new(INSTRUCTIONS[1])).child(
686 Button::new("anthropic_console", ANTHROPIC_CONSOLE_URL)
687 .style(ButtonStyle::Subtle)
688 .icon(IconName::ArrowUpRight)
689 .icon_size(IconSize::XSmall)
690 .icon_color(Color::Muted)
691 .on_click(move |_, _, cx| cx.open_url(ANTHROPIC_CONSOLE_URL))
692 )
693 )
694 .child(Label::new(INSTRUCTIONS[2]))
695 .child(
696 h_flex()
697 .w_full()
698 .my_2()
699 .px_2()
700 .py_1()
701 .bg(cx.theme().colors().editor_background)
702 .border_1()
703 .border_color(cx.theme().colors().border_variant)
704 .rounded_md()
705 .child(self.render_api_key_editor(cx)),
706 )
707 .child(
708 Label::new(
709 format!("You can also assign the {ANTHROPIC_API_KEY_VAR} environment variable and restart Zed."),
710 )
711 .size(LabelSize::Small),
712 )
713 .into_any()
714 } else {
715 h_flex()
716 .size_full()
717 .justify_between()
718 .child(
719 h_flex()
720 .gap_1()
721 .child(Icon::new(IconName::Check).color(Color::Success))
722 .child(Label::new(if env_var_set {
723 format!("API key set in {ANTHROPIC_API_KEY_VAR} environment variable.")
724 } else {
725 "API key configured.".to_string()
726 })),
727 )
728 .child(
729 Button::new("reset-key", "Reset key")
730 .icon(Some(IconName::Trash))
731 .icon_size(IconSize::Small)
732 .icon_position(IconPosition::Start)
733 .disabled(env_var_set)
734 .when(env_var_set, |this| {
735 this.tooltip(Tooltip::text(format!("To reset your API key, unset the {ANTHROPIC_API_KEY_VAR} environment variable.")))
736 })
737 .on_click(cx.listener(|this, _, window, cx| this.reset_api_key(window, cx))),
738 )
739 .into_any()
740 }
741 }
742}