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