Advanced RAG: Query Decomposition and Reasoning

by Tuana Celik (LI, Twitter)

This is part one of the Advanced Use Cases series:

1️⃣ Extract Metadata from Queries to Improve Retrieval cookbook & full article

2️⃣ Query Expansion cookbook & full article

3️⃣ Query Decomposition & the full article

4️⃣ Automated Metadata Enrichment

Query decomposition is a technique we can use to decompose complex queries into simpler steps, answering each sub-question, and getting an LLM to reason about the final answer based on the answers to the sub-questions.

In this recipe, we're using the structured output functionality (currently in beta) by OpenAI to construct Questions which lists the sub-questions based on the original question, as well as keeping track of the intermediate answers to each question.

📺 Code Along

Load the Dataset

For this demo, we're using the Tuana/game-of-thrones dataset on Hugging Face, and we are using a Cohere embedding model to embed the contents.

🧪 Experimental Addition to the OpenAIGenerator for Structured Output Support

🚀 This step is a completely optional advanced step

Let's extend the OpenAIGeneraotor to be able to make use of the strctured output option by OpenAI. Below, we extend the class to call self.client.beta.chat.completions.parse if the user has provides a respose_format in generation_kwargs. This will allow us to provifde a Pydantic Model to the gnerator and request our generator to respond with structured outputs that adhere to this Pydantic schema.

Define the Pydantic Model

For query expansion, we want to keep track of intermediate questions that can be answered independently. So, below we define a Questions schema. Each Question is made up a question and an answer

Define the Query Decomposition Prompt

The first step in our application will be to decompose a question. We define our first splitter_prompt instructing an LLM to take the question step by steo and produce multiple sub questions.

Define a Multi Text Embedder and Retriever

Since we will likely be needing to answer multiple question, let's define a Milti Question Embedder and Retriever. This way, we can have 1 component that can accept multiple questsions, embed them, and another component that can retireve documents that relate to each question too.

Define the Prompt to Answer Multiple Questions

Once we have our decomposed questions, we need to instruct an LLM to answer each question based on the context for each question.

Define the Prompt to Reason About the Final Answer

Our final step will be to instruct the LLM to reason about the final answer based on the decomposed questions and answers to each.

Final Step: Construct the Pipeline