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