通过llamindex实现ReAct RAG Agent
# 加载电商财报数据
from llama_index.core import SimpleDirectoryReader
from dotenv import load_dotenv
load_dotenv()
# 加载 PDF 文件 A.pdf
A_docs = SimpleDirectoryReader(
input_files="../A.pdf" # 修正括号问题
).load_data()
# 加载 PDF 文件 B.pdf
B_docs = SimpleDirectoryReader(
input_files="../B.pdf" # 修正括号问题
).load_data()
# 从文档中创建索引
from llama_index.core import VectorStoreIndex
A_index = VectorStoreIndex.from_documents(A_docs)
B_index = VectorStoreIndex.from_documents(B_docs)
# 持久化索引(保存到本地)
from llama_index.core import StorageContext
# 持久化 A_index 到本地目录 "../storage/A"
A_index.storage_context.persist(persist_dir="../storage/A")
# 持久化 B_index 到本地目录 "../storage/B"
B_index.storage_context.persist(persist_dir="../storage/B")
# 从本地读取索引
from llama_index.core import load_index_from_storage, StorageContext
try:
# 加载 A_index
storage_context = StorageContext.from_defaults(persist_dir="../storage/A")
A_index = load_index_from_storage(storage_context)
# 加载 B_index
storage_context = StorageContext.from_defaults(persist_dir="../storage/B")
B_index = load_index_from_storage(storage_context)
index_loaded = True
except Exception as e:
print(f"加载索引失败: {e}")
index_loaded = False
# 创建查询引擎
A_engine = A_index.as_query_engine(similarity_top_k=3)
B_engine = B_index.as_query_engine(similarity_top_k=3)
# 配置查询工具
from llama_index.core.tools import QueryEngineTool, ToolMetadata
query_engine_tools = [
QueryEngineTool(
query_engine=A_engine,
metadata=ToolMetadata(
name="A_Finance",
description=("用于提供A公司的财务信息"),
),
),
QueryEngineTool(
query_engine=B_engine,
metadata=ToolMetadata(
name="B_Finance",
description=("用于提供B公司的财务信息"),
),
),
]
# 配置大模型
from llama_index.llms.openai import OpenAI
# 初始化 OpenAI 模型
llm = OpenAI(model="gpt-4")
# 创建 React Agent
from llama_index.core.agent import ReactAgent
# 使用查询工具和大模型创建 React Agent
agent = ReactAgent.from_tools(query_engine_tools, llm=llm, verbose=True)
# 让 Agent 完成任务
print(agent.chat("Compare the sales of the two companies"))