AI Agent短期记忆实现：从滑动窗口到向量检索的工程实践

1. 项目概述为什么AI Agent需要“短期记忆”如果你最近在折腾AI Agent大概率会遇到一个让人头疼的问题你问它“我刚才说了什么”它要么答非所问要么直接告诉你“作为AI我没有记忆功能”。这感觉就像在和一个金鱼聊天每次对话都从零开始。今天要聊的就是给这个“金鱼”装上短期记忆让它能记住最近几分钟甚至几轮对话的内容让交互变得更像和一个“人”在交流。这个需求在构建任何有连续对话能力的Agent时都是刚需。无论是客服机器人、编程助手还是游戏里的NPC如果它记不住你上一步的操作、上一句的指令体验就会非常割裂。比如你让Agent帮你写代码先说“创建一个用户模型”然后补充“给这个模型加上邮箱验证字段”如果Agent没有短期记忆它很可能就只执行第二条指令而不知道“这个模型”指的是谁。从技术上看实现短期记忆本质上是在管理对话的“上下文”Context。大模型本身比如通过OpenAI API调用的GPT系列并不具备记忆能力它只是一个根据输入的文本预测下一个词的概率机器。我们每次调用API都需要把完整的“历史对话当前问题”打包成一个提示Prompt送进去。短期记忆系统就是帮我们高效、智能地管理这个不断增长的“历史对话包”确保它在不超出模型上下文长度限制的前提下保留最相关、最重要的信息。2. 短期记忆的核心设计思路与方案选型给Agent加记忆听起来简单做起来有几个关键问题要解决记什么记多久怎么记记满了怎么办这直接决定了Agent的“智商”和“情商”。2.1 记忆的粒度与分类首先我们得定义记忆里存什么。一股脑地把所有对话历史都存下来是最笨的办法不仅效率低而且会让关键信息被淹没。通常我们把短期记忆分为几个层次原始对话记录最基础的按顺序存储用户和Agent的每一轮问答。这是记忆的“原材料”。实体记忆从对话中提取的关键实体信息比如用户的名字“小明”、项目名称“XX系统”、提到的特定参数“max_retries3”。这些是对话的“关键词”。摘要记忆当对话轮数变多时用大模型对之前的对话历史进行总结生成一段精炼的摘要。比如将关于“设计数据库表结构”的10轮讨论总结为“用户决定采用三张核心表User, Order, Product并明确了关联关系”。摘要记忆能极大地压缩信息量。意图/主题记忆记录当前对话围绕的核心主题或用户的意图流。例如从“帮我推荐电影”到“要科幻片”再到“不要近十年的”这是一个连续的意图链。对于“短期记忆”这个场景我们主要聚焦于原始对话记录和实体记忆的即时管理与应用。摘要记忆更偏向于中长期记忆的维护。2.2 记忆的存储与检索记忆存好了用的时候怎么快速找到这里涉及到存储后端和检索策略。存储后端的选择很简单内存存储最简单用Python的List或Dict存在程序运行内存里。优点是快零延迟。缺点是程序重启记忆就清零且无法分布式扩展。适合原型开发、单次会话的简单Agent。向量数据库这是目前的主流方案。将每一条记忆比如一句话或一个实体转换成向量Embedding存入像Chroma、Weaviate、Pinecone这样的向量数据库。检索时将当前问题也转换成向量通过计算余弦相似度来找出最相关的历史记忆。优点是能实现基于语义的相似性搜索即使问题表述不同也能找到相关记忆。缺点是多了一次向量化的开销架构变复杂。传统数据库用SQLite、PostgreSQL等按时间戳、会话ID、实体类型等字段结构化存储。检索时用SQL查询。优点是可精确查询结构清晰。缺点是无法做语义相似度匹配。对于短期记忆如果对话轮次不多比如少于20轮内存存储配合简单的最近N轮滑动窗口是最高效实用的方案。当对话轮次增加需要从长篇历史中精准提取相关片段时才需要考虑引入向量数据库。检索策略是记忆系统的灵魂。你不能每次都将全部记忆喂给模型那会浪费Token并可能超出限制。常见的策略有最近N轮只取最近N轮对话。简单粗暴保证记忆新鲜但可能丢失重要的早期信息。关键词匹配从当前问题中提取关键词去记忆库里做字符串匹配。速度快但不够智能。向量相似度检索如上所述用向量数据库找到语义最相关的几条记忆。最灵活智能但成本最高。混合检索结合多种策略。例如总是带上最近2轮对话保证连贯性同时通过向量检索找出与当前问题最相关的3条历史记忆保证相关性。2.3 记忆的更新与遗忘记忆不是只进不出的。一个健壮的系统必须有“遗忘”机制也就是记忆的更新策略。滑动窗口这是最常用的短期记忆管理方式。只保留一个固定大小的最近对话窗口例如最近10轮。新的对话进来最老的对话就被“挤出去”遗忘掉。这模拟了人类注意力聚焦于最近事件的特点。重要性衰减为每条记忆赋予一个“重要性”分数随着时间或新记忆的加入分数逐渐衰减。当分数低于阈值或需要腾空间时优先移除低分记忆。这个分数可以由规则设定例如包含用户明确指令的记忆分数高也可以尝试用大模型来打分。主动总结与压缩当记忆积累到一定量时触发一个总结动作用大模型生成一段摘要然后用这个摘要替代原来的一大段原始记忆。这样既保留了核心信息又大大节省了空间。在我们的实现中将主要采用内存存储 滑动窗口作为核心这是平衡复杂度和效果的最佳起点。同时我们会设计一个可扩展的架构以便未来可以轻松接入向量数据库或更复杂的检索策略。3. 基于ReAct框架与滑动窗口的短期记忆实现现在我们进入实战环节。我将以一个基于OpenAI API和ReActReasoning and Acting框架的简单Agent为例演示如何为其注入短期记忆能力。ReAct框架让Agent通过“思考Thought-行动Action-观察Observation”的循环来完成任务这正是记忆大显身手的地方。3.1 基础Agent与记忆管理类搭建首先我们定义记忆的数据结构。一条记忆单元Memory Unit至少包含角色、内容和时间戳。from datetime import datetime from typing import List, Dict, Any from pydantic import BaseModel class MemoryUnit(BaseModel): 记忆单元 role: str # user 或 assistant content: str # 对话内容 timestamp: datetime datetime.now() class ShortTermMemory: 短期记忆管理类 def __init__(self, window_size: int 10): self.memories: List[MemoryUnit] [] self.window_size window_size # 滑动窗口大小 def add(self, role: str, content: str): 添加一条新记忆 new_memory MemoryUnit(rolerole, contentcontent) self.memories.append(new_memory) # 应用滑动窗口如果超出限制移除最旧的记忆 if len(self.memories) self.window_size: self.memories.pop(0) def get_recent(self, n: int None) - List[MemoryUnit]: 获取最近N条记忆默认返回全部但受窗口限制 if n is None or n len(self.memories): return self.memories.copy() return self.memories[-n:].copy() def get_formatted_context(self, num_last_turns: int 5) - str: 将记忆格式化为字符串上下文用于拼接到Prompt中 recent_memories self.get_recent(num_last_turns) context_lines [] for mem in recent_memories: prefix 用户 if mem.role user else 助手 context_lines.append(f{prefix}: {mem.content}) return \n.join(context_lines) def clear(self): 清空记忆例如开始一个新会话 self.memories.clear()这个ShortTermMemory类就是Agent的“大脑皮层”。window_size控制着记忆的容量add方法在每次对话后更新记忆并执行滑动窗口淘汰get_formatted_context方法则负责在调用大模型前把记忆打包成模型能理解的文本格式。3.2 将记忆整合进ReAct Agent的循环接下来我们改造一个简单的ReAct Agent让它能在“思考”时参考短期记忆。import openai import re class ReActAgentWithMemory: def __init__(self, api_key: str, memory: ShortTermMemory): openai.api_key api_key self.memory memory # 一个简单的工具集合示例 self.tools { search_web: 在互联网上搜索信息。输入查询字符串。, calculate: 执行数学计算。输入数学表达式。, get_time: 获取当前时间。无需输入。 } def run_step(self, user_input: str) - str: # 1. 将用户输入存入记忆 self.memory.add(user, user_input) # 2. 构建包含记忆的Prompt prompt f 你是一个有帮助的助手可以调用工具。你有以下短期记忆 {self.memory.get_formatted_context(5)} 当前对话 用户{user_input} 你可以使用的工具 {self._format_tools()} 请按照以下格式回应 思考[你对用户请求的分析以及是否需要使用工具、使用哪个工具] 行动工具名称如果不需要工具则写“无”| 工具输入如果没有输入则写“无” # 3. 调用大模型获取“思考”和“行动” response openai.ChatCompletion.create( modelgpt-3.5-turbo, messages[{role: user, content: prompt}], temperature0 ) full_response response.choices[0].message.content # 4. 解析模型输出 thought_match re.search(r思考(.), full_response) action_match re.search(r行动(.)\|(.), full_response) thought thought_match.group(1) if thought_match else action_name action_match.group(1).strip() if action_match else 无 action_input action_match.group(2).strip() if action_match else 无 final_answer # 5. 执行行动工具调用 if action_name ! 无 and action_name in self.tools: observation self._execute_action(action_name, action_input) # 将工具执行结果作为“观察”存入记忆并可能进行下一轮思考简化起见这里只执行一步 self.memory.add(assistant, f[调用工具 {action_name}输入{action_input}]) final_answer f我已通过工具{action_name}处理了你的请求。结果{observation} else: # 无需工具直接提取模型生成的回答这里需要更复杂的解析为简化我们假设模型在思考后直接给出了答案 # 实际上更鲁棒的做法是让模型在“思考”后直接输出最终答案。 answer_match re.search(r最终答案(.), full_response, re.DOTALL) final_answer answer_match.group(1).strip() if answer_match else 我无法处理这个请求。 # 6. 将助手的最终回应存入记忆 self.memory.add(assistant, final_answer) return final_answer def _format_tools(self): return \n.join([f- {name}: {desc} for name, desc in self.tools.items()]) def _execute_action(self, name: str, input_str: str): # 模拟工具执行 if name get_time: from datetime import datetime return datetime.now().strftime(%Y-%m-%d %H:%M:%S) elif name calculate: try: return str(eval(input_str)) except: return 计算错误 elif name search_web: return f模拟搜索关键词 {input_str} 的结果。 return 未知工具这个Agent的工作流程清晰体现了记忆的参与记忆写入用户输入和助手回应都会被即时存入ShortTermMemory。记忆读取在每一步的Prompt构建中通过get_formatted_context方法将最近的记忆作为上下文注入。这样模型在“思考”时就能看到类似“用户我之前让你查了今天的天气。助手今天晴天25度。”这样的历史从而理解“那明天呢”中的“那”指代的是天气。记忆驱动决策历史记忆能帮助模型更好地决定是否需要调用工具、调用哪个工具。例如如果用户刚问完“北京的房价”接着说“上海的呢”模型结合记忆就知道应该再次调用搜索工具但关键词换成了“上海房价”。3.3 关键参数配置与优化实现只是第一步调优才能让记忆系统好用。以下几个参数需要仔细考量滑动窗口大小 (window_size)这是最重要的参数。设得太小如3Agent容易遗忘重要前提设得太大如50会浪费大量Token在无关历史信息上还可能挤占当前问题的空间。一个经验值是8到15轮。对于任务型对话如客服、编码可以小一些保证聚焦对于开放聊天可以大一些。实操心得不要静态设置。可以尝试动态窗口在检测到用户提及“之前”、“刚才”、“上面说的”等指代词时临时扩大检索范围或优先保留那些包含明确指令如“请记住”、“重点是”的记忆。注入记忆的轮数 (num_last_turnsinget_formatted_context)不一定需要把窗口内所有记忆都喂给模型。通常注入最近5-7轮对话已经能覆盖绝大多数短期指代需求。这能有效节省Token。记忆的格式化方式上面例子用了“用户xxx”的简单格式。更优的做法是使用模型训练时熟悉的格式比如OpenAI的ChatML格式|im_start|user\n...|im_end|\n|im_start|assistant\n...|im_end|。这能略微提升模型对上下文的理解。Token计数与截断必须时刻警惕上下文长度限制如GPT-3.5-turbo的4096 Token。在get_formatted_context方法中应该加入Token计数逻辑当格式化后的上下文当前问题超过阈值时主动进行截断。截断策略可以是丢弃最老的记忆或者触发一次记忆总结。import tiktoken # OpenAI的Token计数库 def get_formatted_context_with_token_limit(self, max_tokens: int 2000) - str: 考虑Token限制的上下文格式化 encoder tiktoken.encoding_for_model(gpt-3.5-turbo) formatted_parts [] total_tokens 0 # 从最新记忆开始反向添加 for mem in reversed(self.memories): text f{mem.role}: {mem.content} token_count len(encoder.encode(text)) if total_tokens token_count max_tokens: break formatted_parts.insert(0, text) # 保持时间顺序 total_tokens token_count return \n.join(formatted_parts)4. 高级技巧从短期记忆到“工作记忆”基础的滑动窗口记忆已经能让Agent的对话连贯性大幅提升。但我们可以更进一步模仿人类的“工作记忆”让Agent不仅能记住还能主动提取和利用关键信息。4.1 关键实体提取与记忆与其平等对待所有对话不如让Agent学会抓住重点。我们可以在add记忆时同步运行一个简单的实体提取流程。import spacy # 一个流行的NLP库 class EnhancedShortTermMemory(ShortTermMemory): def __init__(self, window_size: int 10): super().__init__(window_size) # 加载一个小型NLP模型来提取实体 try: self.nlp spacy.load(zh_core_web_sm) # 中文模型 except: self.nlp None self.entities {} # 存储提取的实体格式{实体文本: 出现次数/相关信息} def add_and_extract(self, role: str, content: str): 添加记忆并提取实体 self.add(role, content) if self.nlp: doc self.nlp(content) for ent in doc.ents: # 提取命名实体人名、地名、机构等 entity_text ent.text if entity_text in self.entities: self.entities[entity_text][count] 1 # 可以记录最近出现的位置或上下文 else: self.entities[entity_text] {count: 1, type: ent.label_} # 也可以添加基于规则的提取比如匹配“叫XXX”、“设置为YYY”这样的模式 def get_relevant_entities(self, query: str) - List[str]: 根据查询找出相关实体简单实现字符串包含 relevant [] for entity in self.entities: if entity in query: relevant.append(entity) return relevant这样当用户说“帮我订一张去北京的机票”Agent不仅记下这句话还会提取出“北京”作为一个关键实体地点。当用户后来说“那改成上海吧”Agent可以通过get_relevant_entities发现“上海”与记忆中的实体“北京”属于同一类型地点从而更准确地理解用户是要“更改目的地”。4.2 基于向量检索的记忆增强当对话历史很长滑动窗口无法覆盖所有重要信息时就需要向量检索登场了。我们可以在ShortTermMemory类中集成一个轻量级向量数据库。# 假设使用ChromaDB作为内存向量数据库 import chromadb from chromadb.config import Settings from sentence_transformers import SentenceTransformer # 用于生成向量 class VectorEnhancedMemory(ShortTermMemory): def __init__(self, window_size: int 10, embedding_model_name: str paraphrase-multilingual-MiniLM-L12-v2): super().__init__(window_size) self.embedder SentenceTransformer(embedding_model_name) chroma_client chromadb.Client(Settings(chroma_db_implduckdbparquet, persist_directory./chroma_db)) # 创建一个集合来存储记忆向量 self.memory_collection chroma_client.create_collection(nameconversation_memories) self._current_id 0 def add(self, role: str, content: str): super().add(role, content) # 同时将记忆向量化存储 memory_text f{role}: {content} embedding self.embedder.encode(memory_text).tolist() self.memory_collection.add( embeddings[embedding], documents[memory_text], metadatas[{role: role, timestamp: datetime.now().isoformat()}], ids[str(self._current_id)] ) self._current_id 1 def retrieve_relevant(self, query: str, k: int 3) - List[str]: 基于当前查询检索语义最相关的K条历史记忆 query_embedding self.embedder.encode(query).tolist() results self.memory_collection.query( query_embeddings[query_embedding], n_resultsmin(k, len(self.memories)) ) return results[documents][0] if results[documents] else []在构建Prompt时我们就可以结合两种记忆def get_context_for_query(self, user_query: str) - str: # 1. 获取最近几轮对话保证连贯性 recent_context self.get_formatted_context(3) # 2. 获取语义相关的记忆弥补窗口外的信息 relevant_memories self.retrieve_relevant(user_query, 2) combined [] if recent_context: combined.append(最近对话) combined.append(recent_context) if relevant_memories: combined.append(相关历史) combined.append(\n.join(relevant_memories)) return \n.join(combined)这种“滑动窗口 向量检索”的混合模式既能保证对话的局部连贯性又能从更长的历史中召回相关信息是实现强大短期记忆的实用架构。4.3 记忆的持久化与会话管理对于需要长期运行的Agent服务内存记忆显然不够。我们需要将会话记忆持久化到数据库如Redis、SQLite并以session_id为键进行管理。import json import redis # 示例使用Redis class PersistentMemoryManager: def __init__(self, redis_client): self.redis redis_client self.ttl 3600 # 记忆存活时间例如1小时 def save_session_memory(self, session_id: str, memory: ShortTermMemory): 保存整个记忆对象到Redis # 将MemoryUnit列表转换为可序列化的字典列表 mem_data [mem.dict() for mem in memory.memories] data_str json.dumps({ memories: mem_data, window_size: memory.window_size }) self.redis.setex(fagent:memory:{session_id}, self.ttl, data_str) def load_session_memory(self, session_id: str) - ShortTermMemory: 从Redis加载记忆对象 data_str self.redis.get(fagent:memory:{session_id}) if not data_str: return ShortTermMemory() # 返回新的空记忆 data json.loads(data_str) memory ShortTermMemory(window_sizedata[window_size]) # 将字典列表恢复为MemoryUnit对象列表 memory.memories [MemoryUnit(**mem) for mem in data[memories]] return memory在Web服务或聊天机器人中每次收到用户请求先根据用户ID或会话ID从PersistentMemoryManager中加载记忆Agent处理完后再保存回去。这样就实现了跨请求的短期记忆保持。5. 常见问题、调试技巧与效果评估在实际实现和调试短期记忆系统时你肯定会遇到一些典型问题。下面是我踩过坑后总结的一些经验。5.1 常见问题与排查表问题现象可能原因排查步骤与解决方案Agent完全“忘记”之前说过的话。1. 记忆根本没有被添加到存储中。2.get_formatted_context返回了空字符串。3. 构建Prompt时记忆上下文被错误地放置或格式错误导致模型忽略。1.打日志在add和get_formatted_context方法中加入日志确认记忆的读写流程。2.检查Prompt将最终发送给API的完整Prompt打印出来肉眼检查历史对话是否在其中格式是否清晰如是否有正确的角色标识。3.简化测试先固定注入2-3轮历史看Agent是否能正确引用。Agent引用了错误的历史信息张冠李戴。1. 滑动窗口过大注入了过多无关历史造成干扰。2. 记忆的格式化方式混乱导致模型分不清说话者。3. 向量检索返回了相似但不相关的记忆。1.减小窗口尝试将num_last_turns从10减到5或3。2.标准化格式严格使用“用户”和“助手”这样的前缀并确保每段记忆独立一行。3.优化检索检查向量检索的相似度阈值或尝试在检索时加入元数据过滤如只检索roleuser的记忆。对话变得冗长或重复。1. 记忆中包含了过去冗长的模型输出导致新的回答也模仿了冗长的风格。2. Agent将自己的上一次回答也作为历史读入可能导致循环或重复。1.过滤记忆内容可以考虑不将助手过长的工具调用结果全文存入记忆而是存一个摘要如“用户询问了天气助手回复了晴天”。2.调整记忆策略尝试只存储用户的输入和最终助手答案不存储中间链式思考Thought过程。Token使用量激增成本过高或触发长度限制。1. 记忆窗口太大注入的上下文过长。2. 没有对长文本记忆进行压缩或总结。1.实施Token计数与截断如上文所述在get_formatted_context中加入Token计数和截断逻辑。2.启用摘要记忆当对话轮次超过阈值如15轮时触发一个后台任务用大模型将早期对话总结成一段话替换掉原始的多轮记忆。在多轮复杂任务中Agent迷失方向。记忆虽然存在但缺乏对核心任务目标的突出强调。引入“焦点记忆”在记忆系统中单独维护一个“当前任务目标”字符串。在每一轮Prompt中都显式强调它。例如“当前核心任务为用户预订机票。已收集信息目的地北京时间明天。”5.2 调试技巧与实操心得可视化你的记忆写一个简单的调试端点或函数实时输出当前记忆的内容和实体列表。这比看日志直观得多。进行A/B测试对于关键参数如window_size可以设计一组标准化的多轮对话测试集分别用不同参数运行Agent人工或自动评估其回答的连贯性和准确性。选择效果最好的那个值。关注“指代消解”能力这是检验短期记忆是否有效的金标准。设计包含“它”、“这个”、“那个”、“他”、“她”等指代词的对话看Agent能否正确理解所指。例如用户请介绍Python的列表。 助手Python列表是一种可变的有序集合... 用户它和元组的主要区别是什么 一个拥有良好短期记忆的Agent应该能知道“它”指的是“Python列表”。从简单开始逐步复杂化不要一开始就追求完美的向量检索混合架构。先用一个固定大小的滑动窗口实现最基本的功能确保它工作正常。然后在此基础上逐步添加实体提取、Token截断、持久化等功能。每加一个特性都充分测试。记忆的“毒性”问题要注意并非所有历史信息都是有益的。如果历史对话中包含了错误信息、被用户纠正的内容或者Agent之前失败的尝试这些“有毒”记忆可能会干扰后续决策。一种策略是当用户明确说“不对”、“我指的是”进行纠正时不仅要添加新的正确记忆最好能主动移除或标记之前相关的错误记忆。实现短期记忆是让你的AI Agent从“一问一答的机器”迈向“连续对话的伙伴”的关键一步。它没有想象中那么复杂核心就是一个精心管理的上下文窗口。但把它做好却需要仔细考虑交互设计、资源限制和实际效果。