Sub Question Query Engine

In this tutorial, we showcase how to use a sub question query engine to tackle the problem of answering a complex query using multiple data sources.
It first breaks down the complex query into sub questions for each relevant data source, then gather all the intermediate reponses and synthesizes a final response.

Preparation

If you’re opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.

!pip install llama-index

import os
import openai

os.environ["OPENAI_API_KEY"] = "sk-..."
openai.api_key = os.environ["OPENAI_API_KEY"]

import nest_asyncio

nest_asyncio.apply()

from llama_index import VectorStoreIndex, SimpleDirectoryReader
from llama_index.tools import QueryEngineTool, ToolMetadata
from llama_index.query_engine import SubQuestionQueryEngine
from llama_index.callbacks import CallbackManager, LlamaDebugHandler
from llama_index import ServiceContext

# Using the LlamaDebugHandler to print the trace of the sub questions
# captured by the SUB_QUESTION callback event type
llama_debug = LlamaDebugHandler(print_trace_on_end=True)
callback_manager = CallbackManager([llama_debug])
service_context = ServiceContext.from_defaults(
    callback_manager=callback_manager
)

Download Data

!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'

# load data
pg_essay = SimpleDirectoryReader(input_dir="./data/paul_graham/").load_data()

# build index and query engine
vector_query_engine = VectorStoreIndex.from_documents(
    pg_essay, use_async=True, service_context=service_context
).as_query_engine()

**********
Trace: index_construction
    |_node_parsing ->  0.394271 seconds
      |_chunking ->  0.393344 seconds
    |_embedding ->  0.753133 seconds
    |_embedding ->  0.749828 seconds
**********

Setup sub question query engine

# setup base query engine as tool
query_engine_tools = [
    QueryEngineTool(
        query_engine=vector_query_engine,
        metadata=ToolMetadata(
            name="pg_essay",
            description="Paul Graham essay on What I Worked On",
        ),
    ),
]

query_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=query_engine_tools,
    service_context=service_context,
    use_async=True,
)

Run queries

response = query_engine.query(
    "How was Paul Grahams life different before, during, and after YC?"
)

Generated 3 sub questions.
[pg_essay] Q: What did Paul Graham do before YC?
[pg_essay] Q: What did Paul Graham do during YC?
[pg_essay] Q: What did Paul Graham do after YC?
[pg_essay] A: 
Before YC, Paul Graham was a hacker, writer, and worked on Arc, a programming language. He also wrote essays and worked on other projects.
[pg_essay] A: 
Paul Graham stopped working on YC in March 2014 and began painting. He spent most of the rest of the year painting and then in November he ran out of steam and stopped. He then began writing essays again and in March 2015 he started working on Lisp again.
[pg_essay] A: 
Paul Graham worked on YC in a variety of ways. He wrote essays, worked on internal software in Arc, and created Hacker News. He also helped select and support founders, dealt with disputes between cofounders, and fought with people who maltreated the startups. He worked hard even at the parts he didn't like, and was determined to make YC a success. In 2010, he was offered unsolicited advice to make sure YC wasn't the last cool thing he did, which set him thinking about his future. In 2012, he decided to hand YC over to someone else and recruited Sam Altman to take over. He worked on YC until March 2014, when his mother passed away, and then he checked out completely.
**********
Trace: query
    |_query ->  13.064431 seconds
      |_llm ->  2.499768 seconds
      |_sub_question ->  2.05934 seconds
        |_query ->  2.059142 seconds
          |_retrieve ->  0.278184 seconds
            |_embedding ->  0.274593 seconds
          |_synthesize ->  1.780895 seconds
            |_llm ->  1.740488 seconds
      |_sub_question ->  5.364061 seconds
        |_query ->  5.363695 seconds
          |_retrieve ->  0.230257 seconds
            |_embedding ->  0.226763 seconds
          |_synthesize ->  5.133343 seconds
            |_llm ->  5.091069 seconds
      |_sub_question ->  2.148964 seconds
        |_query ->  2.14889 seconds
          |_retrieve ->  0.323438 seconds
            |_embedding ->  0.319841 seconds
          |_synthesize ->  1.825401 seconds
            |_llm ->  1.783064 seconds
      |_synthesize ->  5.198214 seconds
        |_llm ->  5.175849 seconds
**********

print(response)

Before YC, Paul Graham was a hacker, writer, and worked on Arc, a programming language. During YC, he wrote essays, worked on internal software in Arc, and created Hacker News. He also helped select and support founders, dealt with disputes between cofounders, and fought with people who maltreated the startups. After YC, Paul Graham stopped working on YC and began painting. He then began writing essays again and in March 2015 he started working on Lisp again. Paul Graham's life was different before, during, and after YC in that he changed his focus from programming and writing to painting and then back to programming and writing.

# iterate through sub_question items captured in SUB_QUESTION event
from llama_index.callbacks.schema import CBEventType, EventPayload

for i, (start_event, end_event) in enumerate(
    llama_debug.get_event_pairs(CBEventType.SUB_QUESTION)
):
    qa_pair = end_event.payload[EventPayload.SUB_QUESTION]
    print("Sub Question " + str(i) + ": " + qa_pair.sub_q.sub_question.strip())
    print("Answer: " + qa_pair.answer.strip())
    print("====================================")

Sub Question 0: What did Paul Graham do before YC?
Answer: Before YC, Paul Graham was a hacker, writer, and worked on Arc, a programming language. He also wrote essays and worked on other projects.
====================================
Sub Question 1: What did Paul Graham do during YC?
Answer: Paul Graham worked on YC in a variety of ways. He wrote essays, worked on internal software in Arc, and created Hacker News. He also helped select and support founders, dealt with disputes between cofounders, and fought with people who maltreated the startups. He worked hard even at the parts he didn't like, and was determined to make YC a success. In 2010, he was offered unsolicited advice to make sure YC wasn't the last cool thing he did, which set him thinking about his future. In 2012, he decided to hand YC over to someone else and recruited Sam Altman to take over. He worked on YC until March 2014, when his mother passed away, and then he checked out completely.
====================================
Sub Question 2: What did Paul Graham do after YC?
Answer: Paul Graham stopped working on YC in March 2014 and began painting. He spent most of the rest of the year painting and then in November he ran out of steam and stopped. He then began writing essays again and in March 2015 he started working on Lisp again.
====================================