Retriever Router Query Engine

In this tutorial, we define a router query engine based on a retriever. The retriever will select a set of nodes, and we will in turn select the right QueryEngine.

NOTE: This is a beta feature. The interface for retrieving query engines may change.

Setup

# NOTE: This is ONLY necessary in jupyter notebook.
# Details: Jupyter runs an event-loop behind the scenes. 
#          This results in nested event-loops when we start an event-loop to make async queries.
#          This is normally not allowed, we use nest_asyncio to allow it for convenience.  
import nest_asyncio
nest_asyncio.apply()

import logging
import sys

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

from llama_index import (
    GPTVectorStoreIndex,
    GPTListIndex,
    SimpleDirectoryReader,
    ServiceContext,
    StorageContext
)
from llama_index.data_structs import Node

INFO:numexpr.utils:Note: NumExpr detected 12 cores but "NUMEXPR_MAX_THREADS" not set, so enforcing safe limit of 8.
Note: NumExpr detected 12 cores but "NUMEXPR_MAX_THREADS" not set, so enforcing safe limit of 8.
INFO:numexpr.utils:NumExpr defaulting to 8 threads.
NumExpr defaulting to 8 threads.

/Users/jerryliu/Programming/gpt_index/.venv/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html
  from .autonotebook import tqdm as notebook_tqdm

Load Data

We first show how to convert a Document into a set of Nodes, and insert into a DocumentStore.

# load documents
documents = SimpleDirectoryReader('../paul_graham_essay/data').load_data()

# initialize service context (set chunk size)
service_context = ServiceContext.from_defaults(chunk_size_limit=1024)
nodes = service_context.node_parser.get_nodes_from_documents(documents)

# initialize storage context (by default it's in-memory)
storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(nodes)

Define List Index and Vector Index over Same Data

list_index = GPTListIndex(nodes, storage_context=storage_context)
vector_index = GPTVectorStoreIndex(nodes, storage_context=storage_context)

INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens
> [build_index_from_nodes] Total LLM token usage: 0 tokens
INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 0 tokens
> [build_index_from_nodes] Total embedding token usage: 0 tokens
INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens
> [build_index_from_nodes] Total LLM token usage: 0 tokens
INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 20729 tokens
> [build_index_from_nodes] Total embedding token usage: 20729 tokens

Define Node/Query Engine for these Indices

We define a Node and Query Engine for each Index. We then define an outer “tool” index to store these Nodes, which can be treated as metadata.

list_index_node = Node(
    "Useful for summarization questions related to Paul Graham eassy on What I Worked On.",
    doc_id="list_index"
)
list_query_engine = list_index.as_query_engine(
    response_mode="tree_summarize", use_async=True
)
vector_index_node = Node(
    "Useful for questions around the author's education, from Paul Graham essay on What I Worked On.",
    doc_id="vector_index"
)
vector_query_engine = vector_index.as_query_engine(
    response_mode="tree_summarize", use_async=True
)

Define a Vector Index Retriever for these Nodes

Define a vector index on top of these Nodes which in turn correspond to the underlying query engines.

# create an outer "tool" index to store the underlying index information
tool_index = GPTVectorStoreIndex([list_index_node, vector_index_node])
# get retriever
tool_retriever = tool_index.as_retriever(similarity_top_k=1)

INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total LLM token usage: 0 tokens
> [build_index_from_nodes] Total LLM token usage: 0 tokens
INFO:llama_index.token_counter.token_counter:> [build_index_from_nodes] Total embedding token usage: 40 tokens
> [build_index_from_nodes] Total embedding token usage: 40 tokens

Define Router Query Engine

We define a router query engine using the vector index retriever as input. This retriever will be used to retrieve “Nodes” which contain metadata for query engines. We also take as input a function that maps a Node to a query engine.

def node_to_query_engine(node: Node):
    """Convert node to query engine."""
    # NOTE: hardcode mapping in this case
    mapping = {
        "list_index": list_query_engine,
        "vector_index": vector_query_engine
    }
    return mapping[node.get_doc_id()]

from llama_index.query_engine.router_query_engine import RetrieverRouterQueryEngine


query_engine = RetrieverRouterQueryEngine(
    tool_retriever,
    node_to_query_engine
)

response = query_engine.query('What is the summary of the document?')

print(str(response))

The document is a narrative about the author's journey from writing short stories and programming on an IBM 1401 in high school to studying philosophy in college, discovering AI and Lisp, leaving his startup Viaweb to pursue painting, and eventually working on a new Lisp language called Bel. It discusses the advantages of being independent-minded in fields affected by rapid change, and looks at the history of Y Combinator and how it has helped to create more startups than would have otherwise existed. It also examines the concept of invented versus discovered, using the example of space aliens and McCarthy's Lisp to illustrate the idea that discoveredness can be preserved.

response = query_engine.query('What did Paul Graham do during his time in college?')

print(str(response))

Paul Graham did a variety of things during his time in college. He was a PhD student in computer science, but he also worked on other projects such as writing essays, working on Y Combinator, and writing a new version of Arc. He also took art classes at Harvard and applied to art schools such as RISD and the Accademia di Belli Arti in Florence. He eventually passed the entrance exam for the Accademia and left his PhD program to pursue art.