bedrock.rs

   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_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXX";
 776    const PLACEHOLDER_REGION: &'static str = "us-east-1";
 777
 778    fn new(state: gpui::Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
 779        cx.observe(&state, |_, _, cx| {
 780            cx.notify();
 781        })
 782        .detach();
 783
 784        let load_credentials_task = Some(cx.spawn({
 785            let state = state.clone();
 786            |this, mut cx| async move {
 787                if let Some(task) = state
 788                    .update(&mut cx, |state, cx| state.authenticate(cx))
 789                    .log_err()
 790                {
 791                    // We don't log an error, because "not signed in" is also an error.
 792                    let _ = task.await;
 793                }
 794                this.update(&mut cx, |this, cx| {
 795                    this.load_credentials_task = None;
 796                    cx.notify();
 797                })
 798                .log_err();
 799            }
 800        }));
 801
 802        Self {
 803            access_key_id_editor: cx.new(|cx| {
 804                let mut editor = Editor::single_line(window, cx);
 805                editor.set_placeholder_text(Self::PLACEHOLDER_TEXT, cx);
 806                editor
 807            }),
 808            secret_access_key_editor: cx.new(|cx| {
 809                let mut editor = Editor::single_line(window, cx);
 810                editor.set_placeholder_text(Self::PLACEHOLDER_TEXT, cx);
 811                editor
 812            }),
 813            region_editor: cx.new(|cx| {
 814                let mut editor = Editor::single_line(window, cx);
 815                editor.set_placeholder_text(Self::PLACEHOLDER_REGION, cx);
 816                editor
 817            }),
 818            state,
 819            load_credentials_task,
 820        }
 821    }
 822
 823    fn save_credentials(
 824        &mut self,
 825        _: &menu::Confirm,
 826        _window: &mut Window,
 827        cx: &mut Context<Self>,
 828    ) {
 829        let access_key_id = self
 830            .access_key_id_editor
 831            .read(cx)
 832            .text(cx)
 833            .to_string()
 834            .trim()
 835            .to_string();
 836        let secret_access_key = self
 837            .secret_access_key_editor
 838            .read(cx)
 839            .text(cx)
 840            .to_string()
 841            .trim()
 842            .to_string();
 843        let region = self
 844            .region_editor
 845            .read(cx)
 846            .text(cx)
 847            .to_string()
 848            .trim()
 849            .to_string();
 850
 851        let state = self.state.clone();
 852        cx.spawn(|_, mut cx| async move {
 853            state
 854                .update(&mut cx, |state, cx| {
 855                    let credentials: BedrockCredentials = BedrockCredentials {
 856                        access_key_id: access_key_id.clone(),
 857                        secret_access_key: secret_access_key.clone(),
 858                        region: region.clone(),
 859                    };
 860
 861                    state.set_credentials(credentials, cx)
 862                })?
 863                .await
 864        })
 865        .detach_and_log_err(cx);
 866    }
 867
 868    fn reset_credentials(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 869        self.access_key_id_editor
 870            .update(cx, |editor, cx| editor.set_text("", window, cx));
 871        self.secret_access_key_editor
 872            .update(cx, |editor, cx| editor.set_text("", window, cx));
 873        self.region_editor
 874            .update(cx, |editor, cx| editor.set_text("", window, cx));
 875
 876        let state = self.state.clone();
 877        cx.spawn(|_, mut cx| async move {
 878            state
 879                .update(&mut cx, |state, cx| state.reset_credentials(cx))?
 880                .await
 881        })
 882        .detach_and_log_err(cx);
 883    }
 884
 885    fn make_text_style(&self, cx: &Context<Self>) -> TextStyle {
 886        let settings = ThemeSettings::get_global(cx);
 887        TextStyle {
 888            color: cx.theme().colors().text,
 889            font_family: settings.ui_font.family.clone(),
 890            font_features: settings.ui_font.features.clone(),
 891            font_fallbacks: settings.ui_font.fallbacks.clone(),
 892            font_size: rems(0.875).into(),
 893            font_weight: settings.ui_font.weight,
 894            font_style: FontStyle::Normal,
 895            line_height: relative(1.3),
 896            background_color: None,
 897            underline: None,
 898            strikethrough: None,
 899            white_space: WhiteSpace::Normal,
 900            text_overflow: None,
 901            text_align: Default::default(),
 902            line_clamp: None,
 903        }
 904    }
 905
 906    fn render_aa_id_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
 907        let text_style = self.make_text_style(cx);
 908
 909        EditorElement::new(
 910            &self.access_key_id_editor,
 911            EditorStyle {
 912                background: cx.theme().colors().editor_background,
 913                local_player: cx.theme().players().local(),
 914                text: text_style,
 915                ..Default::default()
 916            },
 917        )
 918    }
 919
 920    fn render_sk_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
 921        let text_style = self.make_text_style(cx);
 922
 923        EditorElement::new(
 924            &self.secret_access_key_editor,
 925            EditorStyle {
 926                background: cx.theme().colors().editor_background,
 927                local_player: cx.theme().players().local(),
 928                text: text_style,
 929                ..Default::default()
 930            },
 931        )
 932    }
 933
 934    fn render_region_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
 935        let text_style = self.make_text_style(cx);
 936
 937        EditorElement::new(
 938            &self.region_editor,
 939            EditorStyle {
 940                background: cx.theme().colors().editor_background,
 941                local_player: cx.theme().players().local(),
 942                text: text_style,
 943                ..Default::default()
 944            },
 945        )
 946    }
 947
 948    fn should_render_editor(&self, cx: &mut Context<Self>) -> bool {
 949        !self.state.read(cx).is_authenticated()
 950    }
 951}
 952
 953impl Render for ConfigurationView {
 954    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
 955        const IAM_CONSOLE_URL: &str = "https://us-east-1.console.aws.amazon.com/iam/home";
 956        const INSTRUCTIONS: [&str; 3] = [
 957            "To use Zed's assistant with Bedrock, you need to add the Access Key ID, Secret Access Key and AWS Region. Follow these steps:",
 958            "- Create a pair at:",
 959            "- Paste your Access Key ID, Secret Key, and Region below and hit enter to use the assistant:",
 960        ];
 961        let env_var_set = self.state.read(cx).credentials_from_env;
 962
 963        let bg_color = cx.theme().colors().editor_background;
 964        let border_color = cx.theme().colors().border_variant;
 965        let input_base_styles = || {
 966            h_flex()
 967                .w_full()
 968                .px_2()
 969                .py_1()
 970                .bg(bg_color)
 971                .border_1()
 972                .border_color(border_color)
 973                .rounded_md()
 974        };
 975
 976        if self.load_credentials_task.is_some() {
 977            div().child(Label::new("Loading credentials...")).into_any()
 978        } else if self.should_render_editor(cx) {
 979            v_flex()
 980                .size_full()
 981                .on_action(cx.listener(ConfigurationView::save_credentials))
 982                .child(Label::new(INSTRUCTIONS[0]))
 983                .child(h_flex().child(Label::new(INSTRUCTIONS[1])).child(
 984                    Button::new("iam_console", IAM_CONSOLE_URL)
 985                        .style(ButtonStyle::Subtle)
 986                        .icon(IconName::ArrowUpRight)
 987                        .icon_size(IconSize::XSmall)
 988                        .icon_color(Color::Muted)
 989                        .on_click(move |_, _window, cx| cx.open_url(IAM_CONSOLE_URL))
 990                )
 991                )
 992                .child(Label::new(INSTRUCTIONS[2]))
 993                .child(
 994                    v_flex()
 995                        .my_2()
 996                        .gap_1()
 997                        .child(input_base_styles().child(self.render_aa_id_editor(cx)))
 998                        .child(input_base_styles().child(self.render_sk_editor(cx)))
 999                        .child(input_base_styles().child(self.render_region_editor(cx)))
1000                )
1001                .child(
1002                    Label::new(
1003                        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."),
1004                    )
1005                        .size(LabelSize::Small),
1006                )
1007                .into_any()
1008        } else {
1009            h_flex()
1010                .size_full()
1011                .justify_between()
1012                .child(
1013                    h_flex()
1014                        .gap_1()
1015                        .child(Icon::new(IconName::Check).color(Color::Success))
1016                        .child(Label::new(if env_var_set {
1017                            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.")
1018                        } else {
1019                            "Credentials configured.".to_string()
1020                        })),
1021                )
1022                .child(
1023                    Button::new("reset-key", "Reset key")
1024                        .icon(Some(IconName::Trash))
1025                        .icon_size(IconSize::Small)
1026                        .icon_position(IconPosition::Start)
1027                        .disabled(env_var_set)
1028                        .when(env_var_set, |this| {
1029                            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.")))
1030                        })
1031                        .on_click(cx.listener(|this, _, window, cx| this.reset_credentials(window, cx))),
1032                )
1033                .into_any()
1034        }
1035    }
1036}