optimize_knowledge_using_ragas_evaluation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Knowledge Base Parameter Optimization using RAGAS\n",
    "This notebook implements evaluation and optimization of Amazon Bedrock Knowledge bases parameters using the RAGAS framework"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      ]
    }
   ],
   "source": [
    "!pip install ragas\n",
    "!pip install datasets\n",
    "!pip install pandas\n",
    "!pip install boto3\n",
    "!pip install langchain\n",
    "!pip install langchain-aws\n",
    "!pip install nltk\n",
    "!pip install seaborn"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import required libraries\n",
    "import pandas as pd\n",
    "import boto3\n",
    "from datetime import datetime\n",
    "from langchain_aws import ChatBedrockConverse, BedrockEmbeddings\n",
    "from datasets import Dataset\n",
    "from ragas import evaluate\n",
    "from ragas.metrics import SemanticSimilarity\n",
    "from ragas.llms import LangchainLLMWrapper\n",
    "from ragas.embeddings import LangchainEmbeddingsWrapper\n",
    "import itertools\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define parameters and permutations to optimize over\n",
    "We'll optimize over the number of results to retrieve from the knowledge base, and the model we'll use to generate the answer.\n",
    "More query time hypter-parameters to try attempt to optimize over include:\n",
    "- custom prompt\n",
    "- Semantic vs Hybrid serach\n",
    "\n",
    "Additionally there are index time parameters you could optimize, for example:\n",
    "- Chunking strategy\n",
    "- Embedding model. \n",
    "\n",
    "To change these you'll need to recreate the KB data source, or the entire knowledge base. This can be done via API in a loop. See example code [here](https://github.com/aws-samples/amazon-bedrock-workshop/blob/main/02_KnowledgeBases_and_RAG/0_create_ingest_documents_test_kb.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[(1, 'mistral.mistral-small-2402-v1:0'),\n",
       " (1, 'ai21.jamba-1-5-mini-v1:0'),\n",
       " (1, 'cohere.command-r-v1:0'),\n",
       " (3, 'mistral.mistral-small-2402-v1:0'),\n",
       " (3, 'ai21.jamba-1-5-mini-v1:0'),\n",
       " (3, 'cohere.command-r-v1:0'),\n",
       " (7, 'mistral.mistral-small-2402-v1:0'),\n",
       " (7, 'ai21.jamba-1-5-mini-v1:0'),\n",
       " (7, 'cohere.command-r-v1:0')]"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# This is a knowledge base created via the AWS console with all defaults and has an S3 datasource that indexes this\n",
    "# filie https://docs.aws.amazon.com/pdfs/whitepapers/latest/aws-overview/aws-overview.pdf\n",
    "BEDROCK_KNOWLEDGE_BASE_ID = \"GXHTSVCWZI\"\n",
    "REGION_NAME = \"us-east-1\"\n",
    "DEFAULT_LLM_MODEL_ID = 'anthropic.claude-3-haiku-20240307-v1:0'\n",
    "\n",
    "# Parameters to optimize\n",
    "NUM_RESULTS_OPTIONS = [1, 3, 7]\n",
    "MODEL_OPTIONS = [\n",
    "    'mistral.mistral-small-2402-v1:0',\n",
    "    'ai21.jamba-1-5-mini-v1:0',\n",
    "    'cohere.command-r-v1:0',\n",
    "]\n",
    "\n",
    "# Generate all possible parameter combinations\n",
    "parameter_combinations = list(itertools.product(NUM_RESULTS_OPTIONS, MODEL_OPTIONS))\n",
    "parameter_combinations\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The questions we'll be asking the knowledge base and the answers we expect to get back (\"ground truth\")\n",
    "test_data_aws_services = [\n",
    "    {\n",
    "        \"question\": \"What is AWS Lambda and how does it work?\",\n",
    "        \"ground_truth\": \"AWS Lambda is a serverless compute service that runs code in response to events without managing servers. It automatically scales and only charges for actual compute time used.\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What is Amazon S3's durability guarantee?\",\n",
    "        \"ground_truth\": \"Amazon S3 provides 99.999999999% (11 9's) durability for objects stored in all S3 storage classes across multiple Availability Zones.\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"How does AWS Direct Connect differ from VPN?\",\n",
    "        \"ground_truth\": \"AWS Direct Connect provides dedicated physical connections to AWS, while VPN creates encrypted tunnels over the public internet. Direct Connect offers more consistent network performance and lower latency.\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What is Amazon Aurora and its key benefits?\",\n",
    "        \"ground_truth\": \"Amazon Aurora is a MySQL/PostgreSQL-compatible database offering up to 5x performance of MySQL and 3x of PostgreSQL, with automated scaling, backup, and fault tolerance built-in.\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"How does AWS Shield protect against DDoS attacks?\",\n",
    "        \"ground_truth\": \"AWS Shield provides automatic DDoS protection for all AWS customers at the network/transport layer (Standard) and additional protection with advanced monitoring for higher-level attacks (Advanced).\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What is Amazon EKS and its primary use case?\",\n",
    "        \"ground_truth\": \"Amazon Elastic Kubernetes Service (EKS) is a managed Kubernetes service for running containerized applications at scale, eliminating the need to manage the Kubernetes control plane.\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"How does AWS CloudFormation enable Infrastructure as Code?\",\n",
    "        \"ground_truth\": \"AWS CloudFormation allows you to define infrastructure using templates (JSON/YAML), enabling automated, version-controlled deployment and management of AWS resources.\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"Which AWS service should I used to store my applicative passwords?\",\n",
    "        \"ground_truth\": \"For storing application passwords securely in AWS use AWS Secrets Manager.\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"How do replace a spare tire?\",\n",
    "        \"ground_truth\": \"Park on flat surface, loosen lug nuts, jack up car, remove flat tire, mount spare, tighten lug nuts in star pattern, lower car, verify lug nut tightness.\"\n",
    "    },\n",
    "    {\n",
    "        \"question\": \"What is Amazon SageMaker's core functionality?\",\n",
    "        \"ground_truth\": \"Amazon SageMaker generates animations of flying shawarmas using serverless technology.\"\n",
    "    }\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "''' \n",
    "This function return Langchain LLM and Embedding wrapper with Bedrock LLMs and Embeddings.\n",
    "'''\n",
    "def get_bedrock_llm_and_embeddings_for_ragas(llm_model = DEFAULT_LLM_MODEL_ID):\n",
    "    config = {\n",
    "        \"region_name\": REGION_NAME,\n",
    "        \"llm\": llm_model,\n",
    "        \"embeddings\": \"amazon.titan-embed-text-v1\",\n",
    "        \"temperature\": 0.1,\n",
    "    }\n",
    "\n",
    "    bedrock_llm = ChatBedrockConverse(\n",
    "        region_name=config[\"region_name\"],\n",
    "        model=config[\"llm\"],\n",
    "        temperature=config[\"temperature\"],\n",
    "    )\n",
    "\n",
    "    bedrock_embeddings = BedrockEmbeddings(\n",
    "        region_name=config[\"region_name\"],\n",
    "        model_id=config[\"embeddings\"],\n",
    "    )\n",
    "\n",
    "    return LangchainLLMWrapper(bedrock_llm), LangchainEmbeddingsWrapper(bedrock_embeddings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [],
   "source": [
    "bedrock_runtime = boto3.client(\n",
    "    service_name = 'bedrock-agent-runtime',\n",
    "    region_name = REGION_NAME\n",
    ")\n",
    "\n",
    "def query_knowledge_base(question :str, model_arn :str, number_of_results :int): \n",
    "    try:\n",
    "        response = bedrock_runtime.retrieve_and_generate(\n",
    "            input={'text': question},\n",
    "            retrieveAndGenerateConfiguration={\n",
    "                'type': 'KNOWLEDGE_BASE',\n",
    "                'knowledgeBaseConfiguration': {\n",
    "                    'knowledgeBaseId': BEDROCK_KNOWLEDGE_BASE_ID,\n",
    "                    'modelArn': model_arn,\n",
    "                    'retrievalConfiguration':{\n",
    "                        'vectorSearchConfiguration': {\n",
    "                            'numberOfResults': number_of_results\n",
    "                        }\n",
    "                    }\n",
    "                }\n",
    "            }\n",
    "        )\n",
    "        \n",
    "        return {\n",
    "            \"output\": response[\"output\"][\"text\"],\n",
    "            \"citations\": [ref['content']['text'] for citation in response.get('citations', [])\n",
    "                         for ref in citation.get('retrievedReferences', [])\n",
    "                         if ref.get('content', {}).get('text')]\n",
    "        }\n",
    "    except Exception as e:\n",
    "        print(f\"Error: {str(e)}\")\n",
    "        return None"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_answers(test_data :str, model_arn : str = DEFAULT_LLM_MODEL_ID, number_of_results : int = 3):\n",
    "    print('Generating answers')\n",
    "    answers = []\n",
    "    for item in test_data:\n",
    "        response = query_knowledge_base(\n",
    "            question = item[\"question\"], \n",
    "            model_arn = model_arn, \n",
    "            number_of_results = number_of_results)\n",
    "                    \n",
    "        if response:\n",
    "            answers.append({\n",
    "                \"question\": item[\"question\"],\n",
    "                \"answer\": response[\"output\"],\n",
    "                \"ground_truth\": item[\"ground_truth\"],\n",
    "                \"retrieved_contexts\": response[\"citations\"]\n",
    "            })\n",
    "    return answers\n",
    "\n",
    "\n",
    "def evaluate_knowledge_base(answers):\n",
    "    dataset = Dataset.from_pandas(pd.DataFrame(answers))\n",
    "    \n",
    "    metrics = [\n",
    "        SemanticSimilarity(),\n",
    "    ]\n",
    "\n",
    "    llm, embeddings = get_bedrock_llm_and_embeddings_for_ragas()\n",
    "    print('Evaluating answers')\n",
    "    results = evaluate(\n",
    "        dataset=dataset,\n",
    "        metrics=metrics,\n",
    "        llm=llm,\n",
    "        embeddings=embeddings\n",
    "    )\n",
    "\n",
    "    return results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [],
   "source": [
    "import statistics\n",
    "def evaluate_parameter_combination(num_results, model_id, test_data):\n",
    "    \"\"\"Evaluate a specific parameter combination\"\"\"\n",
    "    answers = generate_answers(test_data, model_id, num_results)\n",
    "    results = evaluate_knowledge_base(answers)\n",
    "    \n",
    "    return {\n",
    "        'num_results': num_results,\n",
    "        'model_id': model_id,\n",
    "        'semantic_similarity': statistics.mean(results['semantic_similarity']),\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimize_parameters(test_data):\n",
    "    \"\"\"Run optimization across all parameter combinations\"\"\"\n",
    "    results = []\n",
    "    \n",
    "    for num_results, model_id in parameter_combinations:\n",
    "        print(f\"Testing combination: {num_results} results, {model_id}\")\n",
    "        result = evaluate_parameter_combination(num_results, model_id, test_data)\n",
    "        results.append(result)\n",
    "    \n",
    "    # Convert results to DataFrame for analysis\n",
    "    df_results = pd.DataFrame(results)\n",
    "    \n",
    "    # Find optimal combination based on semantic similarity\n",
    "    optimal_row = df_results.loc[df_results['semantic_similarity'].idxmax()]\n",
    "    \n",
    "    return df_results, optimal_row"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing combination: 1 results, mistral.mistral-small-2402-v1:0\n",
      "Generating answers\n",
      "Evaluating answers\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating: 100%|██████████| 10/10 [00:00<00:00, 40.19it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing combination: 1 results, ai21.jamba-1-5-mini-v1:0\n",
      "Generating answers\n",
      "Evaluating answers\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating: 100%|██████████| 10/10 [00:00<00:00, 46.66it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing combination: 1 results, cohere.command-r-v1:0\n",
      "Generating answers\n",
      "Evaluating answers\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating: 100%|██████████| 10/10 [00:00<00:00, 43.43it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing combination: 3 results, mistral.mistral-small-2402-v1:0\n",
      "Generating answers\n",
      "Evaluating answers\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating: 100%|██████████| 10/10 [00:00<00:00, 38.09it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing combination: 3 results, ai21.jamba-1-5-mini-v1:0\n",
      "Generating answers\n",
      "Evaluating answers\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating: 100%|██████████| 10/10 [00:00<00:00, 47.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing combination: 3 results, cohere.command-r-v1:0\n",
      "Generating answers\n",
      "Evaluating answers\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating: 100%|██████████| 10/10 [00:00<00:00, 20.52it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing combination: 7 results, mistral.mistral-small-2402-v1:0\n",
      "Generating answers\n",
      "Evaluating answers\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating: 100%|██████████| 10/10 [00:00<00:00, 47.25it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing combination: 7 results, ai21.jamba-1-5-mini-v1:0\n",
      "Generating answers\n",
      "Evaluating answers\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating: 100%|██████████| 10/10 [00:00<00:00, 45.91it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing combination: 7 results, cohere.command-r-v1:0\n",
      "Generating answers\n",
      "Evaluating answers\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating: 100%|██████████| 10/10 [00:00<00:00, 44.54it/s]\n"
     ]
    }
   ],
   "source": [
    "# Run optimization\n",
    "results_df, optimal_params = optimize_parameters(test_data_aws_services)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Analyze evaluation results\n",
    "### Show overall results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Optimal Parameters:\n",
      "Number of results: 7\n",
      "Model ID: cohere.command-r-v1:0\n",
      "Semantic Similarity Score: 0.8668\n",
      "\n",
      "All Results (sorted by semantic similarity):\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>num_results</th>\n",
       "      <th>model_id</th>\n",
       "      <th>semantic_similarity</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>7</td>\n",
       "      <td>cohere.command-r-v1:0</td>\n",
       "      <td>0.866808</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1</td>\n",
       "      <td>cohere.command-r-v1:0</td>\n",
       "      <td>0.826370</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>3</td>\n",
       "      <td>mistral.mistral-small-2402-v1:0</td>\n",
       "      <td>0.810855</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>mistral.mistral-small-2402-v1:0</td>\n",
       "      <td>0.809441</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>7</td>\n",
       "      <td>mistral.mistral-small-2402-v1:0</td>\n",
       "      <td>0.809415</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>3</td>\n",
       "      <td>cohere.command-r-v1:0</td>\n",
       "      <td>0.808495</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>7</td>\n",
       "      <td>ai21.jamba-1-5-mini-v1:0</td>\n",
       "      <td>0.382790</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1</td>\n",
       "      <td>ai21.jamba-1-5-mini-v1:0</td>\n",
       "      <td>0.343071</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>3</td>\n",
       "      <td>ai21.jamba-1-5-mini-v1:0</td>\n",
       "      <td>0.329163</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   num_results                         model_id  semantic_similarity\n",
       "8            7            cohere.command-r-v1:0             0.866808\n",
       "2            1            cohere.command-r-v1:0             0.826370\n",
       "3            3  mistral.mistral-small-2402-v1:0             0.810855\n",
       "0            1  mistral.mistral-small-2402-v1:0             0.809441\n",
       "6            7  mistral.mistral-small-2402-v1:0             0.809415\n",
       "5            3            cohere.command-r-v1:0             0.808495\n",
       "7            7         ai21.jamba-1-5-mini-v1:0             0.382790\n",
       "1            1         ai21.jamba-1-5-mini-v1:0             0.343071\n",
       "4            3         ai21.jamba-1-5-mini-v1:0             0.329163"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "print(\"\\nOptimal Parameters:\")\n",
    "print(f\"Number of results: {optimal_params['num_results']}\")\n",
    "print(f\"Model ID: {optimal_params['model_id']}\")\n",
    "print(f\"Semantic Similarity Score: {optimal_params['semantic_similarity']:.4f}\")\n",
    "\n",
    "# Display full results sorted by semantic similarity\n",
    "print(\"\\nAll Results (sorted by semantic similarity):\")\n",
    "display(results_df.sort_values('semantic_similarity', ascending=False))"
   ]
  }
 ],
 "metadata": {
  "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}