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