diff --git a/integrations/aws_bedrock/agent.py b/integrations/aws_bedrock/agent.py new file mode 100644 index 0000000..0e2b023 --- /dev/null +++ b/integrations/aws_bedrock/agent.py @@ -0,0 +1,313 @@ +import boto3 +import json +import time +import zipfile +from io import BytesIO + +iam_client = boto3.client('iam') +sts_client = boto3.client('sts') +session = boto3.session.Session() +region = session.region_name +account_id = sts_client.get_caller_identity()["Account"] +dynamodb_client = boto3.client('dynamodb') +dynamodb_resource = boto3.resource('dynamodb') +lambda_client = boto3.client('lambda') +bedrock_agent_client = boto3.client('bedrock-agent') + + +def create_dynamodb(table_name): + table = dynamodb_resource.create_table( + TableName=table_name, + KeySchema=[ + { + 'AttributeName': 'booking_id', + 'KeyType': 'HASH' + } + ], + AttributeDefinitions=[ + { + 'AttributeName': 'booking_id', + 'AttributeType': 'S' + } + ], + BillingMode='PAY_PER_REQUEST' # Use on-demand capacity mode + ) + + # Wait for the table to be created + print(f'Creating table {table_name}...') + table.wait_until_exists() + print(f'Table {table_name} created successfully!') + return + + +def create_lambda(lambda_function_name, lambda_iam_role): + # add to function + + # Package up the lambda function code + s = BytesIO() + z = zipfile.ZipFile(s, 'w') + z.write("lambda_function.py") + z.close() + zip_content = s.getvalue() + + # Create Lambda Function + lambda_function = lambda_client.create_function( + FunctionName=lambda_function_name, + Runtime='python3.12', + Timeout=60, + Role=lambda_iam_role['Role']['Arn'], + Code={'ZipFile': zip_content}, + Handler='lambda_function.lambda_handler' + ) + return lambda_function + + +def create_lambda_role(agent_name, dynamodb_table_name): + lambda_function_role = f'{agent_name}-lambda-role' + dynamodb_access_policy_name = f'{agent_name}-dynamodb-policy' + # Create IAM Role for the Lambda function + try: + assume_role_policy_document = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Principal": { + "Service": "lambda.amazonaws.com" + }, + "Action": "sts:AssumeRole" + } + ] + } + + assume_role_policy_document_json = json.dumps(assume_role_policy_document) + + lambda_iam_role = iam_client.create_role( + RoleName=lambda_function_role, + AssumeRolePolicyDocument=assume_role_policy_document_json + ) + + # Pause to make sure role is created + time.sleep(10) + except iam_client.exceptions.EntityAlreadyExistsException: + lambda_iam_role = iam_client.get_role(RoleName=lambda_function_role) + + # Attach the AWSLambdaBasicExecutionRole policy + iam_client.attach_role_policy( + RoleName=lambda_function_role, + PolicyArn='arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole' + ) + + # Create a policy to grant access to the DynamoDB table + dynamodb_access_policy = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "dynamodb:GetItem", + "dynamodb:PutItem", + "dynamodb:DeleteItem" + ], + "Resource": "arn:aws:dynamodb:{}:{}:table/{}".format( + region, account_id, dynamodb_table_name + ) + } + ] + } + + # Create the policy + dynamodb_access_policy_json = json.dumps(dynamodb_access_policy) + dynamodb_access_policy_response = iam_client.create_policy( + PolicyName=dynamodb_access_policy_name, + PolicyDocument=dynamodb_access_policy_json + ) + + # Attach the policy to the Lambda function's role + iam_client.attach_role_policy( + RoleName=lambda_function_role, + PolicyArn=dynamodb_access_policy_response['Policy']['Arn'] + ) + return lambda_iam_role + + +def create_agent_role_and_policies(agent_name, agent_foundation_model, kb_id=None): + agent_bedrock_allow_policy_name = f"{agent_name}-ba" + agent_role_name = f'AmazonBedrockExecutionRoleForAgents_{agent_name}' + # Create IAM policies for agent + statements = [ + { + "Sid": "AmazonBedrockAgentBedrockFoundationModelPolicy", + "Effect": "Allow", + "Action": "bedrock:InvokeModel", + "Resource": [ + f"arn:aws:bedrock:{region}::foundation-model/{agent_foundation_model}" + ] + } + ] + # add Knowledge Base retrieve and retrieve and generate permissions if agent has KB attached to it + if kb_id: + statements.append( + { + "Sid": "QueryKB", + "Effect": "Allow", + "Action": [ + "bedrock:Retrieve", + "bedrock:RetrieveAndGenerate" + ], + "Resource": [ + f"arn:aws:bedrock:{region}:{account_id}:knowledge-base/{kb_id}" + ] + } + ) + + bedrock_agent_bedrock_allow_policy_statement = { + "Version": "2012-10-17", + "Statement": statements + } + + bedrock_policy_json = json.dumps(bedrock_agent_bedrock_allow_policy_statement) + + agent_bedrock_policy = iam_client.create_policy( + PolicyName=agent_bedrock_allow_policy_name, + PolicyDocument=bedrock_policy_json + ) + + # Create IAM Role for the agent and attach IAM policies + assume_role_policy_document = { + "Version": "2012-10-17", + "Statement": [{ + "Effect": "Allow", + "Principal": { + "Service": "bedrock.amazonaws.com" + }, + "Action": "sts:AssumeRole" + }] + } + + assume_role_policy_document_json = json.dumps(assume_role_policy_document) + agent_role = iam_client.create_role( + RoleName=agent_role_name, + AssumeRolePolicyDocument=assume_role_policy_document_json + ) + + # Pause to make sure role is created + time.sleep(10) + + iam_client.attach_role_policy( + RoleName=agent_role_name, + PolicyArn=agent_bedrock_policy['Policy']['Arn'] + ) + return agent_role + + +def delete_agent_roles_and_policies(agent_name): + agent_bedrock_allow_policy_name = f"{agent_name}-ba" + agent_role_name = f'AmazonBedrockExecutionRoleForAgents_{agent_name}' + dynamodb_access_policy_name = f'{agent_name}-dynamodb-policy' + lambda_function_role = f'{agent_name}-lambda-role' + + for policy in [agent_bedrock_allow_policy_name]: + try: + iam_client.detach_role_policy( + RoleName=agent_role_name, + PolicyArn=f'arn:aws:iam::{account_id}:policy/{policy}' + ) + except Exception as e: + print(f"Could not detach {policy} from {agent_role_name}") + print(e) + + for policy in [dynamodb_access_policy_name]: + try: + iam_client.detach_role_policy( + RoleName=lambda_function_role, + PolicyArn=f'arn:aws:iam::{account_id}:policy/{policy}' + ) + except Exception as e: + print(f"Could not detach {policy} from {lambda_function_role}") + print(e) + + try: + iam_client.detach_role_policy( + RoleName=lambda_function_role, + PolicyArn='arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole' + ) + except Exception as e: + print(f"Could not detach AWSLambdaBasicExecutionRole from {lambda_function_role}") + print(e) + + for role_name in [agent_role_name, lambda_function_role]: + try: + iam_client.delete_role( + RoleName=role_name + ) + except Exception as e: + print(f"Could not delete role {role_name}") + print(e) + + for policy in [agent_bedrock_allow_policy_name, dynamodb_access_policy_name]: + try: + iam_client.delete_policy( + PolicyArn=f'arn:aws:iam::{account_id}:policy/{policy}' + ) + except Exception as e: + print(f"Could not delete policy {policy}") + print(e) + + +def clean_up_resources( + table_name, lambda_function, lambda_function_name, agent_action_group_response, agent_functions, + agent_id, kb_id, alias_id +): + action_group_id = agent_action_group_response['agentActionGroup']['actionGroupId'] + action_group_name = agent_action_group_response['agentActionGroup']['actionGroupName'] + # Delete Agent Action Group, Agent Alias, and Agent + try: + bedrock_agent_client.update_agent_action_group( + agentId=agent_id, + agentVersion='DRAFT', + actionGroupId= action_group_id, + actionGroupName=action_group_name, + actionGroupExecutor={ + 'lambda': lambda_function['FunctionArn'] + }, + functionSchema={ + 'functions': agent_functions + }, + actionGroupState='DISABLED', + ) + bedrock_agent_client.disassociate_agent_knowledge_base( + agentId=agent_id, + agentVersion='DRAFT', + knowledgeBaseId=kb_id + ) + bedrock_agent_client.delete_agent_action_group( + agentId=agent_id, + agentVersion='DRAFT', + actionGroupId=action_group_id + ) + bedrock_agent_client.delete_agent_alias( + agentAliasId=alias_id, + agentId=agent_id + ) + bedrock_agent_client.delete_agent(agentId=agent_id) + print(f"Agent {agent_id}, Agent Alias {alias_id}, and Action Group have been deleted.") + except Exception as e: + print(f"Error deleting Agent resources: {e}") + + # Delete Lambda function + try: + lambda_client.delete_function(FunctionName=lambda_function_name) + print(f"Lambda function {lambda_function_name} has been deleted.") + except Exception as e: + print(f"Error deleting Lambda function {lambda_function_name}: {e}") + + # Delete DynamoDB table + try: + dynamodb_client.delete_table(TableName=table_name) + print(f"Table {table_name} is being deleted...") + waiter = dynamodb_client.get_waiter('table_not_exists') + waiter.wait(TableName=table_name) + print(f"Table {table_name} has been deleted.") + except Exception as e: + print(f"Error deleting table {table_name}: {e}") \ No newline at end of file diff --git a/integrations/aws_bedrock/awsBedrock_x_ragas.ipynb b/integrations/aws_bedrock/awsBedrock_x_ragas.ipynb new file mode 100644 index 0000000..28fe479 --- /dev/null +++ b/integrations/aws_bedrock/awsBedrock_x_ragas.ipynb @@ -0,0 +1,2610 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a0bb5c39-2fde-4336-8127-8debe7cb2741", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "# Create and Evaluate an Agent Integrated with Bedrock Knowledge Bases and Attached Action Group\n", + "\n", + "In this notebook, you will learn how to evaluate an Amazon Bedrock Agent. The agent we'll evaluate is a restaurant agent whose tasks include providing clients with information about adult and children's menus and managing the table booking system. This agent is inspired by a [features example notebooks](https://github.com/aws-samples/amazon-bedrock-samples/tree/main/agents-and-function-calling/bedrock-agents/features-examples/05-create-agent-with-knowledge-base-and-action-group) of [Amazon Bedrock Agents](https://aws.amazon.com/bedrock/agents/) with minor changes. You can learn more about the agent creation process [here](https://github.com/aws-samples/amazon-bedrock-samples/tree/main/agents-and-function-calling/bedrock-agents/features-examples/05-create-agent-with-knowledge-base-and-action-group).\n", + "\n", + "The architecture is illustrated below:\n", + "\n", + "\n", + "
\n", + "\n", + "The steps covered in this notebook include:\n", + "\n", + "1. Importing necessary libraries\n", + "2. Creating the agent\n", + "3. Defining the Ragas metrics\n", + "4. Evaluating the agent\n", + "5. Cleaning up the created resources" + ] + }, + { + "cell_type": "markdown", + "id": "076a5aba-9735-4e98-8a53-0daccd7e94b0", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "## 1. Import the needed libraries" + ] + }, + { + "cell_type": "markdown", + "id": "4fa67d7a", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "First step is to install the pre-requisites packages" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ac05c073-d45b-4d85-9bf8-ae10aa78be8d", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "%pip install --upgrade -q boto3 opensearch-py botocore awscli retrying ragas" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "e8ad6ec2-b283-4c5d-879f-e397e46568c0", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "import os\n", + "import time\n", + "import boto3\n", + "import logging\n", + "import pprint\n", + "import json\n", + "\n", + "from knowledge_base import BedrockKnowledgeBase\n", + "from agent import (\n", + " create_agent_role_and_policies,\n", + " create_lambda_role,\n", + " delete_agent_roles_and_policies,\n", + " create_dynamodb,\n", + " create_lambda,\n", + " clean_up_resources,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d2b2d607-c1f2-4cbb-9f89-d935676e0101", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "# Clients\n", + "s3_client = boto3.client(\"s3\")\n", + "sts_client = boto3.client(\"sts\")\n", + "session = boto3.session.Session()\n", + "region = session.region_name\n", + "account_id = sts_client.get_caller_identity()[\"Account\"]\n", + "bedrock_agent_client = boto3.client(\"bedrock-agent\")\n", + "bedrock_agent_runtime_client = boto3.client(\"bedrock-agent-runtime\")\n", + "logging.basicConfig(\n", + " format=\"[%(asctime)s] p%(process)s {%(filename)s:%(lineno)d} %(levelname)s - %(message)s\",\n", + " level=logging.INFO,\n", + ")\n", + "logger = logging.getLogger(__name__)\n", + "region, account_id" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "d647d2a3", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "suffix = f\"{region}-{account_id}\"\n", + "agent_name = \"booking-agent\"\n", + "knowledge_base_name = f\"{agent_name}-kb\"\n", + "knowledge_base_description = (\n", + " \"Knowledge Base containing the restaurant menu's collection\"\n", + ")\n", + "agent_alias_name = \"booking-agent-alias\"\n", + "bucket_name = f\"{agent_name}-{suffix}\"\n", + "agent_bedrock_allow_policy_name = f\"{agent_name}-ba\"\n", + "agent_role_name = f\"AmazonBedrockExecutionRoleForAgents_{agent_name}\"\n", + "agent_foundation_model = \"amazon.nova-pro-v1:0\"\n", + "\n", + "agent_description = \"Agent in charge of a restaurants table bookings\"\n", + "agent_instruction = \"\"\"\n", + "You are a restaurant agent responsible for managing clients’ bookings (retrieving, creating, or canceling reservations) and assisting with menu inquiries. When handling menu requests, provide detailed information about the requested items. Offer recommendations only when:\n", + "\n", + "1. The customer explicitly asks for a recommendation, even if the item is available (include complementary dishes).\n", + "2. The requested item is unavailable—inform the customer and suggest suitable alternatives.\n", + "3. For general menu inquiries, provide the full menu and add a recommendation only if the customer asks for one.\n", + "\n", + "In all cases, ensure that any recommended items are present in the menu.\n", + "\n", + "Ensure all responses are clear, contextually relevant, and enhance the customer's experience.\n", + "\"\"\"\n", + "\n", + "agent_action_group_description = \"\"\"\n", + "Actions for getting table booking information, create a new booking or delete an existing booking\"\"\"\n", + "\n", + "agent_action_group_name = \"TableBookingsActionGroup\"" + ] + }, + { + "cell_type": "markdown", + "id": "b88af964", + "metadata": {}, + "source": [ + "## 2. Setting up Agent" + ] + }, + { + "cell_type": "markdown", + "id": "38c38fcb-9b87-414e-a644-04c263eea5c9", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "### 2.1 Create Knowledge Base for Amazon Bedrock\n", + "\n", + "Let's start by creating a [Knowledge Base for Amazon Bedrock](https://aws.amazon.com/bedrock/knowledge-bases/) to store the restaurant menus. For this example, we will integrate the knowledge base with Amazon OpenSearch Serverless." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2c39c2d9-7965-4c22-a74b-65d1961c4166", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "knowledge_base = BedrockKnowledgeBase(\n", + " kb_name=knowledge_base_name,\n", + " kb_description=knowledge_base_description,\n", + " data_bucket_name=bucket_name,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "cd8dab17", + "metadata": {}, + "source": [ + "### 2.2 Upload the Dataset to Amazon S3\n", + "\n", + "Now that we have created the knowledge base, let’s populate it with the restaurant menus dataset. In this example, we will use the [boto3 abstraction](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent/client/start_ingestion_job.html) of the API, via our helper classe. \n", + "\n", + "Let’s first upload the menu data available in the dataset folder to Amazon S3." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c74385d9", + "metadata": {}, + "outputs": [], + "source": [ + "def upload_directory(path, bucket_name):\n", + " for root, dirs, files in os.walk(path):\n", + " for file in files:\n", + " file_to_upload = os.path.join(root, file)\n", + " print(f\"uploading file {file_to_upload} to {bucket_name}\")\n", + " s3_client.upload_file(file_to_upload, bucket_name, file)\n", + "\n", + "\n", + "upload_directory(\"dataset\", bucket_name)" + ] + }, + { + "cell_type": "markdown", + "id": "48e46273-0267-4ecb-a686-2b1694b5a604", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "Now we start the ingestion job" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1bb7a9bc-5bba-4066-8fa3-a4c5a1385e95", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "# ensure that the kb is available\n", + "time.sleep(30)\n", + "# sync knowledge base\n", + "knowledge_base.start_ingestion_job()" + ] + }, + { + "cell_type": "markdown", + "id": "9c85cbee-9359-4927-ae16-fa7af42cf981", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "Finally we collect the Knowledge Base Id to integrate it with our Agent later on." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f71287af-0b7f-44d7-99ed-fa292ede4001", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "kb_id = knowledge_base.get_knowledge_base_id()" + ] + }, + { + "cell_type": "markdown", + "id": "315e0df7-4008-4fb4-b28d-a4df6ff446f6", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "#### Testing Knowledge Base with Retrieve and Generate API\n", + "\n", + "First, let’s test the knowledge base using the Retrieve and Generate API to ensure that the knowledge base is functioning correctly." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "f199e822-6e06-4bac-9dbf-40ff0be98598", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "response = bedrock_agent_runtime_client.retrieve_and_generate(\n", + " input={\"text\": \"Which are the mains available in the childrens menu?\"},\n", + " retrieveAndGenerateConfiguration={\n", + " \"type\": \"KNOWLEDGE_BASE\",\n", + " \"knowledgeBaseConfiguration\": {\n", + " \"knowledgeBaseId\": kb_id,\n", + " \"modelArn\": \"arn:aws:bedrock:{}::foundation-model/{}\".format(\n", + " region, agent_foundation_model\n", + " ),\n", + " \"retrievalConfiguration\": {\n", + " \"vectorSearchConfiguration\": {\"numberOfResults\": 5}\n", + " },\n", + " },\n", + " },\n", + ")\n", + "\n", + "print(response[\"output\"][\"text\"], end=\"\\n\" * 2)" + ] + }, + { + "cell_type": "markdown", + "id": "c5edf3bd-3214-4fe3-a9bc-df927e30b8b4", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "### 2.3 Create the DynamoDB Table\n", + "\n", + "We will create a DynamoDB table that contains restaurant booking information." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "18ff65c0-0207-4e34-932f-6d291f4d5c8d", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "table_name = \"restaurant_bookings\"\n", + "create_dynamodb(table_name)" + ] + }, + { + "cell_type": "markdown", + "id": "8460c785-e13f-4182-84d7-3bf7aafd3842", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "### 2.4 Create the Lambda Function\n", + "\n", + "We will now create a Lambda function that interacts with the DynamoDB table." + ] + }, + { + "cell_type": "markdown", + "id": "68871530-2953-4a37-854e-10323cabf095", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "#### Create the Function Code\n", + "\n", + "Create the Lambda function that implements the functions for `get_booking_details`, `create_booking`, and `delete_booking`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12271f1f-4739-472b-8909-cc81c148a941", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "%%writefile lambda_function.py\n", + "import json\n", + "import uuid\n", + "import boto3\n", + "\n", + "dynamodb = boto3.resource('dynamodb')\n", + "table = dynamodb.Table('restaurant_bookings')\n", + "\n", + "def get_named_parameter(event, name):\n", + " \"\"\"\n", + " Get a parameter from the lambda event\n", + " \"\"\"\n", + " return next(item for item in event['parameters'] if item['name'] == name)['value']\n", + "\n", + "\n", + "def get_booking_details(booking_id):\n", + " \"\"\"\n", + " Retrieve details of a restaurant booking\n", + " \n", + " Args:\n", + " booking_id (string): The ID of the booking to retrieve\n", + " \"\"\"\n", + " try:\n", + " response = table.get_item(Key={'booking_id': booking_id})\n", + " if 'Item' in response:\n", + " return response['Item']\n", + " else:\n", + " return {'message': f'No booking found with ID {booking_id}'}\n", + " except Exception as e:\n", + " return {'error': str(e)}\n", + "\n", + "\n", + "def create_booking(date, name, hour, num_guests):\n", + " \"\"\"\n", + " Create a new restaurant booking\n", + " \n", + " Args:\n", + " date (string): The date of the booking\n", + " name (string): Name to idenfity your reservation\n", + " hour (string): The hour of the booking\n", + " num_guests (integer): The number of guests for the booking\n", + " \"\"\"\n", + " try:\n", + " booking_id = str(uuid.uuid4())[:8]\n", + " table.put_item(\n", + " Item={\n", + " 'booking_id': booking_id,\n", + " 'date': date,\n", + " 'name': name,\n", + " 'hour': hour,\n", + " 'num_guests': num_guests\n", + " }\n", + " )\n", + " return {'booking_id': booking_id}\n", + " except Exception as e:\n", + " return {'error': str(e)}\n", + "\n", + "\n", + "def delete_booking(booking_id):\n", + " \"\"\"\n", + " Delete an existing restaurant booking\n", + " \n", + " Args:\n", + " booking_id (str): The ID of the booking to delete\n", + " \"\"\"\n", + " try:\n", + " response = table.delete_item(Key={'booking_id': booking_id})\n", + " if response['ResponseMetadata']['HTTPStatusCode'] == 200:\n", + " return {'message': f'Booking with ID {booking_id} deleted successfully'}\n", + " else:\n", + " return {'message': f'Failed to delete booking with ID {booking_id}'}\n", + " except Exception as e:\n", + " return {'error': str(e)}\n", + " \n", + "\n", + "def lambda_handler(event, context):\n", + " # get the action group used during the invocation of the lambda function\n", + " actionGroup = event.get('actionGroup', '')\n", + " \n", + " # name of the function that should be invoked\n", + " function = event.get('function', '')\n", + " \n", + " # parameters to invoke function with\n", + " parameters = event.get('parameters', [])\n", + "\n", + " if function == 'get_booking_details':\n", + " booking_id = get_named_parameter(event, \"booking_id\")\n", + " if booking_id:\n", + " response = str(get_booking_details(booking_id))\n", + " responseBody = {'TEXT': {'body': json.dumps(response)}}\n", + " else:\n", + " responseBody = {'TEXT': {'body': 'Missing booking_id parameter'}}\n", + "\n", + " elif function == 'create_booking':\n", + " date = get_named_parameter(event, \"date\")\n", + " name = get_named_parameter(event, \"name\")\n", + " hour = get_named_parameter(event, \"hour\")\n", + " num_guests = get_named_parameter(event, \"num_guests\")\n", + "\n", + " if date and hour and num_guests:\n", + " response = str(create_booking(date, name, hour, num_guests))\n", + " responseBody = {'TEXT': {'body': json.dumps(response)}}\n", + " else:\n", + " responseBody = {'TEXT': {'body': 'Missing required parameters'}}\n", + "\n", + " elif function == 'delete_booking':\n", + " booking_id = get_named_parameter(event, \"booking_id\")\n", + " if booking_id:\n", + " response = str(delete_booking(booking_id))\n", + " responseBody = {'TEXT': {'body': json.dumps(response)}}\n", + " else:\n", + " responseBody = {'TEXT': {'body': 'Missing booking_id parameter'}}\n", + "\n", + " else:\n", + " responseBody = {'TEXT': {'body': 'Invalid function'}}\n", + "\n", + " action_response = {\n", + " 'actionGroup': actionGroup,\n", + " 'function': function,\n", + " 'functionResponse': {\n", + " 'responseBody': responseBody\n", + " }\n", + " }\n", + "\n", + " function_response = {'response': action_response, 'messageVersion': event['messageVersion']}\n", + " print(\"Response: {}\".format(function_response))\n", + "\n", + " return function_response" + ] + }, + { + "cell_type": "markdown", + "id": "b8a4eefa-0ee8-4505-bba4-b62f5ba54a79", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "#### Create the required permissions" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "490addc9-16a0-48e1-bd8f-02abb56bb520", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "lambda_iam_role = create_lambda_role(agent_name, table_name)" + ] + }, + { + "cell_type": "markdown", + "id": "895032f3", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "#### Create the function" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "682ea2b4-5654-4e2f-b3dd-bc763a3c5918", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "lambda_function_name = f\"{agent_name}-lambda\"\n", + "lambda_function = create_lambda(lambda_function_name, lambda_iam_role)" + ] + }, + { + "cell_type": "markdown", + "id": "aef843aa-9f0f-473f-ab4e-1c4827974b87", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "### 2.5 Create the IAM Policies Needed for the Agent\n", + "\n", + "Now that we have created the Knowledge Base, our DynamoDB table, and the Lambda function to execute the tasks for our Agent, let’s start creating our Agent." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "aa92a151-36ca-4099-8bc6-156b49b0c3e4", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "agent_role = create_agent_role_and_policies(\n", + " agent_name, agent_foundation_model, kb_id=kb_id\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "595bf767-3a26-4f92-b629-ba2908b90d81", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "### 2.6 Create the Agent\n", + "\n", + "Now that we have created the necessary IAM role, we can use the [`create_agent`](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent/client/create_agent.html) API from boto3 to create a new agent." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7ccaf856-7eff-4bd1-ac47-c89b6fbf7e8f", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "response = bedrock_agent_client.create_agent(\n", + " agentName=agent_name,\n", + " agentResourceRoleArn=agent_role[\"Role\"][\"Arn\"],\n", + " description=agent_description,\n", + " idleSessionTTLInSeconds=1800,\n", + " foundationModel=agent_foundation_model,\n", + " instruction=agent_instruction,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "b6b5884e", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "Let's get our Agent ID. It will be important to perform operations with our agent" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "40c48d3c-ed67-450b-b840-23ec1a99fee7", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "agent_id = response[\"agent\"][\"agentId\"]\n", + "print(\"The agent id is:\", agent_id)" + ] + }, + { + "cell_type": "markdown", + "id": "3f30f874-eb63-4ec3-8c19-36fab0f598df", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "### 2.7 Create the Agent Action Group\n", + "\n", + "We will now create an Agent Action Group that uses the Lambda function created earlier. To inform the agent about the capabilities of the action group, we will provide a description outlining its functionalities.\n", + "\n", + "To define the functions using a function schema, you need to provide the name, description, and parameters for each function." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "91c4c380-ca58-467a-ab67-b475605a0274", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "agent_functions = [\n", + " {\n", + " \"name\": \"get_booking_details\",\n", + " \"description\": \"Retrieve details of a restaurant booking\",\n", + " \"parameters\": {\n", + " \"booking_id\": {\n", + " \"description\": \"The ID of the booking to retrieve\",\n", + " \"required\": True,\n", + " \"type\": \"string\",\n", + " }\n", + " },\n", + " },\n", + " {\n", + " \"name\": \"create_booking\",\n", + " \"description\": \"Create a new restaurant booking\",\n", + " \"parameters\": {\n", + " \"date\": {\n", + " \"description\": \"The date of the booking\",\n", + " \"required\": True,\n", + " \"type\": \"string\",\n", + " },\n", + " \"name\": {\n", + " \"description\": \"Name to idenfity your reservation\",\n", + " \"required\": True,\n", + " \"type\": \"string\",\n", + " },\n", + " \"hour\": {\n", + " \"description\": \"The hour of the booking\",\n", + " \"required\": True,\n", + " \"type\": \"string\",\n", + " },\n", + " \"num_guests\": {\n", + " \"description\": \"The number of guests for the booking\",\n", + " \"required\": True,\n", + " \"type\": \"integer\",\n", + " },\n", + " },\n", + " },\n", + " {\n", + " \"name\": \"delete_booking\",\n", + " \"description\": \"Delete an existing restaurant booking\",\n", + " \"parameters\": {\n", + " \"booking_id\": {\n", + " \"description\": \"The ID of the booking to delete\",\n", + " \"required\": True,\n", + " \"type\": \"string\",\n", + " }\n", + " },\n", + " },\n", + "]" + ] + }, + { + "cell_type": "markdown", + "id": "60d174f6", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "We now use the function schema to create the agent action group using the [`create_agent_action_group`](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent/client/create_agent_action_group.html) API" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "24559f1c-e081-4f03-959f-a82f9444aa49", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "# Pause to make sure agent is created\n", + "time.sleep(30)\n", + "\n", + "# Now, we can configure and create an action group here:\n", + "agent_action_group_response = bedrock_agent_client.create_agent_action_group(\n", + " agentId=agent_id,\n", + " agentVersion=\"DRAFT\",\n", + " actionGroupExecutor={\"lambda\": lambda_function[\"FunctionArn\"]},\n", + " actionGroupName=agent_action_group_name,\n", + " functionSchema={\"functions\": agent_functions},\n", + " description=agent_action_group_description,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "b7d27e07-7d4c-4c59-987d-41a606c6af65", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "### 2.8 Allow the Agent to invoke the Action Group Lambda" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3262f6dd-fb05-4351-904d-8097c0f52134", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "# Create allow to invoke permission on lambda\n", + "lambda_client = boto3.client(\"lambda\")\n", + "response = lambda_client.add_permission(\n", + " FunctionName=lambda_function_name,\n", + " StatementId=\"allow_bedrock\",\n", + " Action=\"lambda:InvokeFunction\",\n", + " Principal=\"bedrock.amazonaws.com\",\n", + " SourceArn=f\"arn:aws:bedrock:{region}:{account_id}:agent/{agent_id}\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "2c8536a4-e625-464c-bae2-d28dba82c556", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "### 2.9 Associate the Knowledge Base to the agent" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "95c7b781-7e82-47b6-a268-237fa97ab58d", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "response = bedrock_agent_client.associate_agent_knowledge_base(\n", + " agentId=agent_id,\n", + " agentVersion=\"DRAFT\",\n", + " description=\"Access the knowledge base when customers ask about the plates in the menu.\",\n", + " knowledgeBaseId=kb_id,\n", + " knowledgeBaseState=\"ENABLED\",\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "b42a8caf-4c76-4760-a7cd-4c6b189f9f5b", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "### 2.10 Prepare the Agent and create an alias\n", + "\n", + "Let's create a DRAFT version of the agent that can be used for internal testing." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4013c967-9a14-4689-b0b5-61aa4f75748c", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "response = bedrock_agent_client.prepare_agent(agentId=agent_id)\n", + "print(response)\n", + "# Pause to make sure agent is prepared\n", + "time.sleep(30)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b9889aaf-9238-479d-b9f1-2ab13b1ec9c0", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "response = bedrock_agent_client.create_agent_alias(\n", + " agentAliasName=\"TestAlias\",\n", + " agentId=agent_id,\n", + " description=\"Test alias\",\n", + ")\n", + "\n", + "alias_id = response[\"agentAlias\"][\"agentAliasId\"]\n", + "print(\"The Agent alias is:\", alias_id)\n", + "time.sleep(30)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8ee4a530", + "metadata": {}, + "outputs": [], + "source": [ + "def invokeAgent(query, session_id, enable_trace=True, session_state=dict()):\n", + " end_session: bool = False\n", + "\n", + " # invoke the agent API\n", + " agentResponse = bedrock_agent_runtime_client.invoke_agent(\n", + " inputText=query,\n", + " agentId=agent_id,\n", + " agentAliasId=alias_id,\n", + " sessionId=session_id,\n", + " enableTrace=enable_trace,\n", + " endSession=end_session,\n", + " sessionState=session_state,\n", + " )\n", + "\n", + " event_stream = agentResponse[\"completion\"]\n", + " try:\n", + " traces = []\n", + " for event in event_stream:\n", + " if \"chunk\" in event:\n", + " data = event[\"chunk\"][\"bytes\"]\n", + " agent_answer = data.decode(\"utf8\")\n", + " end_event_received = True\n", + " return agent_answer, traces\n", + " # End event indicates that the request finished successfully\n", + " elif \"trace\" in event:\n", + " if enable_trace:\n", + " traces.append(event[\"trace\"])\n", + " else:\n", + " raise Exception(\"unexpected event.\", event)\n", + " return agent_answer, traces\n", + " except Exception as e:\n", + " raise Exception(\"unexpected event.\", e)" + ] + }, + { + "cell_type": "markdown", + "id": "9ff4105c", + "metadata": {}, + "source": [ + "## 3. Defining the Ragas metrics" + ] + }, + { + "cell_type": "markdown", + "id": "9cfc3bb2", + "metadata": {}, + "source": [ + "Evaluating agents is different from testing traditional software, where you can simply verify whether the output matches expected results. These agents perform complex tasks that often have multiple valid approaches.\n", + "\n", + "Given their inherent autonomy, evaluating agents is essential to ensure they function properly.\n", + "\n", + "#### Choosing What to Evaluate in Your Agent\n", + "\n", + "Selecting evaluation metrics depends entirely on your use case. A good rule of thumb is to select metrics directly tied to user needs or metrics that clearly drive business value. In the restaurant agent example above, we want the agent to fulfill user requests without unnecessary repetition, provide helpful recommendations when appropriate to enhance customer experience, and maintain consistency with the brand tone.\n", + "\n", + "We’ll define metrics to evaluate these priorities. Ragas provides several user-defined metrics for evaluations.\n", + "\n", + "When defining evaluation criteria, focus on binary decisions or discrete classification scores rather than ambiguous scores. Binary or clear classifications compel you to explicitly define success criteria. Avoid metrics yielding scores between 0 and 100 without clear interpretation, as distinguishing between close scores like 87 and 91 can be challenging, especially when evaluations occur independently.\n", + "\n", + "Ragas includes metrics suited to such evaluations, and we will explore some of them in action:\n", + "- [**Aspect Critic Metric**](): Evaluates whether a submission follows user-defined criteria by leveraging LLM judgments to yield a binary outcome.\n", + "- [**Rubric Score Metric**](): Assesses responses against detailed, user-defined rubrics to consistently assign scores reflecting quality." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "93666bbf", + "metadata": {}, + "outputs": [], + "source": [ + "import getpass\n", + "import os\n", + "\n", + "if \"OPENAI_API_KEY\" not in os.environ:\n", + " os.environ[\"OPENAI_API_KEY\"] = getpass.getpass(\"Enter your OpenAI API key: \")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "91331315", + "metadata": {}, + "outputs": [], + "source": [ + "from ragas.llms import LangchainLLMWrapper\n", + "from ragas.embeddings import LangchainEmbeddingsWrapper\n", + "from langchain_openai import ChatOpenAI, OpenAIEmbeddings\n", + "\n", + "\n", + "evaluator_llm = LangchainLLMWrapper(ChatOpenAI(model=\"gpt-4o-mini\"))\n", + "evaluator_embeddings = LangchainEmbeddingsWrapper(\n", + " OpenAIEmbeddings(model=\"text-embedding-3-small\")\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 100, + "id": "7e13a331", + "metadata": {}, + "outputs": [], + "source": [ + "from ragas.metrics import AspectCritic, RubricsScore\n", + "from ragas.dataset_schema import SingleTurnSample, MultiTurnSample, EvaluationDataset\n", + "from ragas import evaluate\n", + "\n", + "rubrics = {\n", + " \"score-1_description\": (\n", + " \"The item requested by the customer is not present in the menu and no recommendations were made.\"\n", + " ),\n", + " \"score0_description\": (\n", + " \"Either the item requested by the customer is present in the menu, or the conversation does not include any food or menu inquiry (e.g., booking, cancellation), \"\n", + " \"regardless of whether any recommendation was provided.\"\n", + " ),\n", + " \"score1_description\": (\n", + " \"The item requested by the customer is not present in the menu and a recommendation was provided.\"\n", + " ),\n", + "}\n", + "\n", + "recommendations = RubricsScore(rubrics=rubrics, llm=evaluator_llm, name=\"Recommendations\")\n", + "\n", + "\n", + "# Metric to evaluate if the AI fulfills all human requests completely.\n", + "request_completeness = AspectCritic(\n", + " name=\"Request Completeness\",\n", + " llm=evaluator_llm,\n", + " definition=(\n", + " \"Return 1 The agent completely fulfills all the user requests with no omissions. \"\n", + " \"otherwise, return 0.\"\n", + " ),\n", + ")\n", + "\n", + "# Metric to assess if the AI's communication aligns with the desired brand voice.\n", + "brand_tone = AspectCritic(\n", + " name=\"Brand Voice Metric\",\n", + " llm=evaluator_llm,\n", + " definition=(\n", + " \"Return 1 if the AI's communication is friendly, approachable, helpful, clear, and concise; \"\n", + " \"otherwise, return 0.\"\n", + " ),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "ca5b4a45", + "metadata": {}, + "source": [ + "## 4. Evaluating Agent with Ragas" + ] + }, + { + "cell_type": "markdown", + "id": "5def6241", + "metadata": {}, + "source": [ + "In order to perform evaluations using Ragas, the traces need to be converted into the format recognized by Ragas. To convert an AWS Bedrock agent trace into a format suitable for Ragas evaluation, Ragas provides the function [convert_to_ragas_messages][ragas.integrations.swarm.convert_to_ragas_messages], which can be used to transform AWS Bedrock messages into the format expected by Ragas. You can read more about it [here]()." + ] + }, + { + "cell_type": "code", + "execution_count": 93, + "id": "1c3b4f99", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "\n", + "Your booking for 2 people at 7pm on the 5th of May 2025 has been successfully created. Your booking ID is ca2fab70.\n", + "\n", + "\n", + "\n", + "\n", + "CPU times: user 38.3 ms, sys: 46.6 ms, total: 84.9 ms\n", + "Wall time: 10.6 s\n" + ] + } + ], + "source": [ + "%%time\n", + "import uuid\n", + "session_id:str = str(uuid.uuid1())\n", + "query = \"If you have children food then book a table for 2 people at 7pm on the 5th of May 2025.\"\n", + "agent_answer, traces_1 = invokeAgent(query, session_id)\n", + "\n", + "print(agent_answer)" + ] + }, + { + "cell_type": "code", + "execution_count": 94, + "id": "3ac474b8", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Your reservation was found and has been successfully canceled.\n", + "\n" + ] + } + ], + "source": [ + "query = \"Can you check if my previous booking? can you please delete the booking\"\n", + "agent_answer, traces_2 = invokeAgent(query, session_id)\n", + "\n", + "print(agent_answer)" + ] + }, + { + "cell_type": "code", + "execution_count": 95, + "id": "3d181de4", + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "710c7aeb02884fdf8a33a3abf0e425a7", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Evaluating: 0%| | 0/4 [00:00\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
user_inputRequest CompletenessBrand Voice Metric
0[{'content': '[{text=If you have children food...11
1[{'content': '[{text=If you have children food...11
\n", + "" + ], + "text/plain": [ + " user_input Request Completeness \\\n", + "0 [{'content': '[{text=If you have children food... 1 \n", + "1 [{'content': '[{text=If you have children food... 1 \n", + "\n", + " Brand Voice Metric \n", + "0 1 \n", + "1 1 " + ] + }, + "execution_count": 95, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from aws_bedrock import convert_to_ragas_messages\n", + "\n", + "# Convert AWS traces to messages accepted by RAGAS.\n", + "# The convert_to_ragas_messages function transforms AWS-specific trace data \n", + "# into a format that RAGAS can process as conversation messages.\n", + "ragas_messages_trace_1 = convert_to_ragas_messages(traces_1)\n", + "ragas_messages_trace_2 = convert_to_ragas_messages(traces_2)\n", + "\n", + "# Initialize MultiTurnSample objects.\n", + "# MultiTurnSample is a data type defined in RAGAS that encapsulates conversation\n", + "# data for multi-turn evaluation. This conversion is necessary to perform evaluations.\n", + "sample_1 = MultiTurnSample(user_input=ragas_messages_trace_1)\n", + "sample_2 = MultiTurnSample(user_input=ragas_messages_trace_2)\n", + "\n", + "result = evaluate(\n", + " # Create an evaluation dataset from the multi-turn samples\n", + " dataset=EvaluationDataset(samples=[sample_1, sample_2]),\n", + " metrics=[request_completeness, brand_tone],\n", + ")\n", + "\n", + "result.to_pandas()" + ] + }, + { + "cell_type": "markdown", + "id": "af8cd835", + "metadata": {}, + "source": [ + "The scores of 1 were awarded because the agent fully met all user requests without any omissions (completeness) and communicated in a friendly, approachable, helpful, clear, and concise manner (brand voice) for both the conversations." + ] + }, + { + "cell_type": "code", + "execution_count": 96, + "id": "496fcb43", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Yes, we serve Chicken Wings. Here are the details:\n", + "- **Buffalo Chicken Wings**: Classic buffalo wings served with celery sticks and blue cheese dressing. Allergens: Dairy (in blue cheese dressing), Gluten (in the coating), possible Soy (in the sauce).\n", + "\n", + "CPU times: user 34.3 ms, sys: 40.3 ms, total: 74.6 ms\n", + "Wall time: 8.14 s\n" + ] + } + ], + "source": [ + "%%time\n", + "import uuid\n", + "\n", + "session_id:str = str(uuid.uuid1())\n", + "query = \"Do you serve Chicken Wings?\"\n", + "\n", + "agent_answer, traces_3 = invokeAgent(query, session_id)\n", + "print(agent_answer)" + ] + }, + { + "cell_type": "code", + "execution_count": 97, + "id": "581d481d", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "I'm sorry, but we do not have chocolate truffle cake on our dessert menu. However, we have several delicious alternatives you might enjoy:\n", + "\n", + "1. **Classic New York Cheesecake** - Creamy cheesecake with a graham cracker crust, topped with a choice of fruit compote or chocolate ganache.\n", + "2. **Apple Pie à la Mode** - Warm apple pie with a flaky crust, served with a scoop of vanilla ice cream and a drizzle of caramel sauce.\n", + "3. **Chocolate Lava Cake** - Rich and gooey chocolate cake with a molten center, dusted with powdered sugar and served with a scoop of raspberry sorbet.\n", + "4. **Pecan Pie Bars** - Buttery shortbread crust topped with a gooey pecan filling, cut into bars for easy serving.\n", + "5. **Banana Pudding Parfait** - Layers of vanilla pudding, sliced bananas, and vanilla wafers, topped with whipped cream and a sprinkle of crushed nuts.\n", + "\n", + "May I recommend the **Chocolate Lava Cake** for a decadent treat?\n", + "CPU times: user 27.9 ms, sys: 30.7 ms, total: 58.5 ms\n", + "Wall time: 10.9 s\n" + ] + } + ], + "source": [ + "%%time\n", + "session_id:str = str(uuid.uuid1())\n", + "query = \"For desserts, do you have chocolate truffle cake?\"\n", + "agent_answer, traces_4 = invokeAgent(query, session_id)\n", + "print(agent_answer)" + ] + }, + { + "cell_type": "code", + "execution_count": 98, + "id": "a1471ddc", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "I could not find Indian food on our menu. However, we offer a variety of other cuisines including American, Italian, and vegetarian options. Would you like to know more about these options? \n", + "CPU times: user 24.1 ms, sys: 20 ms, total: 44.1 ms\n", + "Wall time: 6.55 s\n" + ] + } + ], + "source": [ + "%%time\n", + "from datetime import datetime\n", + "today = datetime.today().strftime('%b-%d-%Y')\n", + "\n", + "session_id:str = str(uuid.uuid1())\n", + "query = \"Do you have indian food?\"\n", + "session_state = {\n", + " \"promptSessionAttributes\": {\n", + " \"name\": \"John\",\n", + " \"today\": today\n", + " }\n", + "}\n", + "\n", + "agent_answer, traces_5 = invokeAgent(query, session_id, session_state=session_state)\n", + "print(agent_answer)" + ] + }, + { + "cell_type": "code", + "execution_count": 101, + "id": "15590e19", + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "a6b520d376f345caa86ced0d9c3e2367", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Evaluating: 0%| | 0/3 [00:00\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
user_inputRecommendations
0[{'content': '[{text=Do you serve Chicken Wing...0
1[{'content': '[{text=For desserts, do you have...1
2[{'content': '[{text=Do you have indian food?}...1
\n", + "" + ], + "text/plain": [ + " user_input Recommendations\n", + "0 [{'content': '[{text=Do you serve Chicken Wing... 0\n", + "1 [{'content': '[{text=For desserts, do you have... 1\n", + "2 [{'content': '[{text=Do you have indian food?}... 1" + ] + }, + "execution_count": 101, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from aws_bedrock import convert_to_ragas_messages\n", + "\n", + "ragas_messages_trace_3 = convert_to_ragas_messages(traces_3)\n", + "ragas_messages_trace_4 = convert_to_ragas_messages(traces_4)\n", + "ragas_messages_trace_5 = convert_to_ragas_messages(traces_5)\n", + "\n", + "sample_3 = MultiTurnSample(user_input=ragas_messages_trace_3)\n", + "sample_4 = MultiTurnSample(user_input=ragas_messages_trace_4)\n", + "sample_5 = MultiTurnSample(user_input=ragas_messages_trace_5)\n", + "\n", + "result = evaluate(\n", + " dataset=EvaluationDataset(samples=[sample_3, sample_4, sample_5]),\n", + " metrics=[recommendations],\n", + ")\n", + "\n", + "result.to_pandas()" + ] + }, + { + "cell_type": "markdown", + "id": "851b63b1", + "metadata": {}, + "source": [ + "For the Recommendation metric, the chicken wings inquiry scored 0 since the item was available, while both the chocolate truffle cake and Indian food inquiries scored 1 because the requested items were not on the menu and alternative recommendations were provided." + ] + }, + { + "cell_type": "markdown", + "id": "9351a76f", + "metadata": {}, + "source": [ + "To evaluate how well our agent utilizes information retrieved from the knowledge base, we use the RAG evaluation metrics provided by Ragas. You can learn more about these metrics [here]().\n", + "\n", + "In this tutorial, we will use the following RAG metrics:\n", + "- [**ContextRelevance**](): Measures how well the retrieved contexts address the user’s query by evaluating their pertinence through dual LLM judgments.\n", + "- [**Faithfulness**](): Assesses the factual consistency of the response by determining whether all its claims can be supported by the provided retrieved contexts.\n", + "- [**ResponseGroundedness**](): Determines the extent to which each claim in the response is directly supported or “grounded” in the provided contexts." + ] + }, + { + "cell_type": "code", + "execution_count": 102, + "id": "457a10c4", + "metadata": {}, + "outputs": [], + "source": [ + "from ragas.metrics import ContextRelevance, Faithfulness, ResponseGroundedness\n", + "\n", + "metrics = [\n", + " ContextRelevance(llm=evaluator_llm),\n", + " Faithfulness(llm=evaluator_llm),\n", + " ResponseGroundedness(llm=evaluator_llm),\n", + "]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "283cf8d4", + "metadata": {}, + "outputs": [], + "source": [ + "from aws_bedrock import extract_kb_trace\n", + "\n", + "kb_trace_3 = extract_kb_trace(traces_3)\n", + "kb_trace_4 = extract_kb_trace(traces_4)\n", + "\n", + "trace_3_single_turn_sample = SingleTurnSample(\n", + " user_input=kb_trace_3[0].get(\"user_input\"),\n", + " retrieved_contexts=kb_trace_3[0].get(\"retrieved_contexts\"),\n", + " response=kb_trace_3[0].get(\"response\"),\n", + " reference=\"Yes, we do serve chicken wings prepared in Buffalo style, chicken wing that’s typically deep-fried and then tossed in a tangy, spicy Buffalo sauce.\",\n", + ")\n", + "\n", + "trace_4_single_turn_sample = SingleTurnSample(\n", + " user_input=kb_trace_4[0].get(\"user_input\"),\n", + " retrieved_contexts=kb_trace_4[0].get(\"retrieved_contexts\"),\n", + " response=kb_trace_4[0].get(\"response\"),\n", + " reference=\"The desserts on the adult menu are:\\n1. Classic New York Cheesecake\\n2. Apple Pie à la Mode\\n3. Chocolate Lava Cake\\4. Pecan Pie Bars\\n5. Banana Pudding Parfait\",\n", + ")\n", + "\n", + "single_turn_samples = [trace_3_single_turn_sample, trace_4_single_turn_sample]\n", + "\n", + "dataset = EvaluationDataset(samples=single_turn_samples)" + ] + }, + { + "cell_type": "code", + "execution_count": 104, + "id": "1b25e0a1", + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "ec1b31fd7e0d4938ad2cc38b8da268fd", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Evaluating: 0%| | 0/6 [00:00\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
user_inputretrieved_contextsresponsereferencenv_context_relevancefaithfulnessnv_response_groundedness
0Chicken Wings[The Regrettable Experience -- Dinner Menu Ent...Yes, we serve Chicken Wings. Here are the deta...Yes, we do serve chicken wings prepared in Buf...1.01.001.0
1chocolate truffle cake[Allergens: Gluten (in the breading). 3. B...I'm sorry, but we do not have chocolate truffl...The desserts on the adult menu are:\\n1. Classi...0.00.750.5
\n", + "" + ], + "text/plain": [ + " user_input retrieved_contexts \\\n", + "0 Chicken Wings [The Regrettable Experience -- Dinner Menu Ent... \n", + "1 chocolate truffle cake [Allergens: Gluten (in the breading). 3. B... \n", + "\n", + " response \\\n", + "0 Yes, we serve Chicken Wings. Here are the deta... \n", + "1 I'm sorry, but we do not have chocolate truffl... \n", + "\n", + " reference nv_context_relevance \\\n", + "0 Yes, we do serve chicken wings prepared in Buf... 1.0 \n", + "1 The desserts on the adult menu are:\\n1. Classi... 0.0 \n", + "\n", + " faithfulness nv_response_groundedness \n", + "0 1.00 1.0 \n", + "1 0.75 0.5 " + ] + }, + "execution_count": 104, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "kb_results = evaluate(dataset=dataset, metrics=metrics)\n", + "kb_results.to_pandas()" + ] + }, + { + "cell_type": "markdown", + "id": "63fde378", + "metadata": {}, + "source": [ + "Corrected Snippet:\n", + "\n", + "To evaluate whether the agent is able to achieve its goal, we can use the following metrics:\n", + "- [**AgentGoalAccuracyWithReference**](): Determines if the AI achieved the user’s goal by comparing its final outcome against an annotated ideal outcome, yielding a binary result.\n", + "- [**AgentGoalAccuracyWithoutReference**](): Infers whether the AI met the user’s goal solely based on conversational interactions, providing a binary success indicator without an explicit reference." + ] + }, + { + "cell_type": "code", + "execution_count": 105, + "id": "6d141ec5", + "metadata": {}, + "outputs": [], + "source": [ + "from ragas.metrics import (\n", + " AgentGoalAccuracyWithoutReference,\n", + " AgentGoalAccuracyWithReference,\n", + ")\n", + "\n", + "goal_accuracy_with_reference = AgentGoalAccuracyWithReference(llm=evaluator_llm)\n", + "goal_accuracy_without_reference = AgentGoalAccuracyWithoutReference(llm=evaluator_llm)" + ] + }, + { + "cell_type": "code", + "execution_count": 109, + "id": "d80f5d80", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Here are the entrees available for children:\n", + "1. CHICKEN NUGGETS - Crispy chicken nuggets served with a side of ketchup or ranch dressing. Allergens: Gluten (in the coating), possible Soy. Suitable for Vegetarians: No\n", + "2. MACARONI AND CHEESE - Classic macaroni pasta smothered in creamy cheese sauce. Allergens: Dairy, Gluten. Suitable for Vegetarians: Yes\n", + "3. MINI CHEESE QUESADILLAS - Small flour tortillas filled with melted cheese, served with a mild salsa. Allergens: Dairy, Gluten. Suitable for Vegetarians: Yes\n", + "4. PEANUT BUTTER AND BANANA SANDWICH - Peanut butter and banana slices on whole wheat bread. Allergens: Nuts (peanut), Gluten. Suitable for Vegetarians: Yes (if using vegetarian peanut butter)\n", + "5. VEGGIE PITA POCKETS - Mini whole wheat pita pockets filled with hummus, cucumber, and cherry tomatoes. Allergens: Gluten, possible Soy. Suitable for Vegetarians: Yes\n", + "\n", + "\n", + "\n", + "CPU times: user 21.6 ms, sys: 28.1 ms, total: 49.7 ms\n", + "Wall time: 11.6 s\n" + ] + } + ], + "source": [ + "%%time\n", + "import uuid\n", + "\n", + "session_id:str = str(uuid.uuid1())\n", + "query = \"What entrees do you have for children?\"\n", + "\n", + "agent_answer, traces_6 = invokeAgent(query, session_id)\n", + "print(agent_answer)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "ce08fe21", + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "cece52e7cc054b399b41e083df347788", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Evaluating: 0%| | 0/1 [00:00\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
user_inputreferenceagent_goal_accuracy
0[{'content': '[{text=What entrees do you have ...The final outcome provides child-friendly entr...1.0
\n", + "" + ], + "text/plain": [ + " user_input \\\n", + "0 [{'content': '[{text=What entrees do you have ... \n", + "\n", + " reference agent_goal_accuracy \n", + "0 The final outcome provides child-friendly entr... 1.0 " + ] + }, + "execution_count": 107, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "from aws_bedrock import convert_to_ragas_messages\n", + "\n", + "ragas_messages_trace_6 = convert_to_ragas_messages(traces_6)\n", + "\n", + "sample_6 = MultiTurnSample(\n", + " user_input=ragas_messages_trace_6,\n", + " reference=\"Response contains entrees food items for the children.\",\n", + ")\n", + "\n", + "result = evaluate(\n", + " dataset=EvaluationDataset(samples=[sample_6]),\n", + " metrics=[goal_accuracy_with_reference],\n", + ")\n", + "\n", + "result.to_pandas()" + ] + }, + { + "cell_type": "code", + "execution_count": 108, + "id": "a2b12d76", + "metadata": {}, + "outputs": [ + { + "data": { + "application/vnd.jupyter.widget-view+json": { + "model_id": "b400d8aeecc14c4ab0cde031afa72c5e", + "version_major": 2, + "version_minor": 0 + }, + "text/plain": [ + "Evaluating: 0%| | 0/1 [00:00\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
user_inputagent_goal_accuracy
0[{'content': '[{text=What entrees do you have ...1.0
\n", + "" + ], + "text/plain": [ + " user_input agent_goal_accuracy\n", + "0 [{'content': '[{text=What entrees do you have ... 1.0" + ] + }, + "execution_count": 108, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "sample_6 = MultiTurnSample(user_input=ragas_messages_trace_6)\n", + "\n", + "result = evaluate(\n", + " dataset=EvaluationDataset(samples=[sample_6]),\n", + " metrics=[goal_accuracy_without_reference],\n", + ")\n", + "\n", + "result.to_pandas()" + ] + }, + { + "cell_type": "markdown", + "id": "8bb5d818", + "metadata": {}, + "source": [ + "In both scenarios, the agent earned a score of 1 by comprehensively providing all available options—whether listing all children’s entrees." + ] + }, + { + "cell_type": "markdown", + "id": "8ebf4438-1f48-4642-a57c-530a16815064", + "metadata": { + "pycharm": { + "name": "#%% md\n" + } + }, + "source": [ + "## 5. Clean-up \n", + "Let's delete all the associated resources created to avoid unnecessary costs. " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4ed0cb4f-31ab-4535-a1d7-93e33c706b32", + "metadata": { + "pycharm": { + "name": "#%%\n" + }, + "tags": [] + }, + "outputs": [], + "source": [ + "clean_up_resources(\n", + " table_name,\n", + " lambda_function,\n", + " lambda_function_name,\n", + " agent_action_group_response,\n", + " agent_functions,\n", + " agent_id,\n", + " kb_id,\n", + " alias_id,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 111, + "id": "6a9db157", + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "# Delete the agent roles and policies\n", + "delete_agent_roles_and_policies(agent_name)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "17ec0e70", + "metadata": { + "pycharm": { + "name": "#%%\n" + } + }, + "outputs": [], + "source": [ + "# delete KB\n", + "knowledge_base.delete_kb(delete_s3_bucket=True, delete_iam_roles_and_policies=True)" + ] + } + ], + "metadata": { + "availableInstances": [ + { + "_defaultOrder": 0, + "_isFastLaunch": true, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 4, + "name": "ml.t3.medium", + "vcpuNum": 2 + }, + { + "_defaultOrder": 1, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 8, + "name": "ml.t3.large", + "vcpuNum": 2 + }, + { + "_defaultOrder": 2, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 16, + "name": "ml.t3.xlarge", + "vcpuNum": 4 + }, + { + "_defaultOrder": 3, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 32, + "name": "ml.t3.2xlarge", + "vcpuNum": 8 + }, + { + "_defaultOrder": 4, + "_isFastLaunch": true, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 8, + "name": "ml.m5.large", + "vcpuNum": 2 + }, + { + "_defaultOrder": 5, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 16, + "name": "ml.m5.xlarge", + "vcpuNum": 4 + }, + { + "_defaultOrder": 6, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 32, + "name": "ml.m5.2xlarge", + "vcpuNum": 8 + }, + { + "_defaultOrder": 7, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 64, + "name": "ml.m5.4xlarge", + "vcpuNum": 16 + }, + { + "_defaultOrder": 8, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 128, + "name": "ml.m5.8xlarge", + "vcpuNum": 32 + }, + { + "_defaultOrder": 9, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 192, + "name": "ml.m5.12xlarge", + "vcpuNum": 48 + }, + { + "_defaultOrder": 10, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 256, + "name": "ml.m5.16xlarge", + "vcpuNum": 64 + }, + { + "_defaultOrder": 11, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 384, + "name": "ml.m5.24xlarge", + "vcpuNum": 96 + }, + { + "_defaultOrder": 12, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 8, + "name": "ml.m5d.large", + "vcpuNum": 2 + }, + { + "_defaultOrder": 13, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 16, + "name": "ml.m5d.xlarge", + "vcpuNum": 4 + }, + { + "_defaultOrder": 14, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 32, + "name": "ml.m5d.2xlarge", + "vcpuNum": 8 + }, + { + "_defaultOrder": 15, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 64, + "name": "ml.m5d.4xlarge", + "vcpuNum": 16 + }, + { + "_defaultOrder": 16, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 128, + "name": "ml.m5d.8xlarge", + "vcpuNum": 32 + }, + { + "_defaultOrder": 17, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 192, + "name": "ml.m5d.12xlarge", + "vcpuNum": 48 + }, + { + "_defaultOrder": 18, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 256, + "name": "ml.m5d.16xlarge", + "vcpuNum": 64 + }, + { + "_defaultOrder": 19, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 384, + "name": "ml.m5d.24xlarge", + "vcpuNum": 96 + }, + { + "_defaultOrder": 20, + "_isFastLaunch": false, + "category": "General purpose", + "gpuNum": 0, + "hideHardwareSpecs": true, + "memoryGiB": 0, + "name": "ml.geospatial.interactive", + "supportedImageNames": [ + "sagemaker-geospatial-v1-0" + ], + "vcpuNum": 0 + }, + { + "_defaultOrder": 21, + "_isFastLaunch": true, + "category": "Compute optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 4, + "name": "ml.c5.large", + "vcpuNum": 2 + }, + { + "_defaultOrder": 22, + "_isFastLaunch": false, + "category": "Compute optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 8, + "name": "ml.c5.xlarge", + "vcpuNum": 4 + }, + { + "_defaultOrder": 23, + "_isFastLaunch": false, + "category": "Compute optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 16, + "name": "ml.c5.2xlarge", + "vcpuNum": 8 + }, + { + "_defaultOrder": 24, + "_isFastLaunch": false, + "category": "Compute optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 32, + "name": "ml.c5.4xlarge", + "vcpuNum": 16 + }, + { + "_defaultOrder": 25, + "_isFastLaunch": false, + "category": "Compute optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 72, + "name": "ml.c5.9xlarge", + "vcpuNum": 36 + }, + { + "_defaultOrder": 26, + "_isFastLaunch": false, + "category": "Compute optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 96, + "name": "ml.c5.12xlarge", + "vcpuNum": 48 + }, + { + "_defaultOrder": 27, + "_isFastLaunch": false, + "category": "Compute optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 144, + "name": "ml.c5.18xlarge", + "vcpuNum": 72 + }, + { + "_defaultOrder": 28, + "_isFastLaunch": false, + "category": "Compute optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 192, + "name": "ml.c5.24xlarge", + "vcpuNum": 96 + }, + { + "_defaultOrder": 29, + "_isFastLaunch": true, + "category": "Accelerated computing", + "gpuNum": 1, + "hideHardwareSpecs": false, + "memoryGiB": 16, + "name": "ml.g4dn.xlarge", + "vcpuNum": 4 + }, + { + "_defaultOrder": 30, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 1, + "hideHardwareSpecs": false, + "memoryGiB": 32, + "name": "ml.g4dn.2xlarge", + "vcpuNum": 8 + }, + { + "_defaultOrder": 31, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 1, + "hideHardwareSpecs": false, + "memoryGiB": 64, + "name": "ml.g4dn.4xlarge", + "vcpuNum": 16 + }, + { + "_defaultOrder": 32, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 1, + "hideHardwareSpecs": false, + "memoryGiB": 128, + "name": "ml.g4dn.8xlarge", + "vcpuNum": 32 + }, + { + "_defaultOrder": 33, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 4, + "hideHardwareSpecs": false, + "memoryGiB": 192, + "name": "ml.g4dn.12xlarge", + "vcpuNum": 48 + }, + { + "_defaultOrder": 34, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 1, + "hideHardwareSpecs": false, + "memoryGiB": 256, + "name": "ml.g4dn.16xlarge", + "vcpuNum": 64 + }, + { + "_defaultOrder": 35, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 1, + "hideHardwareSpecs": false, + "memoryGiB": 61, + "name": "ml.p3.2xlarge", + "vcpuNum": 8 + }, + { + "_defaultOrder": 36, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 4, + "hideHardwareSpecs": false, + "memoryGiB": 244, + "name": "ml.p3.8xlarge", + "vcpuNum": 32 + }, + { + "_defaultOrder": 37, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 8, + "hideHardwareSpecs": false, + "memoryGiB": 488, + "name": "ml.p3.16xlarge", + "vcpuNum": 64 + }, + { + "_defaultOrder": 38, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 8, + "hideHardwareSpecs": false, + "memoryGiB": 768, + "name": "ml.p3dn.24xlarge", + "vcpuNum": 96 + }, + { + "_defaultOrder": 39, + "_isFastLaunch": false, + "category": "Memory Optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 16, + "name": "ml.r5.large", + "vcpuNum": 2 + }, + { + "_defaultOrder": 40, + "_isFastLaunch": false, + "category": "Memory Optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 32, + "name": "ml.r5.xlarge", + "vcpuNum": 4 + }, + { + "_defaultOrder": 41, + "_isFastLaunch": false, + "category": "Memory Optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 64, + "name": "ml.r5.2xlarge", + "vcpuNum": 8 + }, + { + "_defaultOrder": 42, + "_isFastLaunch": false, + "category": "Memory Optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 128, + "name": "ml.r5.4xlarge", + "vcpuNum": 16 + }, + { + "_defaultOrder": 43, + "_isFastLaunch": false, + "category": "Memory Optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 256, + "name": "ml.r5.8xlarge", + "vcpuNum": 32 + }, + { + "_defaultOrder": 44, + "_isFastLaunch": false, + "category": "Memory Optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 384, + "name": "ml.r5.12xlarge", + "vcpuNum": 48 + }, + { + "_defaultOrder": 45, + "_isFastLaunch": false, + "category": "Memory Optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 512, + "name": "ml.r5.16xlarge", + "vcpuNum": 64 + }, + { + "_defaultOrder": 46, + "_isFastLaunch": false, + "category": "Memory Optimized", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 768, + "name": "ml.r5.24xlarge", + "vcpuNum": 96 + }, + { + "_defaultOrder": 47, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 1, + "hideHardwareSpecs": false, + "memoryGiB": 16, + "name": "ml.g5.xlarge", + "vcpuNum": 4 + }, + { + "_defaultOrder": 48, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 1, + "hideHardwareSpecs": false, + "memoryGiB": 32, + "name": "ml.g5.2xlarge", + "vcpuNum": 8 + }, + { + "_defaultOrder": 49, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 1, + "hideHardwareSpecs": false, + "memoryGiB": 64, + "name": "ml.g5.4xlarge", + "vcpuNum": 16 + }, + { + "_defaultOrder": 50, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 1, + "hideHardwareSpecs": false, + "memoryGiB": 128, + "name": "ml.g5.8xlarge", + "vcpuNum": 32 + }, + { + "_defaultOrder": 51, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 1, + "hideHardwareSpecs": false, + "memoryGiB": 256, + "name": "ml.g5.16xlarge", + "vcpuNum": 64 + }, + { + "_defaultOrder": 52, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 4, + "hideHardwareSpecs": false, + "memoryGiB": 192, + "name": "ml.g5.12xlarge", + "vcpuNum": 48 + }, + { + "_defaultOrder": 53, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 4, + "hideHardwareSpecs": false, + "memoryGiB": 384, + "name": "ml.g5.24xlarge", + "vcpuNum": 96 + }, + { + "_defaultOrder": 54, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 8, + "hideHardwareSpecs": false, + "memoryGiB": 768, + "name": "ml.g5.48xlarge", + "vcpuNum": 192 + }, + { + "_defaultOrder": 55, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 8, + "hideHardwareSpecs": false, + "memoryGiB": 1152, + "name": "ml.p4d.24xlarge", + "vcpuNum": 96 + }, + { + "_defaultOrder": 56, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 8, + "hideHardwareSpecs": false, + "memoryGiB": 1152, + "name": "ml.p4de.24xlarge", + "vcpuNum": 96 + }, + { + "_defaultOrder": 57, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 32, + "name": "ml.trn1.2xlarge", + "vcpuNum": 8 + }, + { + "_defaultOrder": 58, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 512, + "name": "ml.trn1.32xlarge", + "vcpuNum": 128 + }, + { + "_defaultOrder": 59, + "_isFastLaunch": false, + "category": "Accelerated computing", + "gpuNum": 0, + "hideHardwareSpecs": false, + "memoryGiB": 512, + "name": "ml.trn1n.32xlarge", + "vcpuNum": 128 + } + ], + "instance_type": "ml.t3.medium", + "kernelspec": { + "display_name": "agents", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.15" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/integrations/aws_bedrock/aws_bedrock.py b/integrations/aws_bedrock/aws_bedrock.py new file mode 100644 index 0000000..393ea92 --- /dev/null +++ b/integrations/aws_bedrock/aws_bedrock.py @@ -0,0 +1,135 @@ +import json +import typing as t + +from ragas.messages import AIMessage, HumanMessage + + +def get_last_orchestration_value(traces: t.List[t.Dict[str, t.Any]], key: str): + """ + Iterates through the traces to find the last occurrence of a specified key + within the orchestrationTrace. + + Returns: + (index, value): Tuple where index is the last index at which the key was found, and value is the corresponding value, or (None, None) if not found. + """ + last_index = -1 + last_value = None + for i, trace in enumerate(traces): + orchestration = trace.get("trace", {}).get("orchestrationTrace", {}) + if key in orchestration: + last_index = i + last_value = orchestration[key] + return last_index, last_value + + +def extract_messages_from_model_invocation(model_inv): + """ + Extracts messages from the 'text' field of the modelInvocationInput. + Ensures that each message's content is cast to a string. + + Returns: + List of messages as HumanMessage or AIMessage objects. + """ + messages = [] + text_json = json.loads(model_inv.get("text", "{}")) + for msg in text_json.get("messages", []): + content_str = str(msg.get("content", "")) + role = msg.get("role") + if role == "user": + messages.append(HumanMessage(content=content_str)) + elif role == "assistant": + messages.append(AIMessage(content=content_str)) + return messages[:-1] + + +def convert_to_ragas_messages(traces: t.List): + """ + Converts a list of trace dictionaries into a list of messages. + It extracts messages from the last modelInvocationInput and appends + the finalResponse from the observation (if it occurs after the model invocation). + + Returns: + List of HumanMessage and AIMessage objects. + """ + result = [] + + # Get the last modelInvocationInput from the traces. + last_model_inv_index, last_model_inv = get_last_orchestration_value( + traces, "modelInvocationInput" + ) + if last_model_inv is not None: + result.extend(extract_messages_from_model_invocation(last_model_inv)) + + # Get the last observation from the traces. + last_obs_index, last_observation = get_last_orchestration_value( + traces, "observation" + ) + if last_observation is not None and last_obs_index > last_model_inv_index: + final_text = str(last_observation.get("finalResponse", {}).get("text", "")) + result.append(AIMessage(content=final_text)) + + return result + + +def extract_kb_trace(traces): + """ + Extracts groups of traces that follow the specific order: + 1. An element with 'trace' -> 'orchestrationTrace' containing an 'invocationInput' + with invocationType == "KNOWLEDGE_BASE" + 2. Followed (later in the list or within the same trace) by an element with an 'observation' + that contains 'knowledgeBaseLookupOutput' + 3. Followed by an element with an 'observation' that contains 'finalResponse' + + Returns a list of dictionaries each with keys: + 'user_input', 'retrieved_contexts', and 'response' + + This version supports multiple knowledge base invocation groups. + """ + results = [] + groups_in_progress = [] # list to keep track of groups in progress + + for trace in traces: + orchestration = trace.get("trace", {}).get("orchestrationTrace", {}) + + # 1. Look for a KB invocation input. + inv_input = orchestration.get("invocationInput") + if inv_input and inv_input.get("invocationType") == "KNOWLEDGE_BASE": + kb_input = inv_input.get("knowledgeBaseLookupInput", {}) + # Start a new group with the user's input text. + groups_in_progress.append({"user_input": kb_input.get("text")}) + + # 2. Process observations. + obs = orchestration.get("observation", {}) + if obs: + # If the observation contains a KB output, assign it to the earliest group + # that does not yet have a 'retrieved_contexts' key. + if "knowledgeBaseLookupOutput" in obs: + for group in groups_in_progress: + if "user_input" in group and "retrieved_contexts" not in group: + kb_output = obs["knowledgeBaseLookupOutput"] + group["retrieved_contexts"] = [ + retrieved.get("content", {}).get("text") + for retrieved in kb_output.get("retrievedReferences", []) + ] + break + + # 3. When we see a final response, assign it to all groups that have already + # received their KB output but still lack a response. + if "finalResponse" in obs: + final_text = obs["finalResponse"].get("text") + completed_groups = [] + for group in groups_in_progress: + if ( + "user_input" in group + and "retrieved_contexts" in group + and "response" not in group + ): + group["response"] = final_text + completed_groups.append(group) + # Remove completed groups from the in-progress list and add to the final results. + groups_in_progress = [ + g for g in groups_in_progress if g not in completed_groups + ] + results.extend(completed_groups) + + return results diff --git a/integrations/aws_bedrock/dataset/Restaurant_Childrens_Menu.pdf b/integrations/aws_bedrock/dataset/Restaurant_Childrens_Menu.pdf new file mode 100644 index 0000000..9c27571 Binary files /dev/null and b/integrations/aws_bedrock/dataset/Restaurant_Childrens_Menu.pdf differ diff --git a/integrations/aws_bedrock/dataset/Restaurant_Dinner_Menu.pdf b/integrations/aws_bedrock/dataset/Restaurant_Dinner_Menu.pdf new file mode 100644 index 0000000..ce31d43 Binary files /dev/null and b/integrations/aws_bedrock/dataset/Restaurant_Dinner_Menu.pdf differ diff --git a/integrations/aws_bedrock/dataset/Restaurant_week_specials.pdf b/integrations/aws_bedrock/dataset/Restaurant_week_specials.pdf new file mode 100644 index 0000000..4efa5ad Binary files /dev/null and b/integrations/aws_bedrock/dataset/Restaurant_week_specials.pdf differ diff --git a/integrations/aws_bedrock/images/architecture.png b/integrations/aws_bedrock/images/architecture.png new file mode 100644 index 0000000..ef9def8 Binary files /dev/null and b/integrations/aws_bedrock/images/architecture.png differ diff --git a/integrations/aws_bedrock/knowledge_base.py b/integrations/aws_bedrock/knowledge_base.py new file mode 100644 index 0000000..f86e058 --- /dev/null +++ b/integrations/aws_bedrock/knowledge_base.py @@ -0,0 +1,632 @@ +import json +import boto3 +import time +from botocore.exceptions import ClientError +from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth, RequestError +import pprint +from retrying import retry + +valid_embedding_models = ["cohere.embed-multilingual-v3", "cohere.embed-english-v3", "amazon.titan-embed-text-v1"] +pp = pprint.PrettyPrinter(indent=2) + + +def interactive_sleep(seconds: int): + """ + Support functionality to induce an artificial 'sleep' to the code in order to wait for resources to be available + Args: + seconds (int): number of seconds to sleep for + """ + dots = '' + for i in range(seconds): + dots += '.' + print(dots, end='\r') + time.sleep(1) + + +class BedrockKnowledgeBase: + """ + Support class that allows for: + - creation (or retrieval) of a Knowledge Base for Amazon Bedrock with all its pre-requisites + (including OSS, IAM roles and Permissions and S3 bucket) + - Ingestion of data into the Knowledge Base + - Deletion of all resources created + """ + def __init__( + self, + kb_name, + kb_description=None, + data_bucket_name=None, + embedding_model="amazon.titan-embed-text-v1" + ): + """ + Class initializer + Args: + kb_name (str): the knowledge base name + kb_description (str): knowledge base description + data_bucket_name (str): name of s3 bucket to connect with knowledge base + embedding_model (str): embedding model to use + """ + boto3_session = boto3.session.Session() + self.region_name = boto3_session.region_name + self.iam_client = boto3_session.client('iam') + self.account_number = boto3.client('sts').get_caller_identity().get('Account') + self.suffix = str(self.account_number)[:4] + self.identity = boto3.client('sts').get_caller_identity()['Arn'] + self.aoss_client = boto3_session.client('opensearchserverless') + self.s3_client = boto3.client('s3') + self.bedrock_agent_client = boto3.client('bedrock-agent') + credentials = boto3.Session().get_credentials() + self.awsauth = AWSV4SignerAuth(credentials, self.region_name, 'aoss') + + self.kb_name = kb_name + self.kb_description = kb_description + if data_bucket_name is not None: + self.bucket_name = data_bucket_name + else: + self.bucket_name = f"{self.kb_name}-{self.suffix}" + if embedding_model not in valid_embedding_models: + valid_embeddings_str = str(valid_embedding_models) + raise ValueError(f"Invalid embedding model. Your embedding model should be one of {valid_embeddings_str}") + self.embedding_model = embedding_model + self.encryption_policy_name = f"bedrock-sample-rag-sp-{self.suffix}" + self.network_policy_name = f"bedrock-sample-rag-np-{self.suffix}" + self.access_policy_name = f'bedrock-sample-rag-ap-{self.suffix}' + self.kb_execution_role_name = f'AmazonBedrockExecutionRoleForKnowledgeBase_{self.suffix}' + self.fm_policy_name = f'AmazonBedrockFoundationModelPolicyForKnowledgeBase_{self.suffix}' + self.s3_policy_name = f'AmazonBedrockS3PolicyForKnowledgeBase_{self.suffix}' + self.oss_policy_name = f'AmazonBedrockOSSPolicyForKnowledgeBase_{self.suffix}' + + self.vector_store_name = f'bedrock-sample-rag-{self.suffix}' + self.index_name = f"bedrock-sample-rag-index-{self.suffix}" + print("========================================================================================") + print(f"Step 1 - Creating or retrieving {self.bucket_name} S3 bucket for Knowledge Base documents") + self.create_s3_bucket() + print("========================================================================================") + print(f"Step 2 - Creating Knowledge Base Execution Role ({self.kb_execution_role_name}) and Policies") + self.bedrock_kb_execution_role = self.create_bedrock_kb_execution_role() + print("========================================================================================") + print(f"Step 3 - Creating OSS encryption, network and data access policies") + self.encryption_policy, self.network_policy, self.access_policy = self.create_policies_in_oss() + print("========================================================================================") + print(f"Step 4 - Creating OSS Collection (this step takes a couple of minutes to complete)") + self.host, self.collection, self.collection_id, self.collection_arn = self.create_oss() + # Build the OpenSearch client + self.oss_client = OpenSearch( + hosts=[{'host': self.host, 'port': 443}], + http_auth=self.awsauth, + use_ssl=True, + verify_certs=True, + connection_class=RequestsHttpConnection, + timeout=300 + ) + print("========================================================================================") + print(f"Step 5 - Creating OSS Vector Index") + self.create_vector_index() + print("========================================================================================") + print(f"Step 6 - Creating Knowledge Base") + self.knowledge_base, self.data_source = self.create_knowledge_base() + print("========================================================================================") + + def create_s3_bucket(self): + """ + Check if bucket exists, and if not create S3 bucket for knowledge base data source + """ + try: + self.s3_client.head_bucket(Bucket=self.bucket_name) + print(f'Bucket {self.bucket_name} already exists - retrieving it!') + except ClientError as e: + print(f'Creating bucket {self.bucket_name}') + if self.region_name == "us-east-1": + self.s3_client.create_bucket( + Bucket=self.bucket_name + ) + else: + self.s3_client.create_bucket( + Bucket=self.bucket_name, + CreateBucketConfiguration={'LocationConstraint': self.region_name} + ) + + def create_bedrock_kb_execution_role(self): + """ + Create Knowledge Base Execution IAM Role and its required policies. + If role and/or policies already exist, retrieve them + Returns: + IAM role + """ + foundation_model_policy_document = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "bedrock:InvokeModel", + ], + "Resource": [ + f"arn:aws:bedrock:{self.region_name}::foundation-model/{self.embedding_model}" + ] + } + ] + } + + s3_policy_document = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "s3:GetObject", + "s3:ListBucket" + ], + "Resource": [ + f"arn:aws:s3:::{self.bucket_name}", + f"arn:aws:s3:::{self.bucket_name}/*" + ], + "Condition": { + "StringEquals": { + "aws:ResourceAccount": f"{self.account_number}" + } + } + } + ] + } + + assume_role_policy_document = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Principal": { + "Service": "bedrock.amazonaws.com" + }, + "Action": "sts:AssumeRole" + } + ] + } + try: + # create policies based on the policy documents + fm_policy = self.iam_client.create_policy( + PolicyName=self.fm_policy_name, + PolicyDocument=json.dumps(foundation_model_policy_document), + Description='Policy for accessing foundation model', + ) + except self.iam_client.exceptions.EntityAlreadyExistsException: + fm_policy = self.iam_client.get_policy( + PolicyArn=f"arn:aws:iam::{self.account_number}:policy/{self.fm_policy_name}" + ) + + try: + s3_policy = self.iam_client.create_policy( + PolicyName=self.s3_policy_name, + PolicyDocument=json.dumps(s3_policy_document), + Description='Policy for reading documents from s3') + except self.iam_client.exceptions.EntityAlreadyExistsException: + s3_policy = self.iam_client.get_policy( + PolicyArn=f"arn:aws:iam::{self.account_number}:policy/{self.s3_policy_name}" + ) + # create bedrock execution role + try: + bedrock_kb_execution_role = self.iam_client.create_role( + RoleName=self.kb_execution_role_name, + AssumeRolePolicyDocument=json.dumps(assume_role_policy_document), + Description='Amazon Bedrock Knowledge Base Execution Role for accessing OSS and S3', + MaxSessionDuration=3600 + ) + except self.iam_client.exceptions.EntityAlreadyExistsException: + bedrock_kb_execution_role = self.iam_client.get_role( + RoleName=self.kb_execution_role_name + ) + # fetch arn of the policies and role created above + s3_policy_arn = s3_policy["Policy"]["Arn"] + fm_policy_arn = fm_policy["Policy"]["Arn"] + + # attach policies to Amazon Bedrock execution role + self.iam_client.attach_role_policy( + RoleName=bedrock_kb_execution_role["Role"]["RoleName"], + PolicyArn=fm_policy_arn + ) + self.iam_client.attach_role_policy( + RoleName=bedrock_kb_execution_role["Role"]["RoleName"], + PolicyArn=s3_policy_arn + ) + return bedrock_kb_execution_role + + def create_oss_policy_attach_bedrock_execution_role(self, collection_id): + """ + Create OpenSearch Serverless policy and attach it to the Knowledge Base Execution role. + If policy already exists, attaches it + """ + # define oss policy document + oss_policy_document = { + "Version": "2012-10-17", + "Statement": [ + { + "Effect": "Allow", + "Action": [ + "aoss:APIAccessAll" + ], + "Resource": [ + f"arn:aws:aoss:{self.region_name}:{self.account_number}:collection/{collection_id}" + ] + } + ] + } + + oss_policy_arn = f"arn:aws:iam::{self.account_number}:policy/{self.oss_policy_name}" + created = False + try: + self.iam_client.create_policy( + PolicyName=self.oss_policy_name, + PolicyDocument=json.dumps(oss_policy_document), + Description='Policy for accessing opensearch serverless', + ) + created = True + except self.iam_client.exceptions.EntityAlreadyExistsException: + print(f"Policy {oss_policy_arn} already exists, skipping creation") + print("Opensearch serverless arn: ", oss_policy_arn) + + self.iam_client.attach_role_policy( + RoleName=self.bedrock_kb_execution_role["Role"]["RoleName"], + PolicyArn=oss_policy_arn + ) + return created + + def create_policies_in_oss(self): + """ + Create OpenSearch Serverless encryption, network and data access policies. + If policies already exist, retrieve them + """ + try: + encryption_policy = self.aoss_client.create_security_policy( + name=self.encryption_policy_name, + policy=json.dumps( + { + 'Rules': [{'Resource': ['collection/' + self.vector_store_name], + 'ResourceType': 'collection'}], + 'AWSOwnedKey': True + }), + type='encryption' + ) + except self.aoss_client.exceptions.ConflictException: + encryption_policy = self.aoss_client.get_security_policy( + name=self.encryption_policy_name, + type='encryption' + ) + + try: + network_policy = self.aoss_client.create_security_policy( + name=self.network_policy_name, + policy=json.dumps( + [ + {'Rules': [{'Resource': ['collection/' + self.vector_store_name], + 'ResourceType': 'collection'}], + 'AllowFromPublic': True} + ]), + type='network' + ) + except self.aoss_client.exceptions.ConflictException: + network_policy = self.aoss_client.get_security_policy( + name=self.network_policy_name, + type='network' + ) + + try: + access_policy = self.aoss_client.create_access_policy( + name=self.access_policy_name, + policy=json.dumps( + [ + { + 'Rules': [ + { + 'Resource': ['collection/' + self.vector_store_name], + 'Permission': [ + 'aoss:CreateCollectionItems', + 'aoss:DeleteCollectionItems', + 'aoss:UpdateCollectionItems', + 'aoss:DescribeCollectionItems'], + 'ResourceType': 'collection' + }, + { + 'Resource': ['index/' + self.vector_store_name + '/*'], + 'Permission': [ + 'aoss:CreateIndex', + 'aoss:DeleteIndex', + 'aoss:UpdateIndex', + 'aoss:DescribeIndex', + 'aoss:ReadDocument', + 'aoss:WriteDocument'], + 'ResourceType': 'index' + }], + 'Principal': [self.identity, self.bedrock_kb_execution_role['Role']['Arn']], + 'Description': 'Easy data policy'} + ]), + type='data' + ) + except self.aoss_client.exceptions.ConflictException: + access_policy = self.aoss_client.get_access_policy( + name=self.access_policy_name, + type='data' + ) + + return encryption_policy, network_policy, access_policy + + def create_oss(self): + """ + Create OpenSearch Serverless Collection. If already existent, retrieve + """ + try: + collection = self.aoss_client.create_collection(name=self.vector_store_name, type='VECTORSEARCH') + collection_id = collection['createCollectionDetail']['id'] + collection_arn = collection['createCollectionDetail']['arn'] + except self.aoss_client.exceptions.ConflictException: + collection = self.aoss_client.batch_get_collection(names=[self.vector_store_name])['collectionDetails'][0] + pp.pprint(collection) + collection_id = collection['id'] + collection_arn = collection['arn'] + pp.pprint(collection) + + # Get the OpenSearch serverless collection URL + host = collection_id + '.' + self.region_name + '.aoss.amazonaws.com' + print(host) + # wait for collection creation + # This can take couple of minutes to finish + response = self.aoss_client.batch_get_collection(names=[self.vector_store_name]) + # Periodically check collection status + while (response['collectionDetails'][0]['status']) == 'CREATING': + print('Creating collection...') + interactive_sleep(30) + response = self.aoss_client.batch_get_collection(names=[self.vector_store_name]) + print('\nCollection successfully created:') + pp.pprint(response["collectionDetails"]) + # create opensearch serverless access policy and attach it to Bedrock execution role + try: + created = self.create_oss_policy_attach_bedrock_execution_role(collection_id) + if created: + # It can take up to a minute for data access rules to be enforced + print("Sleeping for a minute to ensure data access rules have been enforced") + interactive_sleep(60) + return host, collection, collection_id, collection_arn + except Exception as e: + print("Policy already exists") + pp.pprint(e) + + def create_vector_index(self): + """ + Create OpenSearch Serverless vector index. If existent, ignore + """ + body_json = { + "settings": { + "index.knn": "true", + "number_of_shards": 1, + "knn.algo_param.ef_search": 512, + "number_of_replicas": 0, + }, + "mappings": { + "properties": { + "vector": { + "type": "knn_vector", + "dimension": 1536, + "method": { + "name": "hnsw", + "engine": "faiss", + "space_type": "l2" + }, + }, + "text": { + "type": "text" + }, + "text-metadata": { + "type": "text"} + } + } + } + + # Create index + try: + response = self.oss_client.indices.create(index=self.index_name, body=json.dumps(body_json)) + print('\nCreating index:') + pp.pprint(response) + + # index creation can take up to a minute + interactive_sleep(60) + except RequestError as e: + # you can delete the index if its already exists + # oss_client.indices.delete(index=index_name) + print( + f'Error while trying to create the index, with error {e.error}\nyou may unmark the delete above to ' + f'delete, and recreate the index') + + @retry(wait_random_min=1000, wait_random_max=2000, stop_max_attempt_number=7) + def create_knowledge_base(self): + """ + Create Knowledge Base and its Data Source. If existent, retrieve + """ + opensearch_serverless_configuration = { + "collectionArn": self.collection_arn, + "vectorIndexName": self.index_name, + "fieldMapping": { + "vectorField": "vector", + "textField": "text", + "metadataField": "text-metadata" + } + } + + # Ingest strategy - How to ingest data from the data source + chunking_strategy_configuration = { + "chunkingStrategy": "FIXED_SIZE", + "fixedSizeChunkingConfiguration": { + "maxTokens": 512, + "overlapPercentage": 20 + } + } + + # The data source to ingest documents from, into the OpenSearch serverless knowledge base index + s3_configuration = { + "bucketArn": f"arn:aws:s3:::{self.bucket_name}", + # "inclusionPrefixes":["*.*"] # you can use this if you want to create a KB using data within s3 prefixes. + } + + # The embedding model used by Bedrock to embed ingested documents, and realtime prompts + embedding_model_arn = f"arn:aws:bedrock:{self.region_name}::foundation-model/{self.embedding_model}" + try: + create_kb_response = self.bedrock_agent_client.create_knowledge_base( + name=self.kb_name, + description=self.kb_description, + roleArn=self.bedrock_kb_execution_role['Role']['Arn'], + knowledgeBaseConfiguration={ + "type": "VECTOR", + "vectorKnowledgeBaseConfiguration": { + "embeddingModelArn": embedding_model_arn + } + }, + storageConfiguration={ + "type": "OPENSEARCH_SERVERLESS", + "opensearchServerlessConfiguration": opensearch_serverless_configuration + } + ) + kb = create_kb_response["knowledgeBase"] + pp.pprint(kb) + except self.bedrock_agent_client.exceptions.ConflictException: + kbs = self.bedrock_agent_client.list_knowledge_bases( + maxResults=100 + ) + kb_id = None + for kb in kbs['knowledgeBaseSummaries']: + if kb['name'] == self.kb_name: + kb_id = kb['knowledgeBaseId'] + response = self.bedrock_agent_client.get_knowledge_base(knowledgeBaseId=kb_id) + kb = response['knowledgeBase'] + pp.pprint(kb) + + # Create a DataSource in KnowledgeBase + try: + create_ds_response = self.bedrock_agent_client.create_data_source( + name=self.kb_name, + description=self.kb_description, + knowledgeBaseId=kb['knowledgeBaseId'], + dataSourceConfiguration={ + "type": "S3", + "s3Configuration": s3_configuration + }, + vectorIngestionConfiguration={ + "chunkingConfiguration": chunking_strategy_configuration + } + ) + ds = create_ds_response["dataSource"] + pp.pprint(ds) + except self.bedrock_agent_client.exceptions.ConflictException: + ds_id = self.bedrock_agent_client.list_data_sources( + knowledgeBaseId=kb['knowledgeBaseId'], + maxResults=100 + )['dataSourceSummaries'][0]['dataSourceId'] + get_ds_response = self.bedrock_agent_client.get_data_source( + dataSourceId=ds_id, + knowledgeBaseId=kb['knowledgeBaseId'] + ) + ds = get_ds_response["dataSource"] + pp.pprint(ds) + return kb, ds + + def start_ingestion_job(self): + """ + Start an ingestion job to synchronize data from an S3 bucket to the Knowledge Base + """ + # Start an ingestion job + start_job_response = self.bedrock_agent_client.start_ingestion_job( + knowledgeBaseId=self.knowledge_base['knowledgeBaseId'], + dataSourceId=self.data_source["dataSourceId"] + ) + job = start_job_response["ingestionJob"] + pp.pprint(job) + # Get job + while job['status'] != 'COMPLETE': + get_job_response = self.bedrock_agent_client.get_ingestion_job( + knowledgeBaseId=self.knowledge_base['knowledgeBaseId'], + dataSourceId=self.data_source["dataSourceId"], + ingestionJobId=job["ingestionJobId"] + ) + job = get_job_response["ingestionJob"] + pp.pprint(job) + interactive_sleep(40) + + def get_knowledge_base_id(self): + """ + Get Knowledge Base Id + """ + pp.pprint(self.knowledge_base["knowledgeBaseId"]) + return self.knowledge_base["knowledgeBaseId"] + + def get_bucket_name(self): + """ + Get the name of the bucket connected with the Knowledge Base Data Source + """ + pp.pprint(f"Bucket connected with KB: {self.bucket_name}") + return self.bucket_name + + def delete_kb(self, delete_s3_bucket=False, delete_iam_roles_and_policies=True): + """ + Delete the Knowledge Base resources + Args: + delete_s3_bucket (bool): boolean to indicate if s3 bucket should also be deleted + delete_iam_roles_and_policies (bool): boolean to indicate if IAM roles and Policies should also be deleted + """ + self.bedrock_agent_client.delete_data_source( + dataSourceId=self.data_source["dataSourceId"], + knowledgeBaseId=self.knowledge_base['knowledgeBaseId'] + ) + self.bedrock_agent_client.delete_knowledge_base( + knowledgeBaseId=self.knowledge_base['knowledgeBaseId'] + ) + self.oss_client.indices.delete(index=self.index_name) + self.aoss_client.delete_collection(id=self.collection_id) + self.aoss_client.delete_access_policy( + type="data", + name=self.access_policy_name + ) + self.aoss_client.delete_security_policy( + type="network", + name=self.network_policy_name + ) + self.aoss_client.delete_security_policy( + type="encryption", + name=self.encryption_policy_name + ) + if delete_s3_bucket: + self.delete_s3() + if delete_iam_roles_and_policies: + self.delete_iam_roles_and_policies() + + def delete_iam_roles_and_policies(self): + """ + Delete IAM Roles and policies used by the Knowledge Base + """ + fm_policy_arn = f"arn:aws:iam::{self.account_number}:policy/{self.fm_policy_name}" + s3_policy_arn = f"arn:aws:iam::{self.account_number}:policy/{self.s3_policy_name}" + oss_policy_arn = f"arn:aws:iam::{self.account_number}:policy/{self.oss_policy_name}" + self.iam_client.detach_role_policy( + RoleName=self.kb_execution_role_name, + PolicyArn=s3_policy_arn + ) + self.iam_client.detach_role_policy( + RoleName=self.kb_execution_role_name, + PolicyArn=fm_policy_arn + ) + self.iam_client.detach_role_policy( + RoleName=self.kb_execution_role_name, + PolicyArn=oss_policy_arn + ) + self.iam_client.delete_role(RoleName=self.kb_execution_role_name) + self.iam_client.delete_policy(PolicyArn=s3_policy_arn) + self.iam_client.delete_policy(PolicyArn=fm_policy_arn) + self.iam_client.delete_policy(PolicyArn=oss_policy_arn) + return 0 + + def delete_s3(self): + """ + Delete the objects contained in the Knowledge Base S3 bucket. + Once the bucket is empty, delete the bucket + """ + objects = self.s3_client.list_objects(Bucket=self.bucket_name) + if 'Contents' in objects: + for obj in objects['Contents']: + self.s3_client.delete_object(Bucket=self.bucket_name, Key=obj['Key']) + self.s3_client.delete_bucket(Bucket=self.bucket_name) \ No newline at end of file diff --git a/integrations/aws_bedrock/lambda_function.py b/integrations/aws_bedrock/lambda_function.py new file mode 100644 index 0000000..1cee7ed --- /dev/null +++ b/integrations/aws_bedrock/lambda_function.py @@ -0,0 +1,127 @@ +import json +import uuid +import boto3 + +dynamodb = boto3.resource('dynamodb') +table = dynamodb.Table('restaurant_bookings') + +def get_named_parameter(event, name): + """ + Get a parameter from the lambda event + """ + return next(item for item in event['parameters'] if item['name'] == name)['value'] + + +def get_booking_details(booking_id): + """ + Retrieve details of a restaurant booking + + Args: + booking_id (string): The ID of the booking to retrieve + """ + try: + response = table.get_item(Key={'booking_id': booking_id}) + if 'Item' in response: + return response['Item'] + else: + return {'message': f'No booking found with ID {booking_id}'} + except Exception as e: + return {'error': str(e)} + + +def create_booking(date, name, hour, num_guests): + """ + Create a new restaurant booking + + Args: + date (string): The date of the booking + name (string): Name to idenfity your reservation + hour (string): The hour of the booking + num_guests (integer): The number of guests for the booking + """ + try: + booking_id = str(uuid.uuid4())[:8] + table.put_item( + Item={ + 'booking_id': booking_id, + 'date': date, + 'name': name, + 'hour': hour, + 'num_guests': num_guests + } + ) + return {'booking_id': booking_id} + except Exception as e: + return {'error': str(e)} + + +def delete_booking(booking_id): + """ + Delete an existing restaurant booking + + Args: + booking_id (str): The ID of the booking to delete + """ + try: + response = table.delete_item(Key={'booking_id': booking_id}) + if response['ResponseMetadata']['HTTPStatusCode'] == 200: + return {'message': f'Booking with ID {booking_id} deleted successfully'} + else: + return {'message': f'Failed to delete booking with ID {booking_id}'} + except Exception as e: + return {'error': str(e)} + + +def lambda_handler(event, context): + # get the action group used during the invocation of the lambda function + actionGroup = event.get('actionGroup', '') + + # name of the function that should be invoked + function = event.get('function', '') + + # parameters to invoke function with + parameters = event.get('parameters', []) + + if function == 'get_booking_details': + booking_id = get_named_parameter(event, "booking_id") + if booking_id: + response = str(get_booking_details(booking_id)) + responseBody = {'TEXT': {'body': json.dumps(response)}} + else: + responseBody = {'TEXT': {'body': 'Missing booking_id parameter'}} + + elif function == 'create_booking': + date = get_named_parameter(event, "date") + name = get_named_parameter(event, "name") + hour = get_named_parameter(event, "hour") + num_guests = get_named_parameter(event, "num_guests") + + if date and hour and num_guests: + response = str(create_booking(date, name, hour, num_guests)) + responseBody = {'TEXT': {'body': json.dumps(response)}} + else: + responseBody = {'TEXT': {'body': 'Missing required parameters'}} + + elif function == 'delete_booking': + booking_id = get_named_parameter(event, "booking_id") + if booking_id: + response = str(delete_booking(booking_id)) + responseBody = {'TEXT': {'body': json.dumps(response)}} + else: + responseBody = {'TEXT': {'body': 'Missing booking_id parameter'}} + + else: + responseBody = {'TEXT': {'body': 'Invalid function'}} + + action_response = { + 'actionGroup': actionGroup, + 'function': function, + 'functionResponse': { + 'responseBody': responseBody + } + } + + function_response = {'response': action_response, 'messageVersion': event['messageVersion']} + print("Response: {}".format(function_response)) + + return function_response