LangChain内存配置：仅存储AI响应的实现方法

核心思路 #

默认的ConversationBufferMemory会同时存储用户提问和AI响应，要实现仅保留AI响应，只需重写内存的消息保存逻辑，只将AI的输出存入内存即可。

原示例代码（参考） #

from langchain.chains import RetrievalQA
from langchain.memory import ConversationBufferMemory
from langchain.llms import OpenAI
from langchain.vectorstores import Chroma

# 初始化向量存储与检索器（假设已完成嵌入配置）
vectorstore = Chroma(persist_directory="./chroma_db", embedding_function=embeddings)
retriever = vectorstore.as_retriever()

# 默认内存（同时存用户提问和AI响应）
memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)

# 构建RetrievalQA链
qa_chain = RetrievalQA.from_chain_type(
    llm=OpenAI(temperature=0),
    chain_type="stuff",
    retriever=retriever,
    memory=memory,
    chain_type_kwargs={"memory": memory}
)

改造方案：自定义AI-only内存类 #

继承ConversationBufferMemory，重写save_context方法，仅保存AI的响应内容：

from langchain.memory import ConversationBufferMemory
from langchain.schema import AIMessage

class AIOnlyConversationMemory(ConversationBufferMemory):
    def save_context(self, inputs: dict, outputs: dict) -> None:
        # 提取AI响应并添加到内存，忽略用户提问
        ai_response = outputs["text"]
        self.chat_memory.add_message(AIMessage(content=ai_response))

应用自定义内存 #

将原内存替换为自定义类实例，即可实现仅存储AI响应：

# 使用自定义内存
memory = AIOnlyConversationMemory(memory_key="chat_history", return_messages=True)

# 重新构建RetrievalQA链
qa_chain = RetrievalQA.from_chain_type(
    llm=OpenAI(temperature=0),
    chain_type="stuff",
    retriever=retriever,
    memory=memory,
    chain_type_kwargs={"memory": memory}
)

优化：配合提示词避免重复内容 #

为了让模型更明确地避免重复输出，可自定义提示词，引导模型参考历史AI响应：

from langchain.prompts import PromptTemplate

# 自定义提示词，加入历史AI响应参考
prompt_template = """基于以下上下文回答问题，同时参考之前的AI回答，**不要重复已提及的事实**。

上下文信息:
{context}

之前的AI回答:
{chat_history}

用户问题: {question}
最终回答:"""

PROMPT = PromptTemplate(
    template=prompt_template,
    input_variables=["context", "chat_history", "question"]
)

# 构建带自定义提示词的QA链
qa_chain = RetrievalQA.from_chain_type(
    llm=OpenAI(temperature=0),
    chain_type="stuff",
    retriever=retriever,
    memory=memory,
    chain_type_kwargs={"prompt": PROMPT, "memory": memory}
)

效果验证 #

每次调用qa_chain.run()后，查看内存内容：

print(memory.chat_memory.messages)
# 输出仅包含AIMessage类型的AI响应，无用户提问消息

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