docs/docs/examples/prompts/prompts_rag.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "64624d35-393f-44bc-8583-e39fb8e2a24f",
   "metadata": {},
   "source": [
    "<a href=\"https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/docs/examples/prompts/prompts_rag.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55d722a3-e798-40e9-909e-28b30b2c1fbd",
   "metadata": {},
   "source": [
    "# Prompt Engineering for RAG\n",
    "\n",
    "In this notebook we show various prompt techniques you can try to customize your LlamaIndex RAG pipeline.\n",
    "\n",
    "- Getting and setting prompts for query engines, etc.\n",
    "- Defining template variable mappings (e.g. you have an existing QA prompt)\n",
    "- Adding few-shot examples + performing query transformations/rewriting."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f05505a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install llama-index-llms-openai\n",
    "%pip install llama-index-readers-file pymupdf"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a16ad016-69de-4a5e-9100-4b7828a67cb3",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install llama-index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc38b5cd-ccf2-4476-85d3-53a3ea9b302f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import openai"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9cf5d59-ae01-4ddc-9f63-c796494dbdb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ[\"OPENAI_API_KEY\"] = \"sk-...\"\n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75eaa6b1-a9c4-454e-830c-a4901b7f426c",
   "metadata": {},
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "499b51c3-1037-4f5f-96d5-4f383fd40484",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:numexpr.utils:Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "INFO:numexpr.utils:NumExpr defaulting to 8 threads.\n",
      "NumExpr defaulting to 8 threads.\n"
     ]
    }
   ],
   "source": [
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))\n",
    "\n",
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.core import PromptTemplate\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "42d56cd2-1027-4c6b-8deb-e69d8218a962",
   "metadata": {},
   "source": [
    "#### Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e95e4f26-5ff8-42d3-bfe2-21d6d720379f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mkdir: data: File exists\n",
      "--2023-10-28 23:19:38--  https://arxiv.org/pdf/2307.09288.pdf\n",
      "Resolving arxiv.org (arxiv.org)... 128.84.21.199\n",
      "Connecting to arxiv.org (arxiv.org)|128.84.21.199|:443... connected.\n",
      "HTTP request sent, awaiting response... 200 OK\n",
      "Length: 13661300 (13M) [application/pdf]\n",
      "Saving to: ‘data/llama2.pdf’\n",
      "\n",
      "data/llama2.pdf     100%[===================>]  13.03M  1.50MB/s    in 10s     \n",
      "\n",
      "2023-10-28 23:19:49 (1.31 MB/s) - ‘data/llama2.pdf’ saved [13661300/13661300]\n"
     ]
    }
   ],
   "source": [
    "!mkdir data\n",
    "!wget --user-agent \"Mozilla\" \"https://arxiv.org/pdf/2307.09288.pdf\" -O \"data/llama2.pdf\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e5cde429-c3af-46a3-9ecc-5081954eaa46",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "from llama_index.readers.file import PyMuPDFReader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eef36cdf-8969-4802-a373-b69dd56a3c76",
   "metadata": {},
   "outputs": [],
   "source": [
    "loader = PyMuPDFReader()\n",
    "documents = loader.load(file_path=\"./data/llama2.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3565a377-4a30-456f-a5f6-4e120d075250",
   "metadata": {},
   "source": [
    "#### Load into Vector Store"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ad9b432-06b3-45a3-a3d9-b1e0c6746f0d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core import VectorStoreIndex\n",
    "from llama_index.llms.openai import OpenAI\n",
    "\n",
    "gpt35_llm = OpenAI(model=\"gpt-3.5-turbo\")\n",
    "gpt4_llm = OpenAI(model=\"gpt-4\")\n",
    "\n",
    "index = VectorStoreIndex.from_documents(documents)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dcda1300-c289-4334-b848-f7b1142e6b2d",
   "metadata": {},
   "source": [
    "#### Setup Query Engine / Retriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f23e582c-1b02-443a-ae83-76e9281920c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_str = \"What are the potential risks associated with the use of Llama 2 as mentioned in the context?\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a707c654-572f-463f-91da-48fc6c40d67b",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = index.as_query_engine(similarity_top_k=2, llm=gpt35_llm)\n",
    "# use this for testing\n",
    "vector_retriever = index.as_retriever(similarity_top_k=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "958515ad-af10-4163-89bb-0ceb12ca48d8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The potential risks associated with the use of Llama 2, as mentioned in the context, include the generation of misinformation and the retrieval of information about topics such as bioterrorism or cybercrime. The models have been tuned to avoid these topics and diminish any capabilities they might have offered for those use cases. However, there is a possibility that the safety tuning of the models may go too far, resulting in an overly cautious approach where the model declines certain requests or responds with too many safety details. Users of Llama 2 and Llama 2-Chat need to be cautious and take extra steps in tuning and deployment to ensure responsible use.\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(query_str)\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25e672a6-ee16-49f7-a166-2a8a9a298936",
   "metadata": {},
   "source": [
    "## Viewing/Customizing Prompts\n",
    "\n",
    "First, let's take a look at the query engine prompts, and see how we can customize it."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4441d2d4-b719-4737-8dac-7696b0fafba4",
   "metadata": {},
   "source": [
    "### View Prompts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02253e5f-b2fc-4afd-bfc1-f4e55d78352b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define prompt viewing function\n",
    "def display_prompt_dict(prompts_dict):\n",
    "    for k, p in prompts_dict.items():\n",
    "        text_md = f\"**Prompt Key**: {k}<br>\" f\"**Text:** <br>\"\n",
    "        display(Markdown(text_md))\n",
    "        print(p.get_template())\n",
    "        display(Markdown(\"<br><br>\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d88f692-7017-4ab8-a2ef-6dd8d8b132c5",
   "metadata": {},
   "outputs": [],
   "source": [
    "prompts_dict = query_engine.get_prompts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2463322-db25-4a9e-9560-b89a34706a72",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**Prompt Key**: response_synthesizer:text_qa_template<br>**Text:** <br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context information is below.\n",
      "---------------------\n",
      "{context_str}\n",
      "---------------------\n",
      "Given the context information and not prior knowledge, answer the query.\n",
      "Query: {query_str}\n",
      "Answer: \n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "<br><br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Prompt Key**: response_synthesizer:refine_template<br>**Text:** <br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The original query is as follows: {query_str}\n",
      "We have provided an existing answer: {existing_answer}\n",
      "We have the opportunity to refine the existing answer (only if needed) with some more context below.\n",
      "------------\n",
      "{context_msg}\n",
      "------------\n",
      "Given the new context, refine the original answer to better answer the query. If the context isn't useful, return the original answer.\n",
      "Refined Answer: \n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "<br><br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_prompt_dict(prompts_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecf5703b-4075-4abe-85d3-83793801a78b",
   "metadata": {},
   "source": [
    "### Customize Prompts\n",
    "\n",
    "What if we want to do something different than our standard question-answering prompts?\n",
    "\n",
    "Let's try out the RAG prompt from [LangchainHub](https://smith.langchain.com/hub/rlm/rag-prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f552f06-c72a-4cf4-a131-c694f4908d67",
   "metadata": {},
   "outputs": [],
   "source": [
    "# to do this, you need to use the langchain object\n",
    "\n",
    "from langchain import hub\n",
    "\n",
    "langchain_prompt = hub.pull(\"rlm/rag-prompt\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "912d7e0b-1a6d-4b46-ab2c-8e9079e2532f",
   "metadata": {},
   "source": [
    "One catch is that the template variables in the prompt are different than what's expected by our synthesizer in the query engine:\n",
    "- the prompt uses `context` and `question`,\n",
    "- we expect `context_str` and `query_str`\n",
    "\n",
    "This is not a problem! Let's add our template variable mappings to map variables. We use our `LangchainPromptTemplate` to map to LangChain prompts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ef2bfc3-22c1-41c8-92b9-01914c9bb1d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.prompts import LangchainPromptTemplate\n",
    "\n",
    "lc_prompt_tmpl = LangchainPromptTemplate(\n",
    "    template=langchain_prompt,\n",
    "    template_var_mappings={\"query_str\": \"question\", \"context_str\": \"context\"},\n",
    ")\n",
    "\n",
    "query_engine.update_prompts(\n",
    "    {\"response_synthesizer:text_qa_template\": lc_prompt_tmpl}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac909c71-6a63-4923-af0a-406b08ec7731",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**Prompt Key**: response_synthesizer:text_qa_template<br>**Text:** <br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "input_variables=['question', 'context'] messages=[HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['question', 'context'], template=\"You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.\\nQuestion: {question} \\nContext: {context} \\nAnswer:\"))]\n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "<br><br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Prompt Key**: response_synthesizer:refine_template<br>**Text:** <br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The original query is as follows: {query_str}\n",
      "We have provided an existing answer: {existing_answer}\n",
      "We have the opportunity to refine the existing answer (only if needed) with some more context below.\n",
      "------------\n",
      "{context_msg}\n",
      "------------\n",
      "Given the new context, refine the original answer to better answer the query. If the context isn't useful, return the original answer.\n",
      "Refined Answer: \n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "<br><br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "prompts_dict = query_engine.get_prompts()\n",
    "display_prompt_dict(prompts_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c40474be-28bf-4a9a-bbfa-fa902af1722e",
   "metadata": {},
   "source": [
    "### Try It Out\n",
    "\n",
    "Let's re-run our query engine again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58e49923-4108-45aa-9b97-efc755fa9822",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The potential risks associated with the use of Llama 2 mentioned in the context include the generation of misinformation, retrieval of information about topics like bioterrorism or cybercrime, an overly cautious approach by the model, and the need for users to be cautious and take extra steps in tuning and deployment. However, efforts have been made to tune the models to avoid these topics and diminish any capabilities they might have offered for those use cases.\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(query_str)\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8099f8f4-c4aa-4d30-8b56-aa42255a41bc",
   "metadata": {},
   "source": [
    "## Adding Few-Shot Examples\n",
    "\n",
    "Let's try adding few-shot examples to the prompt, which can be dynamically loaded depending on the query! \n",
    "\n",
    "We do this by setting the `function_mapping` variable in our prompt template - this allows us to compute functions (e.g. return few-shot examples) during prompt formatting time.\n",
    "\n",
    "As an example use case, through this we can coerce the model to output results in a structured format,\n",
    "by showing examples of other structured outputs.\n",
    "\n",
    "Let's parse a pre-generated question/answer file. For the sake of focus we'll skip how the file is generated (tl;dr we used a GPT-4 powered function calling RAG pipeline), but the qa pairs look like this:\n",
    "\n",
    "```\n",
    "{\"query\": \"<query>\", \"response\": \"<output_json>\"}\n",
    "```\n",
    "\n",
    "We embed/index these Q/A pairs, and retrieve the top-k.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66a0f538-9907-4245-9daa-888fcc8e6db1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.schema import TextNode\n",
    "\n",
    "few_shot_nodes = []\n",
    "for line in open(\"../llama2_qa_citation_events.jsonl\", \"r\"):\n",
    "    few_shot_nodes.append(TextNode(text=line))\n",
    "\n",
    "few_shot_index = VectorStoreIndex(few_shot_nodes)\n",
    "few_shot_retriever = few_shot_index.as_retriever(similarity_top_k=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "865cbcc7-e0d0-48aa-9950-3abd732beff3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "\n",
    "def few_shot_examples_fn(**kwargs):\n",
    "    query_str = kwargs[\"query_str\"]\n",
    "    retrieved_nodes = few_shot_retriever.retrieve(query_str)\n",
    "    # go through each node, get json object\n",
    "\n",
    "    result_strs = []\n",
    "    for n in retrieved_nodes:\n",
    "        raw_dict = json.loads(n.get_content())\n",
    "        query = raw_dict[\"query\"]\n",
    "        response_dict = json.loads(raw_dict[\"response\"])\n",
    "        result_str = f\"\"\"\\\n",
    "Query: {query}\n",
    "Response: {response_dict}\"\"\"\n",
    "        result_strs.append(result_str)\n",
    "    return \"\\n\\n\".join(result_strs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ef19f0c-e711-4571-a35f-6d6d69fc41e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# write prompt template with functions\n",
    "qa_prompt_tmpl_str = \"\"\"\\\n",
    "Context information is below.\n",
    "---------------------\n",
    "{context_str}\n",
    "---------------------\n",
    "Given the context information and not prior knowledge, \\\n",
    "answer the query asking about citations over different topics.\n",
    "Please provide your answer in the form of a structured JSON format containing \\\n",
    "a list of authors as the citations. Some examples are given below.\n",
    "\n",
    "{few_shot_examples}\n",
    "\n",
    "Query: {query_str}\n",
    "Answer: \\\n",
    "\"\"\"\n",
    "\n",
    "qa_prompt_tmpl = PromptTemplate(\n",
    "    qa_prompt_tmpl_str,\n",
    "    function_mappings={\"few_shot_examples\": few_shot_examples_fn},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "825af1ce-52e3-44f9-9262-ac79dd8be855",
   "metadata": {},
   "outputs": [],
   "source": [
    "citation_query_str = (\n",
    "    \"Which citations are mentioned in the section on Safety RLHF?\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac6b51eb-fb50-495b-855c-e7acd89be365",
   "metadata": {},
   "source": [
    "Let's see what the formatted prompt looks like with the few-shot examples function.\n",
    "(we fill in test context for brevity)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b550563-a481-4cbf-92b8-f788dd60aadb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context information is below.\n",
      "---------------------\n",
      "test_context\n",
      "---------------------\n",
      "Given the context information and not prior knowledge, answer the query asking about citations over different topics.\n",
      "Please provide your answer in the form of a structured JSON format containing a list of authors as the citations. Some examples are given below.\n",
      "\n",
      "Query: Which citation discusses the impact of safety RLHF measured by reward model score distributions?\n",
      "Response: {'citations': [{'author': 'Llama 2: Open Foundation and Fine-Tuned Chat Models', 'year': 24, 'desc': 'Impact of safety RLHF measured by reward model score distributions. Left: safety reward model scores of generations on the Meta Safety test set. The clustering of samples in the top left corner suggests the improvements of model safety. Right: helpfulness reward model scores of generations on the Meta Helpfulness test set.'}]}\n",
      "\n",
      "Query: Which citations are mentioned in the section on RLHF Results?\n",
      "Response: {'citations': [{'author': 'Gilardi et al.', 'year': 2023, 'desc': ''}, {'author': 'Huang et al.', 'year': 2023, 'desc': ''}]}\n",
      "\n",
      "Query: Which citations are mentioned in the section on Safety RLHF?\n",
      "Answer: \n"
     ]
    }
   ],
   "source": [
    "print(\n",
    "    qa_prompt_tmpl.format(\n",
    "        query_str=citation_query_str, context_str=\"test_context\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfd5ea22-a29c-4316-a552-93ec759be2e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine.update_prompts(\n",
    "    {\"response_synthesizer:text_qa_template\": qa_prompt_tmpl}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f541664e-843a-468b-aaca-722f448d179f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/markdown": [
       "**Prompt Key**: response_synthesizer:text_qa_template<br>**Text:** <br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Context information is below.\n",
      "---------------------\n",
      "{context_str}\n",
      "---------------------\n",
      "Given the context information and not prior knowledge, answer the query asking about citations over different topics.\n",
      "Please provide your answer in the form of a structured JSON format containing a list of authors as the citations. Some examples are given below.\n",
      "\n",
      "{few_shot_examples}\n",
      "\n",
      "Query: {query_str}\n",
      "Answer: \n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "<br><br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/markdown": [
       "**Prompt Key**: response_synthesizer:refine_template<br>**Text:** <br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The original query is as follows: {query_str}\n",
      "We have provided an existing answer: {existing_answer}\n",
      "We have the opportunity to refine the existing answer (only if needed) with some more context below.\n",
      "------------\n",
      "{context_msg}\n",
      "------------\n",
      "Given the new context, refine the original answer to better answer the query. If the context isn't useful, return the original answer.\n",
      "Refined Answer: \n"
     ]
    },
    {
     "data": {
      "text/markdown": [
       "<br><br>"
      ],
      "text/plain": [
       "<IPython.core.display.Markdown object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "display_prompt_dict(query_engine.get_prompts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8388952d-db56-4ec9-9d6a-5678a9d8af84",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'citations': [{'author': 'Llama 2: Open Foundation and Fine-Tuned Chat Models', 'year': 24, 'desc': 'Safety RLHF'}, {'author': 'Bai et al.', 'year': 2022a, 'desc': 'RLHF stage'}, {'author': 'Bai et al.', 'year': 2022a, 'desc': 'adversarial prompts'}, {'author': 'Bai et al.', 'year': 2022a, 'desc': 'safety reward model'}, {'author': 'Bai et al.', 'year': 2022a, 'desc': 'helpfulness reward model'}, {'author': 'Bai et al.', 'year': 2022a, 'desc': 'safety tuning with RLHF'}]}\n"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(citation_query_str)\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10ce0ee5-c7bd-47fa-982d-18197591dbdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(response.source_nodes[1].get_content())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "092be8b1-d42a-4dc6-a941-d4d21be212c2",
   "metadata": {},
   "source": [
    "## Context Transformations - PII Example\n",
    "\n",
    "We can also dynamically add context transformations as functions in the prompt variable. In this example we show how we can process the `context_str` before feeding to the context window - specifically in masking out PII (a step towards alleviating concerns around data privacy/security).\n",
    "\n",
    "**NOTE**: You can do these as steps before feeding into the prompt as well, but this gives you flexibility to perform all this on the fly for any QA prompt you define!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc53b339-0e52-4d4d-9564-a5680d073fe8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.core.postprocessor import (\n",
    "    NERPIINodePostprocessor,\n",
    "    SentenceEmbeddingOptimizer,\n",
    ")\n",
    "from llama_index.core import QueryBundle\n",
    "from llama_index.core.schema import NodeWithScore, TextNode"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17194226-df94-4098-aef3-1470db9c7863",
   "metadata": {},
   "outputs": [],
   "source": [
    "pii_processor = NERPIINodePostprocessor(llm=gpt4_llm)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19ef32db-452e-413f-836c-6cfdf3a9d876",
   "metadata": {},
   "outputs": [],
   "source": [
    "def filter_pii_fn(**kwargs):\n",
    "    # run optimizer\n",
    "    query_bundle = QueryBundle(query_str=kwargs[\"query_str\"])\n",
    "\n",
    "    new_nodes = pii_processor.postprocess_nodes(\n",
    "        [NodeWithScore(node=TextNode(text=kwargs[\"context_str\"]))],\n",
    "        query_bundle=query_bundle,\n",
    "    )\n",
    "    new_node = new_nodes[0]\n",
    "    return new_node.get_content()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac90ebfb-f43f-4cf8-958b-5fbb7715f448",
   "metadata": {},
   "outputs": [],
   "source": [
    "qa_prompt_tmpl_str = (\n",
    "    \"Context information is below.\\n\"\n",
    "    \"---------------------\\n\"\n",
    "    \"{context_str}\\n\"\n",
    "    \"---------------------\\n\"\n",
    "    \"Given the context information and not prior knowledge, \"\n",
    "    \"answer the query.\\n\"\n",
    "    \"Query: {query_str}\\n\"\n",
    "    \"Answer: \"\n",
    ")\n",
    "qa_prompt_tmpl = PromptTemplate(\n",
    "    qa_prompt_tmpl_str, function_mappings={\"context_str\": filter_pii_fn}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d26b15d-44d2-456f-9132-d43e93268184",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine.update_prompts(\n",
    "    {\"response_synthesizer:text_qa_template\": qa_prompt_tmpl}\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "80d88d97-b84e-4e6d-9cb9-4e47834175c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# take a look at the prompt\n",
    "retrieved_nodes = vector_retriever.retrieve(query_str)\n",
    "context_str = \"\\n\\n\".join([n.get_content() for n in retrieved_nodes])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eae91811-7299-4bd6-827c-03002efb4007",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(qa_prompt_tmpl.format(query_str=query_str, context_str=context_str))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae4e94b0-886c-4896-9661-e9bc2a9857fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(query_str)\n",
    "print(str(response))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama_index_v2",
   "language": "python",
   "name": "llama_index_v2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}