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    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(|err| anyhow!(err))?.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 = BedrockSpecificTool::builder()
 429            .name(name.clone())
 430            .build()
 431            .log_err()
 432            .map(BedrockToolChoice::Tool);
 433
 434        if let Some(tool) = BedrockTool::builder()
 435            .name(name.clone())
 436            .description(description.clone())
 437            .input_schema(BedrockToolInputSchema::Json(value_to_aws_document(&schema)))
 438            .build()
 439            .log_err()
 440        {
 441            request.tools.push(tool);
 442        }
 443
 444        let handle = self.handler.clone();
 445
 446        let request = self.stream_completion(request, _cx);
 447        self.request_limiter
 448            .run(async move {
 449                let response = request.map_err(|err| anyhow!(err))?.await;
 450                Ok(extract_tool_args_from_events(name, response, handle)
 451                    .await?
 452                    .boxed())
 453            })
 454            .boxed()
 455    }
 456
 457    fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {
 458        None
 459    }
 460}
 461
 462pub fn into_bedrock(
 463    request: LanguageModelRequest,
 464    model: String,
 465    default_temperature: f32,
 466    max_output_tokens: u32,
 467) -> bedrock::Request {
 468    let mut new_messages: Vec<BedrockMessage> = Vec::new();
 469    let mut system_message = String::new();
 470
 471    for message in request.messages {
 472        if message.contents_empty() {
 473            continue;
 474        }
 475
 476        match message.role {
 477            Role::User | Role::Assistant => {
 478                let bedrock_message_content: Vec<BedrockInnerContent> = message
 479                    .content
 480                    .into_iter()
 481                    .filter_map(|content| match content {
 482                        MessageContent::Text(text) => {
 483                            if !text.is_empty() {
 484                                Some(BedrockInnerContent::Text(text))
 485                            } else {
 486                                None
 487                            }
 488                        }
 489                        _ => None,
 490                    })
 491                    .collect();
 492                let bedrock_role = match message.role {
 493                    Role::User => bedrock::BedrockRole::User,
 494                    Role::Assistant => bedrock::BedrockRole::Assistant,
 495                    Role::System => unreachable!("System role should never occur here"),
 496                };
 497                if let Some(last_message) = new_messages.last_mut() {
 498                    if last_message.role == bedrock_role {
 499                        last_message.content.extend(bedrock_message_content);
 500                        continue;
 501                    }
 502                }
 503                new_messages.push(
 504                    BedrockMessage::builder()
 505                        .role(bedrock_role)
 506                        .set_content(Some(bedrock_message_content))
 507                        .build()
 508                        .expect("failed to build Bedrock message"),
 509                );
 510            }
 511            Role::System => {
 512                if !system_message.is_empty() {
 513                    system_message.push_str("\n\n");
 514                }
 515                system_message.push_str(&message.string_contents());
 516            }
 517        }
 518    }
 519
 520    bedrock::Request {
 521        model,
 522        messages: new_messages,
 523        max_tokens: max_output_tokens,
 524        system: Some(system_message),
 525        tools: vec![],
 526        tool_choice: None,
 527        metadata: None,
 528        stop_sequences: Vec::new(),
 529        temperature: request.temperature.or(Some(default_temperature)),
 530        top_k: None,
 531        top_p: None,
 532    }
 533}
 534
 535// TODO: just call the ConverseOutput.usage() method:
 536// https://docs.rs/aws-sdk-bedrockruntime/latest/aws_sdk_bedrockruntime/operation/converse/struct.ConverseOutput.html#method.output
 537pub fn get_bedrock_tokens(
 538    request: LanguageModelRequest,
 539    cx: &App,
 540) -> BoxFuture<'static, Result<usize>> {
 541    cx.background_executor()
 542        .spawn(async move {
 543            let messages = request.messages;
 544            let mut tokens_from_images = 0;
 545            let mut string_messages = Vec::with_capacity(messages.len());
 546
 547            for message in messages {
 548                use language_model::MessageContent;
 549
 550                let mut string_contents = String::new();
 551
 552                for content in message.content {
 553                    match content {
 554                        MessageContent::Text(text) => {
 555                            string_contents.push_str(&text);
 556                        }
 557                        MessageContent::Image(image) => {
 558                            tokens_from_images += image.estimate_tokens();
 559                        }
 560                        MessageContent::ToolUse(_tool_use) => {
 561                            // TODO: Estimate token usage from tool uses.
 562                        }
 563                        MessageContent::ToolResult(tool_result) => {
 564                            string_contents.push_str(&tool_result.content);
 565                        }
 566                    }
 567                }
 568
 569                if !string_contents.is_empty() {
 570                    string_messages.push(tiktoken_rs::ChatCompletionRequestMessage {
 571                        role: match message.role {
 572                            Role::User => "user".into(),
 573                            Role::Assistant => "assistant".into(),
 574                            Role::System => "system".into(),
 575                        },
 576                        content: Some(string_contents),
 577                        name: None,
 578                        function_call: None,
 579                    });
 580                }
 581            }
 582
 583            // Tiktoken doesn't yet support these models, so we manually use the
 584            // same tokenizer as GPT-4.
 585            tiktoken_rs::num_tokens_from_messages("gpt-4", &string_messages)
 586                .map(|tokens| tokens + tokens_from_images)
 587        })
 588        .boxed()
 589}
 590
 591pub async fn extract_tool_args_from_events(
 592    name: String,
 593    mut events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
 594    handle: Handle,
 595) -> Result<impl Send + Stream<Item = Result<String>>> {
 596    handle
 597        .spawn(async move {
 598            let mut tool_use_index = None;
 599            while let Some(event) = events.next().await {
 600                if let BedrockStreamingResponse::ContentBlockStart(ContentBlockStartEvent {
 601                    content_block_index,
 602                    start,
 603                    ..
 604                }) = event?
 605                {
 606                    match start {
 607                        None => {
 608                            continue;
 609                        }
 610                        Some(start) => match start.as_tool_use() {
 611                            Ok(tool_use) => {
 612                                if name == tool_use.name {
 613                                    tool_use_index = Some(content_block_index);
 614                                    break;
 615                                }
 616                            }
 617                            Err(err) => {
 618                                return Err(anyhow!("Failed to parse tool use event: {:?}", err));
 619                            }
 620                        },
 621                    }
 622                }
 623            }
 624
 625            let Some(tool_use_index) = tool_use_index else {
 626                return Err(anyhow!("Tool is not used"));
 627            };
 628
 629            Ok(events.filter_map(move |event| {
 630                let result = match event {
 631                    Err(_err) => None,
 632                    Ok(output) => match output.clone() {
 633                        BedrockStreamingResponse::ContentBlockDelta(inner) => {
 634                            match inner.clone().delta {
 635                                Some(ContentBlockDelta::ToolUse(tool_use)) => {
 636                                    if inner.content_block_index == tool_use_index {
 637                                        Some(Ok(tool_use.input))
 638                                    } else {
 639                                        None
 640                                    }
 641                                }
 642                                _ => None,
 643                            }
 644                        }
 645                        _ => None,
 646                    },
 647                };
 648
 649                async move { result }
 650            }))
 651        })
 652        .await?
 653}
 654
 655pub fn map_to_language_model_completion_events(
 656    events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
 657    handle: Handle,
 658) -> impl Stream<Item = Result<LanguageModelCompletionEvent>> {
 659    struct RawToolUse {
 660        id: String,
 661        name: String,
 662        input_json: String,
 663    }
 664
 665    struct State {
 666        events: Pin<Box<dyn Send + Stream<Item = Result<BedrockStreamingResponse, BedrockError>>>>,
 667        tool_uses_by_index: HashMap<i32, RawToolUse>,
 668    }
 669
 670    futures::stream::unfold(
 671        State {
 672            events,
 673            tool_uses_by_index: HashMap::default(),
 674        },
 675        move |mut state: State| {
 676            let inner_handle = handle.clone();
 677            async move {
 678                inner_handle
 679                    .spawn(async {
 680                        while let Some(event) = state.events.next().await {
 681                            match event {
 682                                Ok(event) => match event {
 683                                    ConverseStreamOutput::ContentBlockDelta(cb_delta) => {
 684                                        if let Some(ContentBlockDelta::Text(text_out)) =
 685                                            cb_delta.delta
 686                                        {
 687                                            return Some((
 688                                                Some(Ok(LanguageModelCompletionEvent::Text(
 689                                                    text_out,
 690                                                ))),
 691                                                state,
 692                                            ));
 693                                        } else if let Some(ContentBlockDelta::ToolUse(text_out)) =
 694                                            cb_delta.delta
 695                                        {
 696                                            if let Some(tool_use) = state
 697                                                .tool_uses_by_index
 698                                                .get_mut(&cb_delta.content_block_index)
 699                                            {
 700                                                tool_use.input_json.push_str(text_out.input());
 701                                                return Some((None, state));
 702                                            };
 703
 704                                            return Some((None, state));
 705                                        } else if cb_delta.delta.is_none() {
 706                                            return Some((None, state));
 707                                        }
 708                                    }
 709                                    ConverseStreamOutput::ContentBlockStart(cb_start) => {
 710                                        if let Some(start) = cb_start.start {
 711                                            match start {
 712                                                ContentBlockStart::ToolUse(text_out) => {
 713                                                    let tool_use = RawToolUse {
 714                                                        id: text_out.tool_use_id,
 715                                                        name: text_out.name,
 716                                                        input_json: String::new(),
 717                                                    };
 718
 719                                                    state.tool_uses_by_index.insert(
 720                                                        cb_start.content_block_index,
 721                                                        tool_use,
 722                                                    );
 723                                                }
 724                                                _ => {}
 725                                            }
 726                                        }
 727                                    }
 728                                    ConverseStreamOutput::ContentBlockStop(cb_stop) => {
 729                                        if let Some(tool_use) = state
 730                                            .tool_uses_by_index
 731                                            .remove(&cb_stop.content_block_index)
 732                                        {
 733                                            return Some((
 734                                                Some(maybe!({
 735                                                    Ok(LanguageModelCompletionEvent::ToolUse(
 736                                                        LanguageModelToolUse {
 737                                                            id: tool_use.id.into(),
 738                                                            name: tool_use.name,
 739                                                            input: if tool_use.input_json.is_empty()
 740                                                            {
 741                                                                Value::Null
 742                                                            } else {
 743                                                                serde_json::Value::from_str(
 744                                                                    &tool_use.input_json,
 745                                                                )
 746                                                                .map_err(|err| anyhow!(err))?
 747                                                            },
 748                                                        },
 749                                                    ))
 750                                                })),
 751                                                state,
 752                                            ));
 753                                        }
 754                                    }
 755                                    _ => {}
 756                                },
 757                                Err(err) => return Some((Some(Err(anyhow!(err))), state)),
 758                            }
 759                        }
 760                        None
 761                    })
 762                    .await
 763                    .log_err()
 764                    .flatten()
 765            }
 766        },
 767    )
 768    .filter_map(|event| async move { event })
 769}
 770
 771struct ConfigurationView {
 772    access_key_id_editor: Entity<Editor>,
 773    secret_access_key_editor: Entity<Editor>,
 774    region_editor: Entity<Editor>,
 775    state: gpui::Entity<State>,
 776    load_credentials_task: Option<Task<()>>,
 777}
 778
 779impl ConfigurationView {
 780    const PLACEHOLDER_TEXT: &'static str = "XXXXXXXXXXXXXXXXXXX";
 781    const PLACEHOLDER_REGION: &'static str = "us-east-1";
 782
 783    fn new(state: gpui::Entity<State>, window: &mut Window, cx: &mut Context<Self>) -> Self {
 784        cx.observe(&state, |_, _, cx| {
 785            cx.notify();
 786        })
 787        .detach();
 788
 789        let load_credentials_task = Some(cx.spawn({
 790            let state = state.clone();
 791            |this, mut cx| async move {
 792                if let Some(task) = state
 793                    .update(&mut cx, |state, cx| state.authenticate(cx))
 794                    .log_err()
 795                {
 796                    // We don't log an error, because "not signed in" is also an error.
 797                    let _ = task.await;
 798                }
 799                this.update(&mut cx, |this, cx| {
 800                    this.load_credentials_task = None;
 801                    cx.notify();
 802                })
 803                .log_err();
 804            }
 805        }));
 806
 807        Self {
 808            access_key_id_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            secret_access_key_editor: cx.new(|cx| {
 814                let mut editor = Editor::single_line(window, cx);
 815                editor.set_placeholder_text(Self::PLACEHOLDER_TEXT, cx);
 816                editor
 817            }),
 818            region_editor: cx.new(|cx| {
 819                let mut editor = Editor::single_line(window, cx);
 820                editor.set_placeholder_text(Self::PLACEHOLDER_REGION, cx);
 821                editor
 822            }),
 823            state,
 824            load_credentials_task,
 825        }
 826    }
 827
 828    fn save_credentials(
 829        &mut self,
 830        _: &menu::Confirm,
 831        _window: &mut Window,
 832        cx: &mut Context<Self>,
 833    ) {
 834        let access_key_id = self
 835            .access_key_id_editor
 836            .read(cx)
 837            .text(cx)
 838            .to_string()
 839            .trim()
 840            .to_string();
 841        let secret_access_key = self
 842            .secret_access_key_editor
 843            .read(cx)
 844            .text(cx)
 845            .to_string()
 846            .trim()
 847            .to_string();
 848        let region = self
 849            .region_editor
 850            .read(cx)
 851            .text(cx)
 852            .to_string()
 853            .trim()
 854            .to_string();
 855
 856        let state = self.state.clone();
 857        cx.spawn(|_, mut cx| async move {
 858            state
 859                .update(&mut cx, |state, cx| {
 860                    let credentials: BedrockCredentials = BedrockCredentials {
 861                        access_key_id: access_key_id.clone(),
 862                        secret_access_key: secret_access_key.clone(),
 863                        region: region.clone(),
 864                    };
 865
 866                    state.set_credentials(credentials, cx)
 867                })?
 868                .await
 869        })
 870        .detach_and_log_err(cx);
 871    }
 872
 873    fn reset_credentials(&mut self, window: &mut Window, cx: &mut Context<Self>) {
 874        self.access_key_id_editor
 875            .update(cx, |editor, cx| editor.set_text("", window, cx));
 876        self.secret_access_key_editor
 877            .update(cx, |editor, cx| editor.set_text("", window, cx));
 878        self.region_editor
 879            .update(cx, |editor, cx| editor.set_text("", window, cx));
 880
 881        let state = self.state.clone();
 882        cx.spawn(|_, mut cx| async move {
 883            state
 884                .update(&mut cx, |state, cx| state.reset_credentials(cx))?
 885                .await
 886        })
 887        .detach_and_log_err(cx);
 888    }
 889
 890    fn make_text_style(&self, cx: &Context<Self>) -> TextStyle {
 891        let settings = ThemeSettings::get_global(cx);
 892        TextStyle {
 893            color: cx.theme().colors().text,
 894            font_family: settings.ui_font.family.clone(),
 895            font_features: settings.ui_font.features.clone(),
 896            font_fallbacks: settings.ui_font.fallbacks.clone(),
 897            font_size: rems(0.875).into(),
 898            font_weight: settings.ui_font.weight,
 899            font_style: FontStyle::Normal,
 900            line_height: relative(1.3),
 901            background_color: None,
 902            underline: None,
 903            strikethrough: None,
 904            white_space: WhiteSpace::Normal,
 905            text_overflow: None,
 906            text_align: Default::default(),
 907            line_clamp: None,
 908        }
 909    }
 910
 911    fn render_aa_id_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
 912        let text_style = self.make_text_style(cx);
 913
 914        EditorElement::new(
 915            &self.access_key_id_editor,
 916            EditorStyle {
 917                background: cx.theme().colors().editor_background,
 918                local_player: cx.theme().players().local(),
 919                text: text_style,
 920                ..Default::default()
 921            },
 922        )
 923    }
 924
 925    fn render_sk_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
 926        let text_style = self.make_text_style(cx);
 927
 928        EditorElement::new(
 929            &self.secret_access_key_editor,
 930            EditorStyle {
 931                background: cx.theme().colors().editor_background,
 932                local_player: cx.theme().players().local(),
 933                text: text_style,
 934                ..Default::default()
 935            },
 936        )
 937    }
 938
 939    fn render_region_editor(&self, cx: &mut Context<Self>) -> impl IntoElement {
 940        let text_style = self.make_text_style(cx);
 941
 942        EditorElement::new(
 943            &self.region_editor,
 944            EditorStyle {
 945                background: cx.theme().colors().editor_background,
 946                local_player: cx.theme().players().local(),
 947                text: text_style,
 948                ..Default::default()
 949            },
 950        )
 951    }
 952
 953    fn should_render_editor(&self, cx: &mut Context<Self>) -> bool {
 954        !self.state.read(cx).is_authenticated()
 955    }
 956}
 957
 958impl Render for ConfigurationView {
 959    fn render(&mut self, _: &mut Window, cx: &mut Context<Self>) -> impl IntoElement {
 960        const IAM_CONSOLE_URL: &str = "https://us-east-1.console.aws.amazon.com/iam/home";
 961        const INSTRUCTIONS: [&str; 3] = [
 962            "To use Zed's assistant with Bedrock, you need to add the Access Key ID, Secret Access Key and AWS Region. Follow these steps:",
 963            "- Create a pair at:",
 964            "- Paste your Access Key ID, Secret Key, and Region below and hit enter to use the assistant:",
 965        ];
 966        let env_var_set = self.state.read(cx).credentials_from_env;
 967
 968        if self.load_credentials_task.is_some() {
 969            div().child(Label::new("Loading credentials...")).into_any()
 970        } else if self.should_render_editor(cx) {
 971            v_flex()
 972                .size_full()
 973                .on_action(cx.listener(Self::save_credentials))
 974                .child(Label::new(INSTRUCTIONS[0]))
 975                .child(h_flex().child(Label::new(INSTRUCTIONS[1])).child(
 976                    Button::new("iam_console", IAM_CONSOLE_URL)
 977                        .style(ButtonStyle::Subtle)
 978                        .icon(IconName::ExternalLink)
 979                        .icon_size(IconSize::XSmall)
 980                        .icon_color(Color::Muted)
 981                        .on_click(move |_, _window, cx| cx.open_url(IAM_CONSOLE_URL))
 982                )
 983                )
 984                .child(Label::new(INSTRUCTIONS[2]))
 985                .child(
 986                    h_flex()
 987                        .gap_1()
 988                        .child(self.render_aa_id_editor(cx))
 989                        .child(self.render_sk_editor(cx))
 990                        .child(self.render_region_editor(cx))
 991                )
 992                .child(
 993                    Label::new(
 994                        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."),
 995                    )
 996                        .size(LabelSize::Small),
 997                )
 998                .into_any()
 999        } else {
1000            h_flex()
1001                .size_full()
1002                .justify_between()
1003                .child(
1004                    h_flex()
1005                        .gap_1()
1006                        .child(Icon::new(IconName::Check).color(Color::Success))
1007                        .child(Label::new(if env_var_set {
1008                            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.")
1009                        } else {
1010                            "Credentials configured.".to_string()
1011                        })),
1012                )
1013                .child(
1014                    Button::new("reset-key", "Reset key")
1015                        .icon(Some(IconName::Trash))
1016                        .icon_size(IconSize::Small)
1017                        .icon_position(IconPosition::Start)
1018                        .disabled(env_var_set)
1019                        .when(env_var_set, |this| {
1020                            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.")))
1021                        })
1022                        .on_click(cx.listener(|this, _, window, cx| this.reset_credentials(window, cx))),
1023                )
1024                .into_any()
1025        }
1026    }
1027}