Rust: add invoke-agent example in bedrock-agent-runtime (#7479)

* feat: add invoke agent in bedrock-agent-runtime example
---------

Co-authored-by: Tejas Ganesh Naik <tejasgn@amazon.com>

Author: mattiapitossi

.doc_gen/metadata/bedrock-agent-runtime_metadata.yaml

@@ -25,6 +25,14 @@ bedrock-agent-runtime_InvokeAgent:
             - description:
               snippet_tags:
                 - bdz.abapv1.invokeagent
+    Rust:
+      versions:
+        - sdk_version: 1
+          github: rustv1/examples/bedrock-agent-runtime
+          excerpts:
+            - description:
+              snippet_files:
+                - rustv1/examples/bedrock-agent-runtime/src/bin/invoke-agent.rs
   services:
     bedrock-agent-runtime: {InvokeAgent}
 

.github/workflows/lint-rust.yml

@@ -27,7 +27,7 @@ jobs:
       - uses: dtolnay/rust-toolchain@stable
         if: steps.changed-files.outputs.any_changed == 'true'
         with:
-          toolchain: "1.83.0"
+          toolchain: "1.87.0"
           components: clippy, rustfmt
       - name: Set Environment
         if: steps.changed-files.outputs.any_changed == 'true'

rustv1/Dockerfile

@@ -1,5 +1,5 @@
 # syntax=docker/dockerfile:1
-ARG MSRV=1.83.0
+ARG MSRV=1.87.0
 FROM rust:$MSRV
 
 # Update image

rustv1/cross_service/rust-toolchain.toml

@@ -1,3 +1,3 @@
 # This should be kept in sync with https://github.com/awslabs/aws-sdk-rust#supported-rust-versions-msrv
 [toolchain]
-channel = "1.83.0"
+channel = "1.87.0"

rustv1/examples/Cargo.toml

@@ -9,7 +9,8 @@ members = [
     "aurora",
     "auto-scaling",
     "autoscalingplans",
-    "batch",
+    "batch", 
+    "bedrock-agent-runtime",
     "bedrock-runtime",
     "cloudformation",
     "cloudwatch",

rustv1/examples/bedrock-agent-runtime/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "bedrock-agent-runtime"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+aws-config = "1.6.3"
+aws-sdk-bedrockagentruntime = "1.98.0"
+aws-smithy-types = "1.3.2"
+mockall = "0.13.1"
+tokio = { version = "1.45.1", features = ["full"] }

rustv1/examples/bedrock-agent-runtime/README.md

@@ -0,0 +1,79 @@
+# Amazon Bedrock Agents Runtime code examples for the SDK for Rust
+
+## Overview
+
+Shows how to use the AWS SDK for Rust to work with Amazon Bedrock Agents Runtime.
+
+<!--custom.overview.start-->
+<!--custom.overview.end-->
+
+_Amazon Bedrock Agents Runtime offers you the ability to run autonomous agents in your application._
+
+## ⚠ Important
+
+* Running this code might result in charges to your AWS account. For more details, see [AWS Pricing](https://aws.amazon.com/pricing/) and [Free Tier](https://aws.amazon.com/free/).
+* Running the tests might result in charges to your AWS account.
+* We recommend that you grant your code least privilege. At most, grant only the minimum permissions required to perform the task. For more information, see [Grant least privilege](https://docs.aws.amazon.com/IAM/latest/UserGuide/best-practices.html#grant-least-privilege).
+* This code is not tested in every AWS Region. For more information, see [AWS Regional Services](https://aws.amazon.com/about-aws/global-infrastructure/regional-product-services).
+
+<!--custom.important.start-->
+<!--custom.important.end-->
+
+## Code examples
+
+### Prerequisites
+
+For prerequisites, see the [README](../../README.md#Prerequisites) in the `rustv1` folder.
+
+
+<!--custom.prerequisites.start-->
+<!--custom.prerequisites.end-->
+
+### Single actions
+
+Code excerpts that show you how to call individual service functions.
+
+- [InvokeAgent](src/bin/invoke-agent.rs)
+
+
+<!--custom.examples.start-->
+<!--custom.examples.end-->
+
+## Run the examples
+
+### Instructions
+
+
+<!--custom.instructions.start-->
+<!--custom.instructions.end-->
+
+
+
+### Tests
+
+⚠ Running tests might result in charges to your AWS account.
+
+
+To find instructions for running these tests, see the [README](../../README.md#Tests)
+in the `rustv1` folder.
+
+
+
+<!--custom.tests.start-->
+<!--custom.tests.end-->
+
+## Additional resources
+
+- [Amazon Bedrock Agents Runtime User Guide](https://docs.aws.amazon.com/bedrock/latest/userguide/agents.html)
+- [Amazon Bedrock Agents Runtime API Reference](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_Operations_Agents_for_Amazon_Bedrock_Runtime.html)
+- [SDK for JavaScript (v3) Amazon Bedrock Agents Runtime reference](https://crates.io/crates/aws-sdk-bedrockagentruntime)
+
+<!--custom.resources.start-->
+<!--custom.resources.end-->
+
+---
+
+Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+SPDX-License-Identifier: Apache-2.0
+

rustv1/examples/bedrock-agent-runtime/src/bin/invoke-agent.rs

@@ -0,0 +1,155 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+use aws_config::{BehaviorVersion, SdkConfig};
+use aws_sdk_bedrockagentruntime::{
+    self as bedrockagentruntime,
+    types::{error::ResponseStreamError, ResponseStream},
+};
+#[allow(unused_imports)]
+use mockall::automock;
+
+const BEDROCK_AGENT_ID: &str = "AJBHXXILZN";
+const BEDROCK_AGENT_ALIAS_ID: &str = "AVKP1ITZAA";
+const BEDROCK_AGENT_REGION: &str = "us-east-1";
+
+#[cfg(not(test))]
+pub use EventReceiverImpl as EventReceiver;
+#[cfg(test)]
+pub use MockEventReceiverImpl as EventReceiver;
+
+pub struct EventReceiverImpl {
+    inner: aws_sdk_bedrockagentruntime::primitives::event_stream::EventReceiver<
+        ResponseStream,
+        ResponseStreamError,
+    >,
+}
+
+#[cfg_attr(test, automock)]
+impl EventReceiverImpl {
+    #[allow(dead_code)]
+    pub fn new(
+        inner: aws_sdk_bedrockagentruntime::primitives::event_stream::EventReceiver<
+            ResponseStream,
+            ResponseStreamError,
+        >,
+    ) -> Self {
+        Self { inner }
+    }
+
+    pub async fn recv(
+        &mut self,
+    ) -> Result<
+        Option<ResponseStream>,
+        aws_sdk_bedrockagentruntime::error::SdkError<
+            ResponseStreamError,
+            aws_smithy_types::event_stream::RawMessage,
+        >,
+    > {
+        self.inner.recv().await
+    }
+}
+
+#[tokio::main]
+async fn main() -> Result<(), Box<bedrockagentruntime::Error>> {
+    let result = invoke_bedrock_agent("I need help.".to_string(), "123".to_string()).await?;
+    println!("{}", result);
+    Ok(())
+}
+
+async fn invoke_bedrock_agent(
+    prompt: String,
+    session_id: String,
+) -> Result<String, bedrockagentruntime::Error> {
+    let sdk_config: SdkConfig = aws_config::defaults(BehaviorVersion::latest())
+        .region(BEDROCK_AGENT_REGION)
+        .load()
+        .await;
+    let bedrock_client = bedrockagentruntime::Client::new(&sdk_config);
+
+    let command_builder = bedrock_client
+        .invoke_agent()
+        .agent_id(BEDROCK_AGENT_ID)
+        .agent_alias_id(BEDROCK_AGENT_ALIAS_ID)
+        .session_id(session_id)
+        .input_text(prompt);
+
+    let response = command_builder.send().await?;
+
+    let response_stream = response.completion;
+
+    let event_receiver = EventReceiver::new(response_stream);
+
+    process_agent_response_stream(event_receiver).await
+}
+
+async fn process_agent_response_stream(
+    mut event_receiver: EventReceiver,
+) -> Result<String, bedrockagentruntime::Error> {
+    let mut full_agent_text_response = String::new();
+
+    while let Some(event_result) = event_receiver.recv().await? {
+        match event_result {
+            ResponseStream::Chunk(chunk) => {
+                if let Some(bytes) = chunk.bytes {
+                    match String::from_utf8(bytes.into_inner()) {
+                        Ok(text_chunk) => {
+                            full_agent_text_response.push_str(&text_chunk);
+                        }
+                        Err(e) => {
+                            eprintln!("UTF-8 decoding error for chunk: {}", e);
+                        }
+                    }
+                }
+            }
+            _ => {
+                panic!("received an unhandled event type from Bedrock stream",);
+            }
+        }
+    }
+    Ok(full_agent_text_response)
+}
+
+#[cfg(test)]
+mod test {
+
+    use super::*;
+
+    #[tokio::test]
+    async fn test_process_agent_response_stream() {
+        let mut mock = MockEventReceiverImpl::default();
+        mock.expect_recv().times(1).returning(|| {
+            Ok(Some(
+                aws_sdk_bedrockagentruntime::types::ResponseStream::Chunk(
+                    aws_sdk_bedrockagentruntime::types::PayloadPart::builder()
+                        .set_bytes(Some(aws_smithy_types::Blob::new(vec![
+                            116, 101, 115, 116, 32, 99, 111, 109, 112, 108, 101, 116, 105, 111, 110,
+                        ])))
+                        .build(),
+                ),
+            ))
+        });
+
+        // end the stream
+        mock.expect_recv().times(1).returning(|| Ok(None));
+
+        let response = process_agent_response_stream(mock).await.unwrap();
+
+        assert_eq!("test completion", response);
+    }
+
+    #[tokio::test]
+    #[should_panic(expected = "received an unhandled event type from Bedrock stream")]
+    async fn test_process_agent_response_stream_error() {
+        let mut mock = MockEventReceiverImpl::default();
+        mock.expect_recv().times(1).returning(|| {
+            Ok(Some(
+                aws_sdk_bedrockagentruntime::types::ResponseStream::Trace(
+                    aws_sdk_bedrockagentruntime::types::TracePart::builder().build(),
+                ),
+            ))
+        });
+
+        let _ = process_agent_response_stream(mock).await.unwrap();
+    }
+}

rustv1/examples/rust-toolchain.toml

@@ -1,3 +1,3 @@
 # This should be kept in sync with https://github.com/awslabs/aws-sdk-rust#supported-rust-versions-msrv
 [toolchain]
-channel = "1.83.0"
+channel = "1.87.0"