{
"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",
"\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",
"\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 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",
"we are searching for may not have a direct match, leading to suboptimal results.\n",
"A common method of using structured output is to extract information from a\n",
"document and use it to answer a question. Directly, we can be creative in how we\n",
"extract, summarize and generate potential questions in order for our embeddings\n",
"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\n",
"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": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{'hypothetical_questions': ['What is the basic concept behind simple RAG?',\n",
" 'How does simple RAG work for information '\n",
" 'retrieval?'],\n",
" 'keywords': ['Simple RAG',\n",
" 'Retrieval-Augmented Generation',\n",
" 'user query',\n",
" 'embedding search',\n",
" 'vector database',\n",
" 'Wikipedia articles',\n",
" 'information retrieval'],\n",
" 'summary': 'The simplest implementation of Retrieval-Augmented Generation '\n",
" '(RAG) involves embedding a user query and conducting a single '\n",
" 'embedding search in a vector database, like a vector store of '\n",
" 'Wikipedia articles, to retrieve relevant information. This method '\n",
" 'may not be ideal for complex queries or varied data sources.',\n",
" 'topic': 'Simple RAG'}\n",
"{'hypothetical_questions': ['What are the drawbacks of using simple RAG '\n",
" 'systems?',\n",
" 'How does query-document mismatch affect the '\n",
" 'performance of RAG?',\n",
" 'Why is a monolithic search backend a limitation '\n",
" 'for RAG?'],\n",
" 'keywords': ['limitations',\n",
" 'query-document mismatch',\n",
" 'simple RAG',\n",
" 'monolithic search backend',\n",
" 'text search',\n",
" 'planning ability',\n",
" 'contextual information'],\n",
" 'summary': 'Key limitations of the simple RAG include query-document '\n",
" 'mismatch, reliance on a single search backend, constraints of '\n",
" 'text search capabilities, and limited planning ability to '\n",
" 'leverage contextual information. These issues can result in '\n",
" 'suboptimal search outcomes and retrieval of irrelevant or broad '\n",
" 'information.',\n",
" 'topic': 'Limitations of Simple RAG'}\n"
]
}
],
"source": [
"from pprint import pprint\n",
"from collections.abc 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=Iterable[Extraction],\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,\n",
"and keywords in a vector database (i.e. in the metadata fields of a vector\n",
"database), you can then use a vector search to find the best matching document\n",
"for a given query. What you'll find is that the results are much better than if\n",
"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 separately 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 rewritten query, a range of published dates, and a list of domains 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 artificial intelligence', published_daterange=DateRange(start=datetime.date(2024, 1, 1), end=datetime.date(2024, 3, 31)))"
]
},
"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='latest advancements in artificial intelligence', published_daterange=DateRange(chain_of_thought='Since the user is asking for recent developments, it would be relevant to look for articles and papers published within the last year. Therefore, setting the start date to a year before today and the end date to today will cover the most recent advancements.', start=datetime.date(2023, 3, 31), end=datetime.date(2024, 3, 31)))"
]
},
"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": [],
"source": [
"import json\n",
"import instructor\n",
"\n",
"from openai import AsyncOpenAI\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",
"# 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": 8,
"metadata": {},
"outputs": [],
"source": [
"# % pip install pandas wandb\n",
"import pandas as pd\n",
"from typing import Any\n",
"\n",
"\n",
"def flatten_dict(\n",
" d: dict[str, Any], parent_key: str = \"\", sep: str = \"_\"\n",
") -> dict[str, Any]:\n",
" \"\"\"\n",
" Flatten a nested dictionary.\n",
"\n",
" :param d: The nested dictionary to flatten.\n",
" :param parent_key: The base key to use for the flattened keys.\n",
" :param sep: Separator to use between keys.\n",
" :return: A flattened dictionary.\n",
" \"\"\"\n",
" items = []\n",
" for k, v in d.items():\n",
" new_key = f\"{parent_key}{sep}{k}\" if parent_key else k\n",
" if isinstance(v, dict):\n",
" items.extend(flatten_dict(v, new_key, sep=sep).items())\n",
" else:\n",
" items.append((new_key, v))\n",
" return dict(items)\n",
"\n",
"\n",
"def dicts_to_df(list_of_dicts: list[dict[str, Any]]) -> pd.DataFrame:\n",
" \"\"\"\n",
" Convert a list of dictionaries to a pandas DataFrame.\n",
"\n",
" :param list_of_dicts: List of dictionaries, potentially nested.\n",
" :return: A pandas DataFrame representing the flattened data.\n",
" \"\"\"\n",
" # Flatten each dictionary and create a DataFrame\n",
" flattened_data = [flatten_dict(d) for d in list_of_dicts]\n",
" return pd.DataFrame(flattened_data)"
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [],
"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",
"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",
"run.log(\n",
" {\n",
" \"results\": wandb.Table(dataframe=df),\n",
" }\n",
")\n",
"\n",
"files = wandb.Artifact(\"data\", type=\"dataset\")\n",
"files.add_file(\"schema.json\")\n",
"files.add_file(\"results.jsonlines\")\n",
"files.add_file(\"results.csv\")\n",
"\n",
"run.log_artifact(files)\n",
"run.finish()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"The output of Weights and Biases would return something like the below table.\n",
"\n",
"| Metric | Value |\n",
"|--------------------------|--------|\n",
"| average duration (s) | 1.5945 |\n",
"| duration (s) | 6.37799|\n",
"| n_queries | 4 |\n",
"| usage_completion_tokens | 376 |\n",
"| usage_prompt_tokens | 780 |\n",
"| usage_total_tokens | 1156 |\n"
]
},
{
"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": 10,
"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": 11,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{\n",
" \"queries\": [\n",
" {\n",
" \"query\": \"work\",\n",
" \"keywords\": [\n",
" \"work\",\n",
" \"today\"\n",
" ],\n",
" \"email\": \"jason@work.com\",\n",
" \"source\": \"gmail\",\n",
" \"date_range\": {\n",
" \"chain_of_thought\": \"Check today's work schedule\",\n",
" \"start\": \"2024-03-31\",\n",
" \"end\": \"2024-03-31\"\n",
" }\n",
" },\n",
" {\n",
" \"query\": \"new emails\",\n",
" \"keywords\": [\n",
" \"email\",\n",
" \"new\"\n",
" ],\n",
" \"email\": \"jason@work.com\",\n",
" \"source\": \"gmail\",\n",
" \"date_range\": {\n",
" \"chain_of_thought\": \"Check for new emails today\",\n",
" \"start\": \"2024-03-31\",\n",
" \"end\": \"2024-03-31\"\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": 12,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"{\n",
" \"queries\": [\n",
" {\n",
" \"query\": \"Jason's meetings\",\n",
" \"keywords\": [\n",
" \"meeting\",\n",
" \"appointment\",\n",
" \"schedule\",\n",
" \"calendar\"\n",
" ],\n",
" \"email\": \"jason@work.com\",\n",
" \"source\": \"calendar\",\n",
" \"date_range\": {\n",
" \"chain_of_thought\": \"Since today's date is 2024-03-31, we should look for meetings scheduled for this exact date.\",\n",
" \"start\": \"2024-03-31\",\n",
" \"end\": \"2024-03-31\"\n",
" }\n",
" }\n",
" ]\n",
"}\n"
]
}
],
"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": 13,
"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 the population of 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 difference between two population numbers?\",\n",
" \"subquestions\": [\n",
" 1,\n",
" 2\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 decomposed 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.8"
}
},
"nbformat": 4,
"nbformat_minor": 4
}