{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Applying Structured Output to RAG applications\n"
   ]
  },
  {
   "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.\n"
   ]
  },
  {
   "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",
    "- **Query-Document Mismatch:** It assumes that the query and document embeddings will align in the vector space, which is often not the case.\n",
    "- **Text Search Limitations:** The model is restricted to simple text queries without the nuances of advanced search features.\n",
    "- **Limited Planning Ability:** It fails to consider additional contextual information that could refine the search results.\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\n"
   ]
  },
  {
   "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": 1,
   "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.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Extraction(BaseModel):\n",
    "    topic: str\n",
    "    summary: str\n",
    "    hypothetical_questions: List[str] = Field(\n",
    "        default_factory=list,\n",
    "        description=\"Hypothetical questions that this document could answer\",\n",
    "    )\n",
    "    keywords: List[str] = Field(\n",
    "        default_factory=list, description=\"Keywords that this document is about\"\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 's' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mNameError\u001b[0m                                 Traceback (most recent call last)",
      "\u001b[1;32m/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb Cell 10\u001b[0m line \u001b[0;36m3\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=1'>2</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mtyping\u001b[39;00m \u001b[39mimport\u001b[39;00m Iterable\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=4'>5</a>\u001b[0m text_chunk \u001b[39m=\u001b[39m \u001b[39m\"\"\"\u001b[39m\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=5'>6</a>\u001b[0m \u001b[39m## Simple RAG\u001b[39m\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=6'>7</a>\u001b[0m \n\u001b[0;32m   (...)\u001b[0m\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=24'>25</a>\u001b[0m \u001b[39m    - Issue: The model might provide general travel advice, ignoring the specific context of first-time travelers or European destinations.\u001b[39m\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=25'>26</a>\u001b[0m \u001b[39m\"\"\"\u001b[39m\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=27'>28</a>\u001b[0m extractions \u001b[39m=\u001b[39m client\u001b[39m.\u001b[39mchat\u001b[39m.\u001b[39mcompletions\u001b[39m.\u001b[39mcreate(\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=28'>29</a>\u001b[0m     model\u001b[39m=\u001b[39m\u001b[39m\"\u001b[39m\u001b[39mgpt-4-1106-preview\u001b[39m\u001b[39m\"\u001b[39m,\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=29'>30</a>\u001b[0m     stream\u001b[39m=\u001b[39m\u001b[39mTrue\u001b[39;00m,\n\u001b[0;32m---> <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=30'>31</a>\u001b[0m     response_model\u001b[39m=\u001b[39ms,\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=31'>32</a>\u001b[0m     messages\u001b[39m=\u001b[39m[\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=32'>33</a>\u001b[0m         {\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=33'>34</a>\u001b[0m             \u001b[39m\"\u001b[39m\u001b[39mrole\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39m\"\u001b[39m\u001b[39msystem\u001b[39m\u001b[39m\"\u001b[39m,\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=34'>35</a>\u001b[0m             \u001b[39m\"\u001b[39m\u001b[39mcontent\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39m\"\u001b[39m\u001b[39mYour role is to extract chunks from the following and create a set of topics.\u001b[39m\u001b[39m\"\u001b[39m,\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=35'>36</a>\u001b[0m         },\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=36'>37</a>\u001b[0m         {\u001b[39m\"\u001b[39m\u001b[39mrole\u001b[39m\u001b[39m\"\u001b[39m: \u001b[39m\"\u001b[39m\u001b[39muser\u001b[39m\u001b[39m\"\u001b[39m, \u001b[39m\"\u001b[39m\u001b[39mcontent\u001b[39m\u001b[39m\"\u001b[39m: text_chunk},\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=37'>38</a>\u001b[0m     ],\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=38'>39</a>\u001b[0m )\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=41'>42</a>\u001b[0m \u001b[39mfor\u001b[39;00m extraction \u001b[39min\u001b[39;00m extractions:\n\u001b[1;32m     <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X12sZmlsZQ%3D%3D?line=42'>43</a>\u001b[0m     pprint(extraction\u001b[39m.\u001b[39mmodel_dump())\n",
      "\u001b[0;31mNameError\u001b[0m: name 's' is not defined"
     ]
    }
   ],
   "source": [
    "from pprint import pprint\n",
    "from typing import Iterable\n",
    "\n",
    "\n",
    "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",
    "extractions = client.chat.completions.create(\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    stream=True,\n",
    "    response_model=s,\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": \"Your role is to extract chunks from the following and create a set of topics.\",\n",
    "        },\n",
    "        {\"role\": \"user\", \"content\": text_chunk},\n",
    "    ],\n",
    ")\n",
    "\n",
    "\n",
    "for extraction in extractions:\n",
    "    pprint(extraction.model_dump())"
   ]
  },
  {
   "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!\n"
   ]
  },
  {
   "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": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import date\n",
    "\n",
    "\n",
    "class DateRange(BaseModel):\n",
    "    start: date\n",
    "    end: date\n",
    "\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.\n"
   ]
  },
  {
   "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.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Query(rewritten_query='recent developments in AI', published_daterange=DateRange(start=datetime.date(2023, 2, 9), end=datetime.date(2024, 2, 9)))"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def expand_query(q) -> Query:\n",
    "    return client.chat.completions.create(\n",
    "        model=\"gpt-3.5-turbo\",\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",
    "            {\"role\": \"user\", \"content\": f\"query: {q}\"},\n",
    "        ],\n",
    "    )\n",
    "\n",
    "\n",
    "query = expand_query(\"What are some recent developments in AI?\")\n",
    "query"
   ]
  },
  {
   "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.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Query(rewritten_query='Recent developments in Artificial Intelligence', published_daterange=DateRange(chain_of_thought=\"Considering 'recent' generally refers to the past few months or up to a year, the best date range to capture the most recent developments in AI would be from one year ago to today's date.\", start=datetime.date(2023, 2, 9), end=datetime.date(2024, 2, 9)))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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",
    "\n",
    "class Query(BaseModel):\n",
    "    rewritten_query: str = Field(\n",
    "        description=\"Rewrite the query to make it more specific\"\n",
    "    )\n",
    "    published_daterange: DateRange = Field(\n",
    "        description=\"Effective date range to search in\"\n",
    "    )\n",
    "\n",
    "\n",
    "def expand_query(q) -> Query:\n",
    "    return 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",
    "            {\"role\": \"user\", \"content\": f\"query: {q}\"},\n",
    "        ],\n",
    "    )\n",
    "\n",
    "\n",
    "expand_query(\"What are some recent developments in AI?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using Weights and Biases to track experiments\n",
    "\n",
    "While running a function like this production is quite simple, a lot of time will be spend on iterating and improving the model. To do this, we can use Weights and Biases to track our experiments.\n",
    "\n",
    "In order to do so we wand manage a few things\n",
    "\n",
    "1. Save input and output pairs for later\n",
    "2. Save the JSON schema for the response_model\n",
    "3. Having snapshots of the model and data allow us to compare results over time, and as we make changes to the model we can see how the results change.\n",
    "\n",
    "This is particularly useful when we might want to blend a mix of synthetic and real data to evaluate our model. We can use the `wandb` library to track our experiments and save the results to a dashboard.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "ename": "ModuleNotFoundError",
     "evalue": "No module named 'helpers'",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "\u001b[1;32m/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb Cell 20\u001b[0m line \u001b[0;36m5\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X25sZmlsZQ%3D%3D?line=1'>2</a>\u001b[0m \u001b[39mimport\u001b[39;00m \u001b[39minstructor\u001b[39;00m\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X25sZmlsZQ%3D%3D?line=3'>4</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mopenai\u001b[39;00m \u001b[39mimport\u001b[39;00m AsyncOpenAI\n\u001b[0;32m----> <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X25sZmlsZQ%3D%3D?line=4'>5</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mhelpers\u001b[39;00m \u001b[39mimport\u001b[39;00m dicts_to_df\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X25sZmlsZQ%3D%3D?line=5'>6</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mdatetime\u001b[39;00m \u001b[39mimport\u001b[39;00m date\n\u001b[1;32m      <a href='vscode-notebook-cell:/Users/jasonliu/dev/instructor/docs/tutorials/3-0-applications-rag.ipynb#X25sZmlsZQ%3D%3D?line=6'>7</a>\u001b[0m \u001b[39mfrom\u001b[39;00m \u001b[39mpydantic\u001b[39;00m \u001b[39mimport\u001b[39;00m BaseModel, Field\n",
      "\u001b[0;31mModuleNotFoundError\u001b[0m: No module named 'helpers'"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import instructor\n",
    "\n",
    "from openai import AsyncOpenAI\n",
    "from helpers import dicts_to_df\n",
    "from datetime import date\n",
    "from pydantic import BaseModel, Field\n",
    "\n",
    "\n",
    "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",
    "\n",
    "class Query(BaseModel):\n",
    "    rewritten_query: str = Field(\n",
    "        description=\"Rewrite the query to make it more specific\"\n",
    "    )\n",
    "    published_daterange: DateRange = Field(\n",
    "        description=\"Effective date range to search in\"\n",
    "    )\n",
    "\n",
    "    def report(self):\n",
    "        dct = self.model_dump()\n",
    "        dct[\"usage\"] = self._raw_response.usage.model_dump()\n",
    "        return dct\n",
    "\n",
    "\n",
    "\n",
    "# We'll use a different client for async calls\n",
    "# To highlight the difference and how we can use both\n",
    "aclient = instructor.patch(AsyncOpenAI())\n",
    "\n",
    "\n",
    "async def expand_query(\n",
    "    q, *, model: str = \"gpt-4-1106-preview\", temp: float = 0\n",
    ") -> Query:\n",
    "    return await aclient.chat.completions.create(\n",
    "        model=model,\n",
    "        temperature=temp,\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",
    "            {\"role\": \"user\", \"content\": f\"query: {q}\"},\n",
    "        ],\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "wandb version 0.16.1 is available!  To upgrade, please run:\n",
       " $ pip install wandb --upgrade"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Tracking run with wandb version 0.16.0"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Run data is saved locally in <code>/Users/jasonliu/dev/instructor/tutorials/wandb/run-20231227_202003-7c9dxnfl</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Syncing run <strong><a href='https://wandb.ai/instructor/query/runs/7c9dxnfl' target=\"_blank\">blooming-firefly-4</a></strong> to <a href='https://wandb.ai/instructor/query' target=\"_blank\">Weights & Biases</a> (<a href='https://wandb.me/run' target=\"_blank\">docs</a>)<br/>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View project at <a href='https://wandb.ai/instructor/query' target=\"_blank\">https://wandb.ai/instructor/query</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run at <a href='https://wandb.ai/instructor/query/runs/7c9dxnfl' target=\"_blank\">https://wandb.ai/instructor/query/runs/7c9dxnfl</a>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d73fb8a832254b32a938572fd27eca62",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "VBox(children=(Label(value='0.011 MB of 0.011 MB uploaded (0.001 MB deduped)\\r'), FloatProgress(value=1.0, max…"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "wandb: WARNING Source type is set to 'repo' but some required information is missing from the environment. A job will not be created from this run. See https://docs.wandb.ai/guides/launch/create-job\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "W&B sync reduced upload amount by 6.6%             "
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<style>\n",
       "    table.wandb td:nth-child(1) { padding: 0 10px; text-align: left ; width: auto;} td:nth-child(2) {text-align: left ; width: 100%}\n",
       "    .wandb-row { display: flex; flex-direction: row; flex-wrap: wrap; justify-content: flex-start; width: 100% }\n",
       "    .wandb-col { display: flex; flex-direction: column; flex-basis: 100%; flex: 1; padding: 10px; }\n",
       "    </style>\n",
       "<div class=\"wandb-row\"><div class=\"wandb-col\"><h3>Run history:</h3><br/><table class=\"wandb\"><tr><td>average duration (s)</td><td>▁</td></tr><tr><td>duration (s)</td><td>▁</td></tr><tr><td>n_queries</td><td>▁</td></tr><tr><td>usage_completion_tokens</td><td>▁</td></tr><tr><td>usage_prompt_tokens</td><td>▁</td></tr><tr><td>usage_total_tokens</td><td>▁</td></tr></table><br/></div><div class=\"wandb-col\"><h3>Run summary:</h3><br/><table class=\"wandb\"><tr><td>average duration (s)</td><td>2.28692</td></tr><tr><td>duration (s)</td><td>9.14768</td></tr><tr><td>n_queries</td><td>4</td></tr><tr><td>usage_completion_tokens</td><td>359</td></tr><tr><td>usage_prompt_tokens</td><td>780</td></tr><tr><td>usage_total_tokens</td><td>1139</td></tr></table><br/></div></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       " View run <strong style=\"color:#cdcd00\">blooming-firefly-4</strong> at: <a href='https://wandb.ai/instructor/query/runs/7c9dxnfl' target=\"_blank\">https://wandb.ai/instructor/query/runs/7c9dxnfl</a><br/>Synced 4 W&B file(s), 2 media file(s), 5 artifact file(s) and 0 other file(s)"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "Find logs at: <code>./wandb/run-20231227_202003-7c9dxnfl/logs</code>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import asyncio\n",
    "import time\n",
    "import pandas as pd\n",
    "import wandb\n",
    "\n",
    "model = \"gpt-4-1106-preview\"\n",
    "temp = 0\n",
    "\n",
    "run = wandb.init(\n",
    "    project=\"query\",\n",
    "    config={\"model\": model, \"temp\": temp},\n",
    ")\n",
    "\n",
    "test_queries = [\n",
    "    \"latest developments in artificial intelligence last 3 weeks\",\n",
    "    \"renewable energy trends past month\",\n",
    "    \"quantum computing advancements last 2 months\",\n",
    "    \"biotechnology updates last 10 days\",\n",
    "]\n",
    "start = time.perf_counter()\n",
    "queries = await asyncio.gather(\n",
    "    *[expand_query(q, model=model, temp=temp) for q in test_queries]\n",
    ")\n",
    "duration = time.perf_counter() - start\n",
    "\n",
    "with open(\"schema.json\", \"w+\") as f:\n",
    "    schema = Query.model_json_schema()\n",
    "    json.dump(schema, f, indent=2)\n",
    "\n",
    "with open(\"results.jsonlines\", \"w+\") as f:\n",
    "    for query in queries:\n",
    "        f.write(query.model_dump_json() + \"\\n\")\n",
    "\n",
    "df = dicts_to_df([q.report() for q in queries])\n",
    "df[\"input\"] = test_queries\n",
    "df.to_csv(\"results.csv\")\n",
    "\n",
    "\n",
    "run.log({\"schema\": wandb.Table(dataframe=pd.DataFrame([{\"schema\": schema}]))})\n",
    "\n",
    "run.log(\n",
    "    {\n",
    "        \"usage_total_tokens\": df[\"usage_total_tokens\"].sum(),\n",
    "        \"usage_completion_tokens\": df[\"usage_completion_tokens\"].sum(),\n",
    "        \"usage_prompt_tokens\": df[\"usage_prompt_tokens\"].sum(),\n",
    "        \"duration (s)\": duration,\n",
    "        \"average duration (s)\": duration / len(queries),\n",
    "        \"n_queries\": len(queries),\n",
    "    }\n",
    ")\n",
    "\n",
    "\n",
    "run.log(\n",
    "    {\n",
    "        \"results\": wandb.Table(dataframe=df),\n",
    "    }\n",
    ")\n",
    "\n",
    "files = wandb.Artifact(\"data\", type=\"dataset\")\n",
    "\n",
    "files.add_file(\"schema.json\")\n",
    "files.add_file(\"results.jsonlines\")\n",
    "files.add_file(\"results.csv\")\n",
    "\n",
    "\n",
    "run.log_artifact(files)\n",
    "run.finish()"
   ]
  },
  {
   "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": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Literal\n",
    "\n",
    "\n",
    "class SearchClient(BaseModel):\n",
    "    query: str = Field(description=\"The search query that will go into the search bar\")\n",
    "    keywords: List[str]\n",
    "    email: str\n",
    "    source: Literal[\"gmail\", \"calendar\"]\n",
    "    date_range: DateRange\n",
    "\n",
    "\n",
    "class Retrieval(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": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "    \"queries\": [\n",
      "        {\n",
      "            \"query\": \"work\",\n",
      "            \"keywords\": [],\n",
      "            \"email\": \"jason@work.com\",\n",
      "            \"source\": \"calendar\",\n",
      "            \"date_range\": {\n",
      "                \"chain_of_thought\": \"today\",\n",
      "                \"start\": \"2024-02-09\",\n",
      "                \"end\": \"2024-02-09\"\n",
      "            }\n",
      "        },\n",
      "        {\n",
      "            \"query\": \"\",\n",
      "            \"keywords\": [],\n",
      "            \"email\": \"jason@work.com\",\n",
      "            \"source\": \"gmail\",\n",
      "            \"date_range\": {\n",
      "                \"chain_of_thought\": \"today\",\n",
      "                \"start\": \"2024-02-09\",\n",
      "                \"end\": \"2024-02-09\"\n",
      "            }\n",
      "        }\n",
      "    ]\n",
      "}\n"
     ]
    }
   ],
   "source": [
    "retrieval = client.chat.completions.create(\n",
    "    model=\"gpt-3.5-turbo\",\n",
    "    response_model=Retrieval,\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": f\"\"\"You are Jason's personal assistant.\n",
    "                He has two emails jason@work.com jason@personal.com \n",
    "                Today is {date.today()}\"\"\",\n",
    "        },\n",
    "        {\"role\": \"user\", \"content\": \"What do I have today for work? any new emails?\"},\n",
    "    ],\n",
    ")\n",
    "print(retrieval.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.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "retrieval = client.chat.completions.create(\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    response_model=Retrieval,\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": f\"\"\"You are Jason's personal assistant.\n",
    "                He has two emails jason@work.com jason@personal.com \n",
    "                Today is {date.today()}\"\"\",\n",
    "        },\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": \"What meetings do I have today and are there any important emails I should be aware of\",\n",
    "        },\n",
    "    ],\n",
    ")\n",
    "print(retrieval.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.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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(\n",
    "        default_factory=list,\n",
    "        description=\"The subquestions that this question is composed of\",\n",
    "    )\n",
    "\n",
    "\n",
    "class QueryPlan(BaseModel):\n",
    "    root_question: str = Field(..., description=\"The root question that the user asked\")\n",
    "    plan: List[Question] = Field(\n",
    "        ..., description=\"The plan to answer the root question and its subquestions\"\n",
    "    )\n",
    "\n",
    "\n",
    "retrieval = client.chat.completions.create(\n",
    "    model=\"gpt-4-1106-preview\",\n",
    "    response_model=QueryPlan,\n",
    "    messages=[\n",
    "        {\n",
    "            \"role\": \"system\",\n",
    "            \"content\": \"You are a query understanding system capable of decomposing a question into subquestions.\",\n",
    "        },\n",
    "        {\n",
    "            \"role\": \"user\",\n",
    "            \"content\": \"What is the difference between the population of jason's home country and canada?\",\n",
    "        },\n",
    "    ],\n",
    ")\n",
    "\n",
    "print(retrieval.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.\n"
   ]
  }
 ],
 "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
}
