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 region: Option<String>,
87 _subscription: Subscription,
88}
89
90impl State {
91 fn reset_credentials(&self, cx: &mut Context<Self>) -> Task<Result<()>> {
92 let credentials_provider = <dyn CredentialsProvider>::global(cx);
93 cx.spawn(|this, mut cx| async move {
94 credentials_provider
95 .delete_credentials(AMAZON_AWS_URL, &cx)
96 .await
97 .log_err();
98 this.update(&mut cx, |this, cx| {
99 this.credentials = None;
100 this.credentials_from_env = false;
101 cx.notify();
102 })
103 })
104 }
105
106 fn set_credentials(
107 &mut self,
108 credentials: BedrockCredentials,
109 cx: &mut Context<Self>,
110 ) -> Task<Result<()>> {
111 let credentials_provider = <dyn CredentialsProvider>::global(cx);
112 cx.spawn(|this, mut cx| async move {
113 credentials_provider
114 .write_credentials(
115 AMAZON_AWS_URL,
116 "Bearer",
117 &serde_json::to_vec(&credentials)?,
118 &cx,
119 )
120 .await?;
121 this.update(&mut cx, |this, cx| {
122 this.credentials = Some(credentials);
123 cx.notify();
124 })
125 })
126 }
127
128 fn is_authenticated(&self) -> bool {
129 self.credentials.is_some()
130 }
131
132 fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
133 if self.is_authenticated() {
134 return Task::ready(Ok(()));
135 }
136
137 let credentials_provider = <dyn CredentialsProvider>::global(cx);
138 cx.spawn(|this, mut cx| async move {
139 let (credentials, from_env) =
140 if let Ok(credentials) = std::env::var(ZED_AWS_CREDENTIALS_VAR) {
141 (credentials, true)
142 } else {
143 let (_, credentials) = credentials_provider
144 .read_credentials(AMAZON_AWS_URL, &cx)
145 .await?
146 .ok_or_else(|| AuthenticateError::CredentialsNotFound)?;
147 (
148 String::from_utf8(credentials)
149 .context("invalid {PROVIDER_NAME} credentials")?,
150 false,
151 )
152 };
153
154 let credentials: BedrockCredentials =
155 serde_json::from_str(&credentials).context("failed to parse credentials")?;
156
157 this.update(&mut cx, |this, cx| {
158 this.credentials = Some(credentials);
159 this.credentials_from_env = from_env;
160 cx.notify();
161 })?;
162
163 Ok(())
164 })
165 }
166}
167
168pub struct BedrockLanguageModelProvider {
169 http_client: AwsHttpClient,
170 handler: tokio::runtime::Handle,
171 state: gpui::Entity<State>,
172}
173
174impl BedrockLanguageModelProvider {
175 pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self {
176 let state = cx.new(|cx| State {
177 credentials: None,
178 region: Some(String::from("us-east-1")),
179 credentials_from_env: false,
180 _subscription: cx.observe_global::<SettingsStore>(|_, cx| {
181 cx.notify();
182 }),
183 });
184
185 let tokio_handle = Tokio::handle(cx);
186
187 let coerced_client = AwsHttpClient::new(http_client.clone(), tokio_handle.clone());
188
189 Self {
190 http_client: coerced_client,
191 handler: tokio_handle.clone(),
192 state,
193 }
194 }
195}
196
197impl LanguageModelProvider for BedrockLanguageModelProvider {
198 fn id(&self) -> LanguageModelProviderId {
199 LanguageModelProviderId(PROVIDER_ID.into())
200 }
201
202 fn name(&self) -> LanguageModelProviderName {
203 LanguageModelProviderName(PROVIDER_NAME.into())
204 }
205
206 fn icon(&self) -> IconName {
207 IconName::AiBedrock
208 }
209
210 fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
211 let model = bedrock::Model::default();
212 Some(Arc::new(BedrockModel {
213 id: LanguageModelId::from(model.id().to_string()),
214 model,
215 http_client: self.http_client.clone(),
216 handler: self.handler.clone(),
217 state: self.state.clone(),
218 request_limiter: RateLimiter::new(4),
219 }))
220 }
221
222 fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
223 let mut models = BTreeMap::default();
224
225 for model in bedrock::Model::iter() {
226 if !matches!(model, bedrock::Model::Custom { .. }) {
227 models.insert(model.id().to_string(), model);
228 }
229 }
230
231 // Override with available models from settings
232 for model in AllLanguageModelSettings::get_global(cx)
233 .bedrock
234 .available_models
235 .iter()
236 {
237 models.insert(
238 model.name.clone(),
239 bedrock::Model::Custom {
240 name: model.name.clone(),
241 display_name: model.display_name.clone(),
242 max_tokens: model.max_tokens,
243 max_output_tokens: model.max_output_tokens,
244 default_temperature: model.default_temperature,
245 },
246 );
247 }
248
249 models
250 .into_values()
251 .map(|model| {
252 Arc::new(BedrockModel {
253 id: LanguageModelId::from(model.id().to_string()),
254 model,
255 http_client: self.http_client.clone(),
256 handler: self.handler.clone(),
257 state: self.state.clone(),
258 request_limiter: RateLimiter::new(4),
259 }) as Arc<dyn LanguageModel>
260 })
261 .collect()
262 }
263
264 fn is_authenticated(&self, cx: &App) -> bool {
265 self.state.read(cx).is_authenticated()
266 }
267
268 fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
269 self.state.update(cx, |state, cx| state.authenticate(cx))
270 }
271
272 fn configuration_view(&self, window: &mut Window, cx: &mut App) -> AnyView {
273 cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx))
274 .into()
275 }
276
277 fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> {
278 self.state
279 .update(cx, |state, cx| state.reset_credentials(cx))
280 }
281}
282
283impl LanguageModelProviderState for BedrockLanguageModelProvider {
284 type ObservableEntity = State;
285
286 fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> {
287 Some(self.state.clone())
288 }
289}
290
291struct BedrockModel {
292 id: LanguageModelId,
293 model: Model,
294 http_client: AwsHttpClient,
295 handler: tokio::runtime::Handle,
296 state: gpui::Entity<State>,
297 request_limiter: RateLimiter,
298}
299
300impl BedrockModel {
301 fn stream_completion(
302 &self,
303 request: bedrock::Request,
304 cx: &AsyncApp,
305 ) -> Result<
306 BoxFuture<'static, BoxStream<'static, Result<BedrockStreamingResponse, BedrockError>>>,
307 > {
308 let Ok(Ok((access_key_id, secret_access_key, region))) =
309 cx.read_entity(&self.state, |state, _cx| {
310 if let Some(credentials) = &state.credentials {
311 Ok((
312 credentials.access_key_id.clone(),
313 credentials.secret_access_key.clone(),
314 state.region.clone(),
315 ))
316 } else {
317 return Err(anyhow!("Failed to read credentials"));
318 }
319 })
320 else {
321 return Err(anyhow!("App state dropped"));
322 };
323
324 let runtime_client = bedrock_client::Client::from_conf(
325 Config::builder()
326 .stalled_stream_protection(StalledStreamProtectionConfig::disabled())
327 .credentials_provider(Credentials::new(
328 access_key_id,
329 secret_access_key,
330 None,
331 None,
332 "Keychain",
333 ))
334 .region(Region::new(region.unwrap()))
335 .http_client(self.http_client.clone())
336 .build(),
337 );
338
339 let owned_handle = self.handler.clone();
340
341 Ok(async move {
342 let request = bedrock::stream_completion(runtime_client, request, owned_handle);
343 request.await.unwrap_or_else(|e| {
344 futures::stream::once(async move { Err(BedrockError::ClientError(e)) }).boxed()
345 })
346 }
347 .boxed())
348 }
349}
350
351impl LanguageModel for BedrockModel {
352 fn id(&self) -> LanguageModelId {
353 self.id.clone()
354 }
355
356 fn name(&self) -> LanguageModelName {
357 LanguageModelName::from(self.model.display_name().to_string())
358 }
359
360 fn provider_id(&self) -> LanguageModelProviderId {
361 LanguageModelProviderId(PROVIDER_ID.into())
362 }
363
364 fn provider_name(&self) -> LanguageModelProviderName {
365 LanguageModelProviderName(PROVIDER_NAME.into())
366 }
367
368 fn telemetry_id(&self) -> String {
369 format!("bedrock/{}", self.model.id())
370 }
371
372 fn max_token_count(&self) -> usize {
373 self.model.max_token_count()
374 }
375
376 fn max_output_tokens(&self) -> Option<u32> {
377 Some(self.model.max_output_tokens())
378 }
379
380 fn count_tokens(
381 &self,
382 request: LanguageModelRequest,
383 cx: &App,
384 ) -> BoxFuture<'static, Result<usize>> {
385 get_bedrock_tokens(request, cx)
386 }
387
388 fn stream_completion(
389 &self,
390 request: LanguageModelRequest,
391 cx: &AsyncApp,
392 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<LanguageModelCompletionEvent>>>> {
393 let request = into_bedrock(
394 request,
395 self.model.id().into(),
396 self.model.default_temperature(),
397 self.model.max_output_tokens(),
398 );
399
400 let owned_handle = self.handler.clone();
401
402 let request = self.stream_completion(request, cx);
403 let future = self.request_limiter.stream(async move {
404 let response = request.map_err(|e| anyhow!(e)).unwrap().await;
405 Ok(map_to_language_model_completion_events(
406 response,
407 owned_handle,
408 ))
409 });
410 async move { Ok(future.await?.boxed()) }.boxed()
411 }
412
413 fn use_any_tool(
414 &self,
415 request: LanguageModelRequest,
416 name: String,
417 description: String,
418 schema: Value,
419 _cx: &AsyncApp,
420 ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {
421 let mut request = into_bedrock(
422 request,
423 self.model.id().into(),
424 self.model.default_temperature(),
425 self.model.max_output_tokens(),
426 );
427
428 request.tool_choice = Some(BedrockToolChoice::Tool(
429 BedrockSpecificTool::builder()
430 .name(name.clone())
431 .build()
432 .unwrap(),
433 ));
434
435 request.tools = vec![BedrockTool::builder()
436 .name(name.clone())
437 .description(description.clone())
438 .input_schema(BedrockToolInputSchema::Json(value_to_aws_document(&schema)))
439 .build()
440 .unwrap()];
441
442 let handle = self.handler.clone();
443
444 let request = self.stream_completion(request, _cx);
445 self.request_limiter
446 .run(async move {
447 let response = request.map_err(|e| anyhow!(e)).unwrap().await;
448 Ok(extract_tool_args_from_events(name, response, handle)
449 .await?
450 .boxed())
451 })
452 .boxed()
453 }
454
455 fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
456 None
457 }
458}
459
460pub fn into_bedrock(
461 request: LanguageModelRequest,
462 model: String,
463 default_temperature: f32,
464 max_output_tokens: u32,
465) -> bedrock::Request {
466 let mut new_messages: Vec<BedrockMessage> = Vec::new();
467 let mut system_message = String::new();
468
469 for message in request.messages {
470 if message.contents_empty() {
471 continue;
472 }
473
474 match message.role {
475 Role::User | Role::Assistant => {
476 let bedrock_message_content: Vec<BedrockInnerContent> = message
477 .content
478 .into_iter()
479 .filter_map(|content| match content {
480 MessageContent::Text(text) => {
481 if !text.is_empty() {
482 Some(BedrockInnerContent::Text(text))
483 } else {
484 None
485 }
486 }
487 _ => None,
488 })
489 .collect();
490 let bedrock_role = match message.role {
491 Role::User => bedrock::BedrockRole::User,
492 Role::Assistant => bedrock::BedrockRole::Assistant,
493 Role::System => unreachable!("System role should never occur here"),
494 };
495 if let Some(last_message) = new_messages.last_mut() {
496 if last_message.role == bedrock_role {
497 last_message.content.extend(bedrock_message_content);
498 continue;
499 }
500 }
501 new_messages.push(
502 BedrockMessage::builder()
503 .role(bedrock_role)
504 .set_content(Some(bedrock_message_content))
505 .build()
506 .expect("failed to build Bedrock message"),
507 );
508 }
509 Role::System => {
510 if !system_message.is_empty() {
511 system_message.push_str("\n\n");
512 }
513 system_message.push_str(&message.string_contents());
514 }
515 }
516 }
517
518 bedrock::Request {
519 model,
520 messages: new_messages,
521 max_tokens: max_output_tokens,
522 system: Some(system_message),
523 tools: vec![],
524 tool_choice: None,
525 metadata: None,
526 stop_sequences: Vec::new(),
527 temperature: request.temperature.or(Some(default_temperature)),
528 top_k: None,
529 top_p: None,
530 }
531}
532
533// TODO: just call the ConverseOutput.usage() method:
534// https://docs.rs/aws-sdk-bedrockruntime/latest/aws_sdk_bedrockruntime/operation/converse/struct.ConverseOutput.html#method.output
535pub fn get_bedrock_tokens(
536 request: LanguageModelRequest,
537 cx: &App,
538) -> BoxFuture<'static, Result<usize>> {
539 cx.background_executor()
540 .spawn(async move {
541 let messages = request.messages;
542 let mut tokens_from_images = 0;
543 let mut string_messages = Vec::with_capacity(messages.len());
544
545 for message in messages {
546 use language_model::MessageContent;
547
548 let mut string_contents = String::new();
549
550 for content in message.content {
551 match content {
552 MessageContent::Text(text) => {
553 string_contents.push_str(&text);
554 }
555 MessageContent::Image(image) => {
556 tokens_from_images += image.estimate_tokens();
557 }
558 MessageContent::ToolUse(_tool_use) => {
559 // TODO: Estimate token usage from tool uses.
560 }
561 MessageContent::ToolResult(tool_result) => {
562 string_contents.push_str(&tool_result.content);
563 }
564 }
565 }
566
567 if !string_contents.is_empty() {
568 string_messages.push(tiktoken_rs::ChatCompletionRequestMessage {
569 role: match message.role {
570 Role::User => "user".into(),
571 Role::Assistant => "assistant".into(),
572 Role::System => "system".into(),
573 },
574 content: Some(string_contents),
575 name: None,
576 function_call: None,
577 });
578 }
579 }
580
581 // Tiktoken doesn't yet support these models, so we manually use the
582 // same tokenizer as GPT-4.
583 tiktoken_rs::num_tokens_from_messages("gpt-4", &string_messages)
584 .map(|tokens| tokens + tokens_from_images)
585 })
586 .boxed()
587}
588
589pub async fn extract_tool_args_from_events(
590 name: String,
591 mut events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
592 handle: Handle,
593) -> Result<impl Send + Stream<Item = Result<String>>> {
594 handle
595 .spawn(async move {
596 let mut tool_use_index = None;
597 while let Some(event) = events.next().await {
598 if let BedrockStreamingResponse::ContentBlockStart(ContentBlockStartEvent {
599 content_block_index,
600 start,
601 ..
602 }) = event?
603 {
604 match start {
605 None => {
606 continue;
607 }
608 Some(start) => match start.as_tool_use() {
609 Ok(tool_use) => {
610 if name == tool_use.name {
611 tool_use_index = Some(content_block_index);
612 break;
613 }
614 }
615 Err(err) => {
616 return Err(anyhow!("Failed to parse tool use event: {:?}", err));
617 }
618 },
619 }
620 }
621 }
622
623 let Some(tool_use_index) = tool_use_index else {
624 return Err(anyhow!("Tool is not used"));
625 };
626
627 Ok(events.filter_map(move |event| {
628 let result = match event {
629 Err(_err) => None,
630 Ok(output) => match output.clone() {
631 BedrockStreamingResponse::ContentBlockDelta(inner) => {
632 match inner.clone().delta {
633 Some(ContentBlockDelta::ToolUse(tool_use)) => {
634 if inner.content_block_index == tool_use_index {
635 Some(Ok(tool_use.input))
636 } else {
637 None
638 }
639 }
640 _ => None,
641 }
642 }
643 _ => None,
644 },
645 };
646
647 async move { result }
648 }))
649 })
650 .await?
651}
652
653pub fn map_to_language_model_completion_events(
654 events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
655 handle: Handle,
656) -> impl Stream<Item = Result<LanguageModelCompletionEvent>> {
657 struct RawToolUse {
658 id: String,
659 name: String,
660 input_json: String,
661 }
662
663 struct State {
664 events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
665 tool_uses_by_index: HashMap<i32, RawToolUse>,
666 }
667
668 futures::stream::unfold(
669 State {
670 events,
671 tool_uses_by_index: HashMap::default(),
672 },
673 move |mut state: State| {
674 let inner_handle = handle.clone();
675 async move {
676 inner_handle
677 .spawn(async {
678 while let Some(event) = state.events.next().await {
679 match event {
680 Ok(event) => match event {
681 ConverseStreamOutput::ContentBlockDelta(cb_delta) => {
682 if let Some(ContentBlockDelta::Text(text_out)) =
683 cb_delta.delta
684 {
685 return Some((
686 Some(Ok(LanguageModelCompletionEvent::Text(
687 text_out,
688 ))),
689 state,
690 ));
691 } else if let Some(ContentBlockDelta::ToolUse(text_out)) =
692 cb_delta.delta
693 {
694 if let Some(tool_use) = state
695 .tool_uses_by_index
696 .get_mut(&cb_delta.content_block_index)
697 {
698 tool_use.input_json.push_str(text_out.input());
699 return Some((None, state));
700 };
701
702 return Some((None, state));
703 } else if cb_delta.delta.is_none() {
704 return Some((None, state));
705 }
706 }
707 ConverseStreamOutput::ContentBlockStart(cb_start) => {
708 if let Some(start) = cb_start.start {
709 match start {
710 ContentBlockStart::ToolUse(text_out) => {
711 let tool_use = RawToolUse {
712 id: text_out.tool_use_id,
713 name: text_out.name,
714 input_json: String::new(),
715 };
716
717 state.tool_uses_by_index.insert(
718 cb_start.content_block_index,
719 tool_use,
720 );
721 }
722 _ => {}
723 }
724 }
725 }
726 ConverseStreamOutput::ContentBlockStop(cb_stop) => {
727 if let Some(tool_use) = state
728 .tool_uses_by_index
729 .remove(&cb_stop.content_block_index)
730 {
731 return Some((
732 Some(maybe!({
733 Ok(LanguageModelCompletionEvent::ToolUse(
734 LanguageModelToolUse {
735 id: tool_use.id.into(),
736 name: tool_use.name,
737 input: if tool_use.input_json.is_empty()
738 {
739 Value::Null
740 } else {
741 serde_json::Value::from_str(
742 &tool_use.input_json,
743 )
744 .map_err(|err| anyhow!(err))?
745 },
746 },
747 ))
748 })),
749 state,
750 ));
751 }
752 }
753 _ => {}
754 },
755 Err(err) => return Some((Some(Err(anyhow!(err))), state)),
756 }
757 }
758 None
759 })
760 .await
761 .unwrap()
762 }
763 },
764 )
765 .filter_map(|event| async move { event })
766}
767
768struct ConfigurationView {
769 access_key_id_editor: Entity<Editor>,
770 secret_access_key_editor: Entity<Editor>,
771 region_editor: Entity<Editor>,
772 state: gpui::Entity<State>,
773 load_credentials_task: Option<Task<()>>,
774}
775
776impl ConfigurationView {
777 const PLACEHOLDER_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXX";
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_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_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 INSTRUCTIONS: [&str; 3] = [
959 "To use Zed's assistant with Bedrock, you need to add the Access Key ID, Secret Access Key and AWS Region. Follow these steps:",
960 "- Create a pair at:",
961 "- Paste your Access Key ID, Secret Key, and Region below and hit enter to use the assistant:",
962 ];
963 let env_var_set = self.state.read(cx).credentials_from_env;
964
965 if self.load_credentials_task.is_some() {
966 div().child(Label::new("Loading credentials...")).into_any()
967 } else if self.should_render_editor(cx) {
968 v_flex()
969 .size_full()
970 .on_action(cx.listener(Self::save_credentials))
971 .child(Label::new(INSTRUCTIONS[0]))
972 .child(h_flex().child(Label::new(INSTRUCTIONS[1])).child(
973 Button::new("iam_console", IAM_CONSOLE_URL)
974 .style(ButtonStyle::Subtle)
975 .icon(IconName::ExternalLink)
976 .icon_size(IconSize::XSmall)
977 .icon_color(Color::Muted)
978 .on_click(move |_, _window, cx| cx.open_url(IAM_CONSOLE_URL))
979 )
980 )
981 .child(Label::new(INSTRUCTIONS[2]))
982 .child(
983 h_flex()
984 .gap_1()
985 .child(self.render_aa_id_editor(cx))
986 .child(self.render_sk_editor(cx))
987 .child(self.render_region_editor(cx))
988 )
989 .child(
990 Label::new(
991 format!("You can also assign the {ZED_BEDROCK_ACCESS_KEY_ID_VAR}, {ZED_BEDROCK_SECRET_ACCESS_KEY_VAR} and {ZED_BEDROCK_REGION_VAR} environment variable and restart Zed."),
992 )
993 .size(LabelSize::Small),
994 )
995 .into_any()
996 } else {
997 h_flex()
998 .size_full()
999 .justify_between()
1000 .child(
1001 h_flex()
1002 .gap_1()
1003 .child(Icon::new(IconName::Check).color(Color::Success))
1004 .child(Label::new(if env_var_set {
1005 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.")
1006 } else {
1007 "Credentials configured.".to_string()
1008 })),
1009 )
1010 .child(
1011 Button::new("reset-key", "Reset key")
1012 .icon(Some(IconName::Trash))
1013 .icon_size(IconSize::Small)
1014 .icon_position(IconPosition::Start)
1015 .disabled(env_var_set)
1016 .when(env_var_set, |this| {
1017 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.")))
1018 })
1019 .on_click(cx.listener(|this, _, window, cx| this.reset_credentials(window, cx))),
1020 )
1021 .into_any()
1022 }
1023 }
1024}