Google Gemini 1M 长下文窗口 + LlamaIndex agent 处理复杂、多样化文档，精准解答

如何有效地从海量异构文档中提取准确的答案是一项重大挑战。

今天，介绍Google Gemini长上下文窗口的演示项目，展示其在处理复杂问题时的强大能力。

https://colab.research.google.com/drive/1GfAZ4dvoqFuorMCS1Lzhnsp-IU4m2gu6#scrollTo=mPByQN97lEt2

https://x.com/llama_index/status/1798049438814081138

Demo简介

Demo使用的数据是一系列复杂的 PDF，详细介绍了旧金山 2016 年至 2023 年的预算。这是一项艰巨的任务——每个 PDF 的格式都不同，每个 PDF 都引用前后几年进行比较，并且有一个缺少必须推断预算数字的年份。

合并后的 PDF 也非常庞大——即使 Gemini 1.5 Pro 有 1M 词元上下文窗口，也不能一次性发送所有文档的全文。非常适合这次测试。

在这个Colab笔记本中，Demo展示了内置于LlamaIndex代理中的Gemini，如何从一组复杂、多样化文档中回答由多部分组成的问题。通过长上下文窗口，Gemini在解决复杂问题时表现出色。

这不仅展示了RAG（Retrieval-Augmented Generation）在缩小海量数据集搜索范围方面的价值，还突显了长上下文窗口在回答复杂问题时的能力，以及代理如何通过反思和多次查询来得出复杂问题的答案。

代码详解

安装依赖项

!pip install llama-index-core
!pip install llama-index-llms-gemini
!pip install llama-index-embeddings-huggingface
!pip install llama_index.readers.file

通过pip命令安装所需的Python包，包括llama-index-core、llama-index-llms-gemini、llama-index-embeddings-huggingface和llama_index.readers.file，为后续的文档处理和查询提供支持。

导入库

from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.callbacks import CallbackManager, TokenCountingHandler
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.agent import ReActAgent
from llama_index.llms.gemini import Gemini
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from google.colab import userdata
import os

导入处理文档、回调管理、工具元数据、反应式代理、语言模型（Gemini）和嵌入模型（HuggingFace）的相关库和模块。

导入数据

!mkdir data
!wget "https://www.dropbox.com/scl/fi/xt3squt47djba0j7emmjb/2016-CSF_Budget_Book_2016_FINAL_WEB_with-cover-page.pdf?
.....

创建一个名为data的目录，并使用wget命令下载了多个旧金山预算文件到该目录中，为后续的数据处理和查询准备数据。

设置 Token 计数

token_counter = TokenCountingHandler(
    verbose=True
)

Settings.callback_manager = CallbackManager([token_counter])

创建了一个TokenCountingHandler对象以记录和报告token使用情况，并将其添加到回调管理器中，以便在模型处理查询时跟踪token使用情况。

嵌入文档

Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-en-v1.5")

设置嵌入模型为HuggingFaceEmbedding，使用模型BAAI/bge-small-en-v1.5，将文档转化为向量以便进行高效的相似性检索。

设置块大小和重叠

Settings.chunk_size = 512
Settings.chunk_overlap = 50

documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(documents)

设置数据块的大小为512个token，重叠为50个token。然后使用SimpleDirectoryReader从./data目录中加载文档，并创建一个向量存储索引，以便进行快速检索。

简单来说，嵌入模型将文本转换为数字向量，而向量存储技术则利用这些向量实现高效的相似性搜索，两者结合可以帮助我们从海量文本数据中快速找到所需的信息。

初始化 Gemini

os.environ["GOOGLE_API_KEY"] = userdata.get('google-api-key')
Settings.llm = Gemini(
    model_name="models/gemini-1.5-pro-latest",
    temperature=0.2
)

设置环境变量GOOGLE_API_KEY，并初始化Gemini模型gemini-1.5-pro-latest，温度参数设置为0.2，用于生成更加准确和可靠的答案。

创建代理

使用10个上下文（similarity_top_k=10）

query_engine_10 = index.as_query_engine(similarity_top_k=10)
query_engine_tools = [
    QueryEngineTool(
        query_engine=query_engine_10,
        metadata=ToolMetadata(
            name="sf_budgets",
            description=(
                "Has information about the budget of San Francisco, with documents for every year from 2016 to 2023."
            ),
        ),
    ),
]

agent_10 = ReActAgent.from_tools(
    query_engine_tools,
    verbose=True,
    max_iterations=100
)

response = agent_10.chat("

What was the budget of San Francisco for each fiscal year from 2016 to 2023?")
print(str(response))

创建一个查询引擎query_engine_10，设置相似度检索的前10个结果。然后创建一个反应式代理agent_10，让它尝试回答用户的问题。由于检索的上下文数量有限，无法涵盖所有需要的信息，导致模型无法给出完整的回答。

反应式代理（ReActAgent）的工作原理

它通过以下几个步骤来回答问题：

1. 接受输入：代理接受用户的问题作为输入。
2. 查询检索：代理使用查询引擎工具，从文档索引中检索最相关的文档片段（上下文）。
3. 生成回答：代理根据检索到的上下文信息，使用语言模型（如Gemini）生成回答。
4. 迭代优化：如果初次回答不完整或不准确，代理可以根据反馈进行多次迭代查询和回答，直到达到预设的最大迭代次数。

使用100个上下文（similarity_top_k=100）

query_engine_100 = index.as_query_engine(similarity_top_k=100)
query_engine_tools = [
    QueryEngineTool(
        query_engine=query_engine_100,
        metadata=ToolMetadata(
            name="sf_budgets",
            description=(
                "Has information about the budget of San Francisco, with documents for every year from 2016 to 2023."
            ),
        ),
    ),
]

agent_100 = ReActAgent.from_tools(
    query_engine_tools,
    verbose=True,
    max_iterations=100
)

response = agent_100.chat("What was the budget of San Francisco for each fiscal year from 2016 to 2023?")
print(str(response))

增加上下文检索的数量到100个。创建新的查询引擎和代理，重复之前的操作。这次，模型能够获取更多的上下文信息，从而更好地回答问题。然而，某些关键信息仍然缺失，回答仍然不完整。

使用1000个上下文（similarity_top_k=1000）

query_engine_1000 = index.as_query_engine(similarity_top_k=1000)
query_engine_tools = [
    QueryEngineTool(
        query_engine=query_engine_1000,
        metadata=ToolMetadata(
            name="sf_budgets",
            description=(
                "Has information about the budget of San Francisco, with documents for every year from 2016 to 2023."
            ),
        ),
    ),
]

agent_1000 = ReActAgent.from_tools(
    query_engine_tools,
    verbose=True,
    max_iterations=100
)

response = agent_1000.chat("What was the budget of San Francisco for each fiscal year from 2016 to 2023?")
print(str(response))

进一步增加上下文检索的数量到1000个。再次创建新的查询引擎和代理，进行相同的查询。由于检索到的上下文信息更全面，模型有更大的概率找到所需的所有关键信息，从而提供更完整和准确的回答。

这次项目的检索器提供了436,396个Token的上下文——接近百万上下文窗口的一半——现在agent 可以正常工作了！它能够提供每年的预算数字。

值得注意的是，agent 策略在这里起到了作用——它在第一次尝试回答问题时没有得到答案，所以换了一种提问方式，在第二次尝试时得到了答案。

项目核心展示点

长上下文窗口的优势

利用长上下文窗口处理复杂查询，提高答案的准确性和完整性。

反应式代理的应用

反应式代理能够通过多次迭代和查询，不断优化回答的过程，展示了其在实际应用中的潜力。

文档嵌入和检索技术

详细展示了如何使用嵌入模型和向量存储技术，提高大规模文档的检索效率。

限制和潜在改进

尽管本项目有效地展示了 Google Gemini 长上下文窗口在处理复杂问题方面的强大能力，但仍有一些限制和潜在的改进空间:

• 计算成本: 处理大型上下文窗口需要更大的计算资源和更长的处理时间。随着上下文窗口的增长，计算成本会迅速增加，这可能会限制其在实时应用或资源受限环境中的实用性。

  在这个项目里，有网友询问花费，Llamaindex 给到了回答，可以看出来花费不少。

  

• 无关信息的引入: 更大的上下文窗口可能包含与查询无关的信息。这可能会导致模型产生偏差，或降低其识别和关注真正相关信息的能力。需要更智能的上下文窗口管理策略来筛选和过滤无关信息。

结语

对开发人员、研究人员和企业技术团队来说，这个项目具有参考价值。

更大的上下文窗口允许模型从文档中访问更多相关信息，减少遗漏关键细节的可能性。

它使模型能够更好地理解文档中不同信息片段之间的关系，从而产生更准确和全面的答案。

欢迎在评论区留言，让我们一起交流进步。