Bedrock Multi-Agent Collaboration

Introduction

AI agents are very good at doing one thing. On the other hand, they are not very good at solving complex problems if they cover different domains.

This is where multi-agent systems come in. With multi-agent systems, we can combine the strengths of multiple agents to solve complex problems. Each agent is responsible for a specific task, and they work together to solve a larger problem. This way, we can use the strengths of each agent to solve a multi-step problem that would be challenging for a single agent to accomplish alone.

In this article, we'll explore how we can use Amazon Bedrock to create sophisticated multi-agent collaborations.

As with our other hands-on articles, we're also providing you with a complete project that you can run and deploy to your own AWS account.

Feel free to clone the repository and play around with the code.

What is Amazon Bedrock?

Amazon Bedrock is a fully managed service that allows us to build and scale generative AI applications. It provides a wide range of pre-trained models from Amazon, including models for natural language processing, computer vision, and more.

It also allows you to access models from other well-known providers.

Key Features and Capabilities

One great advantage of Bedrock is that it abstracts the underlying infrastructure and underlying models. This allows us to easily switch between different models and different infrastructure providers. Everything is handled under the hood.

This means you can access models from Amazon, Anthropic, and many more. As all the billing is handled by Bedrock, everything will be paid through your AWS account. Or if you're using AWS Organizations, everything will be paid through the organization.

In comparison to integrating with individual model providers, this is a huge advantage, as we don't have to worry about different API specifics and different billing. Especially for large enterprises, this comes in handy. You don't need to fight with the purchasing department to get access to the models you need.

Understanding Multi-Agent Systems

So what are multi-agent systems and why do we need them?

What Are Multi-Agent Systems?

Multi-agent systems are systems that consist of multiple agents. Each agent focuses on a specific task, while being specialized in that task.

This works very well for complex tasks that require a combination of different skills.

In the end, the multi-agent system is able to solve the task at hand by combining the results of the individual agents. On the top, the supervisor agent takes the task of coordinating the other agents.

It's able to see the bigger picture and knows which agents are able to solve which part of the task. It's also able to invoke the other agents (even multiple times) to solve the task at hand.

Benefits of Multi-Agent Collaboration

We've talked about this in the beginning of the article: Agents are very good at doing one thing, or better said: one specific task. They work best if they are carefully crafted for a single domain.

Improved Decision-Making

Agents are designed to specialize in specific tasks, optimizing their reasoning and efficiency for a well-defined scope. A single agent handling multiple complex tasks can lead to

• slower response times
• increased computational overhead, and
• suboptimal decision-making

due to conflicting priorities.

Multi-agent collaboration allows each agent to focus on its expertise. By distributing workloads and making use of specialization, multi-agent systems are a better fit for complex workflows.

Building Our Own Multi-Agent System

Now that we're over the basics and theory, let's get our hands dirty and build our own multi-agent system. As this is an AWS blog, we're going to build a multi-agent system that uses AWS Bedrock to power the agents.

As a simple use case, we'll build

Setting Up the Project

For setting up the project, we're going to use SST v3. If you haven't heard of SST, it's a framework that allows us to define our AWS infra as code. It's very easy to use and highly abstracted. This means we don't need to fight with a lot of AWS specifics, but can bootstrap things quickly.

Let's start by creating a new SST project.

npx create-sst@latest bedrock-multi-agent-collaboration

This will create a new SST project with a basic structure.

Let's choose Next.js as our frontend framework. Generally, we wouldn't even need a frontend, but we're going to build a simple UI to interact with our multi-agent system. This way, we can submit prompts to our agents and see the results.

Primarily, our infrastructure will live in the sst.config.ts file. After running the generator, we'll only have a basic structure with a Next.js app. SST will deploy a Next.js app to CloudFront and run the server-side things on Lambda@Edge.

Defining Our Infrastructure

What do we need for our multi-agent system?

There are a few different components and services we'll need:

• A role that our agents can assume and use to make requests to the models provided through Bedrock.
• Multiple agents that can be used to solve different parts of a task.
• A way to coordinate the agents and combine their results: a supervisor agent.

Let's start by defining the role that our agents will assume. With SST, this is super simple and quickly set up.

export const createBedrockRole = (): aws.iam.Role => {
    const role = new aws.iam.Role('bedrock-role', {
        assumeRolePolicy: JSON.stringify({
            Version: '2012-10-17',
            Statement: [
                {
                    Effect: 'Allow',
                    Action: ['sts:AssumeRole'],
                    Principal: {
                        Service: 'bedrock.amazonaws.com',
                    },
                },
            ],
        }),
    });
    new aws.iam.RolePolicyAttachment('bedrock-policy-attachment', {
        role: role.name,
        policyArn: 'arn:aws:iam::aws:policy/AmazonBedrockFullAccess',
    });
    return role;
};

We can use the already existing AWS-managed policy AmazonBedrockFullAccess to grant the necessary permissions. In a real-world scenario, you might want to create a custom policy to grant only the necessary permissions.

Next, we need to define our agents!

Let's quickly discuss which agents we need and which role they will have in our scenario!

1. Plot Creator: This agent is responsible for creating a plot for a story. It's specializing in creating structured outlines.
2. Writer: This agent is responsible for writing a story. Its role is to transform structured plots into well-written stories.
3. Story Creator: This agent is the supervisor. Its role is to coordinate the other agents and combine their results.

For creating an agent, we have a few required options that we need to set:

• agentResourceRoleArn: The role that the agent will assume to make requests to the models.
• foundationModel: The model that will be used to actually run our prompts.
• instruction: The instruction that will be used as a base for the agent.
• agentCollaboration: The collaboration mode that will be used by the agent. It can be either SUPERVISOR, COLLABORATOR or it can be undefined. When nothing is provided, the agent will be a standalone agent.

So what do we need to do?

• We need to create the collaborators Plot Creator and Writer. They need to be assigned their specific roles via the system prompts.
• We need to create the supervisor Story Creator that will be responsible for coordinating the other agents. This one will be the one that will be invoked by the user.
• We need to create the agent collaborator assignments for the Story Creator to use the collaborators. These also need prompts so that the supervisor knows when and how to use the collaborators.

That's basically it!

For creating the two collaborator agents, we'll use the aws.bedrock.AgentAgent resource from Pulumi that is provided by SST. We need to make sure to assign them our previously created role, the model we want to use, the system prompt and the collaborator mode (COLLABORATOR).

Let's have a look at an example for the Plot Creator agent:

new aws.bedrock.AgentAgent('plot-creator', {
    agentResourceRoleArn: createBedrockRole().arn,
    foundationModel: 'anthropic.claude-3-5-sonnet-20240620-v1:0',
    agentCollaboration: 'COLLABORATOR',
    instruction:
        'You are a story creator responsible for producing a short story based on a given genre and premise. ' +
        'To do this, you must:\n' +
        '1. Invoke the "plot-creator" collaborator with the genre and premise to receive a structured plot.\n' +
        '2. Then invoke the "writer" collaborator with the structured plot to receive the final story.\n\n' +
        'You must not create the plot or the story yourself. Delegate all work to the collaborators.',
});

You'll find the details of the Writer and Story Creator agent in the GitHub repository.

Now we've created our agents, including the supervisor. For both of our collaborators, we also need to define an alias. An alias is a specific version of an agent that can be used to invoke the agent.

Let's have a look at how we define the alias for the Plot Creator agent:

new aws.bedrock.AgentAgentAlias('plot-creator-alias', {
    agentAliasName: `${$app.stage}-plot-creator-alias`,
    agentId: plotCreator.agentId,
    description: 'Plot Creator Alias',
});

With an alias, we can create different version of the same agent, e.g. with different prompts or different models.

Next, we need to assign the collaborators to our supervisor agent.

This is done by having a aws.bedrock.AgentAgentCollaborator resource. This resource will be used to assign the collaborators to the supervisor.

Let's have a look at how we set up the Story Creator agent to use the Plot Creator:

new aws.bedrock.AgentAgentCollaborator('plot-creator-collaborator', {
    agentId: plotCreator.agentId,
    collaborationInstruction:
        'Invoke this agent with a genre and premise to receive a structured plot. ' +
        'Use the result to inform the next step (story writing).',
    collaboratorName: 'plot-creator',
    relayConversationHistory: 'TO_COLLABORATOR',
    agentDescriptor: {
        // The alias ARN of the plot creator agent
        aliasArn: plotCreator.aliasArn,
    },
});

This will assign the Plot Creator agent to the Story Creator agent as a collaborator. We've also told the story creator how and when to invoke the plot creator.

We're also telling the story creator to relay the conversation history to the plot creator. This way, the order of the conversation is preserved.

The last step of the process is to create an alias for our supervisor agent, so we can invoke it later on.

Which Models are Available in Bedrock?

A question that might arise is:

Which models are available in Bedrock and how to get the specific model IDs?

We can get them via the Bedrock web console. Just visit the Model catalog section and search for the model you're interested in.

When clicking on a model, you'll see the model details and the model ID in the table.

A Simple Frontend

This step is not really relevant for the multi-agent system, but it's a good way to test our agents.

Let's create a simple frontend to submit prompts to our agents. Also, as the submission can take a while to complete, we'll store the results in a DynamoDB table.

This is done by having a Lambda-based API that will be invoked by the frontend. The API Lambda function will invoke another function asynchronously that will invoke the Story Creator agent.

This way, we can keep the frontend simple and the user experience responsive. As long as there's no result in the DynamoDB table, the frontend will show a loading spinner.

Testing our Multi-Agent System

Now that we've set up our agents, let's test them! To see their actual behavior, we're going to use the web console of Bedrock to invoke our Story Creator agent.

Let's go to the Agents section in the Bedrock console and select our story-creator agent. By clicking on Test on the top left, we can now invoke our agent.

As a simple test prompt, we can use the following (you're free to use any other prompt):

The story has to be about a cat that always dreams about flying.

By clicking on Run, we're submitting the prompt to our agent.

After a few seconds, we'll see the results. They should look like this when clicking on the Show trace button:

As we can see, our Story Creator is coordinating the other agents to create a story about a cat that always dreams about flying. It first creates a plot by using the Plot Creator agent. Then, it uses the Writer agent to write a story based on the plot.

At the end, the Story Creator is combining the results of the other agents to return us the final story.

When looking at the trace, we can also scroll down to see the actual messages exchanged between the agents.

We can also see the wrapped prompts that were used to invoke the agents. This includes the details about the available agents and the functions that the supervisor can use to coordinate and call the other agents.

Let's take a look at this wrapped prompt:

You can interact with the following agents in this environment using the AgentCommunication::sendMessage tool:

<agents>
  <agent name="plot-creator">
    Invoke this agent with a genre and premise to receive a structured plot. Use the result to inform the next step (story writing).
  </agent>
  <agent name="writer">
    Invoke this agent with a structured plot. It will return a full story based on that outline.
  </agent>
</agents>

We can see that the supervisor gets the information about the available agents and how to invoke them.

Possible Issues Along the Way

While implementing this, we ran into a few issues. Let's go over them and how we solved them.

Throttling Exceptions

You might face the following error:

ThrottlingException: Your request rate is too high. Reduce the frequency of requests. Check your Bedrock model invocation quotas to find the acceptable frequency.

This is due to the fact that the default quotas for the on-demand models are very low. For resolving this, you need to create a basic support ticket (not a service quota limit increase, as the on-demand models are not part of the service quota increase).

Bedrock Models are not available in all regions

Not all Bedrock models are available in all regions. This means you need to check which models are available in the region you're interested in. In our example project, we're using the latest Anthropic models and the us-east-1 region.

Supervisor Agent Is Not Invoking the Other Agents

This was not related to a technical issue with Bedrock, but a prompting issue. The prompt of the supervisor agent, as well as the prompts on how to invoke the other agents, needs to be carefully crafted.

This means, it needs to be very clear for the model what the role of the supervisor is. It needs to know which agents are available and when to invoke them.

This means, there can't be any ambiguity in the prompt. Otherwise, the supervisor agent tends to not invoke the other agents, but tries to solve the task on its own.

Let's compare a good and a bad prompt for our supervisor agent:

Bad prompt:

You are a Plot Creator specializing in creating structured story outlines.
Given a genre and a premise, generate a well-structured plot, including:
• Introduction (setting, protagonist, initial conflict)
• Rising action (key events, character development, challenges)
• Climax (turning point or major confrontation)
• Resolution (how the story ends)
Ensure the plot is compelling and logically structured.
Keep it within the given constraints, if any.

Good prompt:

You are a story creator responsible for producing a short story based on a given genre and premise.
To do this, you must:
1. Invoke the plot-creator collaborator with the genre and premise to receive a structured plot.
2. Then invoke the writer collaborator with the structured plot to receive the final story.

You must not create the plot or the story yourself. Delegate all work to the collaborators.

Why is the second prompt better?

The first prompt is very specific and tells the model exactly what to do. However, it also includes a lot of instructions that are not necessary for the model to understand the task.

The second prompt is more general and tells the model what the task is. It also includes instructions on how and when to invoke the other agents.

This way, the model knows exactly what to do and can delegate the work to the other agents.

Best Practices and Considerations

While multi-agent systems are powerful, they come with their own set of challenges. They go beyond the scope of this article, but we'll cover some of the most important ones.

What are the challenges and limitations of multi-agent systems?

• 💸 Cost Management: Each agent interaction counts as a separate API call, which can quickly add up in costs. Monitor your usage carefully and set up appropriate budget alarms. Remember that the supervisor is able to invoke the other agents multiple times. This means, you should be really careful about how you design your prompts.

• ⚡️ Latency: Multiple agent interactions mean longer response times. Consider this in your application design and set appropriate user expectations. If you design something for a real-time use case, a multi-agent approach might not be the best fit.

• 🐛 Error Handling: With multiple agents, error handling becomes more complex. On each individual step, there's a chance that something goes wrong. This means, you need to have a plan on how to handle these errors. The tracing in the web console can help you to carefully test your prompts.

• 🔐 Agent Boundaries: Each agent should have a clearly defined boundary. This means, it needs to be clear for the model what the role of the agent is. It should know which models it can use and which it can't. This also includes the data that it can access.

• 🔒 Data Privacy: Be mindful of the data you're sending to the models. Consider implementing data anonymization or masking where appropriate. Especially when you're using the models for sensitive data within strictly regulated locations, like the European Union, you need to be careful about the data you're sending.

Conclusion

Multi-agent systems are a great way to solve complex problems. They allow us to combine the strengths of multiple agents to solve complex tasks.

With Bedrock, we can easily create multi-agent systems and invoke them with ease. We can also use multiple, different models for different agents. This way, we can use the best model for the specific task at hand.