1use crate::{count_open_ai_tokens, LanguageModelCompletionProvider};
2use crate::{CompletionProvider, LanguageModel, LanguageModelRequest};
3use anthropic::{stream_completion, Model as AnthropicModel, Request, RequestMessage};
4use anyhow::{anyhow, Result};
5use editor::{Editor, EditorElement, EditorStyle};
6use futures::{future::BoxFuture, stream::BoxStream, FutureExt, StreamExt};
7use gpui::{AnyView, AppContext, Task, TextStyle, View};
8use http::HttpClient;
9use language_model::Role;
10use settings::Settings;
11use std::time::Duration;
12use std::{env, sync::Arc};
13use strum::IntoEnumIterator;
14use theme::ThemeSettings;
15use ui::prelude::*;
16use util::ResultExt;
17
18pub struct AnthropicCompletionProvider {
19 api_key: Option<String>,
20 api_url: String,
21 model: AnthropicModel,
22 http_client: Arc<dyn HttpClient>,
23 low_speed_timeout: Option<Duration>,
24 settings_version: usize,
25}
26
27impl LanguageModelCompletionProvider for AnthropicCompletionProvider {
28 fn available_models(&self) -> Vec<LanguageModel> {
29 AnthropicModel::iter()
30 .map(LanguageModel::Anthropic)
31 .collect()
32 }
33
34 fn settings_version(&self) -> usize {
35 self.settings_version
36 }
37
38 fn is_authenticated(&self) -> bool {
39 self.api_key.is_some()
40 }
41
42 fn authenticate(&self, cx: &AppContext) -> Task<Result<()>> {
43 if self.is_authenticated() {
44 Task::ready(Ok(()))
45 } else {
46 let api_url = self.api_url.clone();
47 cx.spawn(|mut cx| async move {
48 let api_key = if let Ok(api_key) = env::var("ANTHROPIC_API_KEY") {
49 api_key
50 } else {
51 let (_, api_key) = cx
52 .update(|cx| cx.read_credentials(&api_url))?
53 .await?
54 .ok_or_else(|| anyhow!("credentials not found"))?;
55 String::from_utf8(api_key)?
56 };
57 cx.update_global::<CompletionProvider, _>(|provider, _cx| {
58 provider.update_current_as::<_, AnthropicCompletionProvider>(|provider| {
59 provider.api_key = Some(api_key);
60 });
61 })
62 })
63 }
64 }
65
66 fn reset_credentials(&self, cx: &AppContext) -> Task<Result<()>> {
67 let delete_credentials = cx.delete_credentials(&self.api_url);
68 cx.spawn(|mut cx| async move {
69 delete_credentials.await.log_err();
70 cx.update_global::<CompletionProvider, _>(|provider, _cx| {
71 provider.update_current_as::<_, AnthropicCompletionProvider>(|provider| {
72 provider.api_key = None;
73 });
74 })
75 })
76 }
77
78 fn authentication_prompt(&self, cx: &mut WindowContext) -> AnyView {
79 cx.new_view(|cx| AuthenticationPrompt::new(self.api_url.clone(), cx))
80 .into()
81 }
82
83 fn model(&self) -> LanguageModel {
84 LanguageModel::Anthropic(self.model.clone())
85 }
86
87 fn count_tokens(
88 &self,
89 request: LanguageModelRequest,
90 cx: &AppContext,
91 ) -> BoxFuture<'static, Result<usize>> {
92 count_open_ai_tokens(request, cx.background_executor())
93 }
94
95 fn stream_completion(
96 &self,
97 request: LanguageModelRequest,
98 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
99 let request = self.to_anthropic_request(request);
100
101 let http_client = self.http_client.clone();
102 let api_key = self.api_key.clone();
103 let api_url = self.api_url.clone();
104 let low_speed_timeout = self.low_speed_timeout;
105 async move {
106 let api_key = api_key.ok_or_else(|| anyhow!("missing api key"))?;
107 let request = stream_completion(
108 http_client.as_ref(),
109 &api_url,
110 &api_key,
111 request,
112 low_speed_timeout,
113 );
114 let response = request.await?;
115 let stream = response
116 .filter_map(|response| async move {
117 match response {
118 Ok(response) => match response {
119 anthropic::ResponseEvent::ContentBlockStart {
120 content_block, ..
121 } => match content_block {
122 anthropic::ContentBlock::Text { text } => Some(Ok(text)),
123 },
124 anthropic::ResponseEvent::ContentBlockDelta { delta, .. } => {
125 match delta {
126 anthropic::TextDelta::TextDelta { text } => Some(Ok(text)),
127 }
128 }
129 _ => None,
130 },
131 Err(error) => Some(Err(error)),
132 }
133 })
134 .boxed();
135 Ok(stream)
136 }
137 .boxed()
138 }
139
140 fn as_any_mut(&mut self) -> &mut dyn std::any::Any {
141 self
142 }
143}
144
145impl AnthropicCompletionProvider {
146 pub fn new(
147 model: AnthropicModel,
148 api_url: String,
149 http_client: Arc<dyn HttpClient>,
150 low_speed_timeout: Option<Duration>,
151 settings_version: usize,
152 ) -> Self {
153 Self {
154 api_key: None,
155 api_url,
156 model,
157 http_client,
158 low_speed_timeout,
159 settings_version,
160 }
161 }
162
163 pub fn update(
164 &mut self,
165 model: AnthropicModel,
166 api_url: String,
167 low_speed_timeout: Option<Duration>,
168 settings_version: usize,
169 ) {
170 self.model = model;
171 self.api_url = api_url;
172 self.low_speed_timeout = low_speed_timeout;
173 self.settings_version = settings_version;
174 }
175
176 fn to_anthropic_request(&self, mut request: LanguageModelRequest) -> Request {
177 request.preprocess_anthropic();
178
179 let model = match request.model {
180 LanguageModel::Anthropic(model) => model,
181 _ => self.model.clone(),
182 };
183
184 let mut system_message = String::new();
185 if request
186 .messages
187 .first()
188 .map_or(false, |message| message.role == Role::System)
189 {
190 system_message = request.messages.remove(0).content;
191 }
192
193 Request {
194 model,
195 messages: request
196 .messages
197 .iter()
198 .map(|msg| RequestMessage {
199 role: match msg.role {
200 Role::User => anthropic::Role::User,
201 Role::Assistant => anthropic::Role::Assistant,
202 Role::System => unreachable!("filtered out by preprocess_request"),
203 },
204 content: msg.content.clone(),
205 })
206 .collect(),
207 stream: true,
208 system: system_message,
209 max_tokens: 4092,
210 }
211 }
212}
213
214struct AuthenticationPrompt {
215 api_key: View<Editor>,
216 api_url: String,
217}
218
219impl AuthenticationPrompt {
220 fn new(api_url: String, cx: &mut WindowContext) -> Self {
221 Self {
222 api_key: cx.new_view(|cx| {
223 let mut editor = Editor::single_line(cx);
224 editor.set_placeholder_text(
225 "sk-000000000000000000000000000000000000000000000000",
226 cx,
227 );
228 editor
229 }),
230 api_url,
231 }
232 }
233
234 fn save_api_key(&mut self, _: &menu::Confirm, cx: &mut ViewContext<Self>) {
235 let api_key = self.api_key.read(cx).text(cx);
236 if api_key.is_empty() {
237 return;
238 }
239
240 let write_credentials = cx.write_credentials(&self.api_url, "Bearer", api_key.as_bytes());
241 cx.spawn(|_, mut cx| async move {
242 write_credentials.await?;
243 cx.update_global::<CompletionProvider, _>(|provider, _cx| {
244 provider.update_current_as::<_, AnthropicCompletionProvider>(|provider| {
245 provider.api_key = Some(api_key);
246 });
247 })
248 })
249 .detach_and_log_err(cx);
250 }
251
252 fn render_api_key_editor(&self, cx: &mut ViewContext<Self>) -> impl IntoElement {
253 let settings = ThemeSettings::get_global(cx);
254 let text_style = TextStyle {
255 color: cx.theme().colors().text,
256 font_family: settings.ui_font.family.clone(),
257 font_features: settings.ui_font.features.clone(),
258 font_size: rems(0.875).into(),
259 font_weight: settings.ui_font.weight,
260 line_height: relative(1.3),
261 ..Default::default()
262 };
263 EditorElement::new(
264 &self.api_key,
265 EditorStyle {
266 background: cx.theme().colors().editor_background,
267 local_player: cx.theme().players().local(),
268 text: text_style,
269 ..Default::default()
270 },
271 )
272 }
273}
274
275impl Render for AuthenticationPrompt {
276 fn render(&mut self, cx: &mut ViewContext<Self>) -> impl IntoElement {
277 const INSTRUCTIONS: [&str; 4] = [
278 "To use the assistant panel or inline assistant, you need to add your Anthropic API key.",
279 "You can create an API key at: https://console.anthropic.com/settings/keys",
280 "",
281 "Paste your Anthropic API key below and hit enter to use the assistant:",
282 ];
283
284 v_flex()
285 .p_4()
286 .size_full()
287 .on_action(cx.listener(Self::save_api_key))
288 .children(
289 INSTRUCTIONS.map(|instruction| Label::new(instruction).size(LabelSize::Small)),
290 )
291 .child(
292 h_flex()
293 .w_full()
294 .my_2()
295 .px_2()
296 .py_1()
297 .bg(cx.theme().colors().editor_background)
298 .rounded_md()
299 .child(self.render_api_key_editor(cx)),
300 )
301 .child(
302 Label::new(
303 "You can also assign the ANTHROPIC_API_KEY environment variable and restart Zed.",
304 )
305 .size(LabelSize::Small),
306 )
307 .child(
308 h_flex()
309 .gap_2()
310 .child(Label::new("Click on").size(LabelSize::Small))
311 .child(Icon::new(IconName::ZedAssistant).size(IconSize::XSmall))
312 .child(
313 Label::new("in the status bar to close this panel.").size(LabelSize::Small),
314 ),
315 )
316 .into_any()
317 }
318}