knowledge-bases/01-rag-concepts/1b_create_ingest_documents_test_kb_multi_ds.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Knowledge Bases for Amazon Bedrock - End to end example using multiple s3 bucket(s) as data source(s)\n",
    "\n",
    "This notebook provides sample code for building an empty OpenSearch Serverless (OSS) index, Amazon Bedrock knowledge base and ingest documents into the index from multiple data sources. Now you can add multiple data sources with a Knowledge base.\n",
    "\n",
    "\n",
    "#### Notebook Walkthrough\n",
    "\n",
    "A data pipeline that ingests documents (typically stored in multiple Amazon S3) into a knowledge base i.e. a vector database such as Amazon OpenSearch Service Serverless (AOSS) so that it is available for lookup when a question is received.\n",
    "\n",
    "- Load the documents into the knowledge base by connecting your multiple s3 bucket (multiple data source). \n",
    "- Ingestion - Knowledge base will split them into smaller chunks (based on the strategy selected), generate embeddings and store it in the associated vectore store.\n",
    "\n",
    "![data_ingestion.png](./images/data_ingestion.png)\n",
    "\n",
    "\n",
    "#### Steps: \n",
    "- Create Amazon Bedrock Knowledge Base execution role with necessary policies for accessing data from multiple S3 and writing embeddings into OSS.\n",
    "- Create an empty OpenSearch serverless index.\n",
    "- Download documents\n",
    "- Create Amazon Bedrock knowledge base\n",
    "- Create data source(s) within knowledge base which will connect to corresponding Amazon S3\n",
    "- For each data source, start ingestion jobs using KB APIs which will read data from s3, chunk it, convert chunks into embeddings using Amazon Titan Embeddings model and then store these embeddings in AOSS. All of this without having to build, deploy and manage the data pipeline.\n",
    "\n",
    "Once the data is available in the Bedrock Knowledge Base then a question answering application can be built using the Knowledge Base APIs provided by Amazon Bedrock.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Pre-requisites\n",
    "This notebook requires permissions to:\n",
    "- create and delete Amazon IAM roles\n",
    "- create, update and delete Amazon S3 buckets\n",
    "- access Amazon Bedrock\n",
    "- access to Amazon OpenSearch Serverless\n",
    "\n",
    "If running on SageMaker Studio, you should add the following managed policies to your role:\n",
    "- IAMFullAccess\n",
    "- AWSLambda_FullAccess\n",
    "- AmazonS3FullAccess\n",
    "- AmazonBedrockFullAccess\n",
    "- Custom policy for Amazon OpenSearch Serverless such as:\n",
    "```\n",
    "{\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": \"aoss:*\",\n",
    "            \"Resource\": \"*\"\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "```\n",
    "<div class=\"alert alert-block alert-info\">\n",
    "<b>Note:</b> Please make sure to enable `Anthropic Claude 3` and `Anthropic Claude Instant` model access in Amazon Bedrock Console, as the notebook will use Anthropic Claude 3 Sonnet and Claude instant models for testing the knowledge base once its created.\n",
    "</div>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "Before running the rest of this notebook, you'll need to run the cells below to (ensure necessary libraries are installed and) connect to Bedrock."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "%pip install -U opensearch-py==2.3.1\n",
    "%pip install -U boto3==1.34.85\n",
    "%pip install -U retrying==1.3.4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# restart kernel\n",
    "from IPython.core.display import HTML\n",
    "HTML(\"<script>Jupyter.notebook.kernel.restart()</script>\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "import os\n",
    "import boto3\n",
    "from botocore.exceptions import ClientError\n",
    "import pprint\n",
    "from utility import create_bedrock_execution_role, create_bedrock_execution_role_multi_ds, create_oss_policy_attach_bedrock_execution_role, create_policies_in_oss, interactive_sleep\n",
    "import random\n",
    "from retrying import retry\n",
    "suffix = random.randrange(200, 900)\n",
    "\n",
    "sts_client = boto3.client('sts')\n",
    "boto3_session = boto3.session.Session()\n",
    "region_name = boto3_session.region_name\n",
    "bedrock_agent_client = boto3_session.client('bedrock-agent', region_name=region_name)\n",
    "service = 'aoss'\n",
    "s3_client = boto3.client('s3')\n",
    "account_id = sts_client.get_caller_identity()[\"Account\"]\n",
    "s3_suffix = f\"{region_name}-{account_id}\"\n",
    "# Now you can add multiple data sources to a Knowledge Base. For this notebook, we'll test Knowledge Base creation with two different data sources\n",
    "bucket_name_1 = f'bedrock-kb-{s3_suffix}-1' # replace it with your first bucket name.\n",
    "bucket_name_2 = f'bedrock-kb-{s3_suffix}-2' # replace it with your second bucket name.\n",
    "bucket_names=[bucket_name_1,bucket_name_2]\n",
    "pp = pprint.PrettyPrinter(indent=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# Check if bucket exists, and if not create S3 bucket for knowledge base data source\n",
    "for bucket_name in bucket_names:\n",
    "    try:\n",
    "        s3_client.head_bucket(Bucket=bucket_name)\n",
    "        print(f'Bucket {bucket_name} Exists')\n",
    "    except ClientError as e:\n",
    "        print(f'Creating bucket {bucket_name}')\n",
    "        s3bucket = s3_client.create_bucket(\n",
    "            Bucket=bucket_name,\n",
    "            # CreateBucketConfiguration={ 'LocationConstraint': region_name }\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create a vector store - OpenSearch Serverless index\n",
    "\n",
    "### Step 1 - Create OSS policies and collection\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."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import time\n",
    "vector_store_name = f'bedrock-sample-rag-{suffix}'\n",
    "index_name = f\"bedrock-sample-rag-index-{suffix}\"\n",
    "aoss_client = boto3_session.client('opensearchserverless')\n",
    "bedrock_kb_execution_role = create_bedrock_execution_role_multi_ds(bucket_names=bucket_names)\n",
    "\n",
    "bedrock_kb_execution_role_arn = bedrock_kb_execution_role['Role']['Arn']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# create security, network and data access policies within OSS\n",
    "encryption_policy, network_policy, access_policy = create_policies_in_oss(vector_store_name=vector_store_name,\n",
    "                       aoss_client=aoss_client,\n",
    "                       bedrock_kb_execution_role_arn=bedrock_kb_execution_role_arn)\n",
    "collection = aoss_client.create_collection(name=vector_store_name,type='VECTORSEARCH')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pp.pprint(collection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the OpenSearch serverless collection URL\n",
    "collection_id = collection['createCollectionDetail']['id']\n",
    "host = collection_id + '.' + region_name + '.aoss.amazonaws.com'\n",
    "print(host)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "import time\n",
    "# wait for collection creation\n",
    "# This can take couple of minutes to finish\n",
    "response = aoss_client.batch_get_collection(names=[vector_store_name])\n",
    "# Periodically check collection status\n",
    "while (response['collectionDetails'][0]['status']) == 'CREATING':\n",
    "    print('Creating collection...')\n",
    "    interactive_sleep(30)\n",
    "    response = aoss_client.batch_get_collection(names=[vector_store_name])\n",
    "print('\\nCollection successfully created:')\n",
    "pp.pprint(response[\"collectionDetails\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# create opensearch serverless access policy and attach it to Bedrock execution role\n",
    "try:\n",
    "    create_oss_policy_attach_bedrock_execution_role(collection_id=collection_id,\n",
    "                                                    bedrock_kb_execution_role=bedrock_kb_execution_role)\n",
    "    # It can take up to a minute for data access rules to be enforced\n",
    "    interactive_sleep(60)\n",
    "except Exception as e:\n",
    "    print(\"Policy already exists\")\n",
    "    pp.pprint(e)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2 - Create vector index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the vector index in Opensearch serverless, with the knn_vector field index mapping, specifying the dimension size, name and engine.\n",
    "from opensearchpy import OpenSearch, RequestsHttpConnection, AWSV4SignerAuth, RequestError\n",
    "credentials = boto3.Session().get_credentials()\n",
    "awsauth = auth = AWSV4SignerAuth(credentials, region_name, service)\n",
    "\n",
    "index_name = f\"bedrock-sample-index-{suffix}\"\n",
    "body_json = {\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",
    "         \"vector\": {\n",
    "            \"type\": \"knn_vector\",\n",
    "            \"dimension\": 1536,\n",
    "             \"method\": {\n",
    "                 \"name\": \"hnsw\",\n",
    "                 \"engine\": \"faiss\",\n",
    "                 \"space_type\": \"l2\"\n",
    "             },\n",
    "         },\n",
    "         \"text\": {\n",
    "            \"type\": \"text\"\n",
    "         },\n",
    "         \"text-metadata\": {\n",
    "            \"type\": \"text\"         }\n",
    "      }\n",
    "   }\n",
    "}\n",
    "\n",
    "# Build the OpenSearch client\n",
    "oss_client = OpenSearch(\n",
    "    hosts=[{'host': host, 'port': 443}],\n",
    "    http_auth=awsauth,\n",
    "    use_ssl=True,\n",
    "    verify_certs=True,\n",
    "    connection_class=RequestsHttpConnection,\n",
    "    timeout=300\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    },
    "pycharm": {
     "is_executing": true
    }
   },
   "outputs": [],
   "source": [
    "# Create index\n",
    "try:\n",
    "    response = oss_client.indices.create(index=index_name, body=json.dumps(body_json))\n",
    "    print('\\nCreating index:')\n",
    "    pp.pprint(response)\n",
    "\n",
    "    # index creation can take up to a minute\n",
    "    interactive_sleep(60)\n",
    "except RequestError as e:\n",
    "    # you can delete the index if its already exists\n",
    "    # oss_client.indices.delete(index=index_name)\n",
    "    print(f'Error while trying to create the index, with error {e.error}\\nyou may unmark the delete above to delete, and recreate the index')\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download data to ingest into our knowledge base"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Download and prepare dataset for first data source\n",
    "!mkdir -p ./data_1\n",
    "\n",
    "from urllib.request import urlretrieve\n",
    "urls = [\n",
    "    'https://s2.q4cdn.com/299287126/files/doc_financials/2023/ar/2022-Shareholder-Letter.pdf',\n",
    "    'https://s2.q4cdn.com/299287126/files/doc_financials/2022/ar/2021-Shareholder-Letter.pdf'\n",
    "]\n",
    "\n",
    "filenames = [\n",
    "    'AMZN-2022-Shareholder-Letter.pdf',\n",
    "    'AMZN-2021-Shareholder-Letter.pdf'\n",
    "]\n",
    "\n",
    "data_root_1 = \"./data_1/\"\n",
    "\n",
    "for idx, url in enumerate(urls):\n",
    "    file_path = data_root_1 + filenames[idx]\n",
    "    urlretrieve(url, file_path)\n",
    "\n",
    "\n",
    "# Download and prepare dataset for second data source\n",
    "!mkdir -p ./data_2\n",
    "\n",
    "from urllib.request import urlretrieve\n",
    "urls = [\n",
    "    'https://s2.q4cdn.com/299287126/files/doc_financials/2021/ar/Amazon-2020-Shareholder-Letter-and-1997-Shareholder-Letter.pdf',\n",
    "    'https://s2.q4cdn.com/299287126/files/doc_financials/2020/ar/2019-Shareholder-Letter.pdf'\n",
    "]\n",
    "\n",
    "filenames = [\n",
    "    'AMZN-2020-Shareholder-Letter.pdf',\n",
    "    'AMZN-2019-Shareholder-Letter.pdf'\n",
    "]\n",
    "\n",
    "data_root_2 = \"./data_2/\"\n",
    "\n",
    "for idx, url in enumerate(urls):\n",
    "    file_path = data_root_2 + filenames[idx]\n",
    "    urlretrieve(url, file_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Upload data to S3 Bucket data source"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "jupyter": {
     "source_hidden": true
    }
   },
   "outputs": [],
   "source": [
    "# Upload data to s3 to the bucket that was configured as a data source to the knowledge base\n",
    "s3_client = boto3.client(\"s3\")\n",
    "def uploadDirectory(path,bucket_name):\n",
    "        for root,dirs,files in os.walk(path):\n",
    "            for file in files:\n",
    "                s3_client.upload_file(os.path.join(root,file),bucket_name,file)\n",
    "\n",
    "uploadDirectory(data_root_1, bucket_name_1)\n",
    "uploadDirectory(data_root_2, bucket_name_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Knowledge Base\n",
    "Steps:\n",
    "- initialize Open search serverless configuration which will include collection ARN, index name, vector field, text field and metadata field.\n",
    "- initialize the Titan embeddings model ARN, as this will be used to create the embeddings for each of the text chunks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "opensearchServerlessConfiguration = {\n",
    "            \"collectionArn\": collection[\"createCollectionDetail\"]['arn'],\n",
    "            \"vectorIndexName\": index_name,\n",
    "            \"fieldMapping\": {\n",
    "                \"vectorField\": \"vector\",\n",
    "                \"textField\": \"text\",\n",
    "                \"metadataField\": \"text-metadata\"\n",
    "            }\n",
    "        }\n",
    "\n",
    "# The embedding model used by Bedrock to embed ingested documents, and realtime prompts\n",
    "embeddingModelArn = f\"arn:aws:bedrock:{region_name}::foundation-model/amazon.titan-embed-text-v1\"\n",
    "\n",
    "name = f\"bedrock-sample-knowledge-base-{suffix}\"\n",
    "description = \"Amazon shareholder letter knowledge base.\"\n",
    "roleArn = bedrock_kb_execution_role_arn\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Provide the above configurations as input to the `create_knowledge_base` method, which will create the Knowledge base."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a KnowledgeBase\n",
    "from retrying import retry\n",
    "\n",
    "@retry(wait_random_min=1000, wait_random_max=2000,stop_max_attempt_number=7)\n",
    "def create_knowledge_base_func():\n",
    "    create_kb_response = bedrock_agent_client.create_knowledge_base(\n",
    "        name = name,\n",
    "        description = description,\n",
    "        roleArn = roleArn,\n",
    "        knowledgeBaseConfiguration = {\n",
    "            \"type\": \"VECTOR\",\n",
    "            \"vectorKnowledgeBaseConfiguration\": {\n",
    "                \"embeddingModelArn\": embeddingModelArn\n",
    "            }\n",
    "        },\n",
    "        storageConfiguration = {\n",
    "            \"type\": \"OPENSEARCH_SERVERLESS\",\n",
    "            \"opensearchServerlessConfiguration\":opensearchServerlessConfiguration\n",
    "        }\n",
    "    )\n",
    "    return create_kb_response[\"knowledgeBase\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    kb = create_knowledge_base_func()\n",
    "except Exception as err:\n",
    "    print(f\"{err=}, {type(err)=}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pp.pprint(kb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get KnowledgeBase \n",
    "get_kb_response = bedrock_agent_client.get_knowledge_base(knowledgeBaseId = kb['knowledgeBaseId'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next we need to create data source(s), which will be associated with the knowledge base created above. Once the data source(s) is ready, we can then start to ingest the documents."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Data Source(s)\n",
    "Steps:\n",
    "- initialize chunking strategy, based on which KB will split the documents into pieces of size equal to the chunk size mentioned in the `chunkingStrategyConfiguration`.\n",
    "- initialize the s3 configuration, which will be used to create the data source object later.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to create KB\n",
    "def create_ds(bucket_names):\n",
    "    ds_list=[]\n",
    "    for bucket_name in bucket_names:\n",
    "        # Ingest strategy - How to ingest data from the data source\n",
    "        chunkingStrategyConfiguration = {\n",
    "            \"chunkingStrategy\": \"FIXED_SIZE\",\n",
    "            \"fixedSizeChunkingConfiguration\": {\n",
    "                \"maxTokens\": 512,\n",
    "                \"overlapPercentage\": 20\n",
    "            }\n",
    "        }\n",
    "        \n",
    "        # The data source to ingest documents from, into the OpenSearch serverless knowledge base index\n",
    "        s3Configuration = {\n",
    "            \"bucketArn\": f\"arn:aws:s3:::{bucket_name}\",\n",
    "            # \"inclusionPrefixes\":[\"*.*\"] # you can use this if you want to create a KB using data within s3 prefixes.\n",
    "            }\n",
    "        # Create a DataSource in KnowledgeBase \n",
    "        create_ds_response = bedrock_agent_client.create_data_source(\n",
    "            name = f'{name}-{bucket_name}',\n",
    "            description = description,\n",
    "            knowledgeBaseId = kb['knowledgeBaseId'],\n",
    "            dataSourceConfiguration = {\n",
    "                \"type\": \"S3\",\n",
    "                \"s3Configuration\":s3Configuration\n",
    "            },\n",
    "            vectorIngestionConfiguration = {\n",
    "                \"chunkingConfiguration\": chunkingStrategyConfiguration\n",
    "            }\n",
    "        )\n",
    "        ds = create_ds_response[\"dataSource\"]\n",
    "        pp.pprint(ds)\n",
    "        ds_list.append(ds)\n",
    "    return ds_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_sources = create_ds(bucket_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get DataSource \n",
    "for idx, ds in enumerate(data_sources):\n",
    "    print(bedrock_agent_client.get_data_source(knowledgeBaseId = kb['knowledgeBaseId'], dataSourceId = ds[\"dataSourceId\"]))\n",
    "    print(\" \")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Start ingestion job\n",
    "Once the KB and data source(s) created, we can start the ingestion job for each data source.\n",
    "During the ingestion job, KB 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 OSS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ingest_jobs=[]\n",
    "# Start an ingestion job\n",
    "for idx, ds in enumerate(data_sources):\n",
    "    try:\n",
    "        start_job_response = bedrock_agent_client.start_ingestion_job(knowledgeBaseId = kb['knowledgeBaseId'], dataSourceId = ds[\"dataSourceId\"])\n",
    "        job = start_job_response[\"ingestionJob\"]\n",
    "        print(f\"job {idx} started successfully\\n\")\n",
    "        \n",
    "        while(job['status']!='COMPLETE' ):\n",
    "            get_job_response = bedrock_agent_client.get_ingestion_job(\n",
    "              knowledgeBaseId = kb['knowledgeBaseId'],\n",
    "                dataSourceId = ds[\"dataSourceId\"],\n",
    "                ingestionJobId = job[\"ingestionJobId\"]\n",
    "          )\n",
    "            job = get_job_response[\"ingestionJob\"]\n",
    "        pp.pprint(job)\n",
    "        interactive_sleep(40)\n",
    "\n",
    "        ingest_jobs.append(job)\n",
    "    except Exception as e:\n",
    "        print(f\"Couldn't start {idx} job.\\n\")\n",
    "        print(e)\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print the knowledge base Id in bedrock, that corresponds to the Opensearch index in the collection we created before, we will use it for the invocation later\n",
    "kb_id = kb[\"knowledgeBaseId\"]\n",
    "pp.pprint(kb_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# keep the kb_id for invocation later in the invoke request\n",
    "%store kb_id"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Test the knowledge base\n",
    "### Note: If you plan to run any following notebooks, you can skip this section\n",
    "### Using RetrieveAndGenerate API\n",
    "Behind the scenes, RetrieveAndGenerate API converts queries into embeddings, searches the knowledge base, and then augments the foundation model prompt with the search results as context information and returns the FM-generated response to the question. For multi-turn conversations, Knowledge Bases manage short-term memory of the conversation to provide more contextual results.\n",
    "\n",
    "The output of the RetrieveAndGenerate API includes the generated response, source attribution as well as the retrieved text chunks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# # try out KB using RetrieveAndGenerate API\n",
    "bedrock_agent_runtime_client = boto3.client(\"bedrock-agent-runtime\", region_name=region_name)\n",
    "# Lets see how different Anthropic models responds to the input text we provide\n",
    "claude_model_ids = [ [\"Claude 3 Sonnet\", \"anthropic.claude-3-sonnet-20240229-v1:0\"]]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "def ask_bedrock_llm_with_knowledge_base(query: str, model_arn: str, kb_id: str) -> str:\n",
    "    response = bedrock_agent_runtime_client.retrieve_and_generate(\n",
    "        input={\n",
    "            'text': query\n",
    "        },\n",
    "        retrieveAndGenerateConfiguration={\n",
    "            'type': 'KNOWLEDGE_BASE',\n",
    "            'knowledgeBaseConfiguration': {\n",
    "                'knowledgeBaseId': kb_id,\n",
    "                'modelArn': model_arn\n",
    "            }\n",
    "        },\n",
    "    )\n",
    "\n",
    "    return response"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What is Amazon's doing in the field of generative AI?\"\n",
    "\n",
    "for model_id in claude_model_ids:\n",
    "    model_arn = f'arn:aws:bedrock:{region_name}::foundation-model/{model_id[1]}'\n",
    "    response = ask_bedrock_llm_with_knowledge_base(query, model_arn, kb_id)\n",
    "    generated_text = response['output']['text']\n",
    "    citations = response[\"citations\"]\n",
    "    contexts = []\n",
    "    for citation in citations:\n",
    "        retrievedReferences = citation[\"retrievedReferences\"]\n",
    "        for reference in retrievedReferences:\n",
    "            contexts.append(reference[\"content\"][\"text\"])\n",
    "    print(f\"---------- Generated using {model_id[0]}:\")\n",
    "    pp.pprint(generated_text )\n",
    "    print(f'---------- The citations for the response generated by {model_id[0]}:')\n",
    "    pp.pprint(contexts)\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Retrieve API\n",
    "Retrieve API converts user queries into embeddings, searches the knowledge base, and returns the relevant results, giving you more control to build custom workﬂows on top of the semantic search results. The output of the Retrieve API includes the the retrieved text chunks, the location type and URI of the source data, as well as the relevance scores of the retrievals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "# retrieve api for fetching only the relevant context.\n",
    "relevant_documents = bedrock_agent_runtime_client.retrieve(\n",
    "    retrievalQuery= {\n",
    "        'text': query\n",
    "    },\n",
    "    knowledgeBaseId=kb_id,\n",
    "    retrievalConfiguration= {\n",
    "        'vectorSearchConfiguration': {\n",
    "            'numberOfResults': 3 # will fetch top 3 documents which matches closely with the query.\n",
    "        }\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "jupyter": {
     "outputs_hidden": false
    }
   },
   "outputs": [],
   "source": [
    "pp.pprint(relevant_documents[\"retrievalResults\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Clean up\n",
    "Please make sure to comment the below section if you are planning to use the Knowledge Base that you created above for building your RAG application. If you only wanted to try out creating the KB using SDK, then please make sure to delete all the resources that were created as you will be incurred cost for storing documents in OSS index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Delete KnowledgeBase\n",
    "for idx, ds in enumerate(data_sources):\n",
    "    bedrock_agent_client.delete_data_source(dataSourceId = ds[\"dataSourceId\"], knowledgeBaseId=kb['knowledgeBaseId'])\n",
    "bedrock_agent_client.delete_knowledge_base(knowledgeBaseId=kb['knowledgeBaseId'])\n",
    "oss_client.indices.delete(index=index_name)\n",
    "aoss_client.delete_collection(id=collection_id)\n",
    "aoss_client.delete_access_policy(type=\"data\", name=access_policy['accessPolicyDetail']['name'])\n",
    "aoss_client.delete_security_policy(type=\"network\", name=network_policy['securityPolicyDetail']['name'])\n",
    "aoss_client.delete_security_policy(type=\"encryption\", name=encryption_policy['securityPolicyDetail']['name'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# delete role and policies\n",
    "from utility import delete_iam_role_and_policies\n",
    "delete_iam_role_and_policies()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}