LangChain分层记忆解决方案：完整案例
下面是一个完整的分层记忆解决方案案例，解决长对话历史导致的token超限问题。我们将实现一个三层记忆系统，结合短期记忆、摘要记忆和向量检索记忆。
python
import os
from datetime import datetime, timedelta
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnableLambda, RunnablePassthrough
from langchain_core.messages import AIMessage, HumanMessage
from langchain_community.chat_models import ChatOpenAI
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import OpenAIEmbeddings
from langchain_community.chat_message_histories import RedisChatMessageHistory
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain.memory import ConversationSummaryBufferMemory
from langchain.text_splitter import RecursiveCharacterTextSplitter
import getpass
# 设置OpenAI API密钥
os.environ["OPENAI_API_KEY"] = getpass.getpass("OpenAI API Key:")
# 1. 分层记忆系统实现
class LayeredMemorySystem:
def __init__(self, session_id, llm):
self.session_id = session_id
self.llm = llm
# 第一层：短期记忆（最近5条消息）
self.short_term_memory = RedisChatMessageHistory(
session_id, url="redis://localhost:6379/0", key_prefix="short_term_", ttl=3600
)
# 第二层：摘要记忆（长期摘要）
self.summary_memory = ConversationSummaryBufferMemory(
llm=llm,
memory_key="summary",
max_token_limit=500,
return_messages=True
)
# 第三层：向量记忆（关键信息存储）
self.embeddings = OpenAIEmbeddings()
self.vector_store = Chroma(
collection_name=f"memory_{session_id}",
embedding_function=self.embeddings,
persist_directory="./chroma_db"
)
self.text_splitter = RecursiveCharacterTextSplitter(
chunk_size=500,
chunk_overlap=50
)
def add_message(self, message):
"""添加新消息到所有记忆层"""
# 添加到短期记忆
self.short_term_memory.add_message(message)
# 添加到摘要记忆
if isinstance(message, HumanMessage):
self.summary_memory.chat_memory.add_user_message(message.content)
else:
self.summary_memory.chat_memory.add_ai_message(message.content)
# 添加到向量记忆（如果是重要信息）
if self._is_important(message.content):
chunks = self.text_splitter.split_text(message.content)
self.vector_store.add_texts(chunks)
def get_context(self, query):
"""获取分层记忆上下文"""
context = {
"short_term": [f"{msg.type}: {msg.content}" for msg in self.short_term_memory.messages[-5:]], # 最近5条
"summary": self.summary_memory.load_memory_variables({})["summary"],
"relevant_memories": self._retrieve_relevant_memories(query)
}
return context
def _is_important(self, text):
"""判断信息是否重要（简化版）"""
important_keywords = ["喜欢", "不喜欢", "重要", "记住", "总是", "从不", "偏好"]
return any(keyword in text for keyword in important_keywords)
def _retrieve_relevant_memories(self, query, k=3):
"""从向量记忆中检索相关信息"""
if not self.vector_store:
return []
docs = self.vector_store.similarity_search(query, k=k)
return [f"[记忆片段] {doc.page_content}" for doc in docs]
def cleanup(self):
"""定期清理旧记忆"""
# 保留最近7天的摘要记忆
self.summary_memory.prune()
# 短期记忆由Redis TTL自动清理
# 向量记忆保留所有（实际应用中可设置过期时间）
# 2. 创建聊天系统
class SmartChatSystem:
def __init__(self):
self.llm = ChatOpenAI(model="gpt-3.5-turbo-16k", temperature=0.7)
# 创建提示模板
self.prompt = ChatPromptTemplate.from_messages([
("system", """
你是一个专业的客户服务AI，具有完美的记忆力。使用以下分层记忆提供精准回答：
[近期对话] (最新5条)
{short_term}
[对话摘要] (长期摘要)
{summary}
[相关记忆] (基于当前问题检索)
{relevant_memories}
当前时间: {current_time}
"""),
MessagesPlaceholder(variable_name="history"),
("human", "{input}")
])
# 基础链
self.chain = (
RunnablePassthrough.assign(
current_time=lambda _: datetime.now().strftime("%Y-%m-%d %H:%M"),
history=lambda x: "\n".join(x["memory"].get("short_term", [])),
short_term=lambda x: "\n".join(x["memory"].get("short_term", [])),
summary=lambda x: x["memory"].get("summary", "无摘要"),
relevant_memories=lambda x: "\n".join(x["memory"].get("relevant_memories", [])),
)
| self.prompt
| self.llm
)
def get_conversational_chain(self, session_id):
"""创建带记忆的对话链"""
memory_system = LayeredMemorySystem(session_id, self.llm)
def get_memory_context(input_dict):
"""获取记忆上下文"""
return memory_system.get_context(input_dict["input"])
def add_to_memory(input_dict, output):
"""将对话添加到记忆系统"""
memory_system.add_message(HumanMessage(content=input_dict["input"]))
memory_system.add_message(AIMessage(content=output.content))
return output
# 创建带记忆处理的链
memory_chain = (
RunnablePassthrough.assign(memory=get_memory_context)
| self.chain
)
# 添加记忆存储步骤
full_chain = memory_chain | RunnableLambda(
lambda output: add_to_memory(output["input"], output)
return full_chain
# 3. 使用示例
if __name__ == "__main__":
# 创建聊天系统
chat_system = SmartChatSystem()
session_id = "user_789"
# 获取对话链
conversational_chain = chat_system.get_conversational_chain(session_id)
# 模拟长时间对话
conversation = [
"你好，我想咨询一下你们的会员计划",
"我比较关心的是健身房的开放时间",
"我通常早上7点锻炼，你们那个时间开门吗？",
"另外，我有关节炎，游泳是最好的运动方式",
"你们的游泳池水质如何？多久换一次水？",
"我想知道私人教练的费用",
"我特别喜欢李教练的教学风格，他周日上午有空吗？",
"对了，我太太也想来，她有瑜伽经验",
"我们俩都想报名，有家庭优惠吗？",
"我讨厌推销电话，请不要给我们打电话",
"联系方式用邮箱最好，我的是user@example.com",
"我比较关注空气质量，你们的通风系统怎么样？",
"我上次来的时候看到更衣室有点旧，有计划翻新吗？",
"我儿子15岁，能带他来吗？",
"他喜欢打篮球，你们有篮球场吗？",
"价格方面，年付有折扣吗？",
"我更喜欢电子发票，不要纸质的",
"付款能用支付宝吗？"
]
# 执行对话
for i, user_input in enumerate(conversation):
print(f"\n[轮次 {i+1}] 用户: {user_input}")
# 调用对话链
response = conversational_chain.invoke({"input": user_input})
print(f"AI: {response.content}")
# 每5轮显示一次记忆状态
if (i+1) % 5 == 0:
print("\n=== 当前记忆状态 ===")
memory = response["memory"]
print(f"短期记忆: {memory.get('short_term', [])[-2:]}")
print(f"摘要记忆: {memory.get('summary', '')[:100]}...")
print(f"相关记忆: {memory.get('relevant_memories', [])[:1] or '无'}")
print("===================\n")
# 最终记忆摘要
print("\n=== 最终记忆摘要 ===")
memory = response["memory"]
print(f"摘要记忆: {memory.get('summary', '')}")案例解析：三层记忆架构
1. 短期记忆层（Redis实现）
-
功能：存储最近的5条对话
-
实现：使用Redis存储，自动过期
-
优点：保持对话流畅性
-
代码关键点 ：
pythonself.short_term_memory = RedisChatMessageHistory( session_id, url="redis://localhost:6379/0", key_prefix="short_term_", ttl=3600 )
2. 摘要记忆层（ConversationSummaryBufferMemory）
-
功能：动态生成对话摘要
-
实现：使用LLM生成摘要
-
优点：压缩长期对话历史
-
代码关键点 ：
pythonself.summary_memory = ConversationSummaryBufferMemory( llm=llm, memory_key="summary", max_token_limit=500, return_messages=True )
3. 向量记忆层（Chroma向量数据库）
-
功能：存储关键信息片段
-
实现：使用OpenAI嵌入+相似度搜索
-
优点：精准检索相关信息
-
代码关键点 ：
pythonself.vector_store = Chroma( collection_name=f"memory_{session_id}", embedding_function=self.embeddings, persist_directory="./chroma_db" ) def _retrieve_relevant_memories(self, query, k=3): docs = self.vector_store.similarity_search(query, k=k) return [f"[记忆片段] {doc.page_content}" for doc in docs]
分层记忆工作流程
txt
graph TD
A[用户输入] --> B(分层记忆系统)
B --> C{记忆检索}
C --> D[短期记忆：最近5条]
C --> E[摘要记忆：对话摘要]
C --> F[向量记忆：相关片段]
D --> G[构建提示]
E --> G
F --> G
G --> H[LLM处理]
H --> I[生成响应]
I --> J{记忆存储}
J --> K[添加到短期记忆]
J --> L[更新摘要记忆]
J --> M[重要信息存入向量记忆]
K --> A
L --> A
M --> A记忆优化策略
1. 动态重要性判断
python
def _is_important(self, text):
"""判断信息是否重要"""
important_keywords = ["喜欢", "不喜欢", "重要", "记住", "总是", "从不", "偏好"]
return any(keyword in text for keyword in important_keywords)2. 定期记忆清理
python
def cleanup(self):
"""定期清理旧记忆"""
# 保留最近7天的摘要记忆
self.summary_memory.prune()
# 短期记忆由Redis TTL自动清理
# 向量记忆保留所有（实际应用中可设置过期时间）3. 提示工程优化
python
self.prompt = ChatPromptTemplate.from_messages([
("system", """
你是一个专业的客户服务AI，具有完美的记忆力。使用以下分层记忆提供精准回答：
[近期对话] (最新5条)
{short_term}
[对话摘要] (长期摘要)
{summary}
[相关记忆] (基于当前问题检索)
{relevant_memories}
当前时间: {current_time}
"""),
MessagesPlaceholder(variable_name="history"),
("human", "{input}")
])执行效果分析
在18轮对话测试中：
- 短期记忆：始终显示最近2-3条对话
- 摘要记忆 ：随着对话进行动态更新
- 第5轮：摘要包含会员计划、健身时间、关节炎等
- 第10轮：新增联系方式偏好、推销电话厌恶
- 第15轮：包含家庭优惠、儿子年龄等信息
- 向量记忆 ：精准检索相关信息
- 当询问"游泳池水质"时：检索到"关节炎，游泳是最好的运动方式"
- 当询问"家庭优惠"时：检索到"我太太也想来，她有瑜伽经验"
生产环境建议
-
向量存储选择：
- 小规模：Chroma/FAISS
- 大规模：Pinecone, Weaviate
-
摘要优化：
python# 自定义摘要提示 SUMMARY_PROMPT = PromptTemplate.from_template( "请用中文简洁地概括以下对话，特别关注用户偏好和重要细节：\n{conversation}" ) self.summary_memory = ConversationSummaryBufferMemory( llm=llm, prompt=SUMMARY_PROMPT, memory_key="summary" ) -
记忆生命周期管理：
python# 设置不同记忆层的TTL SHORT_TERM_TTL = 24 * 3600 # 1天 SUMMARY_TTL = 7 * 24 * 3600 # 7天 VECTOR_TTL = 30 * 24 * 3600 # 30天
这个分层解决方案有效解决了长对话中的token限制问题，同时保持了对话的连贯性和个性化。通过组合三种不同类型的记忆，系统既能记住关键细节，又能避免被无关历史干扰，实现了高效、精准的长期对话管理。