Build Agents on Amazon Bedrock

Jayant Raj
11 min readAug 27, 2024

--

This post demonstrates building agents on the Amazon Bedrock platform using AWS services such as Agents for Amazon Bedrock, Knowledge Base, OpenSearch Serverless and Bedrock Foundation models.

What are Agents for Amazon Bedrock?

Large Language Model (LLM) Agents are software entities that use a foundation model reasoning capabilities to take a user request, analyze and decompose the request into smaller tasks, orchestrate the sequence of tasks and execute actions (call API, interact with applications, databases, etc.) to complete the request. Agents can also interact with the users to gather more information while fulfilling the tasks.

Agents for Amazon Bedrock provides a fully managed service to build and scale generative AI bots that are able to execute complex tasks autonomously using given tools and integrates with enterprise data sources to perform read and write operations.

Bedrock Agents come with the following features out of the box:

  • Managed service to build and deploy agents
  • Integration with knowledge base for RAG capabilities
  • Memory to maintain conversation history across multiple sessions
  • Code interpretation to generate and execute code
  • Action groups to provide mechanism for agents to take actions
  • Customization of agents with more specific prompt templates
  • Guardrails for agents
  • Trace through the chain-of-thought reasoning

Illustrative use case — Clinical Trials Agent

Let us build an example Clinical Trials agent using Agents for Amazon Bedrock to understand the process of building agents. The agent can answer general questions about health conditions (e.g. high cholesterol) and also autonomously connect to a data source (small subset of clinical trials data related to cholesterol) to pull relevant information if required. The agent will be able to analyze and detect whether it should retrieve information from general literature stored in its knowledge base or connect to a more specific data source to formulate a response. The agent will elicit further information from the user if it determines it needs more specifics to answer a question.

The following diagram shows a high level architecture of the agent setup. The general health information documents are uploaded in PDF format to a S3 bucket and we use this bucket as the data source during the knowledge base creation step. The structured clinical trials data sample is uploaded to a DynamoDB table. A Lambda function has the necessary business logic to query the DynamoDB table. The Bedrock Agent uses the user conversation history and its inbuilt ReACT (Reasoning and Acting) templates for decision making and taking actions.

Following is a step by step outline describing how to build the Clinical Trials agent:

1. Enable foundation model access in Amazon Bedrock

The Bedrock Foundation models do not have access granted by default. An IAM user with the correct permissions has to request for model access via the AWS console. For purposes of this sample application, we will request access to the Claude models provided by Anthropic.

Please refer to the documentation here for more details: https://docs.aws.amazon.com/bedrock/latest/userguide/model-access.html

2. Build a Bedrock Knowledge Base with general articles

We will gather our initial set of cholesterol related articles from sources such as MedlinePlus, FDA and NIH. We will export the articles at the following URLs in PDF format. The list of articles is for illustrative purposes and not meant to be comprehensive enough to answer every possible question around the medical condition.

https://medlineplus.gov/cholesterol.html#

https://medlineplus.gov/lab-tests/cholesterol-levels/

https://medlineplus.gov/lab-tests/fasting-for-a-blood-test/

https://medlineplus.gov/howtolowercholesterolwithdiet.html#

https://medlineplus.gov/cholesterolmedicines.html#

https://www.fda.gov/medical-devices/home-use-tests/cholesterol

https://www.nccih.nih.gov/health/tips/high-cholesterol-and-dietary-supplements

Upload the articles in PDF format to a S3 bucket:

- Navigate to the Amazon Bedrock console. Click on Knowledge bases under Builder tools.

- Click on Create knowledge base

- Enter the following information for Step 1 — Provide knowledge base details

> Knowledge base name: health_information_kb

> Knowledge base description: Knowledge base to store health information

> IAM permissions: Create and use a new service role.

> Service role name: AmazonBedrockExecutionRoleForKnowledgeBase_health_information_kb

> Choose data source: Select S3

> Click Next

- Enter the following information for Step 2 — Configure data source

> Data source name: medlineplus

> S3 URI: <your s3 assets bucket created in previous step>

> Chunking and parsing configurations: Default

> Click Next

- Enter the following information for Step 3 — Select embeddings model and configure vector store

> Embeddings model: Titan Embeddings G1 — Text v1.2

> Vector database: Quick create a new vector store

> Click Next

- Enter the following information for Step 4 — Review and create

> Review the details on the screen and click on Create knowledge base

  • Sync the knowledge base: It may take a few minutes to create the knowledge base. Once the knowledge base has been created, select the medlineplus data sources and click Sync.

- Test knowledge base: Before we proceed, let us test the knowledge base by running some representative queries and checking the generated responses. Select the health_information_kb knowledge base and click Test knowledge base. Select the Claude 3 Sonnet model to generate responses.

Following are some sample questions to test our knowledge base.:

  • How do I reduce LDL?
  • What is considered high LDL?
  • How do I get tested for cholesterol?
  • What is considered a high cholesterol range?

We can see from the screenshot below, Bedrock retrieves the relevant chunks from the knowledge base along with metadata such as source-uri, chunk-id, etc and generates responses to the user questions. The sources are cited in the chat UI response within Bedrock.

Now that we have tested the knowledge base for accuracy of the responses, we proceed to the next step of creating a sample data source with clinical trials data.

3. Create a clinical trials data source

Let us create a simple DynamoDB table and populate it with clinical trials that are cholesterol related and have not started recruiting or actively recruiting.

The data can be found at the clinicaltrials.gov website. Use this URL: https://clinicaltrials.gov/search?cond=%22Cholesterol%22&aggFilters=status:not%20rec

Select all the fields and download the data as a csv file.

Create a DynamoDB table clinical_trials with a primary key NCTNumber. Use the following script to insert the exported data into our DynamoDB table.

import boto3
import csv

# Initialize a session using Amazon DynamoDB
session = boto3.Session(
aws_access_key_id='<your access key id>',
aws_secret_access_key='<your secret access key>',
aws_session_token='<your session token>',
region_name='us-east-1'
)

dynamodb = session.resource('dynamodb')
table = dynamodb.Table('clinical_trials')

def load_csv_to_dynamodb(file_name):
with open(file_name, 'r') as csvfile:
csv_reader = csv.DictReader(csvfile)

# Check for empty headers
headers = csv_reader.fieldnames
if any(header is None or header.strip() == "" for header in headers):
raise ValueError(">> CSV file contains empty headers.")

for row in csv_reader:
# Filter out any empty attribute names
item = {k: v for k, v in row.items() if k.strip() != ""}
print(f"Inserting {item}")
table.put_item(Item=item)

if __name__ == '__main__':
load_csv_to_dynamodb('ctg-studies-trim.csv')

Confirm the data was successfully inserted into DynamoDB, by selecting the clinical_trials table and clicking Explore items.

4. Build a Bedrock Agent

We will build a Bedrock agent that will interact with the knowledgebase and the clinical trials data source to respond to user questions.

Click Create.

We are now able to see the complete Agent Builder screen, where we can configure various aspects of the agent such as the LLM (currently Bedrock supports Anthropic and Amazon models), memory, action groups, knowledge bases, guardrails and advanced prompts (to customize the agent prompt templates).

Agent details section:

Agent name: clinical-trials-agent

Agent description: Help respond to user questions about specific health information and relevant clinical trials around the health condition.

Agent resource role: Create and use a new service role

Select model:

Anthropic, Claude 3 Sonnet

Instructions for the agent

You are a clinical trials agent that has access to health related conditions and clinical trials information. You can respond to general health related questions from your knowledgebase. You can retrieve information such as NCT number, study status, study type, recruiting status from this database.

When a user asks a question, first search through the relevant documents in your knowledge base. If the question pertains to clinical trials, query the clinical trials database to retrieve the most up-to-date and relevant information. Provide citations or references to the documents or trials you use in your response.

Additional Settings:

Code Interpreter: Disabled (Enabling this option would allow the agent to handle tasks that require writing, executing and testing code in a secure environment.)

User input: Enabled (Enabling this option allows the agent to ask clarifying questions, if it needs more information to respond.)

Memory: Disabled

Action groups:

Action groups in Bedrock Agent define what actions the agent can take in order to fulfill tasks. We will define an action group (retrieve-clinical-trials-data) to retrieve information from our DynamoDB table and have the following action group functions:

Function: SummarizeTrialByNCTId — Summarize key details of a clinical trial

Parameters: nctid (nctid of the trial), required

Function: Retrieve5TrialsByRecruitmentStatus — Retrieve 5 trials by recruitment status (RECRUITING or NOT_YET_RECRUITING)

Parameters: study_status (recruitment status of the trial), required

Note, if we indicate the parameters as required, the agent will request the user for these parameters while it tries to fulfill the request.

The Action Groups can be defined in the following ways:

JSON objects — The functions and parameters are defined as JSON objects. This simplified approach was introduced recently.

API schemas: Define the API schema and place it in a S3 bucket.

Select Define with function details

Under Action group invocation, select Quick create a new Lambda function

Click Create.

We will update the business logic within the Lambda function that was created for us.

import json
import boto3
from boto3.dynamodb.conditions import Attr
from boto3.dynamodb.conditions import Key

dynamodb = boto3.resource('dynamodb')

def lambda_handler(event, context):
agent = event['agent']
actionGroup = event['actionGroup']
function = event['function']
parameters = event.get('parameters', [])
print(event)
print(actionGroup)
print(function)

# Execute your business logic here. For more information, refer to: https://docs.aws.amazon.com/bedrock/latest/userguide/agents-lambda.html
if actionGroup == 'retrieve-clinical-trials-data' and function == 'SummarizeTrialByNCTId':
print("----- SummarizeTrialByNCTId")

for p in parameters:
if p['name'] == 'nctid' and p['type'] == 'string':
nctid = p['value']

table_name = 'clinical_trials'
table = dynamodb.Table(table_name)
data = table.query(
KeyConditionExpression=Key('NCTNumber').eq(nctid)
)

items = data.get('Items', [])
print(items)

if items:
study_title = items[0].get('StudyTitle', 'Title not found')
brief_summary = items[0].get('BriefSummary', 'Brief summary not found')
sponsor = items[0].get('Sponsor', 'Sponsor not found')
result = {
'nctid': nctid,
'study_title': study_title,
'brief_summary': brief_summary,
'sponsor': sponsor
}
else:
result = {
'nctid': nctid,
'message': 'No study found with this nctid'
}

response ={
'nctid': nctid,
'study_title': study_title,
'brief_summary': brief_summary,
'sponsor': sponsor

}

elif actionGroup == 'retrieve-clinical-trials-data' and function == 'Retrieve5TrialsByRecruitmentStatus':
print("----- Retrieve5TrialsByRecruitmentStatus")

study_status=''

for p in parameters:
if p['name'] == 'study_status' and p['type'] == 'string':
study_status = p['value']

table_name = 'clinical_trials'
table = dynamodb.Table(table_name)

filter_column = 'StudyStatus'
filter_value = study_status

data = table.scan(
FilterExpression=Attr(filter_column).eq(filter_value),
Limit=5
)

items = data.get('Items', [])
print(data)

response ={
'studies': items

}

else:
response = {
"error": "Unknown action group {} or function {}".format(actionGroup, function)
}

responseBody = {
"TEXT": {
"body": json.dumps(response)
}
}
action_response = {
'actionGroup': actionGroup,
'function': function,
'functionResponse': {
'responseBody': responseBody
}

}

dummy_function_response = {'response': action_response, 'messageVersion': event['messageVersion']}
print("Response: {}".format(dummy_function_response))

return dummy_function_response

Also, add the following policy to the lambda role, to ensure the lambda has permissions to scan and query the DynamoDB table.

{
"Version": "2012-10-17",
"Statement": [
{
"Sid": "VisualEditor0",
"Effect": "Allow",
"Action": [
"dynamodb:Scan",
"dynamodb:Query"
],
"Resource": "arn:aws:dynamodb:*:203744211876:table/clinical_trials*"
}
]
}

We will associate the knowledge base previously created with this Agent. The agent can query the knowledge base to provide additional context to its response and also be more efficient in invoking the action groups.

Click Add.

Click Save and Exit.

Click Prepare to test the changes we have made so far.

5. Test agent with queries

Question: What cholesterol range is considered high?

The Agent determines it is able to respond to this question from its knowledge base. We can see the traces of the agent’s thought process in the agent testing interface. The agent identifies the answer in one of the chunks from the Medline article and formulates a response based on the chunk.

Question: Which clinical trials are recruiting patients with high cholesterol?

The agent determines it does not have sufficient data to respond to this question and it will ask the user for more details.

Question: Show me some clinical trials recruiting participants currently.

Question: Summarize trial NCT06287177 for me.

Conclusion

This post demonstrated how to build Agents on the Bedrock platform to interact with your knowledge bases and take actions autonomously. We used the tools provided within the AWS console to build this prototype application. However, for real life applications, the different steps and business logic has to be codified using infrastructure as code and the agent deployed as an API on the Bedrock platform. Production deployment of Bedrock agents will be a topic for a subsequent post.

References

Unlisted

--

--

No responses yet