LlamaIndex-Recursive Retriever + Query Engine

：Recursive Retriever + Query Engine

概念解释

Recursive Retriever（递归检索器）：递归检索器的概念是我们不仅探索直接最相关的节点，还探索节点关系到其他检索器/查询引擎并执行它们。例如，一个节点可能代表一个结构化表格的简洁摘要，并链接到该结构化表格的SQL/Pandas查询引擎。如果检索到该节点，我们还希望查询底层查询引擎以获取答案。

Query Engine（查询引擎）：查询引擎是执行查询的基础。它负责从索引中检索相关文档，对检索到的节点进行后处理，并将查询、相关数据和提示组合发送给LLM以返回响应。

递归检索器和查询引擎的演示

在这个演示中，我们将展示如何在具有层次关系的数据上使用“RecursiveRetriever”模块。

设置说明

我们使用 camelot 从PDF中提取基于文本的表格。

%pip install llama-index-embeddings-openai
%pip install llama-index-readers-file pymupdf
%pip install llama-index-llms-openai
%pip install llama-index-experimental
import camelot

默认设置

import os

os.environ["OPENAI_API_KEY"] = "YOUR_API_KEY"
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.core import Settings

Settings.llm = OpenAI(model="gpt-3.5-turbo")
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

加载文档（和表格）

我们使用 PyMuPDFReader 读取文档的主要文本。

我们还使用 camelot 从文档中提取一些结构化表格。

file_path = "billionaires_page.pdf"
# 初始化PDF阅读器
reader = PyMuPDFReader()
docs = reader.load(file_path)

# 使用camelot解析表格
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])

创建Pandas查询引擎

我们为每个结构化表格创建一个Pandas查询引擎。

这些可以单独执行，以回答有关每个表格的查询。

# 定义这些表格的查询引擎
llm = OpenAI(model="gpt-4")

df_query_engines = [
    PandasQueryEngine(table_df, llm=llm) 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))

response = df_query_engines[1].query(
    "How many billionaires were there in 2009?"
)
print(str(response))

构建向量索引

在分块文档以及链接到表格的额外IndexNode对象上构建向量索引。

from llama_index.core import Settings

doc_nodes = Settings.node_parser.get_nodes_from_documents(docs)

# 定义索引节点
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)
}

# 构建顶级向量索引 + 查询引擎
vector_index = VectorStoreIndex(doc_nodes + df_nodes)
vector_retriever = vector_index.as_retriever(similarity_top_k=1)

在RetrieverQueryEngine中使用RecursiveRetriever

我们定义一个RecursiveRetriever对象来递归地检索/查询节点。然后将其放入RetrieverQueryEngine中，并使用ResponseSynthesizer来合成响应。

我们传入从id到检索器和id到查询引擎的映射。然后传入一个代表我们首先查询的检索器的根id。

# 基准向量索引（不包括额外的df节点）。
# 用于基准测试
vector_index0 = VectorStoreIndex(doc_nodes)
vector_query_engine0 = vector_index0.as_query_engine()

from llama_index.core.retrievers import RecursiveRetriever
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core 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(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?"
)
print(str(response))

response = query_engine.query("How many billionaires were there in 2009?")
print(str(response))

response = vector_query_engine0.query(
    "How many billionaires were there in 2009?"
)
print(response.source_nodes[0].node.get_content())
print(str(response))

response = query_engine.query(
    "Which billionaires are excluded from this list?"
)
print(str(response))

详细解释代码的具体含义

这段代码展示了如何使用LlamaIndex的RecursiveRetriever和QueryEngine来处理具有层次关系的数据，特别是如何递归地检索和查询节点。以下是代码的详细解释：

1. 基准向量索引（不包括额外的df节点）：

   vector_index0 = VectorStoreIndex(doc_nodes)
   vector_query_engine0 = vector_index0.as_query_engine()
   

   这部分代码创建了一个基准向量索引，该索引仅包含文档节点（doc_nodes），不包括额外的表格节点（df_nodes）。这个基准索引用于后续的基准测试。

2. 导入必要的模块：

   from llama_index.core.retrievers import RecursiveRetriever
   from llama_index.core.query_engine import RetrieverQueryEngine
   from llama_index.core import get_response_synthesizer
   

   这些导入语句引入了RecursiveRetriever、RetrieverQueryEngine和get_response_synthesizer模块，这些模块将在后续代码中使用。

3. 创建RecursiveRetriever对象：

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

   这段代码创建了一个RecursiveRetriever对象，该对象使用向量检索器（vector_retriever）来递归地检索节点。query_engine_dict参数指定了从节点ID到查询引擎的映射，verbose=True表示在检索过程中输出详细信息。

4. 创建ResponseSynthesizer对象：

   response_synthesizer = get_response_synthesizer(response_mode="compact")
   

   这段代码创建了一个ResponseSynthesizer对象，用于合成最终的响应。response_mode="compact"表示在每次LLM调用中“压缩”提示，尽可能多地填充节点文本块。

5. 创建RetrieverQueryEngine对象：

   query_engine = RetrieverQueryEngine.from_args(
       recursive_retriever, response_synthesizer=response_synthesizer
   )
   

   这段代码创建了一个RetrieverQueryEngine对象，该对象使用之前创建的RecursiveRetriever和ResponseSynthesizer来执行查询。

6. 执行查询：

   response = query_engine.query(
       "What's the net worth of the second richest billionaire in 2023?"
   )
   print(str(response))
   
   response = query_engine.query("How many billionaires were there in 2009?")
   print(str(response))
   
   response = vector_query_engine0.query(
       "How many billionaires were there in 2009?"
   )
   print(response.source_nodes[0].node.get_content())
   print(str(response))
   
   response = query_engine.query(
       "Which billionaires are excluded from this list?"
   )
   print(str(response))
   

   这部分代码展示了如何使用查询引擎执行不同的查询，并打印查询结果。前两个查询使用RecursiveRetriever和RetrieverQueryEngine来获取答案，最后一个查询使用基准向量索引（vector_query_engine0）来获取答案，以进行比较。

使用方法

1. 加载文档和表格：

   • 使用PyMuPDFReader读取PDF文档的主要文本。
   • 使用camelot从PDF中提取结构化表格。

2. 创建Pandas查询引擎：

   • 为每个结构化表格创建一个Pandas查询引擎，以便单独执行查询。

3. 构建向量索引：

   • 在分块文档以及链接到表格的额外IndexNode对象上构建向量索引。

4. 创建RecursiveRetriever和RetrieverQueryEngine：

   • 定义一个RecursiveRetriever对象来递归地检索/查询节点。
   • 将其放入RetrieverQueryEngine中，并使用ResponseSynthesizer来合成响应。

5. 执行查询：

   • 使用查询引擎执行不同的查询，并打印查询结果。

总结

通过本课程，我们详细讲解了Recursive Retriever和Query Engine的概念及其在LlamaIndex中的应用。我们介绍了如何加载文档和表格、创建Pandas查询引擎、构建向量索引，以及如何在RetrieverQueryEngine中使用RecursiveRetriever。这些内容将帮助学生更好地理解和应用LlamaIndex中的递归检索和查询功能。