Demystifying Advanced RAG Pipelines

Retrieval-Augmented Generation (RAG) pipelines powered by large language models (LLMs) are gaining popularity for building end-to-end question answering systems. Frameworks such as LlamaIndex and Haystack have made significant progress in making RAG pipelines easy to use. While these frameworks provide excellent abstractions for building advanced RAG pipelines, they do so at the cost of transparency. From a user perspective, it's not readily apparent what's going on under the hood, particularly when errors or inconsistencies arise.

In this EvaDB application, we'll shed light on the inner workings of advanced RAG pipelines by examining the mechanics, limitations, and costs that often remain opaque.

<p align="center"> <img width="70%" src="images/intro.png" title="llama working on a laptop to retrieve data" > <br> <b><i>Llama working on a laptop</i> 🙂</b> </p>

Quick start

If you want to jump right in, use the following commands to run the application:

pip install -r requirements.txt

echo OPENAI_API_KEY='yourkey' > .env
python complex_qa.py

RAG Overview

Retrieval-augmented generation (RAG) is a cutting-edge AI paradigm for LLM-based question answering. A RAG pipeline typically contains:

1. Data Warehouse - A collection of data sources (e.g., documents, tables etc.) that contain information relevant to the question answering task.

2. Vector Retrieval - Given a question, find the top K most similar data chunks to the question. This is done using a vector store (e.g., Faiss).

3. Response Generation - Given the top K most similar data chunks, generate a response using a large language model (e.g. GPT-4).

RAG provides two key advantages over traditional LLM-based question answering:

1. Up-to-date information - The data warehouse can be updated in real-time, so the information is always up-to-date.

2. Source tracking - RAG provides clear traceability, enabling users to identify the sources of information, which is crucial for accuracy verification and mitigating LLM hallucinations.

Building advanced RAG Pipelines

To enable answering more complex questions, recent AI frameworks like LlamaIndex have introduced more advanced abstractions such as the Sub-question Query Engine.

In this application, we'll demystify sophisticated RAG pipelines by using the Sub-question Query Engine as an example. We'll examine the inner workings of the Sub-question Query Engine and simplify the abstractions to their core components. We'll also identify some challenges associated with advanced RAG pipelines.

The setup

A data warehouse is a collection of data sources (e.g., documents, tables etc.) that contain information relevant to the question answering task.

In this example, we'll use a simple data warehouse containing multiple Wikipedia articles for different popular cities, inspired by LlamaIndex's illustrative use-case. Each city's wiki is a separate data source. Note that for simplicity, we limit each document's size to fit within the LLM context limit.

Our goal is to build a system that can answer questions like:

1. "What is the population of Chicago?"
2. "Give me a summary of the positive aspects of Atlanta."
3. "Which city has the highest population?"

As you can see, the questions can be simple factoid/summarization questions over a single data source (Q1/Q2) or complex factoid/summarization questions over multiple data sources (Q3).

We have the following retrieval methods at our disposal:

1. vector retrieval - Given a question and a data source, generate an LLM response using the top-K most similar data chunks to the question from the data source as the context. We use the off-the-shelf FAISS vector index from EvaDB for vector retrieval. However, the concepts are applicable to any vector index.

2. summary retrieval - Given a summary question and a data source, generate an LLM response using the entire data source as context.

The secret sauce

Our key insight is that each component in an advanced RAG pipeline is powered by a single LLM call. The entire pipeline is a series of LLM calls with carefully crafted prompt templates. These prompt templates are the secret sauce that enable advanced RAG pipelines to perform complex tasks.

In fact, any advanced RAG pipeline can be broken down into a series of individual LLM calls that follow a universal input pattern:

<!-- LLM input = **Prompt Template** + **Context** + **Question** -->

where:

• Prompt Template - A curated prompt template for the specific task (e.g., sub-question generation, summarization)
• Context - The context to use to perform the task (e.g. top-K most similar data chunks)
• Question - The question to answer

Now, we illustrate this principle by examining the inner workings of the Sub-question Query Engine.

The Sub-question Query Engine has to perform three tasks:

1. Sub-question generation - Given a complex question, break it down into a set of sub-questions, while identifying the appropriate data source and retrieval function for each sub-question.
2. Vector/Summary Retrieval - For each sub-question, use the chosen retrieval function over the corresponding data source to retrieve the relevant information.
3. Response Aggregation - Aggregate the responses from the sub-questions into a final response.

Let's examine each task in detail.

Task 1: Sub-question Generation

Our goal is to break down a complex question into a set of sub-questions, while identifying the appropriate data source and retrieval function for each sub-question. For example, the question "Which city has the highest population?" is broken down into five sub-questions, one for each city, of the form "What is the population of {city}?". The data source for each sub-question has to be the corresponding city's wiki, and the retrieval function has to be vector retrieval.

At first glance, this seems like a daunting task. Specifically, we need to answer the following questions:

1. How do we know which sub-questions to generate?
2. How do we know which data source to use for each sub-question?
3. How do we know which retrieval function to use for each sub-question?

Remarkably, the answer to all three questions is the same - a single LLM call! The entire sub-question query engine is powered by a single LLM call with a carefully crafted prompt template. Let's call this template the Sub-question Prompt Template.

-- Sub-question Prompt Template --

"""
    You are an AI assistant that specializes in breaking down complex questions into simpler, manageable sub-questions.
    When presented with a complex user question, your role is to generate a list of sub-questions that, when answered, will comprehensively address the original question.
    You have at your disposal a pre-defined set of functions and data sources to utilize in answering each sub-question.
    If a user question is straightforward, your task is to return the original question, identifying the appropriate function and data source to use for its solution.
    Please remember that you are limited to the provided functions and data sources, and that each sub-question should be a full question that can be answered using a single function and a single data source.
"""

The context for the LLM call is the names of the data sources and the functions available to the system. The question is the user question. The LLM outputs a list of sub-questions, each with a function and a data source.

For the three example questions, the LLM returns the following output:

<details> <summary> LLM output Table </summary> <table> <thead> <tr> <th>Question</th> <th>Subquestions</th> <th>Retrieval method</th> <th>Data Source</th> </tr> </thead> <tbody> <tr> <td>"What is the population of Chicago?"</td> <td>"What is the population of Chicago?"</td> <td>vector retrieval</td> <td>Chicago</td> </tr> <tr> <td>"Give me a summary of the positive aspects of Atlanta."</td> <td>"Give me a summary of the positive aspects of Atlanta."</td> <td>summary retrieval</td> <td>Atlanta</td> </tr> <tr> <td rowspan=5>"Which city has the highest population?"</td> <td>"What is the population of Toronto?"</td> <td>vector retrieval</td> <td>Toronto</td> </tr> <tr> <td>"What is the population of Chicago?"</td> <td>vector retrieval</td> <td>Chicago</td> </tr> <tr> <td>"What is the population of Houston?"</td> <td>vector retrieval</td> <td>Houston</td> </tr> <tr> <td>"What is the population of Boston?"</td> <td>vector retrieval</td> <td>Boston</td> </tr> <tr> <td>"What is the population of Atlanta?"</td> <td>vector retrieval</td> <td>Atlanta</td> </tr> </tbody> </table> </details>

Task 2: Vector/Summary Retrieval

For each sub-question, we use the chosen retrieval function over the corresponding data source to retrieve the relevant information. For example, for the sub-question "What is the population of Chicago?", we use vector retrieval over the Chicago data source. Similarly, for the sub-question "Give me a summary of the positive aspects of Atlanta.", we use summary retrieval over the Atlanta data source.

For both retrieval methods, we use the same LLM prompt template. In fact, we find that the popular RAG Prompt from [LangchainHub](https://smith.langchain.