如何在LangChain AgentExecutor工具调用完成后、生成Final Answer前查询向量数据库

要实现所有工具调用完成后、输出最终回答前查询向量数据库，可通过以下两种方式实现：

方法一：自定义OpenAI Tools Agent #

继承原OpenAIToolsAgent，在输出Final Answer前插入向量查询并整合上下文：

from langchain.agents.openai_tools.base import OpenAIToolsAgent
from langchain.schema import AgentFinish, AgentAction
from langchain.agents import AgentExecutor
from langchain.chat_models import ChatOpenAI
from langchain import hub

# 已初始化的向量存储
vector_store = ... 

class CustomOpenAIToolsAgent(OpenAIToolsAgent):
    def plan(
        self, intermediate_steps, **kwargs
    ) -> AgentAction | AgentFinish:
        # 执行原Agent的决策逻辑
        original_plan = super().plan(intermediate_steps, **kwargs)
        
        # 若原计划为输出最终回答，执行向量查询
        if isinstance(original_plan, AgentFinish):
            question = kwargs.get("input")
            docs = vector_store.similarity_search(question)
            doc_content = "\n".join([doc.page_content for doc in docs])
            
            # 整合工具结果与向量文档，生成最终回答
            final_prompt = f"""
工具调用结果: {original_plan.return_values['output']}
参考文档: {doc_content}
请结合以上信息最终回答用户问题：{question}
"""
            final_answer = self.llm.predict(final_prompt)
            return AgentFinish(return_values={"output": final_answer}, log=original_plan.log)
        
        # 若为工具调用动作，直接返回
        return original_plan

# 初始化组件
llm = ChatOpenAI()
tools = [tool_1, tool_2, ...]
prompt = hub.pull("hwchase17/openai-tools-agent")

# 创建自定义Agent并执行
my_agent = CustomOpenAIToolsAgent(llm=llm, tools=tools, prompt=prompt)
agent_executor = AgentExecutor(agent=my_agent, tools=tools, verbose=True)

question = "user's question"
ans = agent_executor.invoke({"input": question})

方法二：使用回调拦截 #

通过回调函数在Agent生成Final Answer前插入向量查询逻辑：

from langchain.callbacks.base import BaseCallbackHandler
from langchain.schema import AgentFinish
from langchain.agents import AgentExecutor
from langchain.chat_models import ChatOpenAI
from langchain import hub

# 已初始化的向量存储
vector_store = ... 

class VectorPostProcessCallback(BaseCallbackHandler):
    def __init__(self, vector_store, llm):
        self.vector_store = vector_store
        self.llm = llm
    
    def on_agent_finish(self, finish: AgentFinish, **kwargs) -> None:
        question = kwargs.get("inputs", {}).get("input")
        if not question:
            return
        
        # 查询向量数据库
        docs = self.vector_store.similarity_search(question)
        doc_content = "\n".join([doc.page_content for doc in docs])
        
        # 重新生成最终回答
        final_prompt = f"""
工具调用结果: {finish.return_values['output']}
参考文档: {doc_content}
请结合以上信息最终回答用户问题：{question}
"""
        finish.return_values["output"] = self.llm.predict(final_prompt)

# 初始化组件
llm = ChatOpenAI()
tools = [tool_1, tool_2, ...]
prompt = hub.pull("hwchase17/openai-tools-agent")
my_agent = create_openai_tools_agent(llm, tools, prompt)

# 添加回调并执行
callback = VectorPostProcessCallback(vector_store=vector_store, llm=llm)
agent_executor = AgentExecutor(agent=my_agent, tools=tools, verbose=True, callbacks=[callback])

question = "user's question"
ans = agent_executor.invoke({"input": question})

内容的提问来源于stack exchange，提问作者chenkun