SubQuestionQueryEngine

在实际应用中，我们经常会遇到需要跨越多个文档来回答复杂查询的场景。

在本 Notebook 中，我们将深入探讨如何通过将复杂查询分解为更简单的子查询，并利用 SubQuestionQueryEngine 来生成答案，从而解决跨越多个文档的复杂查询问题。

安装

!pip install llama-index
!pip install llama-index-llms-anthropic
!pip install llama-index-embeddings-huggingface

设置 API 密钥

import os
os.environ['ANTHROPIC_API_KEY'] = 'YOUR ANTHROPIC API KEY'

设置 LLM 和 Embedding 模型

我们将使用 Anthropic 最新发布的 Claude-3 Opus LLM。

from llama_index.llms.anthropic import Anthropic
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

llm = Anthropic(temperature=0.0, model='claude-3-opus-20240229')
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-base-en-v1.5")

from llama_index.core import Settings
Settings.llm = llm
Settings.embed_model = embed_model
Settings.chunk_size = 512

设置日志

# 注意：这仅在 jupyter notebook 中是必需的。
# 详情：Jupyter 在后台运行一个事件循环。
#          这会导致我们在启动事件循环以进行异步查询时出现嵌套事件循环。
#          这通常是不允许的，我们使用 nest_asyncio 来方便地允许它。
import nest_asyncio

nest_asyncio.apply()

import logging
import sys

# 设置根记录器
logger = logging.getLogger()
logger.setLevel(logging.INFO)  # 将记录器级别设置为 INFO

# 清除所有现有的处理程序
logger.handlers = []

# 设置 StreamHandler 以输出到 sys.stdout (Colab 的输出)
handler = logging.StreamHandler(sys.stdout)
handler.setLevel(logging.INFO)  # 将处理程序级别设置为 INFO

# 将处理程序添加到记录器
logger.addHandler(handler)

from IPython.display import display, HTML

下载数据

我们将使用 Uber 和 Lyft 2021 年的 10K SEC 文件。

!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/uber_2021.pdf' -O './uber_2021.pdf'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/lyft_2021.pdf' -O './lyft_2021.pdf'

--2024-03-08 07:07:32--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/uber_2021.pdf
正在解析 raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.109.133, 185.199.108.133, 185.199.110.133, ...
正在连接 raw.githubusercontent.com (raw.githubusercontent.com)|185.199.109.133|:443... 已连接。
已发送 HTTP 请求，等待响应... 200 OK
文件大小：1880483 (1.8M) [application/octet-stream]
正在保存到 ‘./uber_2021.pdf’

./uber_2021.pdf     100%[===================>]   1.79M  --.-KB/s    在 0.02s 内

2024-03-08 07:07:32 (87.4 MB/s) - ‘./uber_2021.pdf’ 已保存 [1880483/1880483]

--2024-03-08 07:07:33--  https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/10k/lyft_2021.pdf
正在解析 raw.githubusercontent.com (raw.githubusercontent.com)... 185.199.110.133, 185.199.111.133, 185.199.109.133, ...
正在连接 raw.githubusercontent.com (raw.githubusercontent.com)|185.199.110.133|:443... 已连接。
已发送 HTTP 请求，等待响应... 200 OK
文件大小：1440303 (1.4M) [application/octet-stream]
正在保存到 ‘./lyft_2021.pdf’

./lyft_2021.pdf     100%[===================>]   1.37M  --.-KB/s    在 0.02s 内

2024-03-08 07:07:33 (74.9 MB/s) - ‘./lyft_2021.pdf’ 已保存 [1440303/1440303]

加载数据

from llama_index.core import SimpleDirectoryReader
lyft_docs = SimpleDirectoryReader(input_files=["lyft_2021.pdf"]).load_data()
uber_docs = SimpleDirectoryReader(input_files=["uber_2021.pdf"]).load_data()

print(f'已加载 Lyft 10-K 文件，包含 {len(lyft_docs)} 页')
print(f'已加载 Uber 10-K 文件，包含 {len(uber_docs)} 页')

已加载 Lyft 10-K 文件，包含 238 页
已加载 Uber 10-K 文件，包含 307 页

索引数据

from llama_index.core import VectorStoreIndex
lyft_index = VectorStoreIndex.from_documents(lyft_docs[:100])
uber_index = VectorStoreIndex.from_documents(uber_docs[:100])

创建查询引擎

lyft_engine = lyft_index.as_query_engine(similarity_top_k=5)

uber_engine = uber_index.as_query_engine(similarity_top_k=5)

查询

response = await lyft_engine.aquery('What is the revenue of Lyft in 2021? Answer in millions with page reference')
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

HTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"

根据第 79 页的合并运营报表，Lyft 在截至 2021 年 12 月 31 日的财年总收入为 32.083 亿美元。

response = await uber_engine.aquery('What is the revenue of Uber in 2021? Answer in millions, with page reference')
display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

HTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"

根据第 77 页的合并运营报表，Uber 在截至 2021 年 12 月 31 日的财年收入为 174.55 亿美元。

创建工具

from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.query_engine import SubQuestionQueryEngine

query_engine_tools = [
    QueryEngineTool(
        query_engine=lyft_engine,
        metadata=ToolMetadata(name='lyft_10k', description='提供 2021 年 Lyft 财务信息')
    ),
    QueryEngineTool(
        query_engine=uber_engine,
        metadata=ToolMetadata(name='uber_10k', description='提供 2021 年 Uber 财务信息')
    ),
]

创建 SubQuestionQueryEngine

sub_question_query_engine = SubQuestionQueryEngine.from_defaults(query_engine_tools=query_engine_tools)

查询

response = await sub_question_query_engine.aquery('Compare revenue growth of Uber and Lyft from 2020 to 2021')

HTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"
生成了 4 个子问题。
 [1;3;38;2;237;90;200m[uber_10k] Q: What was Uber's revenue in 2020?
 [0m [1;3;38;2;90;149;237m[uber_10k] Q: What was Uber's revenue in 2021?
 [0m [1;3;38;2;11;159;203m[lyft_10k] Q: What was Lyft's revenue in 2020?
 [0m [1;3;38;2;155;135;227m[lyft_10k] Q: What was Lyft's revenue in 2021?
 [0mHTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"
 [1;3;38;2;11;159;203m[lyft_10k] A: 根据 Lyft 的合并运营报表数据，Lyft 在 2020 年的总收入为 2,364,681,000 美元。
 [0mHTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"
 [1;3;38;2;90;149;237m[uber_10k] A: 根据 Uber 的合并运营报表，Uber 在 2021 年的收入为 174.55 亿美元。
 [0mHTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"
 [1;3;38;2;237;90;200m[uber_10k] A: 根据 Uber 的合并运营报表，Uber 在 2020 年的收入为 111.39 亿美元。
 [0mHTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"
 [1;3;38;2;155;135;227m[lyft_10k] A: 根据 Lyft 的合并运营报表，Lyft 在截至 2021 年 12 月 31 日的财年总收入为 3,208,323,000 美元。这包括：

- 根据 ASC 606 的客户合同收入为 2,957,979,000 美元
- 根据 ASC 842 的租赁收入为 250,344,000 美元

因此，Lyft 在截至 2021 年 12 月 31 日的财年总收入为 3,208,323,000 美元。
 [0mHTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"

display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

从 2020 年到 2021 年，Uber 和 Lyft 的收入都实现了显著增长： Uber 的收入从 2020 年的 111.39 亿美元增长到 2021 年的 174.55 亿美元，同比增长约 56.7%。 Lyft 的总收入从 2020 年的 2,364,681,000 美元增长到 2021 年的 3,208,323,000 美元，增长了约 35.7%。 因此，虽然 Lyft 的收入基数较低，但 Uber 在 2020 年至 2021 年间的收入增长百分比更高（分别为 56.7% 和 35.7%）。 由于 2020 年新冠疫情对两家公司的叫车业务产生了重大影响，两家公司在 2021 年都出现了强劲反弹。但 Uber 的收入增长速度更快，达到了 175 亿美元，而 Lyft 为 32 亿美元。

response = await sub_question_query_engine.aquery('Compare the investments made by Uber and Lyft')

HTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"
生成了 4 个子问题。
 [1;3;38;2;237;90;200m[uber_10k] Q: What investments did Uber make in 2021
 [0m [1;3;38;2;90;149;237m[uber_10k] Q: What was the total amount invested by Uber in 2021
 [0m [1;3;38;2;11;159;203m[lyft_10k] Q: What investments did Lyft make in 2021
 [0m [1;3;38;2;155;135;227m[lyft_10k] Q: What was the total amount invested by Lyft in 2021
 [0mHTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"
 [1;3;38;2;90;149;237m[uber_10k] A: 根据提供的内容，Uber 在 2021 年进行了以下投资：

- 收购业务 23 亿美元（净额，扣除已收购现金）
- 购买可交易证券 11 亿美元
- 购买非可交易股权证券 9.82 亿美元
- 购买应收票据 2.97 亿美元
- 购买物业和设备 2.98 亿美元

因此，Uber 在 2021 年的总投资约为 50 亿美元，涵盖了业务收购、可交易和非可交易证券、应收票据以及物业和设备购买。
 [0mHTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"
 [1;3;38;2;11;159;203m[lyft_10k] A: 根据提供的内容，Lyft 在 2021 年投资于可交易证券和定期存款：

- Lyft 在 2021 年购买了价值 38 亿美元的可交易证券。这些可交易证券包括投资级别的可供出售的债务证券。

- Lyft 在 2021 年还投资了价值 5 亿美元的定期存款。这些定期存款按成本计价，接近公允价值。

截至 2021 年 12 月 31 日，Lyft 的投资组合的加权平均剩余期限不到一年。Lyft 的投资政策旨在最大限度地降低信用损失风险。
 [0mHTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"
 [1;3;38;2;237;90;200m[uber_10k] A: 根据提供的内容，Uber 在 2021 年的投资旨在：

- 通过激励措施、折扣和促销活动，增加使用其平台的司机、消费者、商户、托运人和运营商的数量
- 在现有市场和新市场进行扩张
- 增加研发支出
- 扩大营销渠道和运营
- 雇佣更多员工
- 在其平台增加新产品和新服务

内容表明，Uber 预计由于持续进行此类投资而导致运营费用大幅增加，因此在短期内会产生亏损。
 [0mHTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"
 [1;3;38;2;155;135;227m[lyft_10k] A: 根据上下文中提供的财务信息，现金流量表显示 Lyft 在截至 2021 年 12 月 31 日的财年，投资活动产生的净现金为 2.67 亿美元。这主要包括：

- 可交易证券的销售和到期所得款项 38 亿美元
- 定期存款到期 6.755 亿美元
- 部分被购买可交易证券 38 亿美元和定期存款 5 亿美元所抵消

因此，虽然销售/到期所得款项总额约为 45 亿美元，但 Lyft 将大部分资金进行了再投资，2021 年的净新增投资约为 2.67 亿美元。
 [0mHTTP Request: POST https://api.anthropic.com/v1/messages "HTTP/1.1 200 OK"

display(HTML(f'<p style="font-size:20px">{response.response}</p>'))

2021 年，Uber 和 Lyft 的投资类型有所不同： Uber 的总投资约为 50 亿美元，其中包括： - 23 亿美元用于收购业务 - 11 亿美元用于可交易证券 - 9.82 亿美元用于非可交易股权证券 - 2.97 亿美元用于应收票据 - 2.98 亿美元用于物业和设备 Lyft 的净投资约为 2.67 亿美元，主要用于： - 可交易证券，购买额为 38 亿美元，但销售和到期额也为 38 亿美元 - 定期存款 5 亿美元，但到期额为 6.755 亿美元 因此，Uber 大量投资于收购业务、证券、票据和物业，而 Lyft 则专注于短期可交易证券和定期存款，并将大部分收益进行了再投资。与 Lyft 的净投资相比，Uber 的总投资额要大得多。