use std::pin::Pin;use std::str::FromStr;use std::sync::Arc; use crate::ui::InstructionListItem;use anyhow::{anyhow, Context as _, Result};use aws_config::stalled_stream_protection::StalledStreamProtectionConfig;use aws_config::Region;use aws_credential_types::Credentials;use aws_http_client::AwsHttpClient;use bedrock::bedrock_client::types::{    ContentBlockDelta, ContentBlockStart, ContentBlockStartEvent, ConverseStreamOutput,};use bedrock::bedrock_client::{self, Config};use bedrock::{    value_to_aws_document, BedrockError, BedrockInnerContent, BedrockMessage, BedrockSpecificTool,    BedrockStreamingResponse, BedrockTool, BedrockToolChoice, BedrockToolInputSchema, Model,};use collections::{BTreeMap, HashMap};use credentials_provider::CredentialsProvider;use editor::{Editor, EditorElement, EditorStyle};use futures::{future::BoxFuture, stream::BoxStream, FutureExt, Stream, StreamExt};use gpui::{    AnyView, App, AsyncApp, Context, Entity, FontStyle, Subscription, Task, TextStyle, WhiteSpace,};use gpui_tokio::Tokio;use http_client::HttpClient;use language_model::{    AuthenticateError, LanguageModel, LanguageModelCacheConfiguration,    LanguageModelCompletionEvent, LanguageModelId, LanguageModelName, LanguageModelProvider,    LanguageModelProviderId, LanguageModelProviderName, LanguageModelProviderState,    LanguageModelRequest, LanguageModelToolUse, MessageContent, RateLimiter, Role,};use schemars::JsonSchema;use serde::{Deserialize, Serialize};use serde_json::Value;use settings::{Settings, SettingsStore};use strum::IntoEnumIterator;use theme::ThemeSettings;use tokio::runtime::Handle;use ui::{prelude::*, Icon, IconName, List, Tooltip};use util::{maybe, ResultExt}; use crate::AllLanguageModelSettings; const PROVIDER_ID: &str = "amazon-bedrock";const PROVIDER_NAME: &str = "Amazon Bedrock"; #[derive(Default, Clone, Deserialize, Serialize, PartialEq, Debug)]pub struct BedrockCredentials {    pub region: String,    pub access_key_id: String,    pub secret_access_key: String,} #[derive(Default, Clone, Debug, PartialEq)]pub struct AmazonBedrockSettings {    pub session_token: Option<String>,    pub available_models: Vec<AvailableModel>,} #[derive(Clone, Debug, PartialEq, Serialize, Deserialize, JsonSchema)]pub struct AvailableModel {    pub name: String,    pub display_name: Option<String>,    pub max_tokens: usize,    pub cache_configuration: Option<LanguageModelCacheConfiguration>,    pub max_output_tokens: Option<u32>,    pub default_temperature: Option<f32>,} /// The URL of the base AWS service.////// Right now we're just using this as the key to store the AWS credentials/// under in the keychain.const AMAZON_AWS_URL: &str = "https://amazonaws.com"; // These environment variables all use a `ZED_` prefix because we don't want to overwrite the user's AWS credentials.const ZED_BEDROCK_ACCESS_KEY_ID_VAR: &str = "ZED_ACCESS_KEY_ID";const ZED_BEDROCK_SECRET_ACCESS_KEY_VAR: &str = "ZED_SECRET_ACCESS_KEY";const ZED_BEDROCK_REGION_VAR: &str = "ZED_AWS_REGION";const ZED_AWS_CREDENTIALS_VAR: &str = "ZED_AWS_CREDENTIALS"; pub struct State {    credentials: Option<BedrockCredentials>,    credentials_from_env: bool,    _subscription: Subscription,} impl State {    fn reset_credentials(&self, cx: &mut Context<Self>) -> Task<Result<()>> {        let credentials_provider = <dyn CredentialsProvider>::global(cx);        cx.spawn(|this, mut cx| async move {            credentials_provider                .delete_credentials(AMAZON_AWS_URL, &cx)                .await                .log_err();            this.update(&mut cx, |this, cx| {                this.credentials = None;                this.credentials_from_env = false;                cx.notify();            })        })    }     fn set_credentials(        &mut self,        credentials: BedrockCredentials,        cx: &mut Context<Self>,    ) -> Task<Result<()>> {        let credentials_provider = <dyn CredentialsProvider>::global(cx);        cx.spawn(|this, mut cx| async move {            credentials_provider                .write_credentials(                    AMAZON_AWS_URL,                    "Bearer",                    &serde_json::to_vec(&credentials)?,                    &cx,                )                .await?;            this.update(&mut cx, |this, cx| {                this.credentials = Some(credentials);                cx.notify();            })        })    }     fn is_authenticated(&self) -> bool {        self.credentials.is_some()    }     fn authenticate(&self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {        if self.is_authenticated() {            return Task::ready(Ok(()));        }         let credentials_provider = <dyn CredentialsProvider>::global(cx);        cx.spawn(|this, mut cx| async move {            let (credentials, from_env) =                if let Ok(credentials) = std::env::var(ZED_AWS_CREDENTIALS_VAR) {                    (credentials, true)                } else {                    let (_, credentials) = credentials_provider                        .read_credentials(AMAZON_AWS_URL, &cx)                        .await?                        .ok_or_else(|| AuthenticateError::CredentialsNotFound)?;                    (                        String::from_utf8(credentials)                            .context("invalid {PROVIDER_NAME} credentials")?,                        false,                    )                };             let credentials: BedrockCredentials =                serde_json::from_str(&credentials).context("failed to parse credentials")?;             this.update(&mut cx, |this, cx| {                this.credentials = Some(credentials);                this.credentials_from_env = from_env;                cx.notify();            })?;             Ok(())        })    }} pub struct BedrockLanguageModelProvider {    http_client: AwsHttpClient,    handler: tokio::runtime::Handle,    state: gpui::Entity<State>,} impl BedrockLanguageModelProvider {    pub fn new(http_client: Arc<dyn HttpClient>, cx: &mut App) -> Self {        let state = cx.new(|cx| State {            credentials: None,            credentials_from_env: false,            _subscription: cx.observe_global::<SettingsStore>(|_, cx| {                cx.notify();            }),        });         let tokio_handle = Tokio::handle(cx);         let coerced_client = AwsHttpClient::new(http_client.clone(), tokio_handle.clone());         Self {            http_client: coerced_client,            handler: tokio_handle.clone(),            state,        }    }} impl LanguageModelProvider for BedrockLanguageModelProvider {    fn id(&self) -> LanguageModelProviderId {        LanguageModelProviderId(PROVIDER_ID.into())    }     fn name(&self) -> LanguageModelProviderName {        LanguageModelProviderName(PROVIDER_NAME.into())    }     fn icon(&self) -> IconName {        IconName::AiBedrock    }     fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {        let model = bedrock::Model::default();        Some(Arc::new(BedrockModel {            id: LanguageModelId::from(model.id().to_string()),            model,            http_client: self.http_client.clone(),            handler: self.handler.clone(),            state: self.state.clone(),            request_limiter: RateLimiter::new(4),        }))    }     fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {        let mut models = BTreeMap::default();         for model in bedrock::Model::iter() {            if !matches!(model, bedrock::Model::Custom { .. }) {                models.insert(model.id().to_string(), model);            }        }         // Override with available models from settings        for model in AllLanguageModelSettings::get_global(cx)            .bedrock            .available_models            .iter()        {            models.insert(                model.name.clone(),                bedrock::Model::Custom {                    name: model.name.clone(),                    display_name: model.display_name.clone(),                    max_tokens: model.max_tokens,                    max_output_tokens: model.max_output_tokens,                    default_temperature: model.default_temperature,                },            );        }         models            .into_values()            .map(|model| {                Arc::new(BedrockModel {                    id: LanguageModelId::from(model.id().to_string()),                    model,                    http_client: self.http_client.clone(),                    handler: self.handler.clone(),                    state: self.state.clone(),                    request_limiter: RateLimiter::new(4),                }) as Arc<dyn LanguageModel>            })            .collect()    }     fn is_authenticated(&self, cx: &App) -> bool {        self.state.read(cx).is_authenticated()    }     fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {        self.state.update(cx, |state, cx| state.authenticate(cx))    }     fn configuration_view(&self, window: &mut Window, cx: &mut App) -> AnyView {        cx.new(|cx| ConfigurationView::new(self.state.clone(), window, cx))            .into()    }     fn reset_credentials(&self, cx: &mut App) -> Task<Result<()>> {        self.state            .update(cx, |state, cx| state.reset_credentials(cx))    }} impl LanguageModelProviderState for BedrockLanguageModelProvider {    type ObservableEntity = State;     fn observable_entity(&self) -> Option<gpui::Entity<Self::ObservableEntity>> {        Some(self.state.clone())    }} struct BedrockModel {    id: LanguageModelId,    model: Model,    http_client: AwsHttpClient,    handler: tokio::runtime::Handle,    state: gpui::Entity<State>,    request_limiter: RateLimiter,} impl BedrockModel {    fn stream_completion(        &self,        request: bedrock::Request,        cx: &AsyncApp,    ) -> Result<        BoxFuture<'static, BoxStream<'static, Result<BedrockStreamingResponse, BedrockError>>>,    > {        let Ok(Ok((access_key_id, secret_access_key, region))) =            cx.read_entity(&self.state, |state, _cx| {                if let Some(credentials) = &state.credentials {                    Ok((                        credentials.access_key_id.clone(),                        credentials.secret_access_key.clone(),                        credentials.region.clone(),                    ))                } else {                    return Err(anyhow!("Failed to read credentials"));                }            })        else {            return Err(anyhow!("App state dropped"));        };         let runtime_client = bedrock_client::Client::from_conf(            Config::builder()                .stalled_stream_protection(StalledStreamProtectionConfig::disabled())                .credentials_provider(Credentials::new(                    access_key_id,                    secret_access_key,                    None,                    None,                    "Keychain",                ))                .region(Region::new(region))                .http_client(self.http_client.clone())                .build(),        );         let owned_handle = self.handler.clone();         Ok(async move {            let request = bedrock::stream_completion(runtime_client, request, owned_handle);            request.await.unwrap_or_else(|e| {                futures::stream::once(async move { Err(BedrockError::ClientError(e)) }).boxed()            })        }        .boxed())    }} impl LanguageModel for BedrockModel {    fn id(&self) -> LanguageModelId {        self.id.clone()    }     fn name(&self) -> LanguageModelName {        LanguageModelName::from(self.model.display_name().to_string())    }     fn provider_id(&self) -> LanguageModelProviderId {        LanguageModelProviderId(PROVIDER_ID.into())    }     fn provider_name(&self) -> LanguageModelProviderName {        LanguageModelProviderName(PROVIDER_NAME.into())    }     fn telemetry_id(&self) -> String {        format!("bedrock/{}", self.model.id())    }     fn max_token_count(&self) -> usize {        self.model.max_token_count()    }     fn max_output_tokens(&self) -> Option<u32> {        Some(self.model.max_output_tokens())    }     fn count_tokens(        &self,        request: LanguageModelRequest,        cx: &App,    ) -> BoxFuture<'static, Result<usize>> {        get_bedrock_tokens(request, cx)    }     fn stream_completion(        &self,        request: LanguageModelRequest,        cx: &AsyncApp,    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<LanguageModelCompletionEvent>>>> {        let request = into_bedrock(            request,            self.model.id().into(),            self.model.default_temperature(),            self.model.max_output_tokens(),        );         let owned_handle = self.handler.clone();         let request = self.stream_completion(request, cx);        let future = self.request_limiter.stream(async move {            let response = request.map_err(|e| anyhow!(e)).unwrap().await;            Ok(map_to_language_model_completion_events(                response,                owned_handle,            ))        });        async move { Ok(future.await?.boxed()) }.boxed()    }     fn use_any_tool(        &self,        request: LanguageModelRequest,        name: String,        description: String,        schema: Value,        _cx: &AsyncApp,    ) -> BoxFuture<'static, Result<BoxStream<'static, Result<String>>>> {        let mut request = into_bedrock(            request,            self.model.id().into(),            self.model.default_temperature(),            self.model.max_output_tokens(),        );         request.tool_choice = Some(BedrockToolChoice::Tool(            BedrockSpecificTool::builder()                .name(name.clone())                .build()                .unwrap(),        ));         request.tools = vec![BedrockTool::builder()            .name(name.clone())            .description(description.clone())            .input_schema(BedrockToolInputSchema::Json(value_to_aws_document(&schema)))            .build()            .unwrap()];         let handle = self.handler.clone();         let request = self.stream_completion(request, _cx);        self.request_limiter            .run(async move {                let response = request.map_err(|e| anyhow!(e)).unwrap().await;                Ok(extract_tool_args_from_events(name, response, handle)                    .await?                    .boxed())            })            .boxed()    }     fn cache_configuration(&self) -> Option<LanguageModelCacheConfiguration> {        None    }} pub fn into_bedrock(    request: LanguageModelRequest,    model: String,    default_temperature: f32,    max_output_tokens: u32,) -> bedrock::Request {    let mut new_messages: Vec<BedrockMessage> = Vec::new();    let mut system_message = String::new();     for message in request.messages {        if message.contents_empty() {            continue;        }         match message.role {            Role::User | Role::Assistant => {                let bedrock_message_content: Vec<BedrockInnerContent> = message                    .content                    .into_iter()                    .filter_map(|content| match content {                        MessageContent::Text(text) => {                            if !text.is_empty() {                                Some(BedrockInnerContent::Text(text))                            } else {                                None                            }                        }                        _ => None,                    })                    .collect();                let bedrock_role = match message.role {                    Role::User => bedrock::BedrockRole::User,                    Role::Assistant => bedrock::BedrockRole::Assistant,                    Role::System => unreachable!("System role should never occur here"),                };                if let Some(last_message) = new_messages.last_mut() {                    if last_message.role == bedrock_role {                        last_message.content.extend(bedrock_message_content);                        continue;                    }                }                new_messages.push(

Lines 1–500 of 1071. Load more