如何结合Gradio、Langchain Agent实现带会话记忆的OpenAI流式调用

需求说明 #

• 通过LangChain Agent流式调用OpenAI API
• 在Gradio中实时流式输出回答
• 为每个Gradio用户分配独立的会话记忆

关键修改点与完整代码 #

1. 流式输出实现：回调函数对接Gradio生成器 #

原代码的回调仅打印token，需修改为将token传递给Gradio的生成逻辑。我们用线程安全队列在回调和生成器之间传递token，同时让Agent的调用逻辑以生成器形式返回每个token。

2. 独立会话记忆：借助Gradio State组件 #

通过Gradio的State组件存储每个用户的ConversationBufferMemory，每次交互时取出记忆初始化Agent，交互完成后更新记忆并回存到State中。

完整代码如下：

from langchain.agents import load_tools
from langchain.agents import initialize_agent
from langchain.agents import AgentType
from langchain.chat_models import ChatOpenAI
from langchain.memory import ConversationBufferMemory
import gradio as gr
from langchain.callbacks.base import BaseCallbackHandler
from dotenv import load_dotenv
from queue import Queue
from threading import Thread

load_dotenv()

# 自定义回调：将LLM生成的token存入队列
class StreamingCallbackHandler(BaseCallbackHandler):
    def __init__(self, queue):
        self.queue = queue

    def on_llm_new_token(self, token, **kwargs) -> None:
        self.queue.put(token)

    def on_llm_end(self, response, **kwargs) -> None:
        # 标记生成结束
        self.queue.put(None)

def agent_streaming_run(input_text, memory):
    queue = Queue()
    callback = StreamingCallbackHandler(queue)
    
    # 初始化带流式和回调的LLM
    llm = ChatOpenAI(streaming=True, callbacks=[callback], temperature=0)
    tools = load_tools(["wikipedia", "llm-math"], llm=llm)
    
    # 初始化带记忆的Agent
    agent = initialize_agent(
        tools, llm, agent=AgentType.CONVERSATIONAL_REACT_DESCRIPTION, 
        verbose=False, memory=memory
    )
    
    # 启动线程运行Agent，避免阻塞生成器
    def run_agent():
        try:
            agent.run(input_text)
        except Exception as e:
            queue.put(f"\n错误：{str(e)}")
            queue.put(None)
    
    Thread(target=run_agent).start()
    
    # 从队列取token，逐个yield
    while True:
        token = queue.get()
        if token is None:
            break
        yield token

with gr.Blocks() as demo:
    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="输入问题")
    clear = gr.Button("清空对话")
    # 用State存储每个用户的会话记忆
    memory_state = gr.State(lambda: ConversationBufferMemory(memory_key="chat_history", return_messages=True))

    def user_input(user_message, history):
        return "", history + [[user_message, None]]

    def bot_response(history, memory):
        user_query = history[-1][0]
        history[-1][1] = ""
        # 逐个接收Agent流式返回的token
        for token in agent_streaming_run(user_query, memory):
            history[-1][1] += token
            yield history, memory

    msg.submit(user_input, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot_response, [chatbot, memory_state], [chatbot, memory_state]
    )
    clear.click(lambda: (None, ConversationBufferMemory(memory_key="chat_history", return_messages=True)), None, [chatbot, memory_state], queue=False)
    
demo.queue()
demo.launch()

代码说明 #

• 流式输出：StreamingCallbackHandler把每个生成的token存入队列，agent_streaming_run启动线程运行Agent，同时从队列取token并yield，Gradio的bot_response函数逐个接收token更新聊天历史。
• 独立记忆：memory_state作为State组件，每个用户会话会维护自己的ConversationBufferMemory，清空对话时重置记忆。

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