agents-for-bedrock/features-examples/04-create-agent-with-single-knowledge-base/04-create-agent-with-single-knowledge-base.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a0bb5c39-2fde-4336-8127-8debe7cb2741",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "# Create an Agent with a single Knowledge Base only\n",
    "\n",
    "In this notebook you will learn how to create an Amazon Bedrock Agent that connects to a single Knowledge Bases for Amazon Bedrock to retrieve company data and complete tasks. \n",
    "\n",
    "The use case for this notebook is the Amazon Bedrock Documentation pages stored as PDFs. It will allow you to ask questions about Amazon Bedrock and get answers based on documents available in the Knowledge Base.\n",
    "\n",
    "The steps to complete this notebook are:\n",
    "\n",
    "1. Import the needed libraries\n",
    "1. Create an S3 bucket and upload the data to it\n",
    "1. Create the Knowledge Base for Amazon Bedrock and sync data to Knowledge Base\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": "code",
   "execution_count": null,
   "id": "ac05c073-d45b-4d85-9bf8-ae10aa78be8d",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!pip install --upgrade -q opensearch-py\n",
    "!pip install --upgrade -q requests-aws4auth\n",
    "!pip install --upgrade -q boto3\n",
    "!pip install --upgrade -q botocore\n",
    "!pip install --upgrade -q awscli"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8ad6ec2-b283-4c5d-879f-e397e46568c0",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "import boto3\n",
    "import time\n",
    "import json\n",
    "import uuid\n",
    "import pprint\n",
    "import os\n",
    "from opensearchpy import OpenSearch, RequestsHttpConnection\n",
    "from requests_aws4auth import AWS4Auth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2b2d607-c1f2-4cbb-9f89-d935676e0101",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# setting logger\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__)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec5d7363-f715-4849-968b-cd4eb5fae3b0",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# getting boto3 clients for required AWS services\n",
    "sts_client = boto3.client('sts')\n",
    "iam_client = boto3.client('iam')\n",
    "s3_client = boto3.client('s3')\n",
    "lambda_client = boto3.client('lambda')\n",
    "bedrock_agent_client = boto3.client('bedrock-agent')\n",
    "bedrock_agent_runtime_client = boto3.client('bedrock-agent-runtime')\n",
    "open_search_serverless_client = boto3.client('opensearchserverless')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5ab4b38-bb05-4771-9c29-6e48602b3244",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "session = boto3.session.Session()\n",
    "region = session.region_name\n",
    "account_id = sts_client.get_caller_identity()[\"Account\"]\n",
    "region, account_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec1c499b-a834-4928-92c9-0091ae0d381b",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Generate random prefix for unique IAM roles, agent name and S3 Bucket and \n",
    "# assign variables\n",
    "suffix = f\"{region}-{account_id}\"\n",
    "agent_name = \"bedrock-docs-kb-agents\"\n",
    "agent_alias_name = \"bedrock-docs-alias\"\n",
    "bucket_name = f'{agent_name}-{suffix}'\n",
    "bucket_arn = f\"arn:aws:s3:::{bucket_name}\"\n",
    "bedrock_agent_bedrock_allow_policy_name = f\"bda-bedrock-allow-{suffix}\"\n",
    "bedrock_agent_s3_allow_policy_name = f\"bda-s3-allow-{suffix}\"\n",
    "bedrock_agent_kb_allow_policy_name = f\"bda-kb-allow-{suffix}\"\n",
    "agent_role_name = f'AmazonBedrockExecutionRoleForAgents_bedrock_docs'\n",
    "kb_name = f'bedrock-docs-kb-{suffix}'\n",
    "data_source_name = f'bedrock-docs-kb-docs-{suffix}'\n",
    "kb_files_path = 'kb_documents'\n",
    "kb_key = 'kb_documents'\n",
    "kb_role_name = f'AmazonBedrockExecutionRoleForKnowledgeBase_bedrock_docs'\n",
    "kb_bedrock_allow_policy_name = f\"bd-kb-bedrock-allow-{suffix}\"\n",
    "kb_aoss_allow_policy_name = f\"bd-kb-aoss-allow-{suffix}\"\n",
    "kb_s3_allow_policy_name = f\"bd-kb-s3-allow-{suffix}\"\n",
    "kb_collection_name = f'bd-kbc-{suffix}'\n",
    "# Select Amazon titan as the embedding model\n",
    "embedding_model_arn = f'arn:aws:bedrock:{region}::foundation-model/amazon.titan-embed-text-v1'\n",
    "kb_vector_index_name = \"bedrock-knowledge-base-index\"\n",
    "kb_metadataField = 'bedrock-knowledge-base-metadata'\n",
    "kb_textField = 'bedrock-knowledge-base-text'\n",
    "kb_vectorField = 'bedrock-knowledge-base-vector'\n",
    "model_id = \"anthropic.claude-3-sonnet-20240229-v1:0\"\n",
    "\n",
    "# agent configuration\n",
    "agent_instruction = \"\"\"\n",
    "You are an agent that support users working with Amazon Bedrock. You have access to Bedrock's documentation in a Knowledge Base\n",
    "and you can Answer questions from this documentation. Only answer questions based on the documentation and reply with \n",
    "\"There is no information about your question on the Amazon Bedrock Documentation at the moment, sorry! Do you want to ask another question?\" \n",
    "If the answer to the question is not available in the documentation\n",
    "\"\"\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6b5819d-93a9-44d3-8eaa-5ed1da6a8c95",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 2. Upload the dataset to Amazon S3\n",
    "Knowledge Bases for Amazon Bedrock, currently require data to reside in an Amazon S3 bucket. In this section we will create an Amazon S3 bucket and the files."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdc17ae4-3dce-43c4-b79d-ed6884beb7a1",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 2.1 Create the Amazon S3 bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "976cc565-4f58-4284-88eb-dd9498c31958",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "if region != 'us-east-1':\n",
    "    s3_client.create_bucket(\n",
    "        Bucket=bucket_name.lower(),\n",
    "        CreateBucketConfiguration={'LocationConstraint': region}\n",
    "    )\n",
    "else:\n",
    "    s3_client.create_bucket(Bucket=bucket_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7c7a720-5586-4e4b-8a29-2ef17a1cf07a",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 2.2 Upload dataset to the Amazon S3 bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "321c127d-a82a-49a5-b939-32d289376125",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Upload Knowledge Base files to this s3 bucket\n",
    "for f in os.listdir(kb_files_path):\n",
    "    if f.endswith(\".pdf\"):\n",
    "        s3_client.upload_file(kb_files_path+'/'+f, bucket_name, kb_key+'/'+f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6ad39aa6-42bf-4146-9f87-fbb18f674326",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 3. Create a Knowledge Base for Amazon Bedrock\n",
    "\n",
    "In this section we will go through all the steps to create and test a Knowledge Base. \n",
    "\n",
    "These are the steps to complete:\n",
    "    \n",
    "1. Create a Knowledge Base Role and its policies\n",
    "1. Create a Vector Database\n",
    "1. Create an OpenSearch Index\n",
    "1. Create a Knowledge Base\n",
    "1. Create a data source and attach to the recently created Knowledge Base\n",
    "1. Ingest data to your knowledge Base"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d3d808e-81ae-4205-9ac1-362bbb45d96e",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 3.1 Create Knowledge Base Role and Policies\n",
    "\n",
    "Let's first create IAM policies to allow our Knowledge Base to access Bedrock Titan Embedding Foundation model, Amazon OpenSearch Serverless and the S3 bucket with the Knowledge Base Files.\n",
    "\n",
    "Once the policies are ready, we will create the Knowledge Base role"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be047c6d-6d35-414e-915c-e2eec9c656eb",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Create IAM policies for KB to invoke embedding model\n",
    "bedrock_kb_allow_fm_model_policy_statement = {\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Sid\": \"AmazonBedrockAgentBedrockFoundationModelPolicy\",\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": \"bedrock:InvokeModel\",\n",
    "            \"Resource\": [\n",
    "                embedding_model_arn\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "\n",
    "kb_bedrock_policy_json = json.dumps(bedrock_kb_allow_fm_model_policy_statement)\n",
    "\n",
    "kb_bedrock_policy = iam_client.create_policy(\n",
    "    PolicyName=kb_bedrock_allow_policy_name,\n",
    "    PolicyDocument=kb_bedrock_policy_json\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f140db46-1383-4a91-8529-e2193a28675d",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Create IAM policies for KB to access OpenSearch Serverless\n",
    "bedrock_kb_allow_aoss_policy_statement = {\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": \"aoss:APIAccessAll\",\n",
    "            \"Resource\": [\n",
    "                f\"arn:aws:aoss:{region}:{account_id}:collection/*\"\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "\n",
    "\n",
    "kb_aoss_policy_json = json.dumps(bedrock_kb_allow_aoss_policy_statement)\n",
    "\n",
    "kb_aoss_policy = iam_client.create_policy(\n",
    "    PolicyName=kb_aoss_allow_policy_name,\n",
    "    PolicyDocument=kb_aoss_policy_json\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ecedf40d-b99d-477d-92b2-2144524b4bdb",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "kb_s3_allow_policy_statement = {\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Sid\": \"AllowKBAccessDocuments\",\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"s3:GetObject\",\n",
    "                \"s3:ListBucket\"\n",
    "            ],\n",
    "            \"Resource\": [\n",
    "                f\"arn:aws:s3:::{bucket_name}/*\",\n",
    "                f\"arn:aws:s3:::{bucket_name}\"\n",
    "            ],\n",
    "            \"Condition\": {\n",
    "                \"StringEquals\": {\n",
    "                    \"aws:ResourceAccount\": f\"{account_id}\"\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "\n",
    "\n",
    "kb_s3_json = json.dumps(kb_s3_allow_policy_statement)\n",
    "kb_s3_policy = iam_client.create_policy(\n",
    "    PolicyName=kb_s3_allow_policy_name,\n",
    "    PolicyDocument=kb_s3_json\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b568014-8a84-4aa7-b385-e0b7cc63a9ec",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Create IAM Role for the agent and attach IAM policies\n",
    "assume_role_policy_document = {\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [{\n",
    "          \"Effect\": \"Allow\",\n",
    "          \"Principal\": {\n",
    "            \"Service\": \"bedrock.amazonaws.com\"\n",
    "          },\n",
    "          \"Action\": \"sts:AssumeRole\"\n",
    "    }]\n",
    "}\n",
    "\n",
    "assume_role_policy_document_json = json.dumps(assume_role_policy_document)\n",
    "kb_role = iam_client.create_role(\n",
    "    RoleName=kb_role_name,\n",
    "    AssumeRolePolicyDocument=assume_role_policy_document_json\n",
    ")\n",
    "\n",
    "# Pause to make sure role is created\n",
    "time.sleep(10)\n",
    "    \n",
    "iam_client.attach_role_policy(\n",
    "    RoleName=kb_role_name,\n",
    "    PolicyArn=kb_bedrock_policy['Policy']['Arn']\n",
    ")\n",
    "\n",
    "iam_client.attach_role_policy(\n",
    "    RoleName=kb_role_name,\n",
    "    PolicyArn=kb_aoss_policy['Policy']['Arn']\n",
    ")\n",
    "\n",
    "iam_client.attach_role_policy(\n",
    "    RoleName=kb_role_name,\n",
    "    PolicyArn=kb_s3_policy['Policy']['Arn']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "402c9ba1-f46b-4f54-aae0-0aefe9713325",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "kb_role_arn = kb_role[\"Role\"][\"Arn\"]\n",
    "kb_role_arn"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65680beb-9562-4ed0-b354-86588a84a5e8",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 3.2 Create Vector Database\n",
    "\n",
    "Firt of all we have to create a vector store. In this section we will use Amazon OpenSerach Serverless.\n",
    "\n",
    "Amazon OpenSearch Serverless is a serverless option in Amazon OpenSearch Service. As a developer, you can use OpenSearch Serverless to run petabyte-scale workloads without configuring, managing, and scaling OpenSearch clusters. You get the same interactive millisecond response times as OpenSearch Service with the simplicity of a serverless environment. Pay only for what you use by automatically scaling resources to provide the right amount of capacity for your application—without impacting data ingestion.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "83ce2ee2-d9f7-463a-ab77-2fda11a72328",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Create OpenSearch Collection\n",
    "security_policy_json = {\n",
    "    \"Rules\": [\n",
    "        {\n",
    "            \"ResourceType\": \"collection\",\n",
    "            \"Resource\":[\n",
    "                f\"collection/{kb_collection_name}\"\n",
    "            ]\n",
    "        }\n",
    "    ],\n",
    "    \"AWSOwnedKey\": True\n",
    "}\n",
    "security_policy = open_search_serverless_client.create_security_policy(\n",
    "    description='security policy of aoss collection',\n",
    "    name=kb_collection_name,\n",
    "    policy=json.dumps(security_policy_json),\n",
    "    type='encryption'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1009572-4481-46ae-a733-eb82e38abc2f",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "network_policy_json = [\n",
    "  {\n",
    "    \"Rules\": [\n",
    "      {\n",
    "        \"Resource\": [\n",
    "          f\"collection/{kb_collection_name}\"\n",
    "        ],\n",
    "        \"ResourceType\": \"dashboard\"\n",
    "      },\n",
    "      {\n",
    "        \"Resource\": [\n",
    "          f\"collection/{kb_collection_name}\"\n",
    "        ],\n",
    "        \"ResourceType\": \"collection\"\n",
    "      }\n",
    "    ],\n",
    "    \"AllowFromPublic\": True\n",
    "  }\n",
    "]\n",
    "\n",
    "network_policy = open_search_serverless_client.create_security_policy(\n",
    "    description='network policy of aoss collection',\n",
    "    name=kb_collection_name,\n",
    "    policy=json.dumps(network_policy_json),\n",
    "    type='network'\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21f99ce9-140f-421e-b75b-9c2038d8ce18",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "response = sts_client.get_caller_identity()\n",
    "current_role = response['Arn']\n",
    "current_role"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de34fcff-94d6-48f2-b309-cb2ac9d89c49",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "data_policy_json = [\n",
    "  {\n",
    "    \"Rules\": [\n",
    "      {\n",
    "        \"Resource\": [\n",
    "          f\"collection/{kb_collection_name}\"\n",
    "        ],\n",
    "        \"Permission\": [\n",
    "          \"aoss:DescribeCollectionItems\",\n",
    "          \"aoss:CreateCollectionItems\",\n",
    "          \"aoss:UpdateCollectionItems\",\n",
    "          \"aoss:DeleteCollectionItems\"\n",
    "        ],\n",
    "        \"ResourceType\": \"collection\"\n",
    "      },\n",
    "      {\n",
    "        \"Resource\": [\n",
    "          f\"index/{kb_collection_name}/*\"\n",
    "        ],\n",
    "        \"Permission\": [\n",
    "            \"aoss:CreateIndex\",\n",
    "            \"aoss:DeleteIndex\",\n",
    "            \"aoss:UpdateIndex\",\n",
    "            \"aoss:DescribeIndex\",\n",
    "            \"aoss:ReadDocument\",\n",
    "            \"aoss:WriteDocument\"\n",
    "        ],\n",
    "        \"ResourceType\": \"index\"\n",
    "      }\n",
    "    ],\n",
    "    \"Principal\": [\n",
    "        kb_role_arn,\n",
    "        f\"arn:aws:sts::{account_id}:assumed-role/Admin/*\",\n",
    "        current_role\n",
    "    ],\n",
    "    \"Description\": \"\"\n",
    "  }\n",
    "]\n",
    "\n",
    "data_policy = open_search_serverless_client.create_access_policy(\n",
    "    description='data access policy for aoss collection',\n",
    "    name=kb_collection_name,\n",
    "    policy=json.dumps(data_policy_json),\n",
    "    type='data'\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15624c92-e657-4905-954b-c085307bcb14",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "opensearch_collection_response = open_search_serverless_client.create_collection(\n",
    "    description='OpenSearch collection for Amazon Bedrock Knowledge Base',\n",
    "    name=kb_collection_name,\n",
    "    standbyReplicas='DISABLED',\n",
    "    type='VECTORSEARCH'\n",
    ")\n",
    "opensearch_collection_response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5acb0d1d-8813-4a36-b4e0-176116504c7f",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "collection_arn = opensearch_collection_response[\"createCollectionDetail\"][\"arn\"]\n",
    "collection_arn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d3b9036-1843-479d-aa36-1756f8e4da92",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# wait for collection creation\n",
    "response = open_search_serverless_client.batch_get_collection(names=[kb_collection_name])\n",
    "# Periodically check collection status\n",
    "while (response['collectionDetails'][0]['status']) == 'CREATING':\n",
    "    print('Creating collection...')\n",
    "    time.sleep(30)\n",
    "    response = open_search_serverless_client.batch_get_collection(names=[kb_collection_name])\n",
    "print('\\nCollection successfully created:')\n",
    "print(response[\"collectionDetails\"])\n",
    "# Extract the collection endpoint from the response\n",
    "host = (response['collectionDetails'][0]['collectionEndpoint'])\n",
    "final_host = host.replace(\"https://\", \"\")\n",
    "final_host"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "455a3b6f-1bca-410d-86fe-5d38fe61470b",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 3.3 - Create OpenSearch Index\n",
    "\n",
    "Let's now create a vector index to index our data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59f5e1fa-2d27-4d57-b288-327a16edd67b",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "credentials = boto3.Session().get_credentials()\n",
    "service = 'aoss'\n",
    "awsauth = AWS4Auth(\n",
    "    credentials.access_key, \n",
    "    credentials.secret_key,\n",
    "    region, \n",
    "    service, \n",
    "    session_token=credentials.token\n",
    ")\n",
    "\n",
    "# Build the OpenSearch client\n",
    "open_search_client = OpenSearch(\n",
    "    hosts=[{'host': final_host, 'port': 443}],\n",
    "    http_auth=awsauth,\n",
    "    use_ssl=True,\n",
    "    verify_certs=True,\n",
    "    connection_class=RequestsHttpConnection,\n",
    "    timeout=300\n",
    ")\n",
    "# It can take up to a minute for data access rules to be enforced\n",
    "time.sleep(45)\n",
    "index_body = {\n",
    "    \"settings\": {\n",
    "        \"index.knn\": True,\n",
    "        \"number_of_shards\": 1,\n",
    "        \"knn.algo_param.ef_search\": 512,\n",
    "        \"number_of_replicas\": 0,\n",
    "    },\n",
    "    \"mappings\": {\n",
    "        \"properties\": {}\n",
    "    }\n",
    "}\n",
    "\n",
    "index_body[\"mappings\"][\"properties\"][kb_vectorField] = {\n",
    "    \"type\": \"knn_vector\",\n",
    "    \"dimension\": 1536,\n",
    "    \"method\": {\n",
    "         \"name\": \"hnsw\",\n",
    "         \"engine\": \"faiss\"\n",
    "    },\n",
    "}\n",
    "\n",
    "index_body[\"mappings\"][\"properties\"][kb_textField] = {\n",
    "    \"type\": \"text\"\n",
    "}\n",
    "\n",
    "index_body[\"mappings\"][\"properties\"][kb_metadataField] = {\n",
    "    \"type\": \"text\"\n",
    "}\n",
    "\n",
    "# Create index\n",
    "response = open_search_client.indices.create(kb_vector_index_name, body=index_body)\n",
    "print('\\nCreating index:')\n",
    "print(response)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a716b75-29d1-47e2-b9a0-24887b0624ea",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 3.5 - Create Knowledge Base\n",
    "Now that we have the Vector database available in OpenSearch Serverless, let's create a Knowledge Base and associate it with the OpenSearch DB"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dce574a8-5eb9-4fbf-af11-08b7a05e5044",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "storage_configuration = {\n",
    "    'opensearchServerlessConfiguration': {\n",
    "        'collectionArn': collection_arn, \n",
    "        'fieldMapping': {\n",
    "            'metadataField': kb_metadataField,\n",
    "            'textField': kb_textField,\n",
    "            'vectorField': kb_vectorField\n",
    "        },\n",
    "        'vectorIndexName': kb_vector_index_name\n",
    "    },\n",
    "    'type': 'OPENSEARCH_SERVERLESS'\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f853d36e-cf19-4241-8cff-adecf14c3f45",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Creating the knowledge base\n",
    "try:\n",
    "    # ensure the index is created and available\n",
    "    time.sleep(45)\n",
    "    kb_obj = bedrock_agent_client.create_knowledge_base(\n",
    "        name=kb_name, \n",
    "        description='KB that contains the bedrock documentation',\n",
    "        roleArn=kb_role_arn,\n",
    "        knowledgeBaseConfiguration={\n",
    "            'type': 'VECTOR',  # Corrected type\n",
    "            'vectorKnowledgeBaseConfiguration': {\n",
    "                'embeddingModelArn': embedding_model_arn\n",
    "            }\n",
    "        },\n",
    "        storageConfiguration=storage_configuration\n",
    "    )\n",
    "\n",
    "    # Pretty print the response\n",
    "    pprint.pprint(kb_obj)\n",
    "\n",
    "except Exception as e:\n",
    "    print(f\"Error occurred: {e}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92a7acc7-23fc-4c68-bc5d-f56c819b4e45",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Define the S3 configuration for your data source\n",
    "s3_configuration = {\n",
    "    'bucketArn': bucket_arn,\n",
    "    'inclusionPrefixes': [kb_key]  \n",
    "}\n",
    "\n",
    "# Define the data source configuration\n",
    "data_source_configuration = {\n",
    "    's3Configuration': s3_configuration,\n",
    "    'type': 'S3'\n",
    "}\n",
    "\n",
    "knowledge_base_id = kb_obj[\"knowledgeBase\"][\"knowledgeBaseId\"]\n",
    "knowledge_base_arn = kb_obj[\"knowledgeBase\"][\"knowledgeBaseArn\"]\n",
    "\n",
    "chunking_strategy_configuration = {\n",
    "    \"chunkingStrategy\": \"FIXED_SIZE\",\n",
    "    \"fixedSizeChunkingConfiguration\": {\n",
    "        \"maxTokens\": 512,\n",
    "        \"overlapPercentage\": 20\n",
    "    }\n",
    "}\n",
    "\n",
    "# Create the data source\n",
    "try:\n",
    "    # ensure that the KB is created and available\n",
    "    time.sleep(45)\n",
    "    data_source_response = bedrock_agent_client.create_data_source(\n",
    "        knowledgeBaseId=knowledge_base_id,\n",
    "        name=data_source_name,\n",
    "        description='DataSource for the bedrock documentation',\n",
    "        dataSourceConfiguration=data_source_configuration,\n",
    "        vectorIngestionConfiguration = {\n",
    "            \"chunkingConfiguration\": chunking_strategy_configuration\n",
    "        }\n",
    "    )\n",
    "\n",
    "    # Pretty print the response\n",
    "    pprint.pprint(data_source_response)\n",
    "\n",
    "except Exception as e:\n",
    "    print(f\"Error occurred: {e}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9791ecb-9949-4a2d-9ff6-ae463be967cd",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 3.6 - Start ingestion job\n",
    "\n",
    "Once the Knowledge Base and Data Source are created, we can start the ingestion job. During the ingestion job, Knowledge Base will fetch the documents in the data source, pre-process it to extract text, chunk it based on the chunking size provided, create embeddings of each chunk and then write it to the vector database, in this case Amazon OpenSource Serverless.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab88c732-d26b-442e-a502-4a2218e65c80",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Start an ingestion job\n",
    "data_source_id = data_source_response[\"dataSource\"][\"dataSourceId\"]\n",
    "start_job_response = bedrock_agent_client.start_ingestion_job(\n",
    "    knowledgeBaseId=knowledge_base_id, \n",
    "    dataSourceId=data_source_id\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a253838-6dd3-4a44-84aa-9ee9798bb9d0",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 4. Create Agent\n",
    "\n",
    "We will now create the Agent and associate the Knowledge Base to it. To do so we need to: \n",
    "1. Create Agent IAM role and policies\n",
    "1. Create Agent\n",
    "1. Associate Agent to Knowledge Base\n",
    "1. Prepare Agent\n",
    "\n",
    "### 4.1 - Create Agent IAM role and policies\n",
    "First we need to create the agent policies that allow bedrock model invocation and Knowledge Base retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d341725-3d9e-4c9a-aa65-ed661c2775b3",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Create IAM policies for agent\n",
    "bedrock_agent_bedrock_allow_policy_statement = {\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Sid\": \"AmazonBedrockAgentBedrockFoundationModelPolicy\",\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": \"bedrock:InvokeModel\",\n",
    "            \"Resource\": [\n",
    "                f\"arn:aws:bedrock:{region}::foundation-model/{model_id}\"\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "\n",
    "bedrock_policy_json = json.dumps(bedrock_agent_bedrock_allow_policy_statement)\n",
    "\n",
    "agent_bedrock_policy = iam_client.create_policy(\n",
    "    PolicyName=bedrock_agent_bedrock_allow_policy_name,\n",
    "    PolicyDocument=bedrock_policy_json\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38c0f7c3-4c98-49a5-bb94-db7935f754b5",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "bedrock_agent_kb_retrival_policy_statement = {\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"bedrock:Retrieve\"\n",
    "            ],\n",
    "            \"Resource\": [\n",
    "                knowledge_base_arn\n",
    "            ]\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "bedrock_agent_kb_json = json.dumps(bedrock_agent_kb_retrival_policy_statement)\n",
    "agent_kb_schema_policy = iam_client.create_policy(\n",
    "    PolicyName=bedrock_agent_kb_allow_policy_name,\n",
    "    Description=f\"Policy to allow agent to retrieve documents from knowledge base.\",\n",
    "    PolicyDocument=bedrock_agent_kb_json\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd178e18-196d-464e-b8ff-fb36090cbdd8",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "# Create IAM Role for the agent and attach IAM policies\n",
    "assume_role_policy_document = {\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [{\n",
    "          \"Effect\": \"Allow\",\n",
    "          \"Principal\": {\n",
    "            \"Service\": \"bedrock.amazonaws.com\"\n",
    "          },\n",
    "          \"Action\": \"sts:AssumeRole\"\n",
    "    }]\n",
    "}\n",
    "\n",
    "assume_role_policy_document_json = json.dumps(assume_role_policy_document)\n",
    "agent_role = iam_client.create_role(\n",
    "    RoleName=agent_role_name,\n",
    "    AssumeRolePolicyDocument=assume_role_policy_document_json\n",
    ")\n",
    "\n",
    "# Pause to make sure role is created\n",
    "time.sleep(10)\n",
    "    \n",
    "iam_client.attach_role_policy(\n",
    "    RoleName=agent_role_name,\n",
    "    PolicyArn=agent_bedrock_policy['Policy']['Arn']\n",
    ")\n",
    "\n",
    "\n",
    "iam_client.attach_role_policy(\n",
    "    RoleName=agent_role_name,\n",
    "    PolicyArn=agent_kb_schema_policy['Policy']['Arn']\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b3f4fa0-9854-4602-9024-b9b1f05c9c20",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 4.2 - Create 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 function. 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": "4ddf7880-dafc-479a-987f-628369b94035",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Create Agent\n",
    "response = bedrock_agent_client.create_agent(\n",
    "    agentName=agent_name,\n",
    "    agentResourceRoleArn=agent_role['Role']['Arn'],\n",
    "    description=\"Agent supporting Amazon Bedrock Developers.\",\n",
    "    idleSessionTTLInSeconds=1800,\n",
    "    foundationModel=model_id,\n",
    "    instruction=agent_instruction,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e32e05be-f686-42eb-afe9-42026c3c301a",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Let's now store the agent id in a local variable to use it on the next steps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47a5e409-acca-47be-9a89-c72091f539de",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "agent_id = response['agent']['agentId']\n",
    "agent_id"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "accf4d20-6644-48f0-a82d-a60a25fab0a5",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 4.3 - Associate agent to the Knowledge Base\n",
    "Next, we need to associate the agent created with the Knowledge Base for the Bedrock documentation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3b36aba-45ba-4e91-a7af-91937348d264",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "agent_kb_description = bedrock_agent_client.associate_agent_knowledge_base(\n",
    "    agentId=agent_id,\n",
    "    agentVersion='DRAFT',\n",
    "    description=f'Use the information in the {kb_name} knowledge base to provide accurate responses to the questions about Amazon Bedrock.',\n",
    "    knowledgeBaseId=knowledge_base_id \n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5776f5ae-ddb6-4241-86ed-0f3bcb84839a",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "### 4.4 - Prepare Agent\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": "962eb7bf-250d-4ed5-af2d-a5ad6e3b7bff",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "agent_prepare = bedrock_agent_client.prepare_agent(agentId=agent_id)\n",
    "agent_prepare"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b603a0e1-cbfc-492a-86c5-304156bafbbc",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 5 - Testing Agent\n",
    "\n",
    "Now that we have our agent, let's invoke it to test if it is providing correct information about Amazon Bedrock. To do so, let's first create an Agent Alias"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a24e0bb8-ba19-407c-86f7-d20bccbf33c9",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Pause to make sure agent is prepared\n",
    "time.sleep(30)\n",
    "agent_alias = bedrock_agent_client.create_agent_alias(\n",
    "    agentId=agent_id,\n",
    "    agentAliasName=agent_alias_name\n",
    ")\n",
    "# Pause to make sure agent alias is ready\n",
    "time.sleep(30)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36a5e2d5-6f23-4dea-a720-8b3a6d3c07ed",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "agent_alias"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e23a8c-e482-41a9-b5f6-82d28eb5bdfc",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "Now that we've created the agent, let's use the bedrock-agent-runtime client to invoke this agent and get the information from the Knowledge base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d3a53720-96f6-45f3-b4a8-77526d3ed5fd",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Extract the agentAliasId from the response\n",
    "agent_alias_id = agent_alias['agentAlias']['agentAliasId']\n",
    "\n",
    "## create a random id for session initiator id\n",
    "session_id:str = str(uuid.uuid1())\n",
    "enable_trace:bool = True\n",
    "end_session:bool = False\n",
    "\n",
    "# invoke the agent API\n",
    "agentResponse = bedrock_agent_runtime_client.invoke_agent(\n",
    "    inputText=\"How can I evaluate models on Bedrock?\",\n",
    "    agentId=agent_id,\n",
    "    agentAliasId=agent_alias_id, \n",
    "    sessionId=session_id,\n",
    "    enableTrace=enable_trace, \n",
    "    endSession= end_session\n",
    ")\n",
    "\n",
    "logger.info(pprint.pprint(agentResponse))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "430f2083-dc61-4a3e-a373-d36660fd7a82",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "event_stream = agentResponse['completion']\n",
    "try:\n",
    "    for event in event_stream:        \n",
    "        if 'chunk' in event:\n",
    "            data = event['chunk']['bytes']\n",
    "            logger.info(f\"Final answer ->\\n{data.decode('utf8')}\")\n",
    "            agent_answer = data.decode('utf8')\n",
    "            end_event_received = True\n",
    "            # End event indicates that the request finished successfully\n",
    "        elif 'trace' in event:\n",
    "            logger.info(json.dumps(event['trace'], indent=2))\n",
    "        else:\n",
    "            raise Exception(\"unexpected event.\", event)\n",
    "except Exception as e:\n",
    "    raise Exception(\"unexpected event.\", e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28fc780b-d700-4d89-bd23-2e183e03875d",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# And here is the response if you just want to see agent's reply\n",
    "print(agent_answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bf20e96",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "def simple_agent_invoke(input_text, agent_id, agent_alias_id, session_id=None, enable_trace=False, end_session=False):\n",
    "    if session_id is None:\n",
    "        session_id:str = str(uuid.uuid1())\n",
    "\n",
    "    agentResponse = bedrock_agent_runtime_client.invoke_agent(\n",
    "        inputText=input_text,\n",
    "        agentId=agent_id,\n",
    "        agentAliasId=agent_alias_id, \n",
    "        sessionId=session_id,\n",
    "        enableTrace=enable_trace, \n",
    "        endSession= end_session\n",
    "    )\n",
    "    logger.info(pprint.pprint(agentResponse))\n",
    "    \n",
    "    agent_answer = ''\n",
    "    event_stream = agentResponse['completion']\n",
    "    try:\n",
    "        for event in event_stream:        \n",
    "            if 'chunk' in event:\n",
    "                data = event['chunk']['bytes']\n",
    "                logger.info(f\"Final answer ->\\n{data.decode('utf8')}\")\n",
    "                agent_answer = data.decode('utf8')\n",
    "                end_event_received = True\n",
    "                # End event indicates that the request finished successfully\n",
    "            elif 'trace' in event:\n",
    "                logger.info(json.dumps(event['trace'], indent=2))\n",
    "            else:\n",
    "                raise Exception(\"unexpected event.\", event)\n",
    "    except Exception as e:\n",
    "        raise Exception(\"unexpected event.\", e)\n",
    "    return agent_answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6cae6e9",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "simple_agent_invoke(\"what is bedrock provisioned throughput?\", agent_id, agent_alias_id, session_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c4b99e1d",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "simple_agent_invoke(\"what are the components of a Bedrock Guardrail?\", agent_id, agent_alias_id, session_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e9a37f1",
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "simple_agent_invoke(\"what are the components of a Bedrock Guardrail?\", agent_id, agent_alias_id, session_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77e4a008-2176-499a-9fad-93a69ae5cfe0",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## 6 - Clean up (Optional)\n",
    "\n",
    "The next steps are optional and demonstrate how to delete our agent. To delete the agent we need to:\n",
    "\n",
    "1. delete agent alias\n",
    "1. delete agent\n",
    "1. delete the knowledge base\n",
    "1. delete the OpenSearch Serverless vector store\n",
    "1. empty created s3 bucket\n",
    "1. delete s3 bucket\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd1c3418-dc37-4185-818a-743485ecb2b7",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "agent_alias_deletion = bedrock_agent_client.delete_agent_alias(\n",
    "    agentId=agent_id,\n",
    "    agentAliasId=agent_alias['agentAlias']['agentAliasId']\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7d3f07d-fee1-485c-aee8-03fe98f55257",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "agent_deletion = bedrock_agent_client.delete_agent(\n",
    "    agentId=agent_id\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6302823-3fb0-43f2-a926-590fb04e7036",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Empty and delete S3 Bucket\n",
    "\n",
    "objects = s3_client.list_objects(Bucket=bucket_name)  \n",
    "if 'Contents' in objects:\n",
    "    for obj in objects['Contents']:\n",
    "        s3_client.delete_object(Bucket=bucket_name, Key=obj['Key']) \n",
    "s3_client.delete_bucket(Bucket=bucket_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "877f8d8b-5dfc-4f43-882b-176a27475363",
   "metadata": {
    "tags": [],
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Delete IAM Roles and policies and Knowledge Base files\n",
    "for policy in [\n",
    "    agent_bedrock_policy, \n",
    "    agent_kb_schema_policy,\n",
    "    kb_bedrock_policy,\n",
    "    kb_aoss_policy,\n",
    "    kb_s3_policy\n",
    "]:\n",
    "    response = iam_client.list_entities_for_policy(\n",
    "        PolicyArn=policy['Policy']['Arn'],\n",
    "        EntityFilter='Role'\n",
    "    )\n",
    "\n",
    "    for role in response['PolicyRoles']:\n",
    "        iam_client.detach_role_policy(\n",
    "            RoleName=role['RoleName'], \n",
    "            PolicyArn=policy['Policy']['Arn']\n",
    "        )\n",
    "\n",
    "    iam_client.delete_policy(\n",
    "        PolicyArn=policy['Policy']['Arn']\n",
    "    )\n",
    "\n",
    "    \n",
    "\n",
    "for role_name in [\n",
    "    agent_role_name, \n",
    "    kb_role_name\n",
    "]:\n",
    "    try: \n",
    "        iam_client.delete_role(\n",
    "            RoleName=role_name\n",
    "        )\n",
    "    except Exception as e:\n",
    "        print(e)\n",
    "        print(\"couldn't delete role\", role_name)\n",
    "        \n",
    "    \n",
    "try:\n",
    "\n",
    "    open_search_serverless_client.delete_collection(\n",
    "        id=opensearch_collection_response[\"createCollectionDetail\"][\"id\"]\n",
    "    )\n",
    "\n",
    "    open_search_serverless_client.delete_access_policy(\n",
    "          name=kb_collection_name,\n",
    "          type='data'\n",
    "    )    \n",
    "\n",
    "    open_search_serverless_client.delete_security_policy(\n",
    "          name=kb_collection_name,\n",
    "          type='network'\n",
    "    )   \n",
    "\n",
    "    open_search_serverless_client.delete_security_policy(\n",
    "          name=kb_collection_name,\n",
    "          type='encryption'\n",
    "    )    \n",
    "    bedrock_agent_client.delete_knowledge_base(\n",
    "        knowledgeBaseId=knowledge_base_id\n",
    "    )\n",
    "except Exception as e:\n",
    "    print(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15cbc834-c71f-40ae-8a12-6bc2b527ad9b",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "## Conclusion\n",
    "\n",
    "We have now experimented with using boto3 SDK to create, invoke and delete an agent having a single KB connected to it.\n",
    "## Take aways\n",
    "\n",
    "Adapt this notebook to create new agents for your application\n",
    "\n",
    "## Thank You"
   ]
  }
 ],
 "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",
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}