1use std::pin::Pin;
2use std::str::FromStr;
3use std::sync::Arc;
4
5use anyhow::{anyhow, Context as _, Result};
6use aws_config::stalled_stream_protection::StalledStreamProtectionConfig;
7use aws_config::Region;
8use aws_credential_types::Credentials;
9use aws_http_client::AwsHttpClient;
10use bedrock::bedrock_client::types::{
11 ContentBlockDelta, ContentBlockStart, ContentBlockStartEvent, ConverseStreamOutput,
12};
13use bedrock::bedrock_client::{self, Config};
14use bedrock::{
15 value_to_aws_document, BedrockError, BedrockInnerContent, BedrockMessage, BedrockSpecificTool,
16 BedrockStreamingResponse, BedrockTool, BedrockToolChoice, BedrockToolInputSchema, Model,
17};
18use collections::{BTreeMap, HashMap};
19use credentials_provider::CredentialsProvider;
20use editor::{Editor, EditorElement, EditorStyle};
21use futures::{future::BoxFuture, stream::BoxStream, FutureExt, Stream, StreamExt};
22use gpui::{
23 AnyView, App, AsyncApp, Context, Entity, FontStyle, Subscription, Task, TextStyle, WhiteSpace,
24};
25use gpui_tokio::Tokio;
26use http_client::HttpClient;
27use language_model::{
28 AuthenticateError, LanguageModel, LanguageModelCacheConfiguration,
29 LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, LanguageModelProvider,
30 LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState,
31 LanguageModelRequest, LanguageModelToolUse, MessageContent, RateLimiter, Role,
32};
33use schemars::JsonSchema;
34use serde::{Deserialize, Serialize};
35use serde_json::Value;
36use settings::{Settings, SettingsStore};
37use strum::IntoEnumIterator;
38use theme::ThemeSettings;
39use tokio::runtime::Handle;
40use ui::{prelude::*, Icon, IconName, Tooltip};
41use util::{maybe, ResultExt};
42
43use crate::AllLanguageModelSettings;
44
45const PROVIDER_ID: &str = "amazon-bedrock";
46const PROVIDER_NAME: &str = "Amazon Bedrock";
47
48#[derive(Default, Clone, Deserialize, Serialize, PartialEq, Debug)]
49pub struct BedrockCredentials {
50 pub region: String,
51 pub access_key_id: String,
52 pub secret_access_key: String,
53}
54
55#[derive(Default, Clone, Debug, PartialEq)]
56pub struct AmazonBedrockSettings {
57 pub session_token: Option<String>,
58 pub available_models: Vec<AvailableModel>,
59}
60
61#[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]
62pub struct AvailableModel {
63 pub name: String,
64 pub display_name: Option<String>,
65 pub max_tokens: usize,
66 pub cache_configuration: Option<LanguageModelCacheConfiguration>,
67 pub max_output_tokens: Option<u32>,
68 pub default_temperature: Option<f32>,
69}
70
71/// The URL of the base AWS service.
72///
73/// Right now we're just using this as the key to store the AWS credentials
74/// under in the keychain.
75const AMAZON_AWS_URL: &str = "https://amazonaws.com";
76
77// These environment variables all use a `ZED_` prefix because we don't want to overwrite the user's AWS credentials.
78const ZED_BEDROCK_ACCESS_KEY_ID_VAR: &str = "ZED_ACCESS_KEY_ID";
79const ZED_BEDROCK_SECRET_ACCESS_KEY_VAR: &str = "ZED_SECRET_ACCESS_KEY";
80const ZED_BEDROCK_REGION_VAR: &str = "ZED_AWS_REGION";
81const ZED_AWS_CREDENTIALS_VAR: &str = "ZED_AWS_CREDENTIALS";
82
83pub struct State {
84 credentials: Option<BedrockCredentials>,
85 credentials_from_env: bool,
86 _subscription: Subscription,
87}
88
89impl State {
90 fn reset_credentials(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
91 let credentials_provider = <dyn CredentialsProvider>::global(cx);
92 cx.spawn(|this, mut cx| async move {
93 credentials_provider
94 .delete_credentials(AMAZON_AWS_URL, &cx)
95 .await
96 .log_err();
97 this.update(&mut cx, |this, cx| {
98 this.credentials = None;
99 this.credentials_from_env = false;
100 cx.notify();
101 })
102 })
103 }
104
105 fn set_credentials(
106 &mut self,
107 credentials: BedrockCredentials,
108 cx: &mut Context<Self>,
109 ) -> Task<Result<()>> {
110 let credentials_provider = <dyn CredentialsProvider>::global(cx);
111 cx.spawn(|this, mut cx| async move {
112 credentials_provider
113 .write_credentials(
114 AMAZON_AWS_URL,
115 "Bearer",
116 &serde_json::to_vec(&credentials)?,
117 &cx,
118 )
119 .await?;
120 this.update(&mut cx, |this, cx| {
121 this.credentials = Some(credentials);
122 cx.notify();
123 })
124 })
125 }
126
127 fn is_authenticated(&self) -> bool {
128 self.credentials.is_some()
129 }
130
131 fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
132 if self.is_authenticated() {
133 return Task::ready(Ok(()));
134 }
135
136 let credentials_provider = <dyn CredentialsProvider>::global(cx);
137 cx.spawn(|this, mut cx| async move {
138 let (credentials, from_env) =
139 if let Ok(credentials) = std::env::var(ZED_AWS_CREDENTIALS_VAR) {
140 (credentials, true)
141 } else {
142 let (_, credentials) = credentials_provider
143 .read_credentials(AMAZON_AWS_URL, &cx)
144 .await?
145 .ok_or_else(|| AuthenticateError::CredentialsNotFound)?;
146 (
147 String::from_utf8(credentials)
148 .context("invalid {PROVIDER_NAME} credentials")?,
149 false,
150 )
151 };
152
153 let credentials: BedrockCredentials =
154 serde_json::from_str(&credentials).context("failed to parse credentials")?;
155
156 this.update(&mut cx, |this, cx| {
157 this.credentials = Some(credentials);
158 this.credentials_from_env = from_env;
159 cx.notify();
160 })?;
161
162 Ok(())
163 })
164 }
165}
166
167pub struct BedrockLanguageModelProvider {
168 http_client: AwsHttpClient,
169 handler: tokio::runtime::Handle,
170 state: gpui::Entity<State>,
171}
172
173impl BedrockLanguageModelProvider {
174 pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self {
175 let state = cx.new(|cx| State {
176 credentials: None,
177 credentials_from_env: false,
178 _subscription: cx.observe_global::<SettingsStore>(|_, cx| {
179 cx.notify();
180 }),
181 });
182
183 let tokio_handle = Tokio::handle(cx);
184
185 let coerced_client = AwsHttpClient::new(http_client.clone(), tokio_handle.clone());
186
187 Self {
188 http_client: coerced_client,
189 handler: tokio_handle.clone(),
190 state,
191 }
192 }
193}
194
195impl LanguageModelProvider for BedrockLanguageModelProvider {
196 fn id(&self) -> LanguageModelProviderId {
197 LanguageModelProviderId(PROVIDER_ID.into())
198 }
199
200 fn name(&self) -> LanguageModelProviderName {
201 LanguageModelProviderName(PROVIDER_NAME.into())
202 }
203
204 fn icon(&self) -> IconName {
205 IconName::AiBedrock
206 }
207
208 fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
209 let model = bedrock::Model::default();
210 Some(Arc::new(BedrockModel {
211 id: LanguageModelId::from(model.id().to_string()),
212 model,
213 http_client: self.http_client.clone(),
214 handler: self.handler.clone(),
215 state: self.state.clone(),
216 request_limiter: RateLimiter::new(4),
217 }))
218 }
219
220 fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
221 let mut models = BTreeMap::default();
222
223 for model in bedrock::Model::iter() {
224 if !matches!(model, bedrock::Model::Custom { .. }) {
225 models.insert(model.id().to_string(), model);
226 }
227 }
228
229 // Override with available models from settings
230 for model in AllLanguageModelSettings::get_global(cx)
231 .bedrock
232 .available_models
233 .iter()
234 {
235 models.insert(
236 model.name.clone(),
237 bedrock::Model::Custom {
238 name: model.name.clone(),
239 display_name: model.display_name.clone(),
240 max_tokens: model.max_tokens,
241 max_output_tokens: model.max_output_tokens,
242 default_temperature: model.default_temperature,
243 },
244 );
245 }
246
247 models
248 .into_values()
249 .map(|model| {
250 Arc::new(BedrockModel {
251 id: LanguageModelId::from(model.id().to_string()),
252 model,
253 http_client: self.http_client.clone(),
254 handler: self.handler.clone(),
255 state: self.state.clone(),
256 request_limiter: RateLimiter::new(4),
257 }) as Arc<dyn LanguageModel>
258 })
259 .collect()
260 }
261
262 fn is_authenticated(&self, cx: &App) -> bool {
263 self.state.read(cx).is_authenticated()
264 }
265
266 fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
267 self.state.update(cx, |state, cx| state.authenticate(cx))
268 }
269
270 fn configuration_view(&self, window: &mut Window, cx: &mut App) -> AnyView {
271 cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx))
272 .into()
273 }
274
275 fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> {
276 self.state
277 .update(cx, |state, cx| state.reset_credentials(cx))
278 }
279}
280
281impl LanguageModelProviderState for BedrockLanguageModelProvider {
282 type ObservableEntity = State;
283
284 fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> {
285 Some(self.state.clone())
286 }
287}
288
289struct BedrockModel {
290 id: LanguageModelId,
291 model: Model,
292 http_client: AwsHttpClient,
293 handler: tokio::runtime::Handle,
294 state: gpui::Entity<State>,
295 request_limiter: RateLimiter,
296}
297
298impl BedrockModel {
299 fn stream_completion(
300 &self,
301 request: bedrock::Request,
302 cx: &AsyncApp,
303 ) -> Result<
304 BoxFuture<'static, BoxStream<'static, Result<BedrockStreamingResponse, BedrockError>>>,
305 > {
306 let Ok(Ok((access_key_id, secret_access_key, region))) =
307 cx.read_entity(&self.state, |state, _cx| {
308 if let Some(credentials) = &state.credentials {
309 Ok((
310 credentials.access_key_id.clone(),
311 credentials.secret_access_key.clone(),
312 credentials.region.clone(),
313 ))
314 } else {
315 return Err(anyhow!("Failed to read credentials"));
316 }
317 })
318 else {
319 return Err(anyhow!("App state dropped"));
320 };
321
322 let runtime_client = bedrock_client::Client::from_conf(
323 Config::builder()
324 .stalled_stream_protection(StalledStreamProtectionConfig::disabled())
325 .credentials_provider(Credentials::new(
326 access_key_id,
327 secret_access_key,
328 None,
329 None,
330 "Keychain",
331 ))
332 .region(Region::new(region))
333 .http_client(self.http_client.clone())
334 .build(),
335 );
336
337 let owned_handle = self.handler.clone();
338
339 Ok(async move {
340 let request = bedrock::stream_completion(runtime_client, request, owned_handle);
341 request.await.unwrap_or_else(|e| {
342 futures::stream::once(async move { Err(BedrockError::ClientError(e)) }).boxed()
343 })
344 }
345 .boxed())
346 }
347}
348
349impl LanguageModel for BedrockModel {
350 fn id(&self) -> LanguageModelId {
351 self.id.clone()
352 }
353
354 fn name(&self) -> LanguageModelName {
355 LanguageModelName::from(self.model.display_name().to_string())
356 }
357
358 fn provider_id(&self) -> LanguageModelProviderId {
359 LanguageModelProviderId(PROVIDER_ID.into())
360 }
361
362 fn provider_name(&self) -> LanguageModelProviderName {
363 LanguageModelProviderName(PROVIDER_NAME.into())
364 }
365
366 fn telemetry_id(&self) -> String {
367 format!("bedrock/{}", self.model.id())
368 }
369
370 fn max_token_count(&self) -> usize {
371 self.model.max_token_count()
372 }
373
374 fn max_output_tokens(&self) -> Option<u32> {
375 Some(self.model.max_output_tokens())
376 }
377
378 fn count_tokens(
379 &self,
380 request: LanguageModelRequest,
381 cx: &App,
382 ) -> BoxFuture<'static, Result<usize>> {
383 get_bedrock_tokens(request, cx)
384 }
385
386 fn stream_completion(
387 &self,
388 request: LanguageModelRequest,
389 cx: &AsyncApp,
390 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<LanguageModelCompletionEvent>>>> {
391 let request = into_bedrock(
392 request,
393 self.model.id().into(),
394 self.model.default_temperature(),
395 self.model.max_output_tokens(),
396 );
397
398 let owned_handle = self.handler.clone();
399
400 let request = self.stream_completion(request, cx);
401 let future = self.request_limiter.stream(async move {
402 let response = request.map_err(|e| anyhow!(e)).unwrap().await;
403 Ok(map_to_language_model_completion_events(
404 response,
405 owned_handle,
406 ))
407 });
408 async move { Ok(future.await?.boxed()) }.boxed()
409 }
410
411 fn use_any_tool(
412 &self,
413 request: LanguageModelRequest,
414 name: String,
415 description: String,
416 schema: Value,
417 _cx: &AsyncApp,
418 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
419 let mut request = into_bedrock(
420 request,
421 self.model.id().into(),
422 self.model.default_temperature(),
423 self.model.max_output_tokens(),
424 );
425
426 request.tool_choice = Some(BedrockToolChoice::Tool(
427 BedrockSpecificTool::builder()
428 .name(name.clone())
429 .build()
430 .unwrap(),
431 ));
432
433 request.tools = vec![BedrockTool::builder()
434 .name(name.clone())
435 .description(description.clone())
436 .input_schema(BedrockToolInputSchema::Json(value_to_aws_document(&schema)))
437 .build()
438 .unwrap()];
439
440 let handle = self.handler.clone();
441
442 let request = self.stream_completion(request, _cx);
443 self.request_limiter
444 .run(async move {
445 let response = request.map_err(|e| anyhow!(e)).unwrap().await;
446 Ok(extract_tool_args_from_events(name, response, handle)
447 .await?
448 .boxed())
449 })
450 .boxed()
451 }
452
453 fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
454 None
455 }
456}
457
458pub fn into_bedrock(
459 request: LanguageModelRequest,
460 model: String,
461 default_temperature: f32,
462 max_output_tokens: u32,
463) -> bedrock::Request {
464 let mut new_messages: Vec<BedrockMessage> = Vec::new();
465 let mut system_message = String::new();
466
467 for message in request.messages {
468 if message.contents_empty() {
469 continue;
470 }
471
472 match message.role {
473 Role::User | Role::Assistant => {
474 let bedrock_message_content: Vec<BedrockInnerContent> = message
475 .content
476 .into_iter()
477 .filter_map(|content| match content {
478 MessageContent::Text(text) => {
479 if !text.is_empty() {
480 Some(BedrockInnerContent::Text(text))
481 } else {
482 None
483 }
484 }
485 _ => None,
486 })
487 .collect();
488 let bedrock_role = match message.role {
489 Role::User => bedrock::BedrockRole::User,
490 Role::Assistant => bedrock::BedrockRole::Assistant,
491 Role::System => unreachable!("System role should never occur here"),
492 };
493 if let Some(last_message) = new_messages.last_mut() {
494 if last_message.role == bedrock_role {
495 last_message.content.extend(bedrock_message_content);
496 continue;
497 }
498 }
499 new_messages.push(
500 BedrockMessage::builder()
501 .role(bedrock_role)
502 .set_content(Some(bedrock_message_content))
503 .build()
504 .expect("failed to build Bedrock message"),
505 );
506 }
507 Role::System => {
508 if !system_message.is_empty() {
509 system_message.push_str("\n\n");
510 }
511 system_message.push_str(&message.string_contents());
512 }
513 }
514 }
515
516 bedrock::Request {
517 model,
518 messages: new_messages,
519 max_tokens: max_output_tokens,
520 system: Some(system_message),
521 tools: vec![],
522 tool_choice: None,
523 metadata: None,
524 stop_sequences: Vec::new(),
525 temperature: request.temperature.or(Some(default_temperature)),
526 top_k: None,
527 top_p: None,
528 }
529}
530
531// TODO: just call the ConverseOutput.usage() method:
532// https://docs.rs/aws-sdk-bedrockruntime/latest/aws_sdk_bedrockruntime/operation/converse/struct.ConverseOutput.html#method.output
533pub fn get_bedrock_tokens(
534 request: LanguageModelRequest,
535 cx: &App,
536) -> BoxFuture<'static, Result<usize>> {
537 cx.background_executor()
538 .spawn(async move {
539 let messages = request.messages;
540 let mut tokens_from_images = 0;
541 let mut string_messages = Vec::with_capacity(messages.len());
542
543 for message in messages {
544 use language_model::MessageContent;
545
546 let mut string_contents = String::new();
547
548 for content in message.content {
549 match content {
550 MessageContent::Text(text) => {
551 string_contents.push_str(&text);
552 }
553 MessageContent::Image(image) => {
554 tokens_from_images += image.estimate_tokens();
555 }
556 MessageContent::ToolUse(_tool_use) => {
557 // TODO: Estimate token usage from tool uses.
558 }
559 MessageContent::ToolResult(tool_result) => {
560 string_contents.push_str(&tool_result.content);
561 }
562 }
563 }
564
565 if !string_contents.is_empty() {
566 string_messages.push(tiktoken_rs::ChatCompletionRequestMessage {
567 role: match message.role {
568 Role::User => "user".into(),
569 Role::Assistant => "assistant".into(),
570 Role::System => "system".into(),
571 },
572 content: Some(string_contents),
573 name: None,
574 function_call: None,
575 });
576 }
577 }
578
579 // Tiktoken doesn't yet support these models, so we manually use the
580 // same tokenizer as GPT-4.
581 tiktoken_rs::num_tokens_from_messages("gpt-4", &string_messages)
582 .map(|tokens| tokens + tokens_from_images)
583 })
584 .boxed()
585}
586
587pub async fn extract_tool_args_from_events(
588 name: String,
589 mut events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
590 handle: Handle,
591) -> Result<impl Send + Stream<Item = Result<String>>> {
592 handle
593 .spawn(async move {
594 let mut tool_use_index = None;
595 while let Some(event) = events.next().await {
596 if let BedrockStreamingResponse::ContentBlockStart(ContentBlockStartEvent {
597 content_block_index,
598 start,
599 ..
600 }) = event?
601 {
602 match start {
603 None => {
604 continue;
605 }
606 Some(start) => match start.as_tool_use() {
607 Ok(tool_use) => {
608 if name == tool_use.name {
609 tool_use_index = Some(content_block_index);
610 break;
611 }
612 }
613 Err(err) => {
614 return Err(anyhow!("Failed to parse tool use event: {:?}", err));
615 }
616 },
617 }
618 }
619 }
620
621 let Some(tool_use_index) = tool_use_index else {
622 return Err(anyhow!("Tool is not used"));
623 };
624
625 Ok(events.filter_map(move |event| {
626 let result = match event {
627 Err(_err) => None,
628 Ok(output) => match output.clone() {
629 BedrockStreamingResponse::ContentBlockDelta(inner) => {
630 match inner.clone().delta {
631 Some(ContentBlockDelta::ToolUse(tool_use)) => {
632 if inner.content_block_index == tool_use_index {
633 Some(Ok(tool_use.input))
634 } else {
635 None
636 }
637 }
638 _ => None,
639 }
640 }
641 _ => None,
642 },
643 };
644
645 async move { result }
646 }))
647 })
648 .await?
649}
650
651pub fn map_to_language_model_completion_events(
652 events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
653 handle: Handle,
654) -> impl Stream<Item = Result<LanguageModelCompletionEvent>> {
655 struct RawToolUse {
656 id: String,
657 name: String,
658 input_json: String,
659 }
660
661 struct State {
662 events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
663 tool_uses_by_index: HashMap<i32, RawToolUse>,
664 }
665
666 futures::stream::unfold(
667 State {
668 events,
669 tool_uses_by_index: HashMap::default(),
670 },
671 move |mut state: State| {
672 let inner_handle = handle.clone();
673 async move {
674 inner_handle
675 .spawn(async {
676 while let Some(event) = state.events.next().await {
677 match event {
678 Ok(event) => match event {
679 ConverseStreamOutput::ContentBlockDelta(cb_delta) => {
680 if let Some(ContentBlockDelta::Text(text_out)) =
681 cb_delta.delta
682 {
683 return Some((
684 Some(Ok(LanguageModelCompletionEvent::Text(
685 text_out,
686 ))),
687 state,
688 ));
689 } else if let Some(ContentBlockDelta::ToolUse(text_out)) =
690 cb_delta.delta
691 {
692 if let Some(tool_use) = state
693 .tool_uses_by_index
694 .get_mut(&cb_delta.content_block_index)
695 {
696 tool_use.input_json.push_str(text_out.input());
697 return Some((None, state));
698 };
699
700 return Some((None, state));
701 } else if cb_delta.delta.is_none() {
702 return Some((None, state));
703 }
704 }
705 ConverseStreamOutput::ContentBlockStart(cb_start) => {
706 if let Some(start) = cb_start.start {
707 match start {
708 ContentBlockStart::ToolUse(text_out) => {
709 let tool_use = RawToolUse {
710 id: text_out.tool_use_id,
711 name: text_out.name,
712 input_json: String::new(),
713 };
714
715 state.tool_uses_by_index.insert(
716 cb_start.content_block_index,
717 tool_use,
718 );
719 }
720 _ => {}
721 }
722 }
723 }
724 ConverseStreamOutput::ContentBlockStop(cb_stop) => {
725 if let Some(tool_use) = state
726 .tool_uses_by_index
727 .remove(&cb_stop.content_block_index)
728 {
729 return Some((
730 Some(maybe!({
731 Ok(LanguageModelCompletionEvent::ToolUse(
732 LanguageModelToolUse {
733 id: tool_use.id.into(),
734 name: tool_use.name.into(),
735 input: if tool_use.input_json.is_empty()
736 {
737 Value::Null
738 } else {
739 serde_json::Value::from_str(
740 &tool_use.input_json,
741 )
742 .map_err(|err| anyhow!(err))?
743 },
744 },
745 ))
746 })),
747 state,
748 ));
749 }
750 }
751 _ => {}
752 },
753 Err(err) => return Some((Some(Err(anyhow!(err))), state)),
754 }
755 }
756 None
757 })
758 .await
759 .unwrap()
760 }
761 },
762 )
763 .filter_map(|event| async move { event })
764}
765
766struct ConfigurationView {
767 access_key_id_editor: Entity<Editor>,
768 secret_access_key_editor: Entity<Editor>,
769 region_editor: Entity<Editor>,
770 state: gpui::Entity<State>,
771 load_credentials_task: Option<Task<()>>,
772}
773
774impl ConfigurationView {
775 const PLACEHOLDER_ACCESS_KEY_ID_TEXT: &'static str = "XXXXXXXXXXXXXXXX";
776 const PLACEHOLDER_SECRET_ACCESS_KEY_TEXT: &'static str =
777 "XXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXXX";
778 const PLACEHOLDER_REGION: &'static str = "us-east-1";
779
780 fn new(state: gpui::Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
781 cx.observe(&state, |_, _, cx| {
782 cx.notify();
783 })
784 .detach();
785
786 let load_credentials_task = Some(cx.spawn({
787 let state = state.clone();
788 |this, mut cx| async move {
789 if let Some(task) = state
790 .update(&mut cx, |state, cx| state.authenticate(cx))
791 .log_err()
792 {
793 // We don't log an error, because "not signed in" is also an error.
794 let _ = task.await;
795 }
796 this.update(&mut cx, |this, cx| {
797 this.load_credentials_task = None;
798 cx.notify();
799 })
800 .log_err();
801 }
802 }));
803
804 Self {
805 access_key_id_editor: cx.new(|cx| {
806 let mut editor = Editor::single_line(window, cx);
807 editor.set_placeholder_text(Self::PLACEHOLDER_ACCESS_KEY_ID_TEXT, cx);
808 editor
809 }),
810 secret_access_key_editor: cx.new(|cx| {
811 let mut editor = Editor::single_line(window, cx);
812 editor.set_placeholder_text(Self::PLACEHOLDER_SECRET_ACCESS_KEY_TEXT, cx);
813 editor
814 }),
815 region_editor: cx.new(|cx| {
816 let mut editor = Editor::single_line(window, cx);
817 editor.set_placeholder_text(Self::PLACEHOLDER_REGION, cx);
818 editor
819 }),
820 state,
821 load_credentials_task,
822 }
823 }
824
825 fn save_credentials(
826 &mut self,
827 _: &menu::Confirm,
828 _window: &mut Window,
829 cx: &mut Context<Self>,
830 ) {
831 let access_key_id = self
832 .access_key_id_editor
833 .read(cx)
834 .text(cx)
835 .to_string()
836 .trim()
837 .to_string();
838 let secret_access_key = self
839 .secret_access_key_editor
840 .read(cx)
841 .text(cx)
842 .to_string()
843 .trim()
844 .to_string();
845 let region = self
846 .region_editor
847 .read(cx)
848 .text(cx)
849 .to_string()
850 .trim()
851 .to_string();
852
853 let state = self.state.clone();
854 cx.spawn(|_, mut cx| async move {
855 state
856 .update(&mut cx, |state, cx| {
857 let credentials: BedrockCredentials = BedrockCredentials {
858 access_key_id: access_key_id.clone(),
859 secret_access_key: secret_access_key.clone(),
860 region: region.clone(),
861 };
862
863 state.set_credentials(credentials, cx)
864 })?
865 .await
866 })
867 .detach_and_log_err(cx);
868 }
869
870 fn reset_credentials(&mut self, window: &mut Window, cx: &mut Context<Self>) {
871 self.access_key_id_editor
872 .update(cx, |editor, cx| editor.set_text("", window, cx));
873 self.secret_access_key_editor
874 .update(cx, |editor, cx| editor.set_text("", window, cx));
875 self.region_editor
876 .update(cx, |editor, cx| editor.set_text("", window, cx));
877
878 let state = self.state.clone();
879 cx.spawn(|_, mut cx| async move {
880 state
881 .update(&mut cx, |state, cx| state.reset_credentials(cx))?
882 .await
883 })
884 .detach_and_log_err(cx);
885 }
886
887 fn make_text_style(&self, cx: &Context<Self>) -> TextStyle {
888 let settings = ThemeSettings::get_global(cx);
889 TextStyle {
890 color: cx.theme().colors().text,
891 font_family: settings.ui_font.family.clone(),
892 font_features: settings.ui_font.features.clone(),
893 font_fallbacks: settings.ui_font.fallbacks.clone(),
894 font_size: rems(0.875).into(),
895 font_weight: settings.ui_font.weight,
896 font_style: FontStyle::Normal,
897 line_height: relative(1.3),
898 background_color: None,
899 underline: None,
900 strikethrough: None,
901 white_space: WhiteSpace::Normal,
902 text_overflow: None,
903 text_align: Default::default(),
904 line_clamp: None,
905 }
906 }
907
908 fn render_aa_id_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
909 let text_style = self.make_text_style(cx);
910
911 EditorElement::new(
912 &self.access_key_id_editor,
913 EditorStyle {
914 background: cx.theme().colors().editor_background,
915 local_player: cx.theme().players().local(),
916 text: text_style,
917 ..Default::default()
918 },
919 )
920 }
921
922 fn render_sk_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
923 let text_style = self.make_text_style(cx);
924
925 EditorElement::new(
926 &self.secret_access_key_editor,
927 EditorStyle {
928 background: cx.theme().colors().editor_background,
929 local_player: cx.theme().players().local(),
930 text: text_style,
931 ..Default::default()
932 },
933 )
934 }
935
936 fn render_region_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
937 let text_style = self.make_text_style(cx);
938
939 EditorElement::new(
940 &self.region_editor,
941 EditorStyle {
942 background: cx.theme().colors().editor_background,
943 local_player: cx.theme().players().local(),
944 text: text_style,
945 ..Default::default()
946 },
947 )
948 }
949
950 fn should_render_editor(&self, cx: &mut Context<Self>) -> bool {
951 !self.state.read(cx).is_authenticated()
952 }
953}
954
955impl Render for ConfigurationView {
956 fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
957 const IAM_CONSOLE_URL: &str = "https://us-east-1.console.aws.amazon.com/iam/home";
958 const BEDROCK_DOCS_URL: &str =
959 "https://docs.aws.amazon.com/bedrock/latest/userguide/inference-prereq.html";
960 const BEDROCK_MODEL_CATALOG: &str =
961 "https://us-east-1.console.aws.amazon.com/bedrock/home?region=us-east-1#/modelaccess";
962 const INSTRUCTIONS: [&str; 5] = [
963 "To use Zed's assistant with Bedrock, you need to add the Access Key ID, Secret Access Key and AWS Region. Follow these steps:",
964 "- Create a user and security credentials here:",
965 "- Grant that user permissions according to this documentation:",
966 "- Go to the console and select the models you would like access to:",
967 "- Fill the fields below and hit enter to use the assistant:",
968 ];
969 const BEDROCK_MODEL_CATALOG_LABEL: &str = "Bedrock Model Catalog";
970 const BEDROCK_IAM_DOCS: &str = "Prerequisites";
971
972 let env_var_set = self.state.read(cx).credentials_from_env;
973
974 let bg_color = cx.theme().colors().editor_background;
975 let border_color = cx.theme().colors().border_variant;
976 let input_base_styles = || {
977 h_flex()
978 .w_full()
979 .px_2()
980 .py_1()
981 .bg(bg_color)
982 .border_1()
983 .border_color(border_color)
984 .rounded_md()
985 };
986
987 if self.load_credentials_task.is_some() {
988 div().child(Label::new("Loading credentials...")).into_any()
989 } else if self.should_render_editor(cx) {
990 v_flex()
991 .size_full()
992 .on_action(cx.listener(ConfigurationView::save_credentials))
993 .child(Label::new(INSTRUCTIONS[0]))
994 .child(h_flex().flex_wrap().child(Label::new(INSTRUCTIONS[1])).child(
995 Button::new("iam_console", IAM_CONSOLE_URL)
996 .style(ButtonStyle::Subtle)
997 .icon(IconName::ArrowUpRight)
998 .icon_size(IconSize::XSmall)
999 .icon_color(Color::Muted)
1000 .on_click(move |_, _window, cx| cx.open_url(IAM_CONSOLE_URL))
1001 )
1002 )
1003 .child(h_flex().flex_wrap().child(Label::new(INSTRUCTIONS[2])).child(
1004 Button::new("bedrock_iam_docs", BEDROCK_IAM_DOCS)
1005 .style(ButtonStyle::Subtle)
1006 .icon(IconName::ArrowUpRight)
1007 .icon_size(IconSize::XSmall)
1008 .icon_color(Color::Muted)
1009 .on_click(move |_, _window, cx| cx.open_url(BEDROCK_DOCS_URL))
1010 ))
1011 .child(h_flex().flex_wrap().child(Label::new(INSTRUCTIONS[3])).child(
1012 Button::new("bedrock_model_catalog", BEDROCK_MODEL_CATALOG_LABEL)
1013 .style(ButtonStyle::Subtle)
1014 .icon(IconName::ArrowUpRight)
1015 .icon_size(IconSize::XSmall)
1016 .icon_color(Color::Muted)
1017 .on_click(move |_, _window, cx| cx.open_url(BEDROCK_MODEL_CATALOG))
1018 ))
1019 .child(Label::new(INSTRUCTIONS[4]))
1020 .child(
1021 v_flex()
1022 .my_2()
1023 .gap_1p5()
1024 .child(
1025 v_flex()
1026 .gap_0p5()
1027 .child(Label::new("Access Key ID").size(LabelSize::Small))
1028 .child(
1029 input_base_styles().child(self.render_aa_id_editor(cx))
1030 )
1031 )
1032 .child(
1033 v_flex()
1034 .gap_0p5()
1035 .child(Label::new("Secret Access Key").size(LabelSize::Small))
1036 .child(
1037 input_base_styles().child(self.render_sk_editor(cx))
1038 )
1039 )
1040 .child(
1041 v_flex()
1042 .gap_0p5()
1043 .child(Label::new("Region").size(LabelSize::Small))
1044 .child(
1045 input_base_styles().child(self.render_region_editor(cx))
1046 )
1047 )
1048 )
1049 .child(
1050 Label::new(
1051 format!("You can also assign the {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR}, and {ZED_BEDROCK_REGION_VAR} environment variables and restart Zed."),
1052 )
1053 .size(LabelSize::Small),
1054 )
1055 .into_any()
1056 } else {
1057 h_flex()
1058 .size_full()
1059 .justify_between()
1060 .child(
1061 h_flex()
1062 .gap_1()
1063 .child(Icon::new(IconName::Check).color(Color::Success))
1064 .child(Label::new(if env_var_set {
1065 format!("Access Key ID is set in {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, Secret Key is set in {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR}, Region is set in {ZED_BEDROCK_REGION_VAR} environment variables.")
1066 } else {
1067 "Credentials configured.".to_string()
1068 })),
1069 )
1070 .child(
1071 Button::new("reset-key", "Reset key")
1072 .icon(Some(IconName::Trash))
1073 .icon_size(IconSize::Small)
1074 .icon_position(IconPosition::Start)
1075 .disabled(env_var_set)
1076 .when(env_var_set, |this| {
1077 this.tooltip(Tooltip::text(format!("To reset your credentials, unset the {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR}, and {ZED_BEDROCK_REGION_VAR} environment variables.")))
1078 })
1079 .on_click(cx.listener(|this, _, window, cx| this.reset_credentials(window, cx))),
1080 )
1081 .into_any()
1082 }
1083 }
1084}