bedrock-data-analyst

Bedrock Data Analyst Agent Example

This project demonstrates how to build a data analyst agent using Amazon Bedrock, governed by MeshGuard.

The agent is designed to answer questions by querying a mock sales database. MeshGuard is used to enforce fine-grained access control, ensuring that only authorized users and roles can query sensitive data tables or execute specific types of queries.

Architecture

The agent's logic is implemented as an AWS Lambda function, which is configured as the action group for a Bedrock Agent. When the agent needs to query data, it invokes the Lambda, which then checks permissions with MeshGuard before executing the database query.

+------------------+      +----------------------+      +--------------------+
|                  |      |                      |      |                    |
|   End User       +----->+  Amazon Bedrock      +----->+   AWS Lambda       |
| (via chat app)   |      |   (Agent)            |      |  (Action Group)    |
|                  |      |                      |      |                    |
+------------------+      +----------------------+      +--------+-----------+
                                                                 |
                                                                 | 1. check()
                                                                 v
+------------------+      +----------------------+      +--------+-----------+
|                  |      |                      |      |                    |
|   PostgreSQL     +<-----+   MeshGuard Policy   <------+   MeshGuard        |
|   (Sales Data)   | 2b. (allow) |  Decision Point      |   (Governance)     |
|                  |      |                      |      |                    |
+------------------+      +----------------------+      +--------------------+
     ^
     | 2a. Query
     | (if allowed)
     +---------------------------------------------------------+

1. A user asks the Bedrock agent a question like, "What were the total sales for Q1 2023?"
2. The Bedrock Agent invokes the Lambda function, passing the inferred query parameters. The user's identity is passed via session attributes.
3. The Lambda function calls meshguard.check() to see if the user is authorized to perform the query on the requested table.
4. If MeshGuard returns allow, the Lambda function executes the SQL query against the database.
5. The results are returned to the Bedrock agent, which synthesizes a natural language response for the user.

Setup

1. Prerequisites

• Docker and Docker Compose
• AWS Account with Bedrock access
• A MeshGuard account and an Agent Token

2. Configure Environment

Create a .env file in this directory with the following content:

# MeshGuard Configuration
MESHGUARD_GATEWAY_URL=https://dashboard.meshguard.app
MESHGUARD_AGENT_TOKEN=mg_agent_your_token_here

# AWS Credentials (for local testing)
AWS_ACCESS_KEY_ID=your_aws_key
AWS_SECRET_ACCESS_KEY=your_aws_secret
AWS_REGION=us-east-1

# Database Credentials
DB_USER=salesuser
DB_PASSWORD=salespassword
DB_HOST=db
DB_PORT=5432
DB_NAME=sales

3. Build and Run

docker-compose up --build

This will start the Lambda container (running the agent.py handler) and a PostgreSQL database seeded with sample data.

4. Deploy to AWS Lambda

To deploy this for use with a real Bedrock Agent, you'll need to package the code and its dependencies into a zip file or a container image and upload it to AWS Lambda.

1. Package for Lambda:

   pip install -r requirements.txt -t ./package
   cp -r agent.py tools policies ./package/
   cd package && zip -r ../lambda_package.zip .

Note: the example policy lives at policies/bedrock-analyst.yaml.

2. Create a new Lambda function in the AWS Console.
3. Upload lambda_package.zip.
4. Set the handler to agent.lambda_handler.
5. Add the necessary environment variables (MESHGUARD_GATEWAY_URL, MESHGUARD_AGENT_TOKEN, DB credentials).

5. Configure Bedrock Agent

1. Create a new agent in the Amazon Bedrock console.
2. Create an action group.
3. For the action group's API schema, provide an OpenAPI spec that defines the query_sales_data function.
4. Point the action group to the ARN of the Lambda function you deployed.

Running Tests

To run the unit tests:

docker-compose run --rm app pytest