search/retrieval-augmented-generation/examples/question_answering.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "359697d5"
      },
      "source": [
        "# Question Answering Over Documents\n",
        "\n",
        "_Augmented with document retrieval from Vertex AI Search and LangChain_\n",
        "\n",
        "<table align=\"left\">\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/search/retrieval-augmented-generation/examples/question_answering.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/colab-logo-32px.png\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fsearch%2Fretrieval-augmented-generation%2Fexamples%2Fquestion_answering.ipynb\">\n",
        "      <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
        "    </a>\n",
        "  </td>    \n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/search/retrieval-augmented-generation/examples/question_answering.ipynb\">\n",
        "      <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
        "    </a>\n",
        "  </td>\n",
        "  <td style=\"text-align: center\">\n",
        "    <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/search/retrieval-augmented-generation/examples/question_answering.ipynb\">\n",
        "      <img src=\"https://cloud.google.com/ml-engine/images/github-logo-32px.png\" alt=\"GitHub logo\"><br> View on GitHub\n",
        "    </a>\n",
        "  </td>\n",
        "</table>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "31280da3e872"
      },
      "source": [
        "---\n",
        "\n",
        "\n",
        "* Authors: Ruchika Kharwar, Tomas Paquete, Tom Pakeman, Holt Skinner\n",
        "* Created: 06/06/2023\n",
        "\n",
        "---\n",
        "\n",
        "## Objective\n",
        "\n",
        "This notebook shows how to ask and answer questions about your data by combining a Vertex AI Search engine with LLMs. In particular, we focus on querying 'unstructured' data such as PDFs and HTML files.\n",
        "\n",
        "These patterns are useful if you have a Vertex AI Search Engine pointed at a store of documents, such as a Google Cloud Storage bucket containing PDFs.\n",
        "\n",
        "In order to run this notebook you must have created an unstructured search engine and ingested PDF or HTML documents into it.\n",
        "\n",
        "Users may wish to:\n",
        "\n",
        "1. Get an answer to a question from document(s)\n",
        "2. Get an answer to a question from document(s) along with the citation/sources\n",
        "3. Get an answer to a question from document(s) and follow it up with additional questions like in a conversation.\n",
        "4. Users might want to change the default prompt of LangChain while using the results of Vertex AI Search\n",
        "\n",
        "In each case we will use an example of the LangChain [retriever](https://python.langchain.com/v0.2/docs/concepts/#retrievers) interface to fetch contextual documents & segments from a Vertex AI Search engine, and the Vertex AI API to call Google's Gemini model for summarization & answering.\n",
        "\n",
        "Please note that the prompts and examples demonstrated here will need to be customized in order to work well with your own data.\n",
        "\n",
        "---\n",
        "\n",
        "In this notebook the following examples will be elaborated:\n",
        "\n",
        "- ✅ Example of using `VertexAISearchRetriever` with LangChain `RetrievalQA`\n",
        "- ✅ Example of using `VertexAISearchRetriever` with LangChain `RetrievalQAWithSourcesChain`\n",
        "- ✅ Example of using `VertexAISearchRetriever` with LangChain `ConversationalRetrievalChain`\n",
        "- ✅ Example of using `VertexAIMultiTurnSearchRetriever` with LangChain `ConversationalRetrievalChain`\n",
        "- ✅ Examples of changing the default prompt with Vertex AI Search and `RetrievalQAWithSourcesChain`\n",
        "\n",
        "---\n",
        "\n",
        "**References:**\n",
        "\n",
        "- [LangChain Conceptual Documentation](https://docs.langchain.com/docs/)\n",
        "- [LangChain Python Documentation](https://python.langchain.com/en/latest/)\n",
        "- [LangChain Vertex AI Search Retriever Documentation](https://python.langchain.com/docs/integrations/retrievers/google_vertex_ai_search)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1e82b735a19c"
      },
      "source": [
        "## High level flow\n",
        "\n",
        "![Vertex_AI_Search_RAG](https://storage.googleapis.com/github-repo/search/examples/VAIS_RAG.png)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZdBQ2O2rL9i5"
      },
      "source": [
        "The following is a high-level overview of the steps in the examples that are going to be demonstrated:\n",
        "\n",
        "- Users enter a prompt of a question\n",
        "- The datastore is then queried and relevant documents are returned\n",
        "- The documents are added as context into the user's prompt/question\n",
        "- The prompt is sent into the LLM\n",
        "- A response is returned using the relevant documents as sources\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zj-AnsYaOSNq"
      },
      "source": [
        "Services used in the notebook:\n",
        "\n",
        "- ✅ Vertex AI Gemini API for LLM\n",
        "- ✅ Vertex AI Search for document search and retrieval\n",
        "- ✅ LangChain in order to chain tasks\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pQbOO-Lc-2-7"
      },
      "source": [
        "## Install pre-requisites\n",
        "\n",
        "If running in Colab install the pre-requisites into the runtime. Otherwise it is assumed that the notebook is running in Vertex AI Workbench. In that case it is recommended to install the pre-requisites from a terminal using the `--user` option.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ohPUPez8imvE"
      },
      "outputs": [],
      "source": [
        "%pip install -q --upgrade --user google-cloud-aiplatform google-cloud-discoveryengine langchain langchain-google-vertexai langchain-google-community"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R5Xep4W9lq-Z"
      },
      "source": [
        "### Restart current runtime\n",
        "\n",
        "To use the newly installed packages in this Jupyter runtime, you must restart the runtime. You can do this by running the cell below, which will restart the current kernel."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XRvKdaPDTznN"
      },
      "outputs": [],
      "source": [
        "# Restart kernel after installs so that your environment can access the new packages\n",
        "import IPython\n",
        "\n",
        "app = IPython.Application.instance()\n",
        "app.kernel.do_shutdown(True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SbmM4z7FOBpM"
      },
      "source": [
        "<div class=\"alert alert-block alert-warning\">\n",
        "<b>⚠️ The kernel is going to restart. Please wait until it is finished before continuing to the next step. ⚠️</b>\n",
        "</div>\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8vEczuYo_r-g"
      },
      "source": [
        "## Authenticate\n",
        "\n",
        "If running in Colab authenticate with `google.colab.google.auth` otherwise assume that running on Vertex AI Workbench.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "loTfn0KniwB2"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "\n",
        "if \"google.colab\" in sys.modules:\n",
        "    from google.colab import auth as google_auth\n",
        "\n",
        "    google_auth.authenticate_user()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_ARUa9gEK74l"
      },
      "source": [
        "## Configure notebook environment\n",
        "\n",
        "To get started using Vertex AI Search, you must have an existing Google Cloud project and [enable the Vertex AI API and Vertex AI Search API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com,discoveryengine.googleapis.com).\n",
        "\n",
        "Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LgXLTL802_tY"
      },
      "outputs": [],
      "source": [
        "PROJECT_ID = \"YOUR_PROJECT_ID\"  # @param {type:\"string\"}\n",
        "LOCATION = \"us-central1\"  # @param {type:\"string\"}\n",
        "\n",
        "import vertexai\n",
        "\n",
        "vertexai.init(project=PROJECT_ID, location=LOCATION)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "66ff7614"
      },
      "source": [
        "### Set the following constants to reflect your environment\n",
        "\n",
        "- The queries used in the examples here relate to a GCS bucket containing Alphabet investor PDFs, but these should be customized to your own data.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "uVxFSrppK8Oy"
      },
      "outputs": [],
      "source": [
        "DATA_STORE_ID = \"YOUR_DATA_STORE_ID\"  # @param {type:\"string\"}\n",
        "DATA_STORE_LOCATION = \"global\"  # @param {type:\"string\"}\n",
        "\n",
        "MODEL = \"gemini-1.5-pro\"  # @param {type:\"string\"}\n",
        "\n",
        "if PROJECT_ID == \"YOUR_PROJECT_ID\" or DATA_STORE_ID == \"YOUR_DATA_STORE_ID\":\n",
        "    raise ValueError(\n",
        "        \"Please set the PROJECT_ID, DATA_STORE_ID constants to reflect your environment.\"\n",
        "    )"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gUoS5-Tv3T4I"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "1N7SrPF33XA6"
      },
      "outputs": [],
      "source": [
        "from langchain.chains import (\n",
        "    ConversationalRetrievalChain,\n",
        "    RetrievalQA,\n",
        "    RetrievalQAWithSourcesChain,\n",
        ")\n",
        "from langchain.memory import ConversationBufferMemory\n",
        "from langchain.prompts import PromptTemplate\n",
        "from langchain_google_community import (\n",
        "    VertexAIMultiTurnSearchRetriever,\n",
        "    VertexAISearchRetriever,\n",
        ")\n",
        "from langchain_google_vertexai import VertexAI"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "I3tBINK7uUep"
      },
      "source": [
        "## LangChain retrieval Q&A chains\n",
        "\n",
        "We will demonstrate using three LangChain retrieval Q&A chains:\n",
        "\n",
        "- `RetrievalQA`\n",
        "- `RetrievalQAWithSourcesChain`\n",
        "- `ConversationalRetrievalChain`\n",
        "\n",
        "We begin by initializing a Vertex AI LLM and a LangChain 'retriever' to fetch documents from our Vertex AI Search engine.\n",
        "\n",
        "For Q&A chains our retriever is passed directly to the chain and can be used automatically without any further configuration.\n",
        "\n",
        "Behind the scenes, first the search query is passed to the retriever which runs a search and returns relevant document chunks. These chunks are then passed to the prompt used by the LLM to be used as context.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fe224af7"
      },
      "outputs": [],
      "source": [
        "llm = VertexAI(model_name=MODEL)\n",
        "\n",
        "retriever = VertexAISearchRetriever(\n",
        "    project_id=PROJECT_ID,\n",
        "    location_id=DATA_STORE_LOCATION,\n",
        "    data_store_id=DATA_STORE_ID,\n",
        "    get_extractive_answers=True,\n",
        "    max_documents=10,\n",
        "    max_extractive_segment_count=1,\n",
        "    max_extractive_answer_count=5,\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JBf9cgsb-ZvK"
      },
      "source": [
        "### [`RetrievalQA` chain](https://python.langchain.com/docs/modules/chains/popular/vector_db_qa)\n",
        "\n",
        "This is the simplest document Q&A chain offered by LangChain.\n",
        "\n",
        "There are several different chain types available, listed [here](https://docs.langchain.com/docs/components/chains/index_related_chains).\n",
        "\n",
        "- In these examples we use the `stuff` type, which simply inserts all of the document chunks into the prompt.\n",
        "- This has the advantage of only making a single LLM call, which is faster and more cost efficient\n",
        "- However, if we have a large number of search results we run the risk of exceeding the token limit in our prompt, or truncating useful information.\n",
        "- Other chain types such as `map_reduce` and `refine` use an iterative process which makes multiple LLM calls, taking individual document chunks at a time and refining the answer iteratively.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "BwAUJxYS-e9g"
      },
      "outputs": [],
      "source": [
        "search_query = \"What was Alphabet's Revenue in Q2 2021?\"  # @param {type:\"string\"}\n",
        "\n",
        "retrieval_qa = RetrievalQA.from_chain_type(\n",
        "    llm=llm, chain_type=\"stuff\", retriever=retriever\n",
        ")\n",
        "retrieval_qa.invoke(search_query)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "AUSB8Mn8Dvou"
      },
      "source": [
        "#### Inspecting the process\n",
        "\n",
        "If we add `return_source_documents=True` we can inspect the document chunks that were returned by the retriever.\n",
        "\n",
        "This is helpful for debugging, as these chunks may not always be relevant to the answer, or their relevance might not be obvious.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "P236bx-KA9TM"
      },
      "outputs": [],
      "source": [
        "retrieval_qa = RetrievalQA.from_chain_type(\n",
        "    llm=llm, chain_type=\"stuff\", retriever=retriever, return_source_documents=True\n",
        ")\n",
        "\n",
        "results = retrieval_qa.invoke(search_query)\n",
        "\n",
        "print(\"*\" * 79)\n",
        "print(results[\"result\"])\n",
        "print(\"*\" * 79)\n",
        "for doc in results[\"source_documents\"]:\n",
        "    print(\"-\" * 79)\n",
        "    print(doc.page_content)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LxuDV3sdFhUk"
      },
      "source": [
        "### `RetrievalQAWithSourcesChain`\n",
        "\n",
        "This variant returns an answer to the question alongside the source documents that were used to generate the answer.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LlhQ-1JGFfMI"
      },
      "outputs": [],
      "source": [
        "retrieval_qa_with_sources = RetrievalQAWithSourcesChain.from_chain_type(\n",
        "    llm=llm, chain_type=\"stuff\", retriever=retriever\n",
        ")\n",
        "\n",
        "retrieval_qa_with_sources.invoke(search_query, return_only_outputs=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "NaxAUt9FhXR6"
      },
      "source": [
        "### [`ConversationalRetrievalChain`](https://python.langchain.com/docs/modules/chains/popular/chat_vector_db) with `VertexAIMultiTurnSearchRetriever`\n",
        "\n",
        "`ConversationalRetrievalChain` remembers and uses previous questions so you can have a chat-like discovery process.\n",
        "\n",
        "To use this chain we must provide a memory class to store and pass the previous messages to the LLM as context. Here we use the `ConversationBufferMemory` class that comes with LangChain.\n",
        "\n",
        "`VertexAIMultiTurnSearchRetriever` uses [multi-turn search](https://cloud.google.com/generative-ai-app-builder/docs/multi-turn-search) (also called conversational search or search with followups) to preserve context between requests.\n",
        "\n",
        "The following example will work with both retrievers, and the multi-turn retriever can be substituted in any of the previous examples."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "mRdu6hK6F6UD"
      },
      "outputs": [],
      "source": [
        "multi_turn_retriever = VertexAIMultiTurnSearchRetriever(\n",
        "    project_id=PROJECT_ID, location_id=DATA_STORE_LOCATION, data_store_id=DATA_STORE_ID\n",
        ")\n",
        "memory = ConversationBufferMemory(memory_key=\"chat_history\", return_messages=True)\n",
        "conversational_retrieval = ConversationalRetrievalChain.from_llm(\n",
        "    llm=llm, retriever=multi_turn_retriever, memory=memory\n",
        ")\n",
        "\n",
        "search_query = \"What were alphabet revenues in 2022?\"\n",
        "\n",
        "result = conversational_retrieval.invoke(search_query)\n",
        "print(result[\"answer\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ff9b72cd-da78-4690-9919-b74b2b460c44"
      },
      "outputs": [],
      "source": [
        "new_query = \"What about costs and expenses?\"\n",
        "result = conversational_retrieval.invoke(new_query)\n",
        "print(result[\"answer\"])"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "o44PdiaGiM7m"
      },
      "outputs": [],
      "source": [
        "new_query = \"Is this more than in 2021?\"\n",
        "\n",
        "result = conversational_retrieval.invoke(new_query)\n",
        "print(result[\"answer\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1eeba1fe-0a4d-4651-a4b1-b91dfc1254bc"
      },
      "source": [
        "## Advanced: Modifying the default langchain prompt\n",
        "\n",
        "In all of the previous examples we used the default prompt that comes with langchain.\n",
        "\n",
        "We can inspect our chain object to discover the wording of the prompt template being used.\n",
        "\n",
        "We may find that this is not suitable for our purposes, and we may wish to customize the prompt, for example to present our results in a different format, or to specify additional constraints.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c1004f16-ed9d-4170-974e-6dd7528fe36b"
      },
      "outputs": [],
      "source": [
        "qa = RetrievalQA.from_chain_type(\n",
        "    llm=llm, chain_type=\"stuff\", retriever=retriever, return_source_documents=True\n",
        ")\n",
        "\n",
        "print(qa.combine_documents_chain.llm_chain.prompt.template)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "b30e2c43-b57b-43f1-9679-41c4a2a9e8ba"
      },
      "source": [
        "---\n",
        "\n",
        "Let's modify the prompt to return an answer in a single word (useful for yes/no questions). We will constrain the LLM to say 'I don't know' if it cannot answer.\n",
        "\n",
        "We create a new prompt_template and pass this in using the `template` argument.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0ceca59d-f75a-4467-9d58-a15b3241ac22"
      },
      "outputs": [],
      "source": [
        "prompt_template = \"\"\"Use the context to answer the question at the end.\n",
        "You must always use the context and context only to answer the question. Never try to make up an answer. If the context is empty or you do not know the answer, just say \"I don't know\".\n",
        "The answer should consist of only 1 word and not a sentence.\n",
        "\n",
        "Context: {context}\n",
        "\n",
        "Question: {question}\n",
        "Helpful Answer:\n",
        "\"\"\"\n",
        "prompt = PromptTemplate(\n",
        "    template=prompt_template, input_variables=[\"context\", \"question\"]\n",
        ")\n",
        "qa_chain = RetrievalQA.from_llm(\n",
        "    llm=llm, prompt=prompt, retriever=retriever, return_source_documents=True\n",
        ")"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "f58539f7bec3"
      },
      "outputs": [],
      "source": [
        "print(qa_chain.combine_documents_chain.llm_chain.prompt.template)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "287d9859-453d-474a-a449-15252a14aa63"
      },
      "outputs": [],
      "source": [
        "search_query = \"Were 2020 EMEA revenues higher than 2020 APAC revenues?\"\n",
        "\n",
        "results = qa_chain.invoke(search_query)\n",
        "\n",
        "print(\"*\" * 79)\n",
        "print(results[\"result\"])\n",
        "print(\"*\" * 79)\n",
        "for doc in results[\"source_documents\"]:\n",
        "    print(\"-\" * 79)\n",
        "    print(doc.page_content)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "name": "question_answering.ipynb",
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}