使用OpenAI Agent构建查询引擎工具_ai agent 查询引擎

在这篇文章中，我们将介绍如何使用OpenAI Agent和查询引擎工具来处理大数据查询任务。本文将详细讲解如何安装必要的依赖库、下载数据、构建索引并实现一个简单的查询引擎。

安装依赖

首先，我们需要安装llama-index库和其OpenAI扩展。可以使用以下命令来安装：

%pip install llama-index-agent-openai
!pip install llama-index
1
2

导入必要的库

from llama_index.core import (
    SimpleDirectoryReader,
    VectorStoreIndex,
    StorageContext,
    load_index_from_storage,
)

from llama_index.core.tools import QueryEngineTool, ToolMetadata
1
2
3
4
5
6
7
8

加载或构建索引

我们将尝试从已有的存储中加载索引，如果加载失败则重新构建索引。

try:
    storage_context = StorageContext.from_defaults(
        persist_dir="./storage/lyft"
    )
    lyft_index = load_index_from_storage(storage_context)

    storage_context = StorageContext.from_defaults(
        persist_dir="./storage/uber"
    )
    uber_index = load_index_from_storage(storage_context)

    index_loaded = True
except:
    index_loaded = False
1
2
3
4
5
6
7
8
9
10
11
12
13
14

下载数据

如果索引未加载成功，我们需要下载并处理数据。

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

构建索引

if not index_loaded:
    # load data
    lyft_docs = SimpleDirectoryReader(
        input_files=["./data/10k/lyft_2021.pdf"]
    ).load_data()
    uber_docs = SimpleDirectoryReader(
        input_files=["./data/10k/uber_2021.pdf"]
    ).load_data()

    # build index
    lyft_index = VectorStoreIndex.from_documents(lyft_docs)
    uber_index = VectorStoreIndex.from_documents(uber_docs)

    # persist index
    lyft_index.storage_context.persist(persist_dir="./storage/lyft")
    uber_index.storage_context.persist(persist_dir="./storage/uber")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

创建查询引擎

lyft_engine = lyft_index.as_query_engine(similarity_top_k=3)
uber_engine = uber_index.as_query_engine(similarity_top_k=3)

query_engine_tools = [
    QueryEngineTool(
        query_engine=lyft_engine,
        metadata=ToolMetadata(
            name="lyft_10k",
            description=(
                "Provides information about Lyft financials for year 2021. "
                "Use a detailed plain text question as input to the tool."
            ),
        ),
    ),
    QueryEngineTool(
        query_engine=uber_engine,
        metadata=ToolMetadata(
            name="uber_10k",
            description=(
                "Provides information about Uber financials for year 2021. "
                "Use a detailed plain text question as input to the tool."
            ),
        ),
    ),
]
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25

设置OpenAI Agent

from llama_index.agent.openai import OpenAIAgent

agent = OpenAIAgent.from_tools(query_engine_tools, verbose=True)
1
2
3

运行查询引擎

agent.chat_repl()
1

输入以下内容进行测试：

=== Calling Function ===
Calling function: lyft_10k with args: {
  "input": "What was Lyft's revenue growth in 2021?"
}
Got output: 
Lyft's revenue growth in 2021 was 36%.
========================
=== Calling Function ===
Calling function: uber_10k with args: {
  "input": "What was Uber's revenue growth in 2021?"
}
Got output: 
Uber's revenue growth in 2021 was 57%.
========================
Assistant: Lyft's revenue growth in 2021 was 36%, while Uber's revenue growth in 2021 was 57%.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

可能遇到的错误

1. 依赖安装失败：

   • 解决方案：检查网络连接或尝试使用国内镜像源进行安装，如pip install -i https://pypi.tuna.tsinghua.edu.cn/simple llama-index-agent-openai。

2. 数据下载失败：

   • 解决方案：可能是网络问题，检查URL是否可访问，或尝试手动下载文件。

3. 索引加载或构建失败：

   • 解决方案：确保存储路径存在且有写权限，检查数据文件是否完整。

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

参考资料：

• LlamaIndex 官方文档
• OpenAI Agent 官方文档