agents-and-function-calling/bedrock-agents/features-examples/05-create-agent-with-knowledge-base-and-action-group/05-create-agent-with-knowledge-base-and-action-group.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a0bb5c39-2fde-4336-8127-8debe7cb2741",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Create an agent integrated with Bedrock Knowledge Bases and attach Action Group\n",
    "\n",
    "In this notebook you will learn how to create an Amazon Bedrock Agent that makes use of Knowledge Bases for Amazon Bedrock to retrieve data about a restaurant's menu. The use case create a restaurant agent, it's tasks will be to give information to the clients about the adults or childrens menu and be in charge of the table booking system. Client's will be able to create, delete or get booking information. The architecture looks as following:\n",
    "\n",
    "<img src=\"images/architecture.png\" style=\"width:70%;display:block;margin: 0 auto;\">\n",
    "<br/>\n",
    "\n",
    "The steps to complete this notebook are:\n",
    "\n",
    "1. Import the needed libraries\n",
    "1. Create the Knowledge Base for Amazon Bedrock\n",
    "1. Upload the dataset to Amazon S3\n",
    "1. Create the Agent for Amazon Bedrock\n",
    "1. Test the Agent\n",
    "1. Clean-up the resources created"
   ]
  },
  {
   "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 -r requirements.txt"
   ]
  },
  {
   "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 create_agent_role_and_policies, create_lambda_role, delete_agent_roles_and_policies\n",
    "from agent import create_dynamodb, create_lambda, clean_up_resources"
   ]
  },
  {
   "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(format='[%(asctime)s] p%(process)s {%(filename)s:%(lineno)d} %(levelname)s - %(message)s', level=logging.INFO)\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 = \"Knowledge Base containing the restaurant menu's collection\"\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 = \"anthropic.claude-3-sonnet-20240229-v1:0\"\n",
    "\n",
    "agent_description = \"Agent in charge of a restaurants table bookings\"\n",
    "agent_instruction = \"\"\"\n",
    "You are a restaurant agent, helping clients retrieve information from their booking, \n",
    "create a new booking or delete an existing booking\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": "38c38fcb-9b87-414e-a644-04c263eea5c9",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 2. Create Knowledge Base for Amazon Bedrock\n",
    "Let's start by creating a [Knowledge Base for Amazon Bedrock](https://aws.amazon.com/bedrock/knowledge-bases/) to store the restaurant menus. Knowledge Bases allow you to integrate with different vector databases including [Amazon OpenSearch Serverless](https://aws.amazon.com/opensearch-service/features/serverless/), [Amazon Aurora](https://aws.amazon.com/rds/aurora/) and [Pinecone](http://app.pinecone.io/bedrock-integration). For this example, we will integrate the knowledge base with Amazon OpenSearch Serverless. To do so, we will use the helper class `BedrockKnowledgeBase` which will create the knowledge base and all of its pre-requisites:\n",
    "1. IAM roles and policies\n",
    "2. S3 bucket\n",
    "3. Amazon OpenSearch Serverless encryption, network and data access policies\n",
    "4. Amazon OpenSearch Serverless collection\n",
    "5. Amazon OpenSearch Serverless vector index\n",
    "6. Knowledge base\n",
    "7. Knowledge base data source"
   ]
  },
  {
   "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": "c6b5819d-93a9-44d3-8eaa-5ed1da6a8c95",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 3. Upload the dataset to Amazon S3\n",
    "Now that we have created the knowledge base, let's populate it with the menu's dataset. The Knowledge Base data source expects the data to be available on the S3 bucket connected to it and changes on the data can be syncronized to the knowledge base using the `StartIngestionJob` API call. 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's data available on the `dataset` folder to s3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ceecf498-e42e-4fdd-ad2e-e646a4ef4adb",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "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",
    "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": [
    "### 3.1 Test the Knowledge Base\n",
    "Now the Knowlegde Base is available we can test it out using the [**retrieve**](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent-runtime/client/retrieve.html) and [**retrieve_and_generate**](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent-runtime/client/retrieve_and_generate.html) functions. \n",
    "\n",
    "#### Testing Knowledge Base with Retrieve and Generate API\n",
    "\n",
    "Let's first test the knowledge base using the retrieve and generate API. With this API, Bedrock takes care of retrieving the necessary references from the knowledge base and generating the final answer using a LLM model from Bedrock"
   ]
  },
  {
   "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={\n",
    "        \"text\": \"Which are the 5 mains available in the childrens menu?\"\n",
    "    },\n",
    "    retrieveAndGenerateConfiguration={\n",
    "        \"type\": \"KNOWLEDGE_BASE\",\n",
    "        \"knowledgeBaseConfiguration\": {\n",
    "            'knowledgeBaseId': kb_id,\n",
    "            \"modelArn\": \"arn:aws:bedrock:{}::foundation-model/{}\".format(region, agent_foundation_model),\n",
    "            \"retrievalConfiguration\": {\n",
    "                \"vectorSearchConfiguration\": {\n",
    "                    \"numberOfResults\":5\n",
    "                } \n",
    "            }\n",
    "        }\n",
    "    }\n",
    ")\n",
    "\n",
    "print(response['output']['text'],end='\\n'*2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6e4235c",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "As you can see, with the retrieve and generate API we get the final response directly and we don't see the different sources used to generate this response. Let's now retrieve the source information from the knowledge base with the retrieve API.\n",
    "\n",
    "#### Testing Knowledge Base with Retrieve API\n",
    "If you need an extra layer of control, you can retrieve the chuncks that best match your query using the retrieve API. In this setup, we can configure the desired number of results and control the final answer with your own application logic. The API then provides you with the matching content, its S3 location, the similarity score and the chunk metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34a3f96b-926a-411c-bef7-38c4689fed7b",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "response_ret = bedrock_agent_runtime_client.retrieve(\n",
    "    knowledgeBaseId=kb_id, \n",
    "    nextToken='string',\n",
    "    retrievalConfiguration={\n",
    "        \"vectorSearchConfiguration\": {\n",
    "            \"numberOfResults\":5,\n",
    "        } \n",
    "    },\n",
    "    retrievalQuery={\n",
    "        'text': 'Which are the 5 mains available in the childrens menu?'\n",
    "    }\n",
    ")\n",
    "\n",
    "def response_print(retrieve_resp):\n",
    "#structure 'retrievalResults': list of contents. Each list has content, location, score, metadata\n",
    "    for num,chunk in enumerate(response_ret['retrievalResults'],1):\n",
    "        print(f'Chunk {num}: ',chunk['content']['text'],end='\\n'*2)\n",
    "        print(f'Chunk {num} Location: ',chunk['location'],end='\\n'*2)\n",
    "        print(f'Chunk {num} Score: ',chunk['score'],end='\\n'*2)\n",
    "        print(f'Chunk {num} Metadata: ',chunk['metadata'],end='\\n'*2)\n",
    "\n",
    "response_print(response_ret)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e032c452-4feb-406d-aa6e-b85a59aded79",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 4. Create the Agent for Amazon Bedrock\n",
    "\n",
    "In this section we will go through all the steps to create an Agent for Amazon Bedrock. \n",
    "\n",
    "These are the steps to complete:\n",
    "    \n",
    "1. Create an Amazon DynamoDB table\n",
    "2. Create an AWS Lambda function\n",
    "3. Create the IAM policies needed for the Agent\n",
    "4. Create the Agent\n",
    "5. Create the Agent Action Group\n",
    "6. Allow the Agent to invoke the Action Group Lambda\n",
    "7. Associate the Knowledge Base to the agent\n",
    "8. Prepare the Agent and create an alias"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5edf3bd-3214-4fe3-a9bc-df927e30b8b4",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 4.1 Create the DynamoDB table\n",
    "We will create a DynamoDB table which contains the restaurant bookings 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": [
    "### 4.2 Create the Lambda Function\n",
    "\n",
    "We will now create a lambda function that interacts with DynamoDB table. To do so we will:\n",
    "\n",
    "1. Create the `lambda_function.py` file which contains the logic for our lambda function\n",
    "2. Create the IAM role for our Lambda function\n",
    "3. Create the lambda function with the required permissions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68871530-2953-4a37-854e-10323cabf095",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Create the function code\n",
    "When creating an Agent for Amazon Bedrock, you can connect a Lambda function to the Action Group in order to execute the functions required by the agent. In this option, your agent is responsible for the execution of your functions. Let's create the lambda function tha 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\n",
    "Now let's also create the lambda role and its required policies. For this case, we need the lambda to be able to access DynamoDB, that is why we also create a DynamoDB policy and attach to our Lambda. To do so, we will use the support function `create_lambda_role`."
   ]
  },
  {
   "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\n",
    "\n",
    "Now that we have the Lambda function code and its execution role, let's package it into a Zip file and create the Lambda resources"
   ]
  },
  {
   "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'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1cfb2214-9a2e-4521-9d69-3d1c6117c2c3",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "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": [
    "### 4.3 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.\n",
    "\n",
    "\n",
    "First need to create the agent policies that allow bedrock model invocation and Knowledge Base query and the agent IAM role with the policy associated to it. We will allow this agent to invoke the Claude Sonnet model. Here we use the `create_agent_role_and_policies` to create the agent role and its required policies"
   ]
  },
  {
   "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(agent_name, agent_foundation_model, kb_id=kb_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd45ef4e-1f2d-4747-a16e-f64fc6f865f5",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent_role"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "595bf767-3a26-4f92-b629-ba2908b90d81",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 4.4 Create the Agent\n",
    "Once the needed IAM role is created, we can use the bedrock agent client to create a new agent. To do so we use the [`create_agent`](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent/client/create_agent.html) api from boto3. It requires an agent name, underline foundation model and instruction. You can also provide an agent description. Note that the agent created is not yet prepared. We will focus on preparing the agent and then using it to invoke actions and use other APIs"
   ]
  },
  {
   "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",
    ")\n",
    "response"
   ]
  },
  {
   "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": [
    "### 4.5 Create the Agent Action Group\n",
    "We will now create an agent action group that uses the lambda function created before. The [`create_agent_action_group`](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent/client/create_agent_action_group.html) function provides this functionality. We will use `DRAFT` as the agent version since we haven't yet created an agent version or alias. To inform the agent about the action group functionalities, we will provide an action group description containing the functionalities of the action group.\n",
    "\n",
    "In this example, we will provide the Action Group functionality using a [`functionSchema`](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-action-function.html).\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={\n",
    "        'lambda': lambda_function['FunctionArn']\n",
    "    },\n",
    "    actionGroupName=agent_action_group_name,\n",
    "    functionSchema={\n",
    "        'functions': agent_functions\n",
    "    },\n",
    "    description=agent_action_group_description\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a0c3eeb-37f3-4ec9-88d4-43c5306c3195",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "agent_action_group_response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7d27e07-7d4c-4c59-987d-41a606c6af65",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 4.6 Allow the Agent to invoke the Action Group Lambda\n",
    "Before using the action group, we need to allow the agent to invoke the lambda function associated with the action group. This is done via [resource-based policy](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-permissions.html#agents-permissions-lambda). Let's add the resource-based policy to the lambda function created"
   ]
  },
  {
   "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",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb3b8c21-f14a-4645-b28e-3ac9a50f1f7f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c8536a4-e625-464c-bae2-d28dba82c556",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 4.7 Associate the Knowledge Base to the agent\n",
    "Now we have created the Agent we can go ahead and associate the Knowledge Base we created earlier. "
   ]
  },
  {
   "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": "code",
   "execution_count": null,
   "id": "0c72b5ee-2479-41b2-93e1-fd7aa66da08d",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "response"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b42a8caf-4c76-4760-a7cd-4c6b189f9f5b",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 4.8 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.\n"
   ]
  },
  {
   "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(\n",
    "    agentId=agent_id\n",
    ")\n",
    "print(response)\n",
    "# Pause to make sure agent is prepared\n",
    "time.sleep(30)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c2f477d-60b5-4135-bfe4-e9af3874d611",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "You can invoke the DRAFT version of your agent using the test alias id `TSTALIASID` or you can create a new alias and a new version for your agent. Here we are also going to create an Agent alias to later on use to invoke it with the alias id created"
   ]
  },
  {
   "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": "markdown",
   "id": "785eaa07-9059-471b-86cc-9506e9c0b47c",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 5. Test the Agent\n",
    "Now that we've created the agent, let's use the `bedrock-agent-runtime` client to invoke this agent and perform some tasks. You can invoke your agent with the [`invoke_agent`](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent-runtime/client/invoke_agent.html) API"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "607efd7f-469a-4875-9c1c-16ab018fc341",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "import json\n",
    "from datetime import datetime\n",
    "\n",
    "class DateTimeEncoder(json.JSONEncoder):\n",
    "    def default(self, obj):\n",
    "        if isinstance(obj, datetime):\n",
    "            return obj.isoformat()\n",
    "        return super().default(obj)\n",
    "\n",
    "\n",
    "\n",
    "def invokeAgent(query, session_id, enable_trace=False, 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",
    "    if enable_trace:\n",
    "        logger.info(pprint.pprint(agentResponse))\n",
    "    \n",
    "    event_stream = agentResponse['completion']\n",
    "    try:\n",
    "        for event in event_stream:        \n",
    "            if 'chunk' in event:\n",
    "                data = event['chunk']['bytes']\n",
    "                if enable_trace:\n",
    "                    logger.info(f\"Final answer ->\\n{data.decode('utf8')}\")\n",
    "                agent_answer = data.decode('utf8')\n",
    "                end_event_received = True\n",
    "                return agent_answer\n",
    "                # End event indicates that the request finished successfully\n",
    "            elif 'trace' in event:\n",
    "                if enable_trace:\n",
    "                    logger.info(json.dumps(event['trace'], indent=2, cls=DateTimeEncoder))\n",
    "            else:\n",
    "                raise Exception(\"unexpected event.\", event)\n",
    "    except Exception as e:\n",
    "        raise Exception(\"unexpected event.\", e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2bd6097d",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "##### Invoke Agent to query Knowledge Base\n",
    "Let's now use our support `invokeAgent` function to query our Knowledge Base with the Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb134630-0c01-4ec7-b761-1f1a1eca4957",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "import uuid\n",
    "session_id:str = str(uuid.uuid1())\n",
    "query = \"What are the starters in the childrens menu?\"\n",
    "response = invokeAgent(query, session_id)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e59b3614",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "##### Invoke Agent to execute function from Action Group\n",
    "Now let's test our Action Group functionality and create a new reservation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e0201c0-abb3-4ad5-a7e7-1750d14a2717",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "query = \"Hi, I am Anna. I want to create a booking for 2 people, at 8pm on the 5th of May 2024.\"\n",
    "response = invokeAgent(query, session_id)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7043ba4",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "##### Invoke Agent with prompt attribute\n",
    "\n",
    "Great! We've used our agent to do the first reservation. However, often when booking tables in restaurants we are already logged in to systesm that know our names. How great would it be if our agent would know it as well?\n",
    "\n",
    "To do so, we can use the session context to provide some attributes to our prompt. In this case we will provide it directly to the prompt using the [`promptSessionAttributes`](https://docs.aws.amazon.com/bedrock/latest/userguide/agents-session-state.html) parameter. Let's also start a new session id so that our agent does not memorize our name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "139c210f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "session_id:str = str(uuid.uuid1())\n",
    "query = \"I want to create a booking for 2 people, at 8pm on the 5th of May 2024.\"\n",
    "session_state = {\n",
    "    \"promptSessionAttributes\": {\n",
    "        \"name\": \"John\"\n",
    "    }\n",
    "}\n",
    "response = invokeAgent(query, session_id, session_state=session_state)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cf0d98c",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "##### Validating prompt attribute\n",
    "Let's now use our session context to validate that the reservation was done under the correct name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfb8612f-6221-4199-b44b-4c341ac08f7f",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "query = \"What was the name used in my last reservation?\"\n",
    "response = invokeAgent(query, session_id)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a69fa42f",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "##### Retrieving information from the database in a new session\n",
    "\n",
    "Next, let's confirm that our reservation system is working correctly. To do so, let's use our previous booking ID and retrieve our reservation details using a new session id\n",
    "\n",
    "\n",
    "**Important**: remember to replace the information in next queries with your generated booking id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "717da5f1",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "session_id:str = str(uuid.uuid1())\n",
    "query = \"I want to get the information for booking 007659d1\"\n",
    "response = invokeAgent(query, session_id)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31726c96",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "##### Canceling reservation\n",
    "\n",
    "As plans change, we would now like to cancel the reservation we just did using our Agent for it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bf770ef-5602-434e-bef2-0394b9136fe7",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "query = \"I want to delete the booking 007659d1\"\n",
    "response = invokeAgent(query, session_id)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4f4545a",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "And let's make sure everything worked out correctly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bdb62e6",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "session_id:str = str(uuid.uuid1())\n",
    "query = \"I want to get the information for booking 007659d1\"\n",
    "response = invokeAgent(query, session_id)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2b6149d",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "##### Handling context with PromptAttributes\n",
    "\n",
    "With real-life applications, context is really important. We want to make reservations considering the current date and the days sorounding it. Amazon Bedrock Agents also allow you to provide temporal context for the agent with the prompt attributes. Let's test it with a reservation for tomorrow"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7664b37d",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# retrieving today\n",
    "from datetime import datetime\n",
    "today = datetime.today().strftime('%b-%d-%Y')\n",
    "today"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "34bd272d",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "# reserving a table for tomorrow\n",
    "session_id:str = str(uuid.uuid1())\n",
    "query = \"I want to create a booking for 2 people, at 8pm tomorrow.\"\n",
    "session_state = {\n",
    "    \"promptSessionAttributes\": {\n",
    "        \"name\": \"John\",\n",
    "        \"today\": today\n",
    "    }\n",
    "}\n",
    "response = invokeAgent(query, session_id, session_state=session_state)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b53ca04b",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "And finally, let's validate our reservation\n",
    "\n",
    "**Important**: remember to replace the booking id with the new one"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a5d720b",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "session_id:str = str(uuid.uuid1())\n",
    "query = \"I want to get the information for booking 98e6464f\"\n",
    "response = invokeAgent(query, session_id)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5dcfdf43",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "##### Invoke Agent with Trace\n",
    "\n",
    "Amazon Bedrock Agents also provides you with the details of steps being orchestrated by the Agent using the [Trace](https://docs.aws.amazon.com/bedrock/latest/userguide/trace-events.html). You can enable the trace during agent invocation. Let's now invoke the agent with the trace enabled"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebcb9737",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "session_id:str = str(uuid.uuid1())\n",
    "query = \"What are the desserts on the adult menu?\"\n",
    "response = invokeAgent(query, session_id, enable_trace=True)\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "073f0489-bcb8-4329-b16f-afafaf6f4819",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Agent Evaluation Framework - Testing the Agent (Optional)\n",
    "\n",
    "The [Agent Evaluation Framework](https://awslabs.github.io/agent-evaluation/) offers a structured approach for assessing the performance, accuracy, and effectiveness of Bedrock Agents.\n",
    "\n",
    "The next steps are optional, and show you how to write test cases and run them against the Bedrock Agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "830155fb-59f0-4f7d-b3c5-3cada968e6ba",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "!python3 -m pip install agent-evaluation==0.2.0 # Installs the agent-evaluation framework tool\n",
    "\n",
    "!which agenteval # Checks the agent evaluation framework is properly installed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fff61ec1-68b9-4d50-8082-09884f979862",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    },
    "tags": []
   },
   "source": [
    "- The following sections prepare the `agenteval.yml` file by providing the Agent ID and Alias ID created with this notebook into line 5. In the `agenteval.yml` file you will find the different test cases defined to test the Agent.\n",
    "- As of writing, only Claude 3 Sonnet is supported as an evaluator. The Claude-3 model specified in the yml file refers to Claude 3 Sonnet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fd95939-2da3-4743-bc4e-fd53b6c9ed63",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "agent_id # Prints the Agent ID for reference here\n",
    "\n",
    "!sed -i 's/{{agent_id}}/{alias_id}/g' agenteval.yml\n",
    "!sed -i 's/none/{agent_id}/' agenteval.yml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54f7f2f7-5369-4b60-86c8-2e0269724d83",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Run the defined test cases that are part of the repo (i.e. the agenteval.yml file)\n",
    "\n",
    "!agenteval run"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96af7a1c-400b-4904-bd39-6069b56c7eda",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "- The output above shows the results of the testing evaluation. You can find a detailed report about the tests evaluation under the following generated file: `agenteval_summary.md.` You can preview it by right-clicking and select open with -> markdown preview\n",
    "  \n",
    "- You will notice that some new files have been generated as well once the test have been executed (e.g. `check_number_of_vacation.json`), where you will find the detailed traces from the test conversation between the test User and the test Agent."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ebf4438-1f48-4642-a57c-530a16815064",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 6. 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, lambda_function, lambda_function_name, agent_action_group_response, agent_functions, \n",
    "    agent_id, kb_id, alias_id\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": "Python 3 (ipykernel)",
   "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}