RAG (Bedrock and Ragas).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "cac5459b-7237-4a57-944e-33c6ca92f9e5",
   "metadata": {},
   "source": [
    "# Evaluating Knowledge Bases for Amazon Bedrock with Ragas\n",
    "\n",
    "> *This notebook has been tested in the Python 3 kernel of SageMaker Studio JupyterLab (Distribution v1.9)*\n",
    "\n",
    "In this notebook we'll explore ways to evaluate the quality of [Retrieval-Augmented Generation (RAG)](https://aws.amazon.com/what-is/retrieval-augmented-generation/) flows managed by [Amazon Bedrock Knowledge Bases](https://docs.aws.amazon.com/bedrock/latest/userguide/knowledge-base.html): Both at solution development time with open-source tools like [Ragas](https://docs.ragas.io/en/latest/), and at run-time with [checks built in to Bedrock](https://aws.amazon.com/blogs/aws/guardrails-for-amazon-bedrock-can-now-detect-hallucinations-and-safeguard-apps-built-using-custom-or-third-party-fms/) itself."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbb8e54d-a3af-4e95-9d79-0a7f69491afc",
   "metadata": {},
   "source": [
    "## Pre-requisites\n",
    "\n",
    "### Additional permissions for Amazon OpenSearch\n",
    "\n",
    "To complete the manual Bedrock Knowledge setup steps in this notebook, your **AWS Console user/role** will need:\n",
    "\n",
    "- [Permissions to work with Amazon OpenSearch vector collections](https://docs.aws.amazon.com/opensearch-service/latest/developerguide/serverless-vector-search.html)\n",
    "- Permission to **create IAM roles** and attach policies to them, including: `iam:AttachRolePolicy`, `iam:CreateRole`, `iam:DetachRolePolicy`, `iam:GetRole`, `iam:PassRole`, `iam:CreatePolicy`, `iam:CreatePolicyVersion`, and `iam:DeletePolicyVersion`.\n",
    "\n",
    "> ℹ️ **Note:** In testing, we saw `NetworkError` issues when attempting to create Bedrock KBs using only the above-linked `aoss` policy statements. This was resolved by granting `aoss:*` on `*` instead, but you should consider reducing these permissions before using in production environments.\n",
    "\n",
    "If you're in an instructor-led workshop using temporary accounts provided by AWS, this setup should already have been completed for you. If not, refer to the [AWS Console for Identity and Access Management (IAM)](https://console.aws.amazon.com/iam/home?#/home) to grant permissions to your user or role."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47fe714b-5458-4656-9666-8ed7760ea3d0",
   "metadata": {},
   "source": [
    "## Imports and setup\n",
    "\n",
    "First, we'll need to install [Ragas](https://docs.ragas.io/en/latest/) as it's not included by default on the SageMaker Studio base notebook kernel. We'll also:\n",
    "\n",
    "1. Check the installed version of boto3 (the AWS SDK for Python) is sufficiently up-to-date to use Bedrock's built-in contextual grounding checks\n",
    "2. Force Pydantic to a specified version range, to avoid https://github.com/explodinggradients/ragas/issues/867"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9994b796-5abf-4ad2-80fb-8ce611307bcf",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%pip install boto3 \"botocore>=1.34.155\" \"langchain-aws>=0.1,<0.2\" \"ragas==0.1.13\" \"pydantic>=2.8,<3\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc2727e3-7ddb-4d8d-89ce-a23355ac9a78",
   "metadata": {},
   "source": [
    "Next, let's import the libraries that'll be used in the rest of the notebook - and set some **configurations** that we'll use later:\n",
    "\n",
    "- An [Amazon S3 `bucket_name`](https://docs.aws.amazon.com/AmazonS3/latest/userguide/Welcome.html) is required to store our document corpus. By default, we'll use the *default bucket for Amazon SageMaker* - but you could change this to any bucket that this notebook's [execution role](https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-roles.html) (if running in SageMaker) or your IAM user/role (if running locally) has access to\n",
    "- A folder prefix under the bucket where artifacts will be stored (to keep things tidy in case the bucket is used for other projects also)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e1a7597-28ee-4ee9-83f2-cdc06feabff3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Python Built-Ins:\n",
    "from concurrent.futures import ThreadPoolExecutor\n",
    "import json\n",
    "\n",
    "# External Dependencies:\n",
    "import boto3  # General Python SDK for AWS (including Bedrock)\n",
    "from datasets import Dataset  # For use with Ragas\n",
    "from langchain_community.embeddings import BedrockEmbeddings as LangChainBedrockEmbed\n",
    "from langchain_aws import ChatBedrock as LangChainBedrock\n",
    "import pandas as pd  # For working with tabular data\n",
    "import ragas\n",
    "import sagemaker  # Just used for looking up default bucket\n",
    "from tqdm.notebook import tqdm  # Progress bars\n",
    "\n",
    "bucket_name = sagemaker.Session().default_bucket()\n",
    "s3_prefix = \"bedrock-rag-eval\"\n",
    "\n",
    "botosess = boto3.Session()\n",
    "region = botosess.region_name\n",
    "bedrock = botosess.client(\"bedrock\")\n",
    "br_agents_runtime = botosess.client(\"bedrock-agent-runtime\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cde364a-4c0c-401c-8180-bb33d73b1892",
   "metadata": {},
   "source": [
    "## Create the knowledge base\n",
    "\n",
    "> ⚠️ ***Watch out:** This section includes steps you'll need to take manually, not just running the code cells!*\n",
    "\n",
    "First, we'll need to upload the sample documents to Amazon S3 - for which you can just run the code cell below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f014d34e-c914-4f43-85d4-048c30355bd9",
   "metadata": {},
   "outputs": [],
   "source": [
    "corpus_s3uri = f\"s3://{bucket_name}/{s3_prefix}/corpus\"\n",
    "\n",
    "print(f\"Syncing corpus to:\\n{corpus_s3uri}/\")\n",
    "\n",
    "!aws s3 sync --quiet ./datasets/question-answering/corpus {corpus_s3uri}/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efdc3dc7-6804-440e-beaf-73d0860dc442",
   "metadata": {},
   "source": [
    "The simplest way to set up the actual Bedrock Knowledge Base for testing will be **manually through the AWS Console**:\n",
    "\n",
    "▶️ First, **open** the [AWS Console for Amazon Bedrock](https://console.aws.amazon.com/bedrock/home?#/knowledge-bases) and select *Orchestration > Knowledge bases* from the left sidebar menu, as shown in the screenshot below:\n",
    "\n",
    "> ⚠️ **Check** you're working in the correct *AWS Region* in the top right corner of the UI\n",
    "\n",
    "![](images/bedrock-kbs/01-bedrock-kb-console.png \"Screenshot of AWS Console for Amazon Bedrock Knowledge Bases, showing 'Create knowledge base' action button\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87714878-b4ab-4d65-b033-cde80593a24e",
   "metadata": {},
   "source": [
    "▶️ **Click** the *Create knowledge base* button to start the workflow. In the screen that opens:\n",
    "\n",
    "- For **knowledge base name**, enter `example-squad-kb`\n",
    "- For **knowledge base description**, you can provide (something like) `Demo knowledge base for question answering evaluation`\n",
    "- Leave the other settings as default (allow creating a new execution role, and no tags)\n",
    "\n",
    "Your configuration should be as shown below:\n",
    "\n",
    "![](images/bedrock-kbs/02a-create-kb-basics.png \"Screenshot of step 1 in Bedrock Knowledge Base creation workflow: with KB name, description, (create new) execution role, and (empty) tags configured. At the end of the form, a 'Next' button is visible.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4092b22f-33e3-4b07-b136-4c2774347a9c",
   "metadata": {},
   "source": [
    "▶️ In the **Next** screen, you'll configure the S3 data source:\n",
    "\n",
    "- For **data source name**, enter `example-squad-corpus`\n",
    "- For **S3 URI**, refer to the previous cell of this notebook where we uploaded the data and output `Syncing corpus to: ...`\n",
    "\n",
    "> ⚠️ **Be sure to include the trailing slash** in your `s3://.../.../` URI. If you omit it, you may find that creating the KB succeeds but it fails to sync any documents later (because the auto-created execution role for Amazon Bedrock will be granted IAM `s3:GetObject` permissions to `.../corpus` instead of to `.../corpus/*`).\n",
    "\n",
    "You should leave all *Advanced settings* as default, but feel free to expand out this section to explore the options available.\n",
    "\n",
    "![](images/bedrock-kbs/02b-create-kb-data-source.png \"Screenshot of S3 data source configuration, with name and S3 URI configured and 'next' button visible\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97d9521b-3b1c-4626-83da-334ceb868998",
   "metadata": {},
   "source": [
    "▶️ In the **Next** screen, you'll configure the vector index:\n",
    "\n",
    "- For **embeddings model**, select `Cohere Embed Multilingual`\n",
    "\n",
    "> ⚠️ **Check** in the [Amazon Bedrock Model Access console](https://console.aws.amazon.com/bedrock/home?#/modelaccess) that you've enabled access to this model in the current region.\n",
    ">\n",
    "> If needed, you should be able to select an alternative embedding model instead... But we haven't tested all options for this walkthrough.\n",
    "\n",
    "- For **Vector database**, select `Quick create a new vector store`\n",
    "\n",
    "You can find more information from this screen or the [Amazon Bedrock Developer Guide](https://docs.aws.amazon.com/bedrock/latest/userguide/knowledge-base-setup.html) about the different vector stores Bedrock Knowledge Bases support. This default option will create a new [Amazon OpenSearch Serverless](https://docs.aws.amazon.com/opensearch-service/latest/developerguide/serverless-overview.html) cluster\n",
    "\n",
    "Leave other settings at their defaults as shown below, and you should be ready to proceed:\n",
    "\n",
    "![](images/bedrock-kbs/02c-create-kb-index.png \"Screenshot of Knowledge Base vector index settings including Cohere Embed Multilingual embedding model, and quick-create vector store. 'Next' button is visible.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61381a1e-7b14-4715-970d-5b008dc1441d",
   "metadata": {},
   "source": [
    "▶️ Click **Next** to review your configuration, and then **Create knowledge base** to complete the process.\n",
    "\n",
    "> ⏰ It might take **a few minutes** for the creation to complete. A progress indicator banner should be visible if you scroll up. Alternatively in a separate tab, you could check the [Amazon OpenSearch Serverless Collections console](https://console.aws.amazon.com/aos/home?#opensearch/collections) - where you should see the underlying vector collection being created.\n",
    "\n",
    "Once your Knowledge Base is completed successfully, you'll be directed to the its detail screen as shown below:\n",
    "\n",
    "![](images/bedrock-kbs/03-kb-detail-page.png \"Detail screen for the created Amazon Bedrock Knowledge Base, showing creation success banner. Includes sections 'Knowledge base overview' (containing the KB ID, name, and other details); 'Tags' (empty); 'Data source' (one Amazon S3 data source listed); 'Embeddings model' (Cohere Embed); and an interactive 'Test knowledge base' chat sidebar on the right with a warning that some data sources have not been synced.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e8fd8dfb-a9c9-45ae-b0fc-ebccafacd8b3",
   "metadata": {},
   "source": [
    "As mentioned in the alert box shown ahead, your new knowledge base will not yet contain your documents until we **sync** the data source:\n",
    "\n",
    "▶️ **Select** your S3 data source using the radio button to the left of it's name in the data sources list, and **click the Sync button** above to start the sync.\n",
    "\n",
    "The sync should only take a few seconds, after which your data source's *Status* will return to `Available`\n",
    "\n",
    "![](images/bedrock-kbs/04a-kb-data-source-after-sync.png \"Screenshot of KB 'data source' section after running sync, with the data source selected and status showing as 'available'\")\n",
    "\n",
    "With the sync completed, your Knowledge Base should be ready to use.\n",
    "\n",
    "Optionally, you can click through to your data source name to check the sync `Added` the 20 files as expected:\n",
    "\n",
    "![](images/bedrock-kbs/04b-kb-data-sync-details.png \"Data source details screen showing sync completed successfully with 20 files detected and added to the index, and 0 files failed\")."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ff171c96-de81-40ae-bc8e-f3fe55e8f3f9",
   "metadata": {},
   "source": [
    "## Try out your Knowledge Base\n",
    "\n",
    "Before we discuss evaluation at scale, let's run a couple of test queries to check the KB is working properly"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58cfd6b7-6ddb-4f9d-a6b6-ccd3700001e7",
   "metadata": {},
   "source": [
    "### ...From the AWS Console\n",
    "\n",
    "Your Knowledge Base's detail screen in the [AWS Console for Amazon Bedrock](https://console.aws.amazon.com/bedrock/home?2#/knowledge-bases) includes an interactive widget on the right sidebar, for trying out queries.\n",
    "\n",
    "▶️ **Click** the orange *Select model* button to get started, and select `Claude 3 Sonnet` with on-demand throughput as shown below:\n",
    "\n",
    "![](images/bedrock-kbs/05-kb-select-model-claude-3-sonnet.png \"Screenshot of model selection interface with Claude 3 Sonnet selected on on-demand throughput\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9928fecf-e46e-46a5-b831-cff97ca709f0",
   "metadata": {},
   "source": [
    "▶️ **Type** an example question into the chat panel on the right of the screen, and click **Run** to try it out. You can ask for example:\n",
    "\n",
    "```\n",
    "In what country is Normandy located?\n",
    "```\n",
    "\n",
    "You should find the system is able to respond, based on the [Normans.txt](datasets/question-answering/corpus/Normans.txt) file we ingested, as shown below:\n",
    "\n",
    "![](images/bedrock-kbs/06a-kb-test.png \"example-squad-kb detail page with the sample question already asked in the interactive try it out sidebar. The model's response, about a paragraph long, correctly identifies Normandy as being in France and includes a link to 'show source details\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55e20b02-f6cb-458c-97be-3f5e8516f85d",
   "metadata": {},
   "source": [
    "▶️ **Explore** the source details, and the configuration menu available in the top left of the 'Test knowledge base' widget. You'll see:\n",
    "\n",
    "1. An extensive range of configuration options are available, covering both source document retrieval and final answer generation\n",
    "2. In this case, the `Normans` source article has been automatically split into two **chunks** during ingestion to the knowledge base - to improve answer relevancy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a312c6d1-07a2-417b-a9f2-8f27c013a515",
   "metadata": {},
   "source": [
    "### ...From Python code\n",
    "\n",
    "To use our Knowledge Base programmatically, we'll need to look up the **knowledge base ID**. This automatically-generated, alphanumeric string is different from the *name* we gave our KB during creation. You can find it [from the AWS Console](https://console.aws.amazon.com/bedrock/home?#/knowledge-bases in the \"Knowledge base overview\" section of your KB's detail screen.\n",
    "\n",
    "▶️ **Replace** the below placeholder with your knowledge base's unique ID, and run the cells below to continue:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b9d1f47-46a6-42aa-987d-bd2eafcfd43f",
   "metadata": {},
   "outputs": [],
   "source": [
    "knowledge_base_id = \"TODO\"  # Something like \"55GUAMQYUT\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d7503dd-a959-497b-ab5c-1533ab15baf5",
   "metadata": {},
   "source": [
    "With the ID identified, you can use the Bedrock runtime [RetrieveAndGenerate API](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_RetrieveAndGenerate.html) (see [corresponding boto3 doc page for Python](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-agent-runtime/client/retrieve_and_generate.html)) to query your knowledge base.\n",
    "\n",
    "As in the manual example, you'll also need to select which text generation model to use - which we've pre-populated below for Claude 3 Sonnet:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff65763a-d729-4b59-a5c5-ff4ad2da121c",
   "metadata": {},
   "outputs": [],
   "source": [
    "rag_resp = br_agents_runtime.retrieve_and_generate(\n",
    "    input={\"text\": \"In what country is Normandy located?\"},\n",
    "    retrieveAndGenerateConfiguration={\n",
    "        \"knowledgeBaseConfiguration\": {\n",
    "            \"knowledgeBaseId\": knowledge_base_id,\n",
    "            \"modelArn\": f\"arn:aws:bedrock:{region}::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0\",\n",
    "        },\n",
    "        \"type\": \"KNOWLEDGE_BASE\",\n",
    "    },\n",
    "    # Optional session ID can help improve results for follow-up questions:\n",
    "    # sessionId='string'\n",
    ")\n",
    "\n",
    "print(\"Plain text response:\")\n",
    "print(\"--------------------\")\n",
    "print(rag_resp[\"output\"][\"text\"], end=\"\\n\\n\\n\")\n",
    "\n",
    "print(\"Full API output:\")\n",
    "print(\"----------------\")\n",
    "rag_resp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4dde089-551f-4568-a503-c461a0280738",
   "metadata": {},
   "source": [
    "As shown in the full API response from the above cell, the `RetrieveAndGenerate` action provides:\n",
    "\n",
    "- The final text answer\n",
    "- The `retrievedReferences` from the search engine\n",
    "- Specific `citations` localizing which references should be cited by different parts of the text answer\n",
    "\n",
    "...Similarly to the information we saw when trying the KB out in the AWS Console.\n",
    "\n",
    "It's also possible to run **only the retrieval** through the API, and skip the generative answer synthesis step - as shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3070579-9785-4d5f-aedb-489c09e72556",
   "metadata": {},
   "outputs": [],
   "source": [
    "retrieve_resp = br_agents_runtime.retrieve(\n",
    "    knowledgeBaseId=knowledge_base_id,\n",
    "    retrievalQuery={\"text\": \"In what country is Normandy located?\"},\n",
    ")\n",
    "print(json.dumps(retrieve_resp[\"retrievalResults\"], indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea0cc84f-2a06-4552-b585-95a25d5e3c4c",
   "metadata": {},
   "source": [
    "## Evaluate Bedrock Knowledge Bases with Ragas\n",
    "\n",
    "Now we have our Bedrock Knowledge Base set up, we'd like to **evaluate** the quality of its results against a test dataset - to help us **optimize** the configuration for high quality and low cost.\n",
    "\n",
    "First, let's load the sample dataset of questions, reference answers, and their source documents, that we [prepared earlier](datasets/Prepare-SQuAD.ipynb):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b12ee252-ffeb-4c46-99ce-7c80b4c07054",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_df = pd.read_json(\"datasets/question-answering/qa.manifest.jsonl\", lines=True)\n",
    "dataset_df.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15585ef1-fcb5-4fbe-be21-2744b1b0e2d8",
   "metadata": {},
   "source": [
    "Records in this dataset include:\n",
    "\n",
    "- (`doc`) The full text of the source document for this example\n",
    "- (`doc_id`) A unique identifier for the source document\n",
    "- (`question`) The user question to be asked\n",
    "- (`question_id`) A unique identifier for the question\n",
    "- (`answers`) A list of (possibly multiple) reference 'correct' answers, supported by the document"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c852e15-2c11-4a50-bfe1-07e6689675cf",
   "metadata": {},
   "source": [
    "### Run the knowledge base against our test set\n",
    "\n",
    "As shown in [Ragas' API Reference](https://docs.ragas.io/en/latest/references/evaluation.html), records in Ragas evaluation datasets typically include:\n",
    "\n",
    "- The `question` that was asked\n",
    "- The `answer` the system generated\n",
    "- The actual text `contexts` the answer was based on (i.e. snippets of document text retrieved by the search engine)\n",
    "- The `ground_truth` answer(s)\n",
    "\n",
    "We can run our example questions through the Bedrock KB RAG as shown below, to fetch the outputs ready to calculate metrics:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48a90920-b197-4a16-b314-73a671409840",
   "metadata": {},
   "outputs": [],
   "source": [
    "from concurrent.futures import ThreadPoolExecutor\n",
    "\n",
    "def retrieve_and_generate(\n",
    "    question: str,\n",
    "    kb_id: str,\n",
    "    generate_model_arn: str = f\"arn:aws:bedrock:{region}::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0\",\n",
    "):\n",
    "    rag_resp = br_agents_runtime.retrieve_and_generate(\n",
    "        input={\"text\": question},\n",
    "        retrieveAndGenerateConfiguration={\n",
    "            \"knowledgeBaseConfiguration\": {\n",
    "                \"knowledgeBaseId\": kb_id,\n",
    "                \"modelArn\": generate_model_arn,\n",
    "            },\n",
    "            \"type\": \"KNOWLEDGE_BASE\",\n",
    "        },\n",
    "    )\n",
    "    # Fetch flat list of references from the nested citations->retrievedReferences:\n",
    "    all_refs = [r for cite in rag_resp[\"citations\"] for r in cite[\"retrievedReferences\"]]\n",
    "    ref_s3uris = [r[\"location\"][\"s3Location\"][\"uri\"] for r in all_refs]\n",
    "    # Map e.g. 's3://.../doc_id.txt' to 'doc_id':\n",
    "    ref_ids = [uri.rpartition(\"/\")[2].rpartition(\".\")[0] for uri in ref_s3uris]\n",
    "    return {\n",
    "        \"answer\": rag_resp[\"output\"][\"text\"],\n",
    "        \"retrieved_doc_ids\": ref_ids,\n",
    "        \"retrieved_doc_texts\": [r[\"content\"][\"text\"] for r in all_refs]\n",
    "    }\n",
    "\n",
    "with ThreadPoolExecutor(max_workers=2) as pool:\n",
    "    rag_futures = [\n",
    "        pool.submit(retrieve_and_generate, question=rec.question, kb_id=knowledge_base_id)\n",
    "        for _, rec in dataset_df.iterrows()\n",
    "    ]\n",
    "    outputs_df = pd.DataFrame([f.result() for f in tqdm(rag_futures, desc=\"Running RAG...\")])\n",
    "    # Combine & clarify the column names for a nice tabular representation:\n",
    "    results_df = pd.concat((dataset_df, outputs_df), axis=1).rename(\n",
    "        columns={\n",
    "            \"answer\": \"model_answer\",\n",
    "            \"answers\": \"gt_answers\",\n",
    "            \"doc_id\": \"gt_doc_id\",\n",
    "            \"doc\": \"gt_doc_text\",\n",
    "        }\n",
    "    )\n",
    "results_df.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b2545bf-8373-4a27-8adf-2da334ebb2c2",
   "metadata": {},
   "source": [
    "Ragas supports a broad [range of metrics](https://docs.ragas.io/en/latest/concepts/metrics/index.html) with the option to configure which ones you calculate - but many of these depend on providing an **evaluator LLM** (or evaluator embedding model) which will be used in scoring.\n",
    "\n",
    "We'll set up Claude 3 Sonnet as the evaluator LLM and Cohere Embedding Multilingual as the evaluator embedding model, to be able to demonstrate the full suite of available metrics.\n",
    "\n",
    "Although Ragas defines its own base classes (see [BaseRagasLLM](https://github.com/explodinggradients/ragas/blob/2d793651f778b6c0da07a834e9ce2765be13cc9f/src/ragas/llms/base.py#L46), [BaseRagasEmbeddings](https://github.com/explodinggradients/ragas/blob/2d793651f778b6c0da07a834e9ce2765be13cc9f/src/ragas/embeddings/base.py#L19)) for these interfaces, integration is typically via LangChain for simplicity - so that's the pattern we'll follow here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7f1dba2-1a9d-4891-b698-54d6b425ab5a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ragas_result = ragas.evaluation.evaluate(\n",
    "    Dataset.from_pandas(results_df),\n",
    "    metrics=[\n",
    "        # A looot of metrics to give a general overview:\n",
    "        ragas.metrics.answer_relevancy,\n",
    "        ragas.metrics.faithfulness,\n",
    "        ragas.metrics.context_precision,\n",
    "        ragas.metrics.context_recall,\n",
    "        ragas.metrics.context_utilization,\n",
    "        ragas.metrics.answer_similarity,\n",
    "        ragas.metrics.answer_correctness,\n",
    "        ragas.metrics.critique.conciseness,\n",
    "    ],\n",
    "    llm=LangChainBedrock(model_id=\"anthropic.claude-3-sonnet-20240229-v1:0\"),\n",
    "    embeddings=LangChainBedrockEmbed(model_id=\"cohere.embed-multilingual-v3\"),\n",
    "    column_map={\n",
    "        \"answer\": \"model_answer\",\n",
    "        \"contexts\": \"retrieved_doc_texts\",\n",
    "        \"ground_truths\": \"gt_answers\",\n",
    "        \"question\": \"question\",\n",
    "    },\n",
    ")\n",
    "\n",
    "print(\"Overall scores\")\n",
    "print(\"--------------\")\n",
    "print(ragas_result, end=\"\\n\\n\")\n",
    "print(\"Details\")\n",
    "print(\"-------\")\n",
    "scores_df = ragas_result.to_pandas()\n",
    "scores_df.head(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5bdcc9e8-e368-4226-a179-d328bbb7930b",
   "metadata": {},
   "source": [
    "Although your figures may differ due to variability of generation, in our test run on the sample dataset we observed:\n",
    "\n",
    "```json\n",
    "{\n",
    "    'answer_relevancy': 0.8380,\n",
    "    'faithfulness': 0.9519,\n",
    "    'context_precision': 1.0000,\n",
    "    'context_recall': 1.0000,\n",
    "    'context_recall': 1.0000,\n",
    "    'answer_similarity': 0.5708,\n",
    "    'answer_correctness': 0.6102,\n",
    "    'conciseness': 1.0000,\n",
    "}\n",
    "```\n",
    "\n",
    "Generally, the system appears to be performing well at retrieving the correct document and generating appropriate answers on this dataset - while many of the **lower** scores relate to the significant different distribution between the reference answers (which just extract specific words or phrases from the source document, without providing any explanation), versus the RAG system responses (which typically provide a bit more contextual background).\n",
    "\n",
    "For full discussion of the various metrics and their interpretation, refer to the [Ragas documentation](https://docs.ragas.io/en/latest/concepts/metrics/index.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00e01977-7a7e-4513-86b2-499a5374376b",
   "metadata": {},
   "source": [
    "## More traditional search engine metrics\n",
    "\n",
    "Some particular points to note about the Ragas evaluation results include:\n",
    "\n",
    "1. Nearly all the metrics (even those like 'conciseness') are based on either **LLM self-critique, or embedding-based scores**\n",
    "    - As a consequence, we should be careful to monitor for potential **bias** between different evaluator LLMs and candidate LLMs - since [LLM evaluators have been shown to recognise and prefer their own generations](https://arxiv.org/abs/2404.13076).\n",
    "    - We should aim to have **humans review** a subset of the same data and rate the same metrics, allowing us to *measure* how well these automated evaluations align with real user preferences - and thus quantify how much trust we should place in the automated metrics when run over bigger datasets that it wouldn't be practical for humans to label.\n",
    "2. Not all the available information is being utilized in this case\n",
    "    - Although metrics like `context_precision` and `context_recall` aim to rate how well the retrieved context relates to the question and the ground-truth answer, in this case we have labelled examples of which document each answer should come from (our `gt_doc_text` and `gt_doc_id` fields), which the metrics are not using\n",
    "\n",
    "If you have labelled data for the *correct source document / chunk* per question in your use-case, you might also be interested to calculate more traditional search engine performance metrics, like:\n",
    "\n",
    "- **Recall at K:** In what percentage of cases did the target document appear in the first *K* snippets returned by the search engine?\n",
    "- **Precision:** What percentage of the snippets returned by the search engine to use in generation, came from 'correct'/relevant documents?\n",
    "- **Normalized Discounted Cumulative Gain (nDCG) at K:** Summarizing how early 'relevant' documents are ranked, in cases where many of your example questions have multiple relevant documents that should be returned.\n",
    "\n",
    "Precision and recall are fairly straightforward to calculate from the available data as shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c31e0dfc-b291-44bf-96f4-4be7a37c60a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_recall_at_1 = [\n",
    "    # Recall@1 is % of the time the first retrieved doc was the target\n",
    "    1.0 if rec.gt_doc_id == rec.retrieved_doc_ids[0] else 0.0\n",
    "    for _, rec in results_df.iterrows()\n",
    "]\n",
    "print(f\"Recall@1: {sum(sample_recall_at_1) / len(sample_recall_at_1):.2%}\")\n",
    "\n",
    "sample_precisions = [\n",
    "    # Precision is count(retrieved_doc == target_doc) / number_of_retrieved_docs\n",
    "    sum((doc_id == rec.gt_doc_id for doc_id in rec.retrieved_doc_ids)) / len(rec.retrieved_doc_ids)\n",
    "    for _, rec in results_df.iterrows()\n",
    "]\n",
    "print(f\"Average Precision: {sum(sample_precisions) / len(sample_precisions):.2%}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "576a5710-69cd-407a-b7ff-72d47d01f624",
   "metadata": {},
   "source": [
    "In our tests, we recorded `100.00%` for both these metrics on the sample dataset.\n",
    "\n",
    "In this dataset with only 20 documents each covering very distinct topics (see for yourself in [datasets/question-answering/corpus](datasets/question-answering/corpus)), it's easy for the retriever to fetch the correct document 100% of the time. In more typical enterprise contexts with large knowledge bases of heavily overlapping content, this retrieval is much more likely to be a challenging bottleneck to overall RAG performance."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d960d58-e476-4f3c-8e31-a5e6fd9f0f90",
   "metadata": {},
   "source": [
    "## Online checks with Guardrails for Amazon Bedrock\n",
    "\n",
    "Bulk evaluations at solution development time are useful for understanding and optimizing response quality - and estimating likely performance in production... But we could further improve robustness by **checking answers in real time** and intervening if it seems like the RAG pipeline failed to find relevant information or generate an answer that's faithful to the source documents.\n",
    "\n",
    "[Guardrails for Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails.html) offers a suite of configurable but fully-managed safety checks - including a [contextual grounding check](https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails-contextual-grounding-check.html) specifically tailored for RAG systems.\n",
    "\n",
    "The contextual grounding check includes two components:\n",
    "\n",
    "- **Grounding** – This checks if the model response is factually accurate based on the source and is grounded in the source. Any new information introduced in the response will be considered un-grounded.\n",
    "- **Relevance** – This checks if the model response is relevant to the user query.\n",
    "\n",
    "Each has a separate **configurable threshold** from 0.99 (very strict) to 0.0 (disabled).\n",
    "\n",
    "Guardrails can be created [through the AWS Console](https://console.aws.amazon.com/bedrock/home?#/guardrails), or through APIs as shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb109ac4-ec31-41a9-aed4-0ac34bfe3d81",
   "metadata": {},
   "outputs": [],
   "source": [
    "guardrail_resp = bedrock.create_guardrail(\n",
    "    name=\"eval-workshop-rag-grounding\",\n",
    "    description=(\n",
    "        \"Example guardrail to demonstrate RAG contextual grounding checks for LLM Evaluation \"\n",
    "        \"workshop\"\n",
    "    ),\n",
    "    # Many additional options available here, such as topicPolicyConfig, contentPolicyConfig, etc\n",
    "    contextualGroundingPolicyConfig={\n",
    "        \"filtersConfig\": [\n",
    "            {\"type\": \"GROUNDING\", \"threshold\": 0.7},\n",
    "            {\"type\": \"RELEVANCE\", \"threshold\": 0.7},\n",
    "        ]\n",
    "    },\n",
    "    blockedInputMessaging=\"(Required but wouldn't be triggered by a grounding-only guardrail)\",\n",
    "    blockedOutputsMessaging=\"Sorry, I couldn't find an answer to your question\",\n",
    ")\n",
    "\n",
    "guardrail_id = guardrail_resp[\"guardrailId\"]\n",
    "guardrail_ver = guardrail_resp[\"version\"]\n",
    "\n",
    "print(f\"Created guardrail {guardrail_id} at version {guardrail_ver}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19739eb7-c4ae-4b22-83f0-ddaa29e51d9d",
   "metadata": {},
   "source": [
    "Note that guardrails are **versioned**, allowing you to test updates before deploying them to live applications - and track the history of a guardrail's configuration. In this example, we'll just point to the `DRAFT` version rather than explicitly calling the [create_bedrock_version API](https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock/client/create_guardrail_version.html).\n",
    "\n",
    "> ℹ️ One guardrail can enable many checks at once: From contextual grounding as shown here, to \"denied topics\", sensitive information/PII filters, [and more](https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails-components.html).\n",
    "\n",
    "Once the guardrail is created, we can *reference* it when invoking an LLM or Knowledge Base in Bedrock (or building an Agent).\n",
    "\n",
    "Let's enable the guardrail and ask a question that's only partially answered by the available knowledge base:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c89e28b3-f9d8-4833-9dd7-558ce56e6654",
   "metadata": {},
   "outputs": [],
   "source": [
    "rag_resp = br_agents_runtime.retrieve_and_generate(\n",
    "    input={\"text\": \"Where are Normandy and Singapore located?\"},\n",
    "    retrieveAndGenerateConfiguration={\n",
    "        \"knowledgeBaseConfiguration\": {\n",
    "            \"knowledgeBaseId\": knowledge_base_id,\n",
    "            \"modelArn\": f\"arn:aws:bedrock:{region}::foundation-model/anthropic.claude-3-sonnet-20240229-v1:0\",\n",
    "\n",
    "            # Use the contextual grounding guardrail configured above:\n",
    "            \"generationConfiguration\": {\n",
    "                \"guardrailConfiguration\": {\n",
    "                    \"guardrailId\": guardrail_id,\n",
    "                    \"guardrailVersion\": guardrail_ver,\n",
    "                },\n",
    "            },\n",
    "        },\n",
    "        \"type\": \"KNOWLEDGE_BASE\",\n",
    "    },\n",
    ")\n",
    "\n",
    "\n",
    "print(\"Guardrail action:\")\n",
    "print(\"--------------------\")\n",
    "print(rag_resp[\"guardrailAction\"], end=\"\\n\\n\\n\")\n",
    "\n",
    "print(\"Plain text response:\")\n",
    "print(\"--------------------\")\n",
    "print(rag_resp[\"output\"][\"text\"], end=\"\\n\\n\\n\")\n",
    "\n",
    "print(\"Full API output:\")\n",
    "print(\"----------------\")\n",
    "rag_resp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d37ea3d7-e97a-4a17-b87b-3d32ed055c4e",
   "metadata": {},
   "source": [
    "In our test, the guardrail action was `INTERVENED` which caused the system to override the response to the guardrail-configured *\"Sorry, I couldn't find an answer to your question\"*.\n",
    "\n",
    "Depending on your question and corpus, you'll likely find cases where the generator model already recognises that the context doesn't contain enough information to answer the question - so generates a response that explains this. Sometimes the guardrail action may be `NONE` in these cases - as the model's response is not un-grounded enough to require intervention\n",
    "\n",
    "▶️ **Try** asking some other questions like `Can I fly to Normandy from Singapore?` to see how the guardrail responds. You could also edit your guardrail thresholds in the [Bedrock Guardrails Console](https://console.aws.amazon.com/bedrock/home?#/guardrails)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33d8c9d3-3378-4b7f-ba1f-908f781faa75",
   "metadata": {},
   "source": [
    "## Clean-Up\n",
    "\n",
    "Bedrock Knowledge Bases store and index data backed by an underlying vector store, so once you're done experimenting you should **delete your knowledge base** to avoid unnecessary ongoing charges from the vector store itself.\n",
    "\n",
    "> ⚠️ **Before you clean up:** The [conversational tests example notebook](conversational-tests/Conversational%20Tests.ipynb) re-uses the Knowledge Base we created here... So if you're going to run through that example, **do it first** before cleaning up the resources below!\n",
    "\n",
    "1. Find your KB in the [Orchestration > Knowledge bases section of the Amazon Bedrock console](https://console.aws.amazon.com/bedrock/home?#/knowledge-bases)\n",
    "2. Check the underlying vector store from the KB details: In this example it should be an Amazon OpenSearch Serverless collection\n",
    "3. **Delete** your KB from the Amazon Bedrock console\n",
    "4. Check in the [Serverless Collections section of the Amazon OpenSearch console](https://console.aws.amazon.com/aos/home?region=#opensearch/collections) (or whichever other service is relevant, if your KB was backed by Aurora Postgres or a different store), that your underlying collection has also been deleted - and delete it manually if not.\n",
    "5. Consider also removing the source data from [Amazon S3](https://console.aws.amazon.com/s3/buckets), to avoid any potential S3 charges.\n",
    "\n",
    "For more information, refer to the pricing pages for [Amazon OpenSearch Serverless](https://aws.amazon.com/opensearch-service/pricing/), [Amazon Bedrock](https://aws.amazon.com/bedrock/pricing/), and [Amazon S3](https://aws.amazon.com/s3/pricing/)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddde3e36-7dca-44ba-bbeb-cc4a282436aa",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this notebook we saw how to create a [Knowledge Base for Amazon Bedrock](https://aws.amazon.com/bedrock/knowledge-bases/) to deploy a fully-managed RAG pipeline on AWS, and then ran some basic result quality evaluations using the [RetrieveAndGenerate API](https://docs.aws.amazon.com/bedrock/latest/APIReference/API_agent-runtime_RetrieveAndGenerate.html) and the open-source library [Ragas](https://docs.ragas.io/en/latest/) for RAG evaluations.\n",
    "\n",
    "While Ragas provides a broad range of pre-implemented metrics that can help summarize performance and pinpoint limitations of RAG systems on validation datasets, it's important to remember that the included metrics are largely LLM-evaluated - and therefore potentially subject to bias, especially if comparing between different candidate and evaluator LLMs.\n",
    "\n",
    "- **Collecting** source attribution data up-front in your validation datasets (i.e. source document, not just question and reference answer), will provide additional options for you to evaluate the performance of the retrieval component itself separately from the answer generation/synthesis - helping to identify which component is the bottleneck for overall result quality\n",
    "- **Comparing** Ragas-reported metrics against human evaluations for the same datasets and metrics, can help to quantify how representative the automated metrics are of real user perceptions, and therefore how much trust should be placed in them.\n",
    "\n",
    "As well as evaluation and optimization at build-time, real-time contextual grounding checks can be deployed into live RAG pipelines to safeguard users at run-time. [Guardrails for Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/guardrails.html) provides a combined, fully-managed solution for this and many other safety checks. Alternatively, you could explore [Prompt flows for Amazon Bedrock](https://docs.aws.amazon.com/bedrock/latest/userguide/flows.html) to build your own custom chains."
   ]
  }
 ],
 "metadata": {
  "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.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}