names = ['c-rag','self-rag','kg-rag']
files = ['../../data/c-rag.pdf','../../data/self-rag.pdf','../../data/kg-rag.pdf']

......省略import部分与准备llm部分......

#采用chroma向量数据库
chroma = chromadb.HttpClient(host="localhost", port=8000)
collection = chroma.get_or_create_collection(name="agentic_rag") 
vector_store = ChromaVectorStore(chroma_collection=collection)

#创建针对某个文档的tool_agent
def create_tool_agent(file,name):

    #文档拆分
    print(f'Starting to create tool agent for 【{name}】...\n')
    docs =SimpleDirectoryReader(input_files = [file]).load_data()
    splitter = SentenceSplitter(chunk_size=500,chunk_overlap=50)
    nodes = splitter.get_nodes_from_documents(docs)
    
    #创建向量索引，并做持久保存
    if not os.path.exists(f"./storage/{name}"):
        print('Creating vector index...\n')
        storage_context = StorageContext.from_defaults(vector_store=vector_store)
        vector_index = VectorStoreIndex(nodes,storage_context=storage_context)
        vector_index.storage_context.persist(persist_dir=f"./storage/{name}")
    else:
        print('Loading vector index...\n')
        storage_context = StorageContext.from_defaults(persist_dir=f"./storage/{name}",vector_store=vector_store)
        vector_index = load_index_from_storage(storage_context=storage_context)

    #创建基于向量的查询引擎
    query_engine = vector_index.as_query_engine(similarity_top_k=5)

    #创建摘要索引与对应的查询引擎
    summary_index = SummaryIndex(nodes)
    summary_engine = summary_index.as_query_engine(response_mode="tree_summarize")

    #将RAG引擎转化为两个tool
    query_tool = QueryEngineTool.from_defaults(query_engine=query_engine,name=f'query_tool',description=f'Use if you want to query details about {name}')
    summary_tool = QueryEngineTool.from_defaults(query_engine=summary_engine,name=f'summary_tool',description=f'Use ONLY IF you want to get a holistic summary of the documents. DO NOT USE if you want to query some details about {name}.')

    #创建一个tool agent
    tool_agent = ReActAgent.from_tools([query_tool,summary_tool],verbose=True,
                                  system_prompt=f"""
                                                          You are a specialized agent designed to answer queries about {name}.You must ALWAYS use at least one of the tools provided when answering a question; DO NOT rely on prior knowledge. DO NOT fabricate answer.
                                                            """)
    return tool_agent

#创建不同文档的agent
print('===============================================\n')
print('Creating tool agents for different documents...\n')
tool_agents_dict = {}
for name, file in zip(names, files):
    tool_agent = create_tool_agent(file, name)
    tool_agents_dict[name] = tool_agent

#首先将Tool Agent进行“工具化”
print('===============================================\n')
print('Creating tools from tool agents...\n')
all_tools = []

for name in names:
    agent_tool = QueryEngineTool.from_defaults(
            #注意，Agent本身也是一种Query Engine，所以直接转为tool
           query_engine=tool_agents_dict[name],

            #这个工具的名字
            name=f"tool_{name.replace("-", "")}",

            #描述这个工具的作用和使用方法
            description=f"Use this tool if you want to answer any questions about {name}."
    )

    all_tools.append(agent_tool)

#创建Top Agent
print('Creating top agent...\n')
top_agent = OpenAIAgent.from_tools(tools=all_tools,verbose=True,system_prompt="""You are an agent designed to answer queries over a set of given papers.Please always use the tools provided to answer a question.Do not rely on prior knowledge.DO NOT fabricate answer""" )

top_agent.chat_repl()

#创建工具检索器
print('===============================================\n')
print('Creating tool retrieve index...\n')
obj_index = ObjectIndex.from_objects(all_tools,index_cls=VectorStoreIndex,)
tool_retriever = obj_index.as_retriever(similarity_top_k=5,verbose=True)

......
top_agent = OpenAIAgent.from_tools(tool_retriever=tool_retriever,
                                   verbose=True,
                                   system_prompt="""You are an agent designed to answer queries over a set of given papers.Please always use the tools provided to answer a question.Do not rely on prior knowledge.""")
.......

tools_needed = tool_retriever.retrieve("What is the Adaptive retrieval in the c-RAG?")
print('Tools needed to answer the question:')
for tool in tools_needed:
    print(tool.metadata.name)