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