instructor/tutorials/3.applications-rag.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Applying Structured Output to RAG applications "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**What is RAG?**\n",
    "\n",
    "Retrieval Augmented Generation (RAG) models are the bridge between large language models and external knowledge databases. They fetch the relevant data for a given query. For example, if you have some documents and want to ask questions related to the content of those documents, RAG models help by retrieving data from those documents and passing it to the LLM in queries.\n",
    "\n",
    "**How do RAG models work?**\n",
    "\n",
    "The typical RAG process involves embedding a user query and searching a vector database to find the most relevant information to supplement the generated response. This approach is particularly effective when the database contains information closely matching the query but not more than that.\n",
    "\n",
    "![Image](https://jxnl.github.io/instructor/blog/img/dumb_rag.png)\n",
    "\n",
    "**Why is there a need for them?**\n",
    "\n",
    "Pre-trained large language models do not learn over time. If you ask them a question they have not been trained on, they will often hallucinate. Therefore, we need to embed our own data to achieve a better output."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simple RAG\n",
    "\n",
    "**What is it?**\n",
    "\n",
    "The simplest implementation of RAG embeds a user query and do a single embedding search in a vector database, like a vector store of Wikipedia articles. However, this approach often falls short when dealing with complex queries and diverse data sources.\n",
    "\n",
    "**What are the limitations?**\n",
    "\n",
    "- **Query-Document Mismatch:** It assumes that the query and document embeddings will align in the vector space, which is often not the case.\n",
    "    - Query: \"Tell me about climate change effects on marine life.\"\n",
    "    - Issue: The model might retrieve documents related to general climate change or marine life, missing the specific intersection of both topics.\n",
    "<!-- blank -->\n",
    "- **Monolithic Search Backend:** It relies on a single search method and backend, reducing flexibility and the ability to handle multiple data sources.\n",
    "    - Query: \"Latest research in quantum computing.\"\n",
    "    - Issue: The model might only search in a general science database, missing out on specialized quantum computing resources.\n",
    "<!-- blank -->\n",
    "- **Text Search Limitations:** The model is restricted to simple text queries without the nuances of advanced search features.\n",
    "    - Query: \"what problems did we fix last week\"\n",
    "    - Issue: cannot be answered by a simple text search since documents that contain problem, last week are going to be present at every week.\n",
    "<!-- blank -->\n",
    "- **Limited Planning Ability:** It fails to consider additional contextual information that could refine the search results.\n",
    "    - Query: \"Tips for first-time Europe travelers.\"\n",
    "    - Issue: The model might provide general travel advice, ignoring the specific context of first-time travelers or European destinations.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Improving the RAG model\n",
    "\n",
    "**What's the solution?**\n",
    "\n",
    "Enhancing RAG requires a more sophisticated approach known as query understanding.\n",
    "\n",
    "This process involves analyzing the user's query and transforming it to better match the backend's search capabilities.\n",
    "\n",
    "By doing so, we can significantly improve both the precision and recall of the search results, providing more accurate and relevant responses.\n",
    "\n",
    "![Image](https://jxnl.github.io/instructor/blog/img/query_understanding.png)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Practical Examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the examples below, we're going to use the [`instructor`](https://github.com/jxnl/instructor) library to simplify the interaction between the programmer and language models via the function-calling API.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "import instructor \n",
    "\n",
    "from openai import OpenAI\n",
    "from typing import List\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "client = instructor.patch(OpenAI())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 1) Improving Extractions\n",
    "\n",
    "One of the big limitations is that often times the query we embed and the text \n",
    "A common method of using structured output is to extract information from a document and use it to answer a question. Directly, we can be creative in how we extract, summarize and generate potential questions in order for our embeddings to do better. \n",
    "\n",
    "For example, instead of using just a text chunk we could try to:\n",
    "\n",
    "1. extract key words and themes\n",
    "2. extract hypothetical questions\n",
    "3. generate a summary of the text\n",
    "\n",
    "In the example below, we use the `instructor` library to extract the key words and themes from a text chunk and use them to answer a question."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Extraction(BaseModel):\n",
    "    summary: str \n",
    "    hypothetical_questions: List[str] = Field(default_factory=list, description=\"Hypothetical questions that this document could answer\")\n",
    "    keywords: List[str] = Field(default_factory=list, description=\"Keywords that this document is about\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"summary\": \"Simple RAG, an approach to query embedding and document retrieval, often falls short due to misalignment between query and document embeddings, reliance on a single search method and backend, text search limitations, and inadequate consideration of contextual information.\",\n",
      "  \"hypothetical_questions\": [\n",
      "    \"How can RAG be improved to better handle complex queries?\",\n",
      "    \"What other methods exist to enhance the alignment of queries and documents in vector spaces?\",\n",
      "    \"How can multiple data sources be incorporated into the search mechanism of RAG?\",\n",
      "    \"In what ways can advanced search features be integrated into RAG?\",\n",
      "    \"What approaches enable RAG to use contextual information for improved search results?\"\n",
      "  ],\n",
      "  \"keywords\": [\n",
      "    \"RAG\",\n",
      "    \"query embedding\",\n",
      "    \"document retrieval\",\n",
      "    \"vector database\",\n",
      "    \"query-document mismatch\",\n",
      "    \"monolithic search backend\",\n",
      "    \"text search limitations\",\n",
      "    \"planning ability\",\n",
      "    \"contextual information\"\n",
      "  ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "text_chunk = \"\"\"\n",
    "## Simple RAG\n",
    "\n",
    "**What is it?**\n",
    "\n",
    "The simplest implementation of RAG embeds a user query and do a single embedding search in a vector database, like a vector store of Wikipedia articles. However, this approach often falls short when dealing with complex queries and diverse data sources.\n",
    "\n",
    "**What are the limitations?**\n",
    "\n",
    "- **Query-Document Mismatch:** It assumes that the query and document embeddings will align in the vector space, which is often not the case.\n",
    "    - Query: \"Tell me about climate change effects on marine life.\"\n",
    "    - Issue: The model might retrieve documents related to general climate change or marine life, missing the specific intersection of both topics.\n",
    "- **Monolithic Search Backend:** It relies on a single search method and backend, reducing flexibility and the ability to handle multiple data sources.\n",
    "    - Query: \"Latest research in quantum computing.\"\n",
    "    - Issue: The model might only search in a general science database, missing out on specialized quantum computing resources.\n",
    "- **Text Search Limitations:** The model is restricted to simple text queries without the nuances of advanced search features.\n",
    "    - Query: \"what problems did we fix last week\"\n",
    "    - Issue: cannot be answered by a simple text search since documents that contain problem, last week are going to be present at every week.\n",
    "- **Limited Planning Ability:** It fails to consider additional contextual information that could refine the search results.\n",
    "    - Query: \"Tips for first-time Europe travelers.\"\n",
    "    - Issue: The model might provide general travel advice, ignoring the specific context of first-time travelers or European destinations.\n",
    "\"\"\"\n",
    "\n",
    "extraction = client.chat.completions.create(\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    response_model=Extraction,\n",
    "    messages=[{\n",
    "        \"role\": \"system\", \n",
    "        \"content\": \"Your role is to extract data from the following text chunk and create a RAG document.\"\n",
    "    }, {\n",
    "        \"role\": \"user\", \n",
    "        \"content\": text_chunk\n",
    "    }])\n",
    "\n",
    "\n",
    "print(extraction.model_dump_json(indent=2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you can imagine if you were to embed the summaries, hypothetical questions, and keywords in a vector database, you can then use a vector search to find the best matching document for a given query. What you'll find is that the results are much better than if you were to just embed the text chunk! "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 2) Understanding 'recent queries' to add temporal context\n",
    "\n",
    "One common application of using structured outputs for query understanding is to identify the intent of a user's query. In this example we're going to use a simple schema to seperately process the query to add additional temporal context.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import date\n",
    "\n",
    "class DateRange(BaseModel):\n",
    "    start: date\n",
    "    end: date\n",
    "\n",
    "class Query(BaseModel):\n",
    "    rewritten_query: str\n",
    "    published_daterange: DateRange"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this example, `DateRange` and `Query` are Pydantic models that structure the user's query with a date range and a list of domains to search within.\n",
    "\n",
    "These models **restructure** the user's query by including a <u>rewritten query</u>, a <u>range of published dates</u>, and a <u>list of domains</u> to search in."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the new restructured query, we can apply this pattern to our function calls to obtain results that are optimized for our backend."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"rewritten_query\": \"recent developments in Artificial Intelligence\",\n",
      "    \"published_daterange\": {\n",
      "        \"start\": \"2023-04-01\",\n",
      "        \"end\": \"2023-11-18\"\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "query = client.chat.completions.create(\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    response_model=Query,\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"system\", \n",
    "            \"content\": f\"You're a query understanding system for the Metafor Systems search engine. Today is {date.today()}. Here are some tips: ...\"\n",
    "        },\n",
    "        {\n",
    "            \"role\": \"user\", \n",
    "            \"content\": \"query: What are some recent developments in AI?\"\n",
    "        }\n",
    "    ],\n",
    ")\n",
    "\n",
    "print(query.model_dump_json(indent=4)) # Printing the Json dump of the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This isn't just about adding some date ranges. We can even use some chain of thought prompting to generate tailored searches that are deeply integrated with our backend. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"rewritten_query\": \"recent developments in artificial intelligence\",\n",
      "    \"published_daterange\": {\n",
      "        \"chain_of_thought\": \"Given that it's currently late 2023, a recent timeframe would ideally be within the last year to ensure the developments are current. Therefore, a suitable date range for recent AI developments could start from late 2022 to the present date in 2023.\",\n",
      "        \"start\": \"2022-11-18\",\n",
      "        \"end\": \"2023-11-18\"\n",
      "    }\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "class DateRange(BaseModel):\n",
    "    chain_of_thought: str = Field(\n",
    "        description=\"Think step by step to plan what is the best time range to search in\"\n",
    "    )\n",
    "    start: date\n",
    "    end: date\n",
    "\n",
    "class Query(BaseModel):\n",
    "    rewritten_query: str\n",
    "    published_daterange: DateRange\n",
    "\n",
    "\n",
    "query = client.chat.completions.create(\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    response_model=Query,\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"system\", \n",
    "            \"content\": f\"You're a query understanding system for a search engine. Today is {date.today()}.\"\n",
    "        },\n",
    "        {\n",
    "            \"role\": \"user\", \n",
    "            \"content\": \"What are some recent developments in AI?\"\n",
    "        }\n",
    "    ],\n",
    ")\n",
    "\n",
    "print(query.model_dump_json(indent=4)) # Printing the Json dump of the model"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 3) Personal Assistants, parallel processing\n",
    "\n",
    "A personal assistant application needs to interpret vague queries and fetch information from multiple backends, such as emails and calendars. By modeling the assistant's capabilities using Pydantic, we can dispatch the query to the correct backend and retrieve a unified response.\n",
    "\n",
    "For instance, when you ask, \"What's on my schedule today?\", the application needs to fetch data from various sources like events, emails, and reminders. This data is stored across different backends, but the goal is to provide a consolidated summary of results.\n",
    "\n",
    "It's important to note that the data from these sources may not be embedded in a search backend. Instead, they could be accessed through different clients like a calendar or email, spanning both personal and professional accounts.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal\n",
    "\n",
    "class SearchClient(BaseModel):\n",
    "    query: str\n",
    "    keywords: List[str]\n",
    "    email: str\n",
    "    source: Literal[\"gmail\", \"calendar\"] \n",
    "    date_range: DateRange\n",
    "\n",
    "class Retrival(BaseModel):\n",
    "    queries: List[SearchClient]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we can utilize this with a straightforward query such as \"What do I have today?\".\n",
    "\n",
    "The system will attempt to asynchronously dispatch the query to the appropriate backend.\n",
    "\n",
    "However, it's still crucial to remember that effectively prompting the language model is still a key aspect.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"queries\": [\n",
      "        {\n",
      "            \"query\": \"schedule\",\n",
      "            \"keywords\": [\n",
      "                \"appointments\",\n",
      "                \"meetings\",\n",
      "                \"schedule\",\n",
      "                \"events\"\n",
      "            ],\n",
      "            \"email\": \"jason.assistant@busybot.com\",\n",
      "            \"source\": \"calendar\",\n",
      "            \"date_range\": {\n",
      "                \"start\": \"2023-11-18\",\n",
      "                \"end\": \"2023-11-18\"\n",
      "            }\n",
      "        }\n",
      "    ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "retrival = client.chat.completions.create(\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    response_model=Retrival,\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\":f\"You are Jason's personal assistant. Today is {date.today()}\"},\n",
    "        {\"role\": \"user\", \"content\": \"What do I have today?\"}\n",
    "    ],\n",
    ")\n",
    "print(retrival.model_dump_json(indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make it more challenging, we will assign it multiple tasks, followed by a list of queries that are routed to various search backends, such as email and calendar. Not only do we dispatch to different backends, over which we have no control, but we are also likely to render them to the user in different ways."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"queries\": [\n",
      "        {\n",
      "            \"query\": \"meetings\",\n",
      "            \"keywords\": [\n",
      "                \"meetings\",\n",
      "                \"appointments\",\n",
      "                \"schedule\",\n",
      "                \"calendar\"\n",
      "            ],\n",
      "            \"email\": \"user@email.com\",\n",
      "            \"source\": \"calendar\",\n",
      "            \"date_range\": {\n",
      "                \"start\": \"2023-11-18\",\n",
      "                \"end\": \"2023-11-18\"\n",
      "            }\n",
      "        },\n",
      "        {\n",
      "            \"query\": \"important emails\",\n",
      "            \"keywords\": [\n",
      "                \"important\",\n",
      "                \"priority\",\n",
      "                \"urgent\",\n",
      "                \"follow-up\"\n",
      "            ],\n",
      "            \"email\": \"user@email.com\",\n",
      "            \"source\": \"gmail\",\n",
      "            \"date_range\": {\n",
      "                \"start\": \"2023-11-18\",\n",
      "                \"end\": \"2023-11-18\"\n",
      "            }\n",
      "        }\n",
      "    ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "retrival = client.chat.completions.create(\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    response_model=Retrival,\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\": f\"You are Jason's personal assistant. Today is {date.today()}\"},\n",
    "        {\"role\": \"user\", \"content\": \"What meetings do I have today and are there any important emails I should be aware of?\"}\n",
    "    ],\n",
    ")\n",
    "print(retrival.model_dump_json(indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example 4) Decomposing questions \n",
    "\n",
    "Lastly, a lightly more complex example of a problem that can be solved with structured output is decomposing questions. Where you ultimately want to decompose a question into a series of sub-questions that can be answered by a search backend. For example \n",
    "\n",
    "\"Whats the difference in populations of jason's home country and canada?\"\n",
    "\n",
    "You'd ultimately need to know a few things\n",
    "\n",
    "1. Jason's home country\n",
    "2. The population of Jason's home country\n",
    "3. The population of Canada\n",
    "4. The difference between the two\n",
    "\n",
    "This would not be done correctly as a single query, nor would it be done in parallel, however there are some opportunities try to be parallel since not all of the sub-questions are dependent on each other."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"root_question\": \"What is the difference between the population of jason's home country and canada?\",\n",
      "    \"plan\": [\n",
      "        {\n",
      "            \"id\": 1,\n",
      "            \"query\": \"What is Jason's home country?\",\n",
      "            \"subquestions\": []\n",
      "        },\n",
      "        {\n",
      "            \"id\": 2,\n",
      "            \"query\": \"What is the population of Canada?\",\n",
      "            \"subquestions\": []\n",
      "        },\n",
      "        {\n",
      "            \"id\": 3,\n",
      "            \"query\": \"What is the population of {Jason's home country}?\",\n",
      "            \"subquestions\": [\n",
      "                1\n",
      "            ]\n",
      "        },\n",
      "        {\n",
      "            \"id\": 4,\n",
      "            \"query\": \"What is the difference between the population of {Jason's home country} and the population of Canada?\",\n",
      "            \"subquestions\": [\n",
      "                2,\n",
      "                3\n",
      "            ]\n",
      "        }\n",
      "    ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "class Question(BaseModel):\n",
    "    id: int = Field(..., description=\"A unique identifier for the question\")\n",
    "    query: str = Field(..., description=\"The question decomposited as much as possible\")\n",
    "    subquestions: List[int] = Field(default_factory=list, description=\"The subquestions that this question is composed of\")\n",
    "\n",
    "\n",
    "class QueryPlan(BaseModel):\n",
    "    root_question: str = Field(..., description=\"The root question that the user asked\")\n",
    "    plan: List[Question] = Field(..., description=\"The plan to answer the root question and its subquestions\")\n",
    "\n",
    "\n",
    "retrival = client.chat.completions.create(\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    response_model=QueryPlan,\n",
    "    messages=[\n",
    "        {\"role\": \"system\", \"content\":\"You are a query understanding system capable of decomposing a question into subquestions.\"},\n",
    "        {\"role\": \"user\", \"content\": \"What is the difference between the population of jason's home country and canada?\"}\n",
    "    ],\n",
    ")\n",
    "\n",
    "print(retrival.model_dump_json(indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "I hope in this section I've exposed you to some ways we can be creative in modeling structured outputs to leverage LLMS in building some lightweight components for our systems."
   ]
  }
 ],
 "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}