![Agents for Amazon Bedrock: Handling return of control in code](https://assets.community.aws/a/2g8mxcnMj4TPIfbANxo57rXdytl/a-ro.webp)
Agents for Amazon Bedrock: Handling return of control in code
Learn how to handle Amazon Bedrock agents return of control in your backend services and simplify your integrations.
Danilo Poccia
Amazon Employee
Published May 8, 2024
Last Modified May 9, 2024
You can push the limits of what you can achieve with generative artificial intelligence (generative AI) by using an agent. An agent can process more advanced tasks because it applies a chain of thoughts to orchestrate multiple calls to a foundation model (FM) and augments the model capabilities by retrieving data from multiple sources.
The latest updates to Agents for Amazon Bedrock introduced the capability to return control to the application invoking the agent instead of handling external requests with an AWS Lambda function. In this way, an application can handle the external interactions and return the response as the agent continues its orchestration.
The option to return control is particularly useful in four main scenarios:
- When it's more convenient to call an API from your existing application rather than building a new Lambda function with the required authentication and networking configurations.
- For handling tasks that exceed Lambda's 15-minute timeout limit and instead use containers, virtual servers, or workflow orchestration tools like AWS Step Functions.
- For time-consuming actions where the agent doesn't need to wait for completion before proceeding to the next step, allowing your application to run these actions asynchronously in the background.
- During development and testing phases, to provide a quick way to mock API interactions without fully implementing them.
Let's dive deeper into how return of control works and explore its implementation details in code. To do that, I'll use the same setup as in the launch post. This time, I'll invoke the agent and process return of control using the AWS SDK for Python (Boto3). The same concepts apply for other AWS SDKs that expose similar functionalities.
Here's a quick summary of the agent setup from the launch post.
The customer-feedback agent is instructed to help reply to customer feedback emails with a solution tailored to the customer account settings. To do that, the agent is provided with two action groups:
- retrieve-customer-settings, to retrieve customer settings including the customer ID. It can do that with the retrieve-customer-settings-from-crm function that retrieves customer settings from a customer relationship management (CRM) system using the customer email in the sender/from fields of the email. This function is configured to use return of control.
- check-login-status, to check customer login status in the login system using the customer ID from settings. To do that, it can use the check-customer-login-status-in-login-system function backed by a Lambda function.
![The configuration of Amazon Bedrock agent (action groups and functions) A description of how the Amazon Bedrock agent is configured.](https://assets.community.aws/a/2g8mGygflQQzWR3CJjKbuVadhN6/bedr.webp)
To invoke an agent, I need an agent alias. In the Amazon Bedrock console, I create an agent version and an alias associated with that version. Now, to invoke the agent, I use the agent ID and the agent alias ID form the console.
Here's the code of the
app.py
script:The session ID is a unique string that identifies the interactions with an agent which are part of the same session, such as my reply after return of control. For the session ID, the script generates a random string every time is run.
The first invocation of the agent, in the main block, uses as input the content of a customer email. Then, the function processing the response looks at the content of the completion.
The completion is an event stream part of the invokeAgent response that can contain three types of events:
- trace events that, when tracing is enabled for the request, give an insight into the chain of thoughts and the orchestration of the agent.
- chunk events, each providing a part of the final response from the agent.
- returnControl events, that provide the information to reply with a result to a function configured with return of control.
In the returnControl event, I look for the action group, the function, and the input parameters for the function. Using those values, I can provide the result for the function. In this case, to handle return of control, I simulate an external API providing the customer ID in JSON format. In general, the result is text and doesn't need to be structured as JSON.
To reply to the returnControl event, I invoke the agent a second time (in the
process_response
function). I use the same sessionId and provide, in the sessionState, the same invocationId as in the returnControl event and a list of function results (returnControlInvocationResults). Each functionResult includes the action group and function they are providing a result for and the responseBody that is passed to the agent.To run the script, I create a Python virtual environment and install the dependencies. Note that the
requirements.txt
file only contains boto3
.Now, let's see what happens when I run the script:
Each time the script runs, the flow and the result are slightly different because they are generated using a foundation model (FM). However, the overall flow and result should be similar to this example.
When started, the script processes the first invocation and finds a returnControl event with the corresponding returnControlInvocationResults.
The script intercepts the returnControl event and replies (with a seconf invocation) providing a customer ID. Then, the agent uses another action group and function to retrieve the status of the login from a CRM system. This second function is baked by a Lambda function whose code is in the launch post.
The result of the Lambda function tells the agent what the problem is, and that information is used for the final completion.
In more advanced scenarios, you can have multiple functions per action group. Each function can be configured to use a Lambda function or return of control. To use return of control in your application, you just need to manage the returnControl events and reply with the information required by the corresponding functions.
To learn more, this notebook creates an agent using the new capabilities for function definition and return of control. The example in the notebook works in code (no need to access the Amazon Bedrock console), and does not require an API Schema file or a Lambda function.
See the Amazon Bedrock User Guide to start using this new capability and let me know what you build with it!
Any opinions in this post are those of the individual author and may not reflect the opinions of AWS.