Recursive Retriever + Query Engine Demo#

In this demo, we walk through a use case of showcasing our “RecursiveRetriever” module over hierarchical data.

The concept of recursive retrieval is that we not only explore the directly most relevant nodes, but also explore node relationships to additional retrievers/query engines and execute them. For instance, a node may represent a concise summary of a structured table, and link to a SQL/Pandas query engine over that structured table. Then if the node is retrieved, we want to also query the underlying query engine for the answer.

This can be especially useful for documents with hierarchical relationships. In this example, we walk through a Wikipedia article about billionaires (in PDF form), which contains both text and a variety of embedded structured tables. We first create a Pandas query engine over each table, but also represent each table by an IndexNode (stores a link to the query engine); this Node is stored along with other Nodes in a vector store.

During query-time, if an IndexNode is fetched, then the underlying query engine/retriever will be queried.

Notes about Setup

We use camelot to extract text-based tables from PDFs.

import camelot
from llama_index import Document, SummaryIndex

# https://en.wikipedia.org/wiki/The_World%27s_Billionaires
from llama_index import VectorStoreIndex, ServiceContext
from llama_index.query_engine import PandasQueryEngine, RetrieverQueryEngine
from llama_index.retrievers import RecursiveRetriever
from llama_index.schema import IndexNode
from llama_index.llms import OpenAI

from llama_hub.file.pymu_pdf.base import PyMuPDFReader
from pathlib import Path
from typing import List

Load in Document (and Tables)#

We use our PyMuPDFReader to read in the main text of the document.

We also use camelot to extract some structured tables from the document

file_path = "billionaires_page.pdf"
# initialize PDF reader
reader = PyMuPDFReader()
docs = reader.load(file_path)
# use camelot to parse tables
def get_tables(path: str, pages: List[int]):
    table_dfs = []
    for page in pages:
        table_list = camelot.read_pdf(path, pages=str(page))
        table_df = table_list[0].df
        table_df = (
            table_df.rename(columns=table_df.iloc[0])
            .drop(table_df.index[0])
            .reset_index(drop=True)
        )
        table_dfs.append(table_df)
    return table_dfs
table_dfs = get_tables(file_path, pages=[3, 25])
# shows list of top billionaires in 2023
table_dfs[0]
No. Name Net worth\n(USD) Age Nationality Primary source(s) of wealth
0 1 Bernard Arnault &\nfamily $211 billion 74 France LVMH
1 2 Elon Musk $180 billion 51 United\nStates Tesla, SpaceX, X Corp.
2 3 Jeff Bezos $114 billion 59 United\nStates Amazon
3 4 Larry Ellison $107 billion 78 United\nStates Oracle Corporation
4 5 Warren Buffett $106 billion 92 United\nStates Berkshire Hathaway
5 6 Bill Gates $104 billion 67 United\nStates Microsoft
6 7 Michael Bloomberg $94.5 billion 81 United\nStates Bloomberg L.P.
7 8 Carlos Slim & family $93 billion 83 Mexico Telmex, América Móvil, Grupo\nCarso
8 9 Mukesh Ambani $83.4 billion 65 India Reliance Industries
9 10 Steve Ballmer $80.7 billion 67 United\nStates Microsoft
# shows list of top billionaires
table_dfs[1]
Year Number of billionaires Group's combined net worth
0 2023[2] 2,640 $12.2 trillion
1 2022[6] 2,668 $12.7 trillion
2 2021[11] 2,755 $13.1 trillion
3 2020 2,095 $8.0 trillion
4 2019 2,153 $8.7 trillion
5 2018 2,208 $9.1 trillion
6 2017 2,043 $7.7 trillion
7 2016 1,810 $6.5 trillion
8 2015[18] 1,826 $7.1 trillion
9 2014[67] 1,645 $6.4 trillion
10 2013[68] 1,426 $5.4 trillion
11 2012 1,226 $4.6 trillion
12 2011 1,210 $4.5 trillion
13 2010 1,011 $3.6 trillion
14 2009 793 $2.4 trillion
15 2008 1,125 $4.4 trillion
16 2007 946 $3.5 trillion
17 2006 793 $2.6 trillion
18 2005 691 $2.2 trillion
19 2004 587 $1.9 trillion
20 2003 476 $1.4 trillion
21 2002 497 $1.5 trillion
22 2001 538 $1.8 trillion
23 2000 470 $898 billion
24 Sources: Forbes.[18][67][66][68]

Create Pandas Query Engines#

We create a pandas query engine over each structured table.

These can be executed on their own to answer queries about each table.

# define query engines over these tables
llm = OpenAI(model="gpt-4")

service_context = ServiceContext.from_defaults(llm=llm)
df_query_engines = [
    PandasQueryEngine(table_df, service_context=service_context)
    for table_df in table_dfs
]
response = df_query_engines[0].query(
    "What's the net worth of the second richest billionaire in 2023?"
)
print(str(response))
$180 billion
response = df_query_engines[1].query(
    "How many billionaires were there in 2009?"
)
print(str(response))
793

Build Vector Index#

Build vector index over the chunked document as well as over the additional IndexNode objects linked to the tables.

llm = OpenAI(temperature=0, model="gpt-4")

service_context = ServiceContext.from_defaults(
    llm=llm,
)
doc_nodes = service_context.node_parser.get_nodes_from_documents(docs)
# define index nodes
summaries = [
    (
        "This node provides information about the world's richest billionaires"
        " in 2023"
    ),
    (
        "This node provides information on the number of billionaires and"
        " their combined net worth from 2000 to 2023."
    ),
]

df_nodes = [
    IndexNode(text=summary, index_id=f"pandas{idx}")
    for idx, summary in enumerate(summaries)
]

df_id_query_engine_mapping = {
    f"pandas{idx}": df_query_engine
    for idx, df_query_engine in enumerate(df_query_engines)
}
# construct top-level vector index + query engine
vector_index = VectorStoreIndex(doc_nodes + df_nodes)
vector_retriever = vector_index.as_retriever(similarity_top_k=1)

Use RecursiveRetriever in our RetrieverQueryEngine#

We define a RecursiveRetriever object to recursively retrieve/query nodes. We then put this in our RetrieverQueryEngine along with a ResponseSynthesizer to synthesize a response.

We pass in mappings from id to retriever and id to query engine. We then pass in a root id representing the retriever we query first.

# baseline vector index (that doesn't include the extra df nodes).
# used to benchmark
vector_index0 = VectorStoreIndex(doc_nodes)
vector_query_engine0 = vector_index0.as_query_engine()
from llama_index.retrievers import RecursiveRetriever
from llama_index.query_engine import RetrieverQueryEngine
from llama_index.response_synthesizers import get_response_synthesizer

recursive_retriever = RecursiveRetriever(
    "vector",
    retriever_dict={"vector": vector_retriever},
    query_engine_dict=df_id_query_engine_mapping,
    verbose=True,
)

response_synthesizer = get_response_synthesizer(
    # service_context=service_context,
    response_mode="compact"
)

query_engine = RetrieverQueryEngine.from_args(
    recursive_retriever, response_synthesizer=response_synthesizer
)
response = query_engine.query(
    "What's the net worth of the second richest billionaire in 2023?"
)
Retrieving with query id None: What's the net worth of the second richest billionaire in 2023?
Retrieved node with id, entering: pandas0
Retrieving with query id pandas0: What's the net worth of the second richest billionaire in 2023?
Got response: $180 billion

response.source_nodes[0].node.get_content()
"Query: What's the net worth of the second richest billionaire in 2023?\nResponse: $180\xa0billion"
str(response)
'$180 billion.'
response = query_engine.query("How many billionaires were there in 2009?")
Retrieving with query id None: How many billionaires were there in 2009?
Retrieved node with id, entering: pandas1
Retrieving with query id pandas1: How many billionaires were there in 2009?
Got response: 793

str(response)
'793'
response = vector_query_engine0.query(
    "How many billionaires were there in 2009?"
)
print(response.source_nodes[0].node.get_content())
print(str(response))
Based on the context information, it is not possible to determine the exact number of billionaires in 2009. The provided information only mentions the number of billionaires in 2013 and 2014.
response.source_nodes[0].node.get_content()
response = query_engine.query(
    "Which billionaires are excluded from this list?"
)
print(str(response))
Royal families and dictators whose wealth is contingent on a position are excluded from this list.