ARTICLE DETAIL

资讯详情

深耕网站视觉设计与运营推广的一线实战洞察。

基于ChromaDB与RAG的AI Agent记忆系统设计与实现

基于ChromaDB与RAG的AI Agent记忆系统设计与实现 1. 项目概述当AI助手开始“健忘”最近在折腾AI Agent项目时我遇到了一个挺典型的问题Agent记不住事。你让它处理一个多轮对话或者执行一个需要上下文信息的复杂任务它经常表现得像个“金鱼”七秒之后就忘了之前聊过什么。比如你让它帮你整理一份会议纪要先给了它会议主题再上传录音文件让它提取要点最后让它根据要点生成总结。结果在最后一步它可能完全忘了会议主题是什么或者把不同发言人的观点张冠李戴。这显然不是我们想要的“智能”助手。这个问题的核心就是Agent的“记忆”Memory能力缺失。一个没有记忆的Agent就像一台没有硬盘的电脑每次对话都是全新的开始无法积累知识也无法进行连贯的思考。为了解决这个问题业界普遍采用“记忆模块”结合“检索增强生成”RAG的方案。简单来说就是给Agent装上一个“外置大脑”让它不仅能记住当前对话的短期上下文工作记忆还能从海量的历史数据或知识库中精准地检索出相关信息长期记忆从而做出更准确、更连贯的决策和回复。我这次实践的目标就是为我的Agent补上这块短板。具体来说我选择了ChromaDB作为向量数据库来存储和检索记忆并设计了一套结合短期记忆Conversation Buffer Memory和长期记忆RAG检索的混合架构。整个过程涉及记忆的写入、存储、检索、融合与应用等多个环节。下面我就把这套方案的完整设计思路、技术选型、实操步骤以及踩过的坑毫无保留地分享出来。2. 核心架构设计短期记忆与长期记忆的融合一个健壮的Agent记忆系统不能是单一维度的。人类的记忆本身就分为短期记忆和长期记忆。短期记忆处理即时信息容量有限但存取快长期记忆存储经验知识容量巨大但需要线索提取。我们的Agent记忆系统也应如此设计。2.1 记忆系统的分层设计我的架构主要分为三层短期记忆层Conversation Buffer Memory作用缓存当前对话轮次例如最近10轮的完整历史。这是Agent的“工作台”保证了对话的即时连贯性。实现通常使用一个简单的键值对或列表在内存中维护。在LangChain等框架中有现成的ConversationBufferMemory类可以直接使用。特点存取速度极快但容量小会话结束或重启后即丢失。长期记忆存储与检索层RAG with Vector DB作用将重要的对话片段、用户偏好、任务结果等经过处理后存入向量数据库形成Agent的“知识库”或“经验库”。当需要时通过检索Recall将其召回。实现核心是向量化和向量检索。使用文本嵌入模型如text-embedding-ada-002,BGE将文本转换为高维向量存入ChromaDB。检索时将当前问题或上下文也转换为向量在数据库中查找最相似的向量即最相关的记忆。特点容量近乎无限可持久化但检索需要计算存在精度和召回率的权衡。记忆融合与决策层作用将短期记忆和检索到的长期记忆进行融合、去重和排序形成一份完整的“上下文背景板”提交给大语言模型LLM进行最终的任务规划和内容生成。实现这是最体现设计功力的地方。简单的做法是直接拼接但更好的做法是引入重排序Re-ranking技术对检索结果进行二次精排确保最相关的信息排在前面。2.2 为什么选择ChromaDB在众多向量数据库如Milvus, Pinecone, Weaviate, Qdrant中我选择了ChromaDB主要基于以下几点考量轻量易用开发友好ChromaDB的设计哲学就是“简单”。它可以直接在Python中以内嵌模式运行几行代码就能启动无需复杂的服务部署。对于快速原型验证和个人项目来说这极大地降低了门槛。与LangChain生态无缝集成我的Agent基础框架基于LangChain而ChromaDB是LangChain官方推荐且集成度最高的向量数据库之一。Chroma类提供了开箱即用的方法用于文档加载、拆分、嵌入、存储和检索大大减少了样板代码。足够的性能与功能对于中小规模的知识库百万级向量以内ChromaDB的性能完全够用。它支持多种距离函数默认是L2也支持余弦相似度和内积并且具备基本的过滤Filtering能力能满足大部分RAG场景的需求。开源且可持久化虽然内存模式方便测试但ChromaDB也支持将数据持久化到磁盘这对于长期运行的Agent应用是必需的。注意如果你的项目面向海量数据亿级向量和高并发生产环境可能需要评估Milvus或Pinecone这类分布式向量数据库。但对于绝大多数Agent应用和RAG场景ChromaDB是一个平衡了易用性与能力的绝佳起点。2.3 整体工作流程一次完整的Agent交互其记忆系统的流程如下用户输入用户提出新的请求或信息。短期记忆读取Agent从ConversationBufferMemory中读取最近的对话历史。长期记忆检索基于用户输入和/或短期记忆生成检索查询Query。将查询向量化在ChromaDB中进行相似性搜索返回Top-K个相关的历史记忆片段。记忆融合将短期记忆和检索到的长期记忆片段按照一定的策略如按时间、按相关性组合成一个完整的上下文提示Prompt。LLM推理与行动将融合后的上下文和用户当前请求一起提交给LLM。LLM基于所有记忆进行思考规划步骤调用工具并生成回复。记忆写入将本轮交互中有价值的信息如LLM的最终结论、工具执行的关键结果写入短期记忆缓冲区。同时判断是否需要将某些信息如达成的共识、用户明确指示的偏好存入ChromaDB作为长期记忆。3. 实操搭建从零构建记忆增强型Agent理论讲完了我们动手搭建。这里我以构建一个“智能学习助手”Agent为例它需要记住用户的学习目标、已掌握的知识点并在后续对话中提供连贯的辅导。3.1 环境准备与依赖安装首先创建一个干净的Python环境推荐3.9并安装核心库。# 创建虚拟环境可选 python -m venv agent_memory_env source agent_memory_env/bin/activate # Linux/Mac # agent_memory_env\Scripts\activate # Windows # 安装核心依赖 pip install langchain langchain-community langchain-openai chromadb # 安装文本嵌入模型这里使用OpenAI的embedding也可用本地模型如BGE pip install openai # 安装一个轻量级LLM用于测试这里用OpenAI GPT也可用Ollama本地部署其他模型 # 确保你已设置好OPENAI_API_KEY环境变量3.2 初始化记忆组件与向量数据库我们首先初始化短期记忆和长期记忆的存储后端。import os from langchain.memory import ConversationBufferMemory from langchain.vectorstores import Chroma from langchain.embeddings import OpenAIEmbeddings from langchain.schema import Document # 1. 初始化短期记忆 short_term_memory ConversationBufferMemory( memory_keychat_history, # 在prompt中引用的变量名 return_messagesTrue # 以消息列表格式返回更适合与ChatModel配合 ) # 2. 初始化嵌入模型和向量数据库长期记忆 # 使用OpenAI的嵌入模型需要设置API Key os.environ[OPENAI_API_KEY] your-openai-api-key embeddings OpenAIEmbeddings(modeltext-embedding-ada-002) # 指定一个持久化路径 persist_directory ./chroma_db # 初始化Chroma向量库。如果路径不存在则新建如果存在则加载。 vectorstore Chroma( collection_nameagent_long_term_memory, # 集合名称类似数据库的表 embedding_functionembeddings, persist_directorypersist_directory ) # 检查一下是否已有数据 print(f当前向量库中的文档数量{vectorstore._collection.count()})3.3 设计记忆的存储格式与切片策略直接存储大段的对话文本到向量数据库效率低下且检索精度差。我们需要对记忆进行结构化处理。存储格式每条记忆不应是原始的“用户说...AI说...”而应该是原子化的、富含信息的事实或指令。例如原始对话“用户我的目标是三个月内学会Python数据分析。AI好的我们可以从Pandas和Matplotlib开始。”结构化记忆1用户目标Document(page_content用户学习目标在三个月内掌握Python数据分析。, metadata{type: goal, user: default, timestamp: 2024-05-27})结构化记忆2AI建议Document(page_content建议学习路径起点Pandas库和Matplotlib库。, metadata{type: advice, source: AI, timestamp: 2024-05-27})切片策略使用RecursiveCharacterTextSplitter将长的文档或对话记录切分成语义相对完整的小片段chunks。这里的关键是chunk_size和chunk_overlap参数。from langchain.text_splitter import RecursiveCharacterTextSplitter text_splitter RecursiveCharacterTextSplitter( chunk_size500, # 每个片段大约500字符 chunk_overlap50, # 片段间重叠50字符避免语义被割裂 length_functionlen, separators[\n\n, \n, 。, , , , , , ] # 中文分隔符 )3.4 实现记忆的写入与检索逻辑接下来我们封装记忆的增删改查操作。class AgentMemoryManager: def __init__(self, vectorstore, short_term_memory): self.vectorstore vectorstore self.short_term_memory short_term_memory self.text_splitter RecursiveCharacterTextSplitter(chunk_size500, chunk_overlap50) def save_to_long_term(self, text, metadataNone): 将重要信息存入长期记忆向量数据库 if metadata is None: metadata {} # 1. 创建基础文档 doc Document(page_contenttext, metadatametadata) # 2. 文本切片如果是长文本 splits self.text_splitter.split_documents([doc]) # 3. 添加到向量库 self.vectorstore.add_documents(splits) # 4. 持久化到磁盘 self.vectorstore.persist() print(f[长期记忆] 已保存 {len(splits)} 个片段。) def retrieve_from_long_term(self, query, k3, filter_dictNone): 从长期记忆中检索相关信息 # 使用相似性搜索 docs_and_scores self.vectorstore.similarity_search_with_relevance_scores(query, kk, filterfilter_dict) # 返回文档列表和相似度分数 retrieved_docs [doc for doc, score in docs_and_scores] scores [score for doc, score in docs_and_scores] return retrieved_docs, scores def get_short_term_context(self): 获取当前的短期记忆对话历史 # 从memory中加载历史变量 history_dict self.short_term_memory.load_memory_variables({}) return history_dict.get(chat_history, []) def update_short_term(self, human_input, ai_output): 更新短期记忆保存一轮对话 self.short_term_memory.save_context({input: human_input}, {output: ai_output}) # 初始化记忆管理器 memory_manager AgentMemoryManager(vectorstore, short_term_memory)3.5 构建融合记忆的Agent执行链这是最核心的一步我们将短期记忆、长期记忆和LLM组合成一个可执行的链。from langchain_openai import ChatOpenAI from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder from langchain.schema.runnable import RunnablePassthrough # 1. 初始化LLM llm ChatOpenAI(modelgpt-3.5-turbo, temperature0.2) # 2. 定义Prompt模板其中预留了位置给历史对话和检索到的记忆 prompt_template ChatPromptTemplate.from_messages([ (system, 你是一个智能学习助手拥有与用户对话的短期记忆和存储知识点的长期记忆。 请充分利用以下背景信息来回答用户的问题或执行任务。 【长期记忆相关知识点】 {long_term_memory} 【短期记忆最近对话】 {short_term_memory} 当前用户输入{input} 请开始你的回答或行动规划), MessagesPlaceholder(variable_namechat_history), # 这里也会自动填充但我们在上面已经显式格式化了 ]) # 3. 构建一个Runnable动态获取记忆 def format_memories(input_dict): 一个自定义函数用于在链运行时获取并格式化记忆 user_input input_dict[input] # 获取短期记忆并格式化成字符串 short_term_msgs memory_manager.get_short_term_context() short_term_str \n.join([f{msg.type}: {msg.content} for msg in short_term_msgs[-6:]]) # 取最近3轮 # 基于用户输入检索长期记忆 retrieved_docs, _ memory_manager.retrieve_from_long_term(user_input, k2) long_term_str \n.join([f- {doc.page_content} (来源{doc.metadata}) for doc in retrieved_docs]) return { input: user_input, short_term_memory: short_term_str, long_term_memory: long_term_str if long_term_str else 暂无直接相关的长期记忆, chat_history: short_term_msgs # 仍然保留原始消息列表给Placeholder如果模板需要 } # 4. 组装执行链 agent_chain ( RunnablePassthrough.assign(**{input: lambda x: x[input]}) # 传递输入 | RunnablePassthrough.assign(**{memories: format_memories}) # 调用函数获取记忆 | RunnablePassthrough.assign(**{ short_term_memory: lambda x: x[memories][short_term_memory], long_term_memory: lambda x: x[memories][long_term_memory], }) # 将格式化后的记忆提取出来 | prompt_template | llm ) # 5. 封装一个方便调用的函数 def run_agent_with_memory(user_input): # 运行链得到LLM回复 response agent_chain.invoke({input: user_input}) ai_output response.content # 更新短期记忆 memory_manager.update_short_term(user_input, ai_output) # 判断是否需要将本轮信息存入长期记忆这里用简单规则如果用户陈述了明确目标或事实 if 目标 in user_input or 记住 in user_input or 偏好 in user_input: memory_manager.save_to_long_term( textf用户陈述{user_input}, metadata{type: user_fact, timestamp: auto_generated} ) print([系统] 已将此轮信息纳入长期记忆。) return ai_output4. 实战演示与效果对比让我们用代码实际运行一下看看效果。print( 第一轮对话设定目标 ) response1 run_agent_with_memory(我的目标是学习Python用于自动化办公特别是处理Excel和PPT。) print(f助手{response1}\n) print( 第二轮对话询问具体建议 ) # 注意这里没有重复目标 response2 run_agent_with_memory(那我应该先学什么库呢) print(f助手{response2}\n) print( 第三轮对话隔了几轮后再次询问 ) # 模拟一些其他对话冲刷短期记忆缓冲区 run_agent_with_memory(今天天气怎么样) run_agent_with_memory(给我讲个笑话。) # 现在再问回学习 response3 run_agent_with_memory(对了我们刚才说到哪了关于学习自动化办公的。) print(f助手{response3})预期输出分析第一轮助手会直接回答关于openpyxl/xlrd和python-pptx库的建议。第二轮由于短期记忆存在助手能连贯地承接上一轮给出更具体的学习路径比如“既然目标是自动化办公那么应该先学openpyxl来处理Excel...”。第三轮短期记忆已经被“天气”和“笑话”冲刷。如果没有长期记忆助手很可能回答“我们刚才有聊过吗”或转向其他话题。但有了长期记忆RAG检索当用户提到“自动化办公”时系统会从ChromaDB中检索到第一轮存入的“用户目标”记忆片段并将其融入上下文。因此助手应该能回答“我们之前讨论过您的目标是使用Python进行自动化办公特别是处理Excel和PPT。当时建议从openpyxl和python-pptx库开始...”这个简单的演示清晰地展示了MemoryRAG如何让Agent克服“健忘症”实现跨会话的、基于历史知识的连贯交互。5. 高级优化与问题排查基础功能跑通后我们会发现很多细节问题。以下是几个关键的优化点和常见坑位。5.1 检索质量优化超越简单的向量搜索单纯的向量相似度搜索如余弦相似度在以下情况会失效词汇不匹配用户问“怎么让表格变好看”但知识库里存的是“Excel样式美化”。语义相似但意图不同“Python循环”和“迭代器”语义相关但当用户问“for循环的语法”时返回“迭代器协议”的文档就没用。多维度过滤只想检索某个特定用户或某个时间段内的记忆。解决方案混合检索与重排序关键词检索如BM25作为补充BM25算法对关键词匹配更敏感可以与向量检索结果融合。LangChain提供了EnsembleRetriever。from langchain.retrievers import BM25Retriever, EnsembleRetriever from langchain.vectorstores import Chroma # 假设我们有文档列表 doc_list bm25_retriever BM25Retriever.from_documents(doc_list) vector_retriever vectorstore.as_retriever(search_kwargs{k: 4}) ensemble_retriever EnsembleRetriever( retrievers[bm25_retriever, vector_retriever], weights[0.3, 0.7] # 调整权重 )元数据过滤ChromaDB支持在检索时传入filter字典精确筛选。# 只检索“目标”类型的记忆且用户是“小明” filter_dict {type: goal, user: 小明} docs vectorstore.similarity_search(query, k3, filterfilter_dict)重排序Re-ranking使用一个更精细的交叉编码器模型对初步检索到的Top-N个结果进行相关性重排提升Top-K的精度。可以使用Cohere或BGE的重排模型。# 伪代码需安装FlagEmbedding等库 from FlagEmbedding import FlagReranker reranker FlagReranker(BAAI/bge-reranker-large, use_fp16True) # 对ensemble_retriever返回的docs进行rerank5.2 记忆的更新、遗忘与冲突解决记忆不是只增不减的。无效的、过时的记忆需要清理。更新机制当用户说“我改变主意了我的新目标是...”我们需要更新长期记忆。一种做法是给记忆条目增加版本号或时间戳检索时优先返回最新的。更复杂的做法是设计一个“记忆修正”工具让Agent主动发起对旧记忆的修正。遗忘机制可以基于时间自动清理N天前的记忆、使用频率清理从未被检索到的记忆或用户显式指令“忘记我刚才说的密码”来实现。冲突解决如果检索到两条矛盾的记忆如“用户喜欢咖啡”和“用户讨厌咖啡”需要在融合时处理。简单的策略是相信时间最近的或者让LLM在上下文中进行判断和取舍。5.3 性能与成本考量嵌入模型成本如果使用OpenAI等付费API大量文本的嵌入向量化会产生费用。对于静态知识库可以预计算并缓存嵌入向量。对于动态生成的记忆需要评估其价值和频率。检索延迟向量检索虽然比全表扫描快但在数据量大时仍有延迟。确保k值返回数量不要过大一般3-5足矣。对检索结果进行缓存也是常用优化手段。上下文长度限制LLM有上下文窗口限制如4K, 8K, 128K。融合后的记忆短期长期可能很长需要进行摘要或选择性纳入。例如只纳入与当前查询最相关的1-2条长期记忆或者对短期记忆进行摘要。5.4 常见问题与排查技巧检索不到相关内容检查嵌入模型确保用于存储和检索的嵌入模型是同一个。不同模型生成的向量空间不同无法直接比较。检查文本预处理存储和查询时的文本预处理如去除空格、标点是否一致不一致会导致向量差异大。调整切片策略chunk_size太大可能包含过多无关信息太小可能割裂语义。尝试调整大小和重叠度。尝试不同的距离函数ChromaDB默认用L2距离对于文本相似度余弦相似度cosine通常是更好的选择。在初始化Chroma时可以通过collection_metadata设置hnsw:space为cosine。vectorstore Chroma.from_documents( documentssplits, embeddingembeddings, persist_directorypersist_directory, collection_metadata{hnsw:space: cosine} # 使用余弦相似度 )记忆混乱或无关信息被引入优化元数据为每条记忆添加丰富、准确的元数据如type,topic,user_id,timestamp并在检索时利用过滤功能。引入相关性分数阈值相似度分数过低的结果很可能不相关。可以设置一个阈值如0.7过滤掉低分结果。docs_and_scores vectorstore.similarity_search_with_relevance_scores(query, k5) filtered_docs [doc for doc, score in docs_and_scores if score 0.7]Agent响应变慢向量库索引优化ChromaDB默认使用HNSW索引对于大规模数据确保索引参数合理。但对于中小规模默认值通常足够。减少检索量降低k值或先使用关键词进行粗筛再用向量进行精筛。异步操作将记忆检索和LLM调用设计为异步流程避免阻塞。6. 从记忆到真正的“智能体”进阶思考补全了记忆我们的Agent才算有了“过去”。但这距离一个真正自主的智能体Agent还有距离。记忆是基础在此之上还需要反思Reflection能力让Agent定期或不定期地回顾自己的记忆和行为总结规律形成更高层次的“经验”或“原则”并存入长期记忆。例如“用户每次在周五晚上都会询问周末学习计划”那么Agent可以在周五主动推送提醒。记忆的主动管理当前的记忆存储是被动的由我们定义的规则触发。更高级的Agent应该能自主判断一段信息是否重要、是否值得存储、以及如何分类存储。多模态记忆目前的记忆是纯文本的。未来的Agent可能需要处理图像、音频等多模态记忆并实现跨模态的检索例如根据描述检索图片或根据图片生成描述文本存入记忆。这次为Agent添加Memory和RAG的实践让我深刻体会到记忆是智能的基石。没有记忆的AI再强大的推理能力也是无根之木。这套方案虽然还有优化空间但已经极大地提升了Agent的实用性和用户体验。它不再是一个“一问一答”的机器而是一个能延续对话、积累知识、个性化服务的初级“伙伴”。
返回列表