使用OpenAI Agent进行高级查询计划

在这篇文章中，我们将探讨如何在OpenAI Agent中添加一个QueryPlanTool，从而让Agent能够通过单一工具进行高级查询计划。QueryPlanTool设计用于与OpenAI函数API很好地结合。工具函数签名包含QueryPlan Pydantic对象，该对象可以包含定义计算图的QueryNode对象的DAG。Agent负责通过调用工具时的函数签名来定义这个图。工具本身则在任何相应的工具上执行DAG。

在这个设置中，我们使用一个熟悉的例子：Uber 2022年3月、6月和9月的10-Q报告。

如果您在Colab上打开这个Notebook，您可能需要安装LlamaIndex 🦙。

%pip install llama-index-agent-openai
%pip install llama-index-llms-openai
!pip install llama-index

# # uncomment to turn on logging
# import logging
# import sys

# logging.basicConfig(stream=sys.stdout, level=logging.INFO)
# logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

%load_ext autoreload
%autoreload 2

from llama_index.core import SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.response.pprint_utils import pprint_response
from llama_index.llms.openai import OpenAI

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

# 下载数据
!mkdir -p 'data/10q/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10q/uber_10q_march_2022.pdf' -O 'data/10q/uber_10q_march_2022.pdf'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10q/uber_10q_june_2022.pdf' -O 'data/10q/uber_10q_june_2022.pdf'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/docs/examples/data/10q/uber_10q_sept_2022.pdf' -O 'data/10q/uber_10q_sept_2022.pdf'

# 加载数据
march_2022 = SimpleDirectoryReader(
    input_files=["./data/10q/uber_10q_march_2022.pdf"]
).load_data()
june_2022 = SimpleDirectoryReader(
    input_files=["./data/10q/uber_10q_june_2022.pdf"]
).load_data()
sept_2022 = SimpleDirectoryReader(
    input_files=["./data/10q/uber_10q_sept_2022.pdf"]
).load_data()

# 构建索引
march_index = VectorStoreIndex.from_documents(march_2022)
june_index = VectorStoreIndex.from_documents(june_2022)
sept_index = VectorStoreIndex.from_documents(sept_2022)

march_engine = march_index.as_query_engine(similarity_top_k=3, llm=llm)
june_engine = june_index.as_query_engine(similarity_top_k=3, llm=llm)
sept_engine = sept_index.as_query_engine(similarity_top_k=3, llm=llm)

# 使用Query Plan Tool
from llama_index.core.tools import QueryEngineTool

query_tool_sept = QueryEngineTool.from_defaults(
    query_engine=sept_engine,
    name="sept_2022",
    description=(
        f"Provides information about Uber quarterly financials ending"
        f" September 2022"
    ),
)
query_tool_june = QueryEngineTool.from_defaults(
    query_engine=june_engine,
    name="june_2022",
    description=(
        f"Provides information about Uber quarterly financials ending June"
        f" 2022"
    ),
)
query_tool_march = QueryEngineTool.from_defaults(
    query_engine=march_engine,
    name="march_2022",
    description=(
        f"Provides information about Uber quarterly financials ending March"
        f" 2022"
    ),
)

# 定义查询计划工具
from llama_index.core.tools import QueryPlanTool
from llama_index.core import get_response_synthesizer

response_synthesizer = get_response_synthesizer()
query_plan_tool = QueryPlanTool.from_defaults(
    query_engine_tools=[query_tool_sept, query_tool_june, query_tool_march],
    response_synthesizer=response_synthesizer,
)

# 转换为OpenAI工具
query_plan_tool.metadata.to_openai_tool()

from llama_index.agent.openai import OpenAIAgent

agent = OpenAIAgent.from_tools(
    [query_plan_tool],
    max_function_calls=10,
    llm=OpenAI(temperature=0, model="gpt-4-0613"),
    verbose=True,
)

response = agent.query("What were the risk factors in sept 2022?")
print(str(response))

# 示例生成的查询计划
from llama_index.core.tools.query_plan import QueryPlan, QueryNode

query_plan = QueryPlan(
    nodes=[
        QueryNode(
            id=1,
            query_str="risk factors",
            tool_name="sept_2022",
            dependencies=[],
        )
    ]
)
print(query_plan.schema())

response = agent.query(
    "Analyze Uber revenue growth in March, June, and September"
)
print(str(response))

上述代码展示了如何使用QueryPlanTool来执行高级查询计划。该工具允许我们将多个查询引擎工具组合成一个复杂的DAG，Agent通过调用这些工具来生成回答。我们还可以看到如何使用LlamaIndex来加载数据、构建索引并进行查询。

可能遇到的错误

1. 无法连接到下载链接：如果数据下载失败，请确保网络连接正常，或尝试手动下载文件并放置在指定目录中。
2. 工具初始化失败：确保所有依赖库已正确安装，并且版本兼容。
3. OpenAI API调用失败：如果无法调用OpenAI API，请确保API密钥正确配置，并且API服务可用。

如果你觉得这篇文章对你有帮助,请点赞,关注我的博客,谢谢!

参考资料:

• OpenAI文档
• LlamaIndex文档