ReAct RAG Agent

通过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"))
posted @ 2025-03-07 17:05  Michael_Z  阅读(73)  评论(0)    收藏  举报