Add Bedrock Cloud Model provider

Cargo.toml

@@ -9,6 +9,7 @@ members = [
     "crates/assistant_tool",
     "crates/audio",
     "crates/auto_update",
+    "crates/bedrock",
     "crates/auto_update_ui",
     "crates/breadcrumbs",
     "crates/call",
@@ -170,7 +171,6 @@ members = [
     #
     # Tooling
     #
-
     "tooling/xtask",
 ]
 default-members = ["crates/zed"]
@@ -191,6 +191,7 @@ assistant_tool = { path = "crates/assistant_tool" }
 audio = { path = "crates/audio" }
 auto_update = { path = "crates/auto_update" }
 auto_update_ui = { path = "crates/auto_update_ui" }
+bedrock = { path = "crates/bedrock" }
 breadcrumbs = { path = "crates/breadcrumbs" }
 call = { path = "crates/call" }
 channel = { path = "crates/channel" }
@@ -335,6 +336,10 @@ async-trait = "0.1"
 async-tungstenite = "0.28"
 async-watch = "0.3.1"
 async_zip = { version = "0.0.17", features = ["deflate", "deflate64"] }
+aws-smithy-runtime-api = { version = "1.7.3" }
+aws-credential-types = { version = "1.2.1", features = ["hardcoded-credentials"] }
+aws-config = { version = "1.1.7", features = ["behavior-version-latest"] }
+aws-sdk-bedrockruntime = { version = "1.57.0", features = ["behavior-version-latest"]}
 base64 = "0.22"
 bitflags = "2.6.0"
 blade-graphics = { git = "https://github.com/kvark/blade", rev = "e142a3a5e678eb6a13e642ad8401b1f3aa38e969" }

crates/assistant/src/assistant_settings.rs

@@ -46,6 +46,11 @@ pub enum AssistantProviderContentV1 {
         default_model: Option<OllamaModel>,
         api_url: Option<String>,
     },
+    #[serde(rename = "bedrock")]
+    Bedrock {
+        default_model: Option<CloudModel>,
+        region: Option<String>,
+    },
 }
 
 #[derive(Debug, Default)]
@@ -415,6 +420,7 @@ pub struct LanguageModelSelection {
 fn providers_schema(_: &mut schemars::gen::SchemaGenerator) -> schemars::schema::Schema {
     schemars::schema::SchemaObject {
         enum_values: Some(vec![
+            "bedrock".into(),
             "anthropic".into(),
             "google".into(),
             "ollama".into(),

crates/bedrock/Cargo.toml

@@ -0,0 +1,33 @@
+[package]
+name = "bedrock"
+version = "0.1.0"
+edition = "2021"
+publish = false
+license = "AGPL-3.0-or-later"
+
+[features]
+default = []
+schemars = ["dep:schemars"]
+
+[lints]
+workspace = true
+
+[lib]
+path = "src/bedrock.rs"
+
+[dependencies]
+anyhow.workspace = true
+chrono.workspace = true
+futures.workspace = true
+http_client.workspace = true
+schemars = { workspace = true, optional = true }
+serde.workspace = true
+serde_json.workspace = true
+strum.workspace = true
+thiserror.workspace = true
+util.workspace = true
+aws-sdk-bedrockruntime = { workspace = true, features = ["behavior-version-latest"]}
+aws-config = {workspace = true, features = ["behavior-version-latest"]}
+
+[dev-dependencies]
+tokio.workspace = true

crates/bedrock/src/bedrock.rs

@@ -0,0 +1,108 @@
+mod models;
+
+use anyhow::{anyhow, Context, Result};
+use aws_sdk_bedrockruntime::config::SharedHttpClient;
+use futures::stream;
+use futures::StreamExt;
+use http_client::HttpClient;
+use serde::{Deserialize, Serialize};
+use std::sync::Arc;
+use thiserror::Error;
+
+use aws_sdk_bedrockruntime as bedrock;
+pub use aws_sdk_bedrockruntime as bedrock_client;
+pub use bedrock::operation::converse_stream::ConverseStreamInput as BedrockStreamingRequest;
+pub use bedrock::types::ContentBlock as BedrockRequestContent;
+pub use bedrock::types::ConversationRole as BedrockRole;
+use bedrock::types::ConverseOutput as Response;
+pub use bedrock::types::ConverseStreamOutput as BedrockStreamingResponse;
+pub use bedrock::types::Message as BedrockMessage;
+pub use bedrock::types::ResponseStream as BedrockResponseStream;
+use futures::stream::BoxStream;
+use strum::Display;
+//TODO: Re-export the Bedrock stuff
+// https://doc.rust-lang.org/rustdoc/write-documentation/re-exports.html
+
+pub use models::*;
+
+pub async fn complete(
+    client: &bedrock::Client,
+    request: Request,
+) -> Result<Response, BedrockError> {
+    let mut response = bedrock::Client::converse(client)
+        .model_id(request.model.clone())
+        .set_messages(request.messages.into())
+        .send()
+        .await
+        .context("Failed to send request to Bedrock");
+
+    match response {
+        Ok(output) => Ok(output.output.unwrap()),
+        Err(err) => Err(BedrockError::Other(err)),
+    }
+}
+
+pub async fn stream_completion(
+    client: &bedrock::Client,
+    request: Request,
+) -> Result<BoxStream<'static, Result<Option<BedrockStreamingResponse>, BedrockError>>, BedrockError> {
+    let response = bedrock::Client::converse_stream(client)
+        .model_id(request.model)
+        .set_messages(request.messages.into())
+        .send()
+        .await;
+
+    match response {
+        Ok(mut output) => {
+            let stream = stream::unfold(output.stream, |mut stream| async move {
+                match stream.recv().await {
+                    Ok(Some(output)) => Some((Ok(Some(output)), stream)),
+                    Ok(None) => Some((Ok(None), stream)),
+                    Err(e) => Some((
+                        Err(BedrockError::ClientError(anyhow!(
+                            "Failed to receive response from Bedrock"
+                        ))),
+                        stream,
+                    )),
+                }
+            })
+            .boxed();
+
+            Ok(stream)
+        }
+        Err(e) => Err(BedrockError::ClientError(anyhow!(e))),
+    }
+}
+
+//TODO: A LOT of these types need to re-export the Bedrock types instead of making custom ones
+
+#[derive(Debug)]
+pub struct Request {
+    pub model: String,
+    pub max_tokens: u32,
+    pub messages: Vec<BedrockMessage>,
+    // #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub system: Option<String>,
+    // #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub metadata: Option<Metadata>,
+    // #[serde(default, skip_serializing_if = "Vec::is_empty")]
+    pub stop_sequences: Vec<String>,
+    // #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub temperature: Option<f32>,
+    // #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub top_k: Option<u32>,
+    // #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub top_p: Option<f32>,
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct Metadata {
+    pub user_id: Option<String>,
+}
+
+#[derive(Error, Debug, Display)]
+pub enum BedrockError {
+    ClientError(anyhow::Error),
+    ExtensionError(anyhow::Error),
+    Other(anyhow::Error),
+}