Observability with LLM Agents
How to trace Amazon Bedrock Agents with Amazon OpenSearch
Agents for Amazon Bedrock enables developers to create fully managed agents with just a few clicks. Agents for Amazon Bedrock accelerate the delivery of generative AI applications that can manage and perform tasks by making API calls to your systems. Agents understand user requests, break down complex tasks into multiple steps, and take actions to fulfill the user request.
Given the non-deterministic nature of LLMs, it can be a challenge at times to trace and understand the exact steps and function calls that an agent decided to take to respond to a given user prompt. While there are many purpose built tools available such as arize, LangSmith, or whylabs, this post explores how far you can get with observability tools that you most likely already have such as OpenTelemetry and Amazon OpenSearch.
OpenSearch is a distributed, community-driven, Apache 2.0-licensed, 100% open-source search and analytics suite used for a broad set of use cases like real-time application monitoring, log analytics, and website search. Not only is OpenSearch frequently used as vector database in RAG applications, but we can also use its observability plugin to instrument and monitor Generative AI (GenAI) applications. The observability plugin provides a unified experience for collecting and monitoring metrics, logs, and traces from common data sources in one place to enable full-stack, end-to-end observability of your entire stack.
This tutorial explores how you can use this plugin in the context of GenAI agents for observability.
In part 1 we cover how to build and deploy an Amazon Bedrock Agent, and in part 2 we will incorporate observability with Amazon OpenSearch.
# Part 1: Build and deploy Amazon Bedrock Agent
In this part (which is inspired by this quickstart repository), we will create a chat application that takes in user questions, and then retrieves financial information from the internet for a given company/stock along with a user’s existing stock portfolio, to provide investment advice leveraging an Amazon Bedrock agent dubbed “Agent Finance”.
This solution architecture uses Amazon ECR to host the container image which encapsulates all of the tools that we want to make available to our Agent. As runtime we leverage AWS Lambda, which is a serverless execution environment and therefore ideal for this use case given its highly volatile usage/workload pattern. And finally we use Amazon Bedrock to leverage LLMs as a service.
Deployment Steps
To deploy the solution, log into your AWS management console. Then click on this link to deploy the required resources via Infrastructure as Code (IaC).
The CloudFormation stack will create following resources in your account:
- Amazon Simple Storage Service (Amazon S3) bucket
- Glue database, crawler, and table for sample dataset
- 3 AWS Lambda functions & Lambda layers
- 2 IAM roles
- Elastic Container Registry (hosts container image for Lambda function)
- Bedrock Knowledge Base & Bedrock Agent
- OpenSearch Serverless Collection as backend for Bedrock KB
- Provisioned OpenSearch cluster for observability
- OpenSearch Ingestion pipeline for observability
- temporary EC2 instance to pull and push Docker image to ECR
To avoid incurring additional charges to your account, stop and delete all of the resources by deleting the CloudFormation template at the end of this tutorial and ensuring that all created resources have been deleted.
Prerequisites
Some familiarly with Python and using services such as AWS Lambda and Elastic Container Registry is helpful. No AI/ML experience is necessary.
This tutorial assumes you are working in an environment with access to [Python 3.11](https://www.python.org/getit/) or later, and [Docker](https://www.docker.com/) as well as sufficient access to AWS Cloud Formation, Amazon Bedrock, and Amazon OpenSearch.
Let’s go over the main components of this solution, while the solution is being deployed.
Solution Overview
1.Code Repository
To test the agent tools locally, clone the following repository.
git clone https://github.com/fhuthmacher/bedrock-agent-finance
cd amazon-bedrock-agents-finance
Then run the below pip install to download the dependencies.
pip install -r requirements.txt
In `tools.py` we have defined the functions which we are going to make available to our agent. To test these tools locally, you can run the below.
python test_tools.py
2. Creation of IAM roles
The Bedrock Agent, as well as the Lambda function that runs our tools, require IAM roles with sufficient permissions to Amazon S3, Amazon Bedrock, and AWS Glue. Refer to the documentation for the specific Agent service role permissions.
To review the two IAM roles that have been created as part of the IaC deployment, go to the IAM console, and search for “AmazonBedrockExecutionRole”.
3. “AgentTools” Lambda Function
For our Agent to use tools, we must encapsulate the logic into a Lambda function. `index.py` has the logic to parse a request from the agent, and then call the corresponding method. The function will then format the response and send it back to the agent.
The function code is packaged as a Docker container using the below Dockerfile.
FROM amazon/aws-lambda-python:3.11
ARG AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION:-"us-east-1"}
ARG AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID:-""}
ARG AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY:-""}
ENV AWS_DEFAULT_REGION=${AWS_DEFAULT_REGION}
ENV AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
ENV AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
RUN yum install unzip aws-cli -y
RUN mkdir -p /opt
RUN curl $(aws lambda get-layer-version-by-arn --arn arn:aws:lambda:us-east-1:901920570463:layer:aws-otel-python-amd64-ver-1-21-0:1 --query 'Content.Location' --output text) --output layer.zip
RUN unzip layer.zip -d /opt
RUN rm layer.zip
# Install deps
COPY requirements.txt ${LAMBDA_TASK_ROOT}
RUN pip install --upgrade pip
RUN pip install -U boto3 botocore
RUN pip install --no-cache-dir -r requirements.txt
# Copy function code
COPY index.py ${LAMBDA_TASK_ROOT}
COPY tools.py ${LAMBDA_TASK_ROOT}
COPY collector.yml ${LAMBDA_TASK_ROOT}
# Set the CMD to your handler
CMD [ "index.handler" ]
To make changes, create a new ECR following the instructions here.
To test the Lambda function, go to the AWS Lambda console, and select the lambda function (e.g. “demo-BedrockAgentToolsLambda”) that has been created via CloudFormation.
Then, verify that the Lambda function has a resource-based policy that allows the Bedrock service to invoke the Lambda function (see example below).
{
"Version": "2012-10-17",
"Id": "default",
"Statement": [
{
"Sid": "agent",
"Effect": "Allow",
"Principal": {
"Service": "bedrock.amazonaws.com"
},
"Action": "lambda:InvokeFunction",
"Resource": "arn:aws:lambda:us-east-1:<your AWS account>:function:bedrock_tools"
}
]
}
Next, click on the “Test” tab near the top of the page. Then configure a test event that matches how the Agent will send a request.
{
"agent": {
"alias": "TSTALIASID",
"name": "Agent-Finance",
"version": "DRAFT",
"id": "WCZVZDPF7M"
},
"sessionId": "975786472213626",
"httpMethod": "GET",
"sessionAttributes": {},
"inputText": "Should I invest in Amazon?",
"promptSessionAttributes": {},
"apiPath": "/retrieve_inv_research",
"messageVersion": "1.0",
"actionGroup": "agent_action_group",
"parameters": [
{
"name": "query",
"type": "string",
"value": "Should I invest in Amazon?"
}
]
}
Click on “Test” to execute the Lambda function. You should see the results of the function invocation, similar to the below example.
4. OpenAPI Spec
The [OpenAPI Specification (OAS)](https://swagger.io/specification/) defines a standard, language-agnostic interface to HTTP APIs which allows both humans and computers to discover and understand the capabilities of the service without access to source code.
Our agent will be able to understand what tool to use based on the user request, and then call the correct endpoint based on the OpenAPI spec.
In our example, we defined two methods with corresponding inputs and outputs as shown in the extract below.
The Bedrock Agent references this schema , which is why the IaC deployment uploaded the `agent_aws_openapi.json` file to an Amazon S3 bucket.
5. Bedrock Agent
You can configure Bedrock Agents via the AWS management console or via the AWS SDK (see also CloudFormation template).
To get started with the agent, open the Bedrock console, select Agents in the left navigation panel, and then choose the Agent that has been created by the CloudFormation stack to review the agent details.
If you want to test the Agent right a way, go to step 6 below.
Alternatively, if you want to create a new agent via the AWS Bedrock Console leveraging the components that have been deployed, go back to the Bedrock Agents console, and click on “create agent”. This starts the agent creation workflow outlined below.
5a. Provide agent details
First we need to provide details such as agent name, description, and whether the agent is allowed to request additional user inputs, and choose the IAM role created earlier.
Here is what I used as description for this Agent.
Agent Finance is an automated, AI-powered agent that helps customers with financial investments.
5b. Select foundation model
Then select a foundation model from Bedrock that fits your use case. Here, you provide an instruction to your agent in natural language. The instruction tells the agent what task it’s supposed to perform and the persona it’s supposed to assume. Here is the instruction I used for this example.
Agent Finance is an automated, AI-powered agent that provides investment advice based on a company's current stock price and financial news
5c. Create an Action Group
Finally we have to create an action group. An action is a task that the agent can perform automatically by making API calls to your company systems. A set of actions is defined in an action group. The OpenAPI schema defines the API for all the actions in the group. The agent will invoke our Lambda function based on request, that will invoke the tools we built.
5d. Click on the `Add` button in Action groups section to create a new Action Group for the agent.
5e. Provide the Action Group Name and Description. Choose the Lambda function and the S3 object for the API spec uploaded earlier.
5f. Select Knowledge Base (Optional)
Knowledge Bases for Amazon Bedrock is a fully managed capability that helps you implement the entire RAG workflow from ingestion to retrieval and prompt augmentation without having to build custom integrations to data sources and manage data flows.
Here the IaC deployment created the vector database (OpenSearch Serverless) for us and vectorized Amazon’s latest 10k report.
To give our Agent access to this knowledge base, we can select it here and then provide instructions for the agent about its content.
This knowledge base contains financial information about companys such as quarterly reports and 10k reports.
Now click on next, and then on `Save and exit` button in the bottom right of the page.
Once the Agent is in `PREPARED` state, you can start testing.
6. Testing
Again we can use our prompt “Should I invest in Amazon?” from earlier to test our new Agent as shown below.
In addition to the chat interface which simplifies testing, Amazon Bedrock also offers tracing for Bedrock Agents, which allows you to better understand which tools have been used for a given prompt.
In part 2, we will add further observability to the Agent to monitor the execution and usage from Amazon OpenSearch.
Feel free to experiment further with other prompts.
Summary
In part 1 of this tutorial we have successfully created an Amazon Bedrock Agent, which we will use in part 2 as a starting point when we instrument the agent with OpenSearch.
Our “Finance Agent” has two functions at its disposal (get_investment_research & get_existing_portfolio) and the LLM (Anthropic’s Claude) will use the ReAct framework to determine if and when to use them.
Refer to this AWS blog for more information on Amazon Bedrock Agents.