88 递归检索器与查询引擎：探索分层数据的强大工具

递归检索器与查询引擎：探索分层数据的强大工具

在编程的世界里，处理分层数据是一项常见但复杂的任务。今天，我们将深入探讨一种名为“递归检索器（RecursiveRetriever）”的模块，它能够帮助我们更高效地处理和查询分层数据。

动机

递归检索的概念不仅仅局限于直接检索最相关的节点，还包括探索节点之间的关系，以进一步检索其他检索器或查询引擎并执行它们。例如，一个节点可能代表一个结构化表格的简洁摘要，并链接到该结构化表格上的SQL/Pandas查询引擎。如果检索到该节点，我们还希望查询底层查询引擎以获取答案。

这种技术对于具有分层关系的文档尤其有用。在本例中，我们将通过一个关于亿万富翁的Wikipedia文章（PDF格式）来演示，该文章包含文本和各种嵌入的结构化表格。我们首先为每个表格创建一个Pandas查询引擎，但也通过一个IndexNode表示每个表格（存储链接到查询引擎）；这个节点与其他节点一起存储在向量存储中。

在查询时，如果检索到IndexNode，则将查询底层查询引擎/检索器。

设置说明

我们使用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])
# 显示2023年顶级亿万富翁列表
table_dfs[0]
# 显示顶级亿万富翁列表
table_dfs[1]

创建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[