Agents for Amazon Bedrock: Handling return of control in code

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:

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:

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:

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!