feat: AWS Bedrock provider (#318)

* feat: aws-bedrock-provider

Added new model provider

* feat: AWS Bedrock provider

- changes as requested in PR comments

* feat: image generation

Author: mkranjac

Cargo.lock

@@ -9,6 +9,8 @@ members = [
     "rig-qdrant",
     "rig-core/rig-core-derive",
     "rig-sqlite",
-    "rig-eternalai", "rig-fastembed",
     "rig-surrealdb",
+    "rig-eternalai",
+    "rig-fastembed",
+    "rig-bedrock",
 ]

Cargo.toml

@@ -0,0 +1,26 @@
+[package]
+name = "rig-bedrock"
+version = "0.1.0"
+edition = "2021"
+license = "MIT"
+readme = "README.md"
+description = "AWS Bedrock model provider for Rig integration."
+
+[dependencies]
+rig-core = { version = "0.11.0", features = ["image"]  }
+rig-derive = { path = "../rig-core/rig-core-derive", version = "0.1.0" }
+serde = { version = "1.0.193", features = ["derive"] }
+serde_json = "1.0.108"
+schemars = "0.8.16"
+tracing = "0.1.40"
+aws-config = { version = "1.6.0", features = ["behavior-version-latest"] }
+aws-sdk-bedrockruntime = "1.77.0"
+aws-smithy-types = "1.3.0"
+base64 = "0.22.1"
+async-stream = "0.3.6"
+
+[dev-dependencies]
+anyhow = "1.0.75"
+tokio = { version = "1.34.0", features = ["full"] }
+tracing-subscriber = "0.3.18"
+reqwest = { version = "0.12.12", features = ["json", "stream"] }

rig-bedrock/Cargo.toml

@@ -0,0 +1,23 @@
+## Rig-Bedrock
+This companion crate integrates AWS Bedrock as model provider with Rig.
+
+## Usage
+
+Add the companion crate to your `Cargo.toml`, along with the rig-core crate:
+
+```toml
+[dependencies]
+rig-bedrock = "0.1.0"
+rig-core = "0.9.1"
+```
+
+You can also run `cargo add rig-bedrock rig-core` to add the most recent versions of the dependencies to your project.
+
+See the [`/examples`](./examples) folder for usage examples.
+
+Make sure to have AWS credentials env vars loaded before starting client such as:
+```shell
+export AWS_DEFAULT_REGION=us-east-1
+export AWS_SECRET_ACCESS_KEY=.......
+export AWS_ACCESS_KEY_ID=......
+```

rig-bedrock/README.md

@@ -0,0 +1,121 @@
+use rig::{agent::AgentBuilder, completion::Prompt, loaders::FileLoader};
+use rig_bedrock::{
+    client::{Client, ClientBuilder},
+    completion::AMAZON_NOVA_LITE,
+};
+use tracing::info;
+
+mod common;
+
+/// Runs 4 agents based on AWS Bedrock (derived from the agent_with_grok example)
+#[tokio::main]
+async fn main() -> Result<(), anyhow::Error> {
+    tracing_subscriber::fmt()
+        .with_max_level(tracing::Level::INFO)
+        .with_target(false)
+        .init();
+
+    info!("Running basic agent");
+    basic().await?;
+
+    info!("\nRunning agent with tools");
+    tools().await?;
+
+    info!("\nRunning agent with loaders");
+    loaders().await?;
+
+    info!("\nRunning agent with context");
+    context().await?;
+
+    info!("\n\nAll agents ran successfully");
+    Ok(())
+}
+
+async fn client() -> Client {
+    ClientBuilder::new().build().await
+}
+
+async fn partial_agent() -> AgentBuilder<rig_bedrock::completion::CompletionModel> {
+    let client = client().await;
+    client.agent(AMAZON_NOVA_LITE)
+}
+
+/// Create an AWS Bedrock agent with a system prompt
+async fn basic() -> Result<(), anyhow::Error> {
+    let agent = partial_agent()
+        .await
+        .preamble("Answer with json format only")
+        .build();
+
+    let response = agent.prompt("Describe solar system").await?;
+    info!("{}", response);
+
+    Ok(())
+}
+
+/// Create an AWS Bedrock with tools
+async fn tools() -> Result<(), anyhow::Error> {
+    let calculator_agent = partial_agent()
+        .await
+        .preamble("You must only do math by using a tool.")
+        .max_tokens(1024)
+        .tool(common::Adder)
+        .build();
+
+    info!(
+        "Calculator Agent: add 400 and 20\nResult: {}",
+        calculator_agent.prompt("add 400 and 20").await?
+    );
+
+    Ok(())
+}
+
+async fn context() -> Result<(), anyhow::Error> {
+    let model = client().await.completion_model(AMAZON_NOVA_LITE);
+
+    // Create an agent with multiple context documents
+    let agent = AgentBuilder::new(model)
+        .preamble("Answer the question")
+        .context("Definition of a *flurbo*: A flurbo is a green alien that lives on cold planets")
+        .context("Definition of a *glarb-glarb*: A glarb-glarb is a ancient tool used by the ancestors of the inhabitants of planet Jiro to farm the land.")
+        .context("Definition of a *linglingdong*: A term used by inhabitants of the far side of the moon to describe humans.")
+        .build();
+
+    // Prompt the agent and print the response
+    let response = agent.prompt("What does \"glarb-glarb\" mean?").await?;
+
+    info!("What does \"glarb-glarb\" mean?\n{}", response);
+
+    Ok(())
+}
+
+/// Based upon the `loaders` example
+///
+/// This example loads in all the rust examples from the rig-core crate and uses them as\\
+///  context for the agent
+async fn loaders() -> Result<(), anyhow::Error> {
+    let model = client().await.completion_model(AMAZON_NOVA_LITE);
+
+    // Load in all the rust examples
+    let examples = FileLoader::with_glob("rig-core/examples/*.rs")?
+        .read_with_path()
+        .ignore_errors()
+        .into_iter();
+
+    // Create an agent with multiple context documents
+    let agent = examples
+        .fold(AgentBuilder::new(model), |builder, (path, content)| {
+            builder.context(format!("Rust Example {:?}:\n{}", path, content).as_str())
+        })
+        .preamble("Answer the question")
+        .build();
+
+    // Prompt the agent and print the response
+    let response = agent
+        .prompt("Which rust example is best suited for the operation 1 + 2")
+        .await?;
+
+    info!("{}", response);
+
+    Ok(())
+}