基于LCEL改造LLMChain实现流式支持：对话记忆不更新的问题求助

问题背景 #

原本使用LLMChain，切换到LCEL实现流式支持后，发现ConversationBufferMemory在每次调用后不会自动更新对话历史；尝试使用RunnableWithMessageHistory也未解决问题。

原LLMChain代码：

llm_chain = LLMChain(
            llm=llm_chat,
            prompt=prompt,
            output_key="answer",
            verbose=False,
        )

首次尝试的LCEL代码（内存不更新）：

retrieved_chat_history = ChatMessageHistory(messages=[]) ## example
conversation_memory = ConversationBufferMemory(
            input_key="question",
            output_key="answer",
            memory_key="chat_history",
            return_messages=True,
            chat_memory=retrieved_chat_history,
        )

lcel_chain = (
    RunnablePassthrough.assign(
        chat_history=lambda x: conversation_memory.load_memory_variables(x)["chat_history"]
    )
    | prompt
    | llm_chat
    | {"answer": RunnablePassthrough()}
)

尝试的RunnableWithMessageHistory代码（未生效）：

from langchain_core.runnables.history import RunnableWithMessageHistory
base_chain = prompt | llm_chat 

# Create a function to get an empty message history
def get_session_history(session_id):
    return ChatMessageHistory()

# Wrap the base chain with RunnableWithMessageHistory
chain_with_memory = RunnableWithMessageHistory(
    base_chain,
    get_session_history,
    input_messages_key="question",
    history_messages_key=memory_key,
    output_messages_key="answer"
)

session_id = "user_123"  # Use a unique identifier for each conversation
result = chain_with_memory.invoke(
    {"question": "Hello, my name is XYZ"},
    config={"configurable": {"session_id": session_id}}
)

修复方案 #

方案一：正确使用RunnableWithMessageHistory（推荐） #

此方式能自动管理多会话的历史，无需手动调用保存方法。

步骤1：确保Prompt正确引用对话历史

Prompt中必须包含与history_messages_key对应的占位符，示例：

from langchain_core.prompts import ChatPromptTemplate

prompt = ChatPromptTemplate.from_messages([
    ("system", "你是一个专业助手，根据对话历史回答问题"),
    ("placeholder", "{chat_history}"),  # 对应history_messages_key
    ("human", "{question}")
])

步骤2：实现会话历史的持久化存储

之前的代码每次调用get_session_history都会新建ChatMessageHistory，导致历史无法保留。改用字典维护不同会话的历史：

from langchain_core.chat_history import ChatMessageHistory

# 全局存储不同会话的历史，生产环境可替换为数据库/缓存
session_histories = {}

def get_session_history(session_id):
    if session_id not in session_histories:
        session_histories[session_id] = ChatMessageHistory()
    return session_histories[session_id]

步骤3：正确构建带历史的LCEL链

from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_core.runnables import RunnablePassthrough

# 基础链：Prompt -> LLM -> 格式化输出为{"answer": 回复内容}
base_chain = prompt | llm_chat | {"answer": RunnablePassthrough()}

memory_key = "chat_history"
chain_with_memory = RunnableWithMessageHistory(
    base_chain,
    get_session_history,
    input_messages_key="question",  # 输入中用户问题的键名
    history_messages_key=memory_key,  # 注入到Prompt的历史变量名
    output_messages_key="answer"  # 输出中AI回复的键名，用于存入历史
)

调用示例

session_id = "user_123"
# 第一次调用
result1 = chain_with_memory.invoke(
    {"question": "你好，我叫小明"},
    config={"configurable": {"session_id": session_id}}
)
print(result1["answer"])  # 输出助手回复

# 第二次调用，验证历史是否生效
result2 = chain_with_memory.invoke(
    {"question": "我叫什么名字？"},
    config={"configurable": {"session_id": session_id}}
)
print(result2["answer"])  # 应能正确回答"小明"

方案二：手动更新ConversationBufferMemory #

如果坚持使用ConversationBufferMemory，需要在链中加入手动保存上下文的步骤：

from langchain_core.runnables import RunnableLambda

retrieved_chat_history = ChatMessageHistory(messages=[])
conversation_memory = ConversationBufferMemory(
    input_key="question",
    output_key="answer",
    memory_key="chat_history",
    return_messages=True,
    chat_memory=retrieved_chat_history,
)

lcel_chain = (
    RunnablePassthrough.assign(
        chat_history=lambda x: conversation_memory.load_memory_variables(x)["chat_history"]
    )
    | prompt
    | llm_chat
    | {"answer": RunnablePassthrough()}
    # 新增：调用save_context保存对话，使用lambda确保返回结果不被覆盖
    | RunnableLambda(lambda x: conversation_memory.save_context({"question": x["question"]}, {"answer": x["answer"]}) or x)
)

错误原因分析 #

1. 会话历史未持久化：之前的get_session_history每次返回新的ChatMessageHistory，导致会话历史无法跨调用保留。
2. Prompt未引用历史变量：如果Prompt中没有对应history_messages_key的占位符，历史无法注入到模型输入中。
3. 未配置自动保存逻辑：直接使用ConversationBufferMemory时，LCEL不会自动调用save_context，需手动添加保存步骤。

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