AI Agent 记忆系统设计实战：从短期记忆到长期知识管理的完整方案

发表于 分类于 ，
本文字数： 11k 阅读时长 ≈ 10 分钟
分享构建AI Agent记忆系统的完整实践经验，从技术选型到落地实现，包含短期记忆、长期记忆、知识图谱等多层次记忆架构的设计思路和代码实现。

AI Agent 记忆系统设计实战：从短期记忆到长期知识管理的完整方案

技术主题：AI Agent（人工智能/工作流）
内容方向：实际使用经验分享（工具/框架选型、项目落地心得）

引言

在构建智能客服AI Agent的过程中，我们发现一个关键问题：Agent缺乏”记忆”能力，每次对话都像第一次见面，无法建立连续性的用户体验。经过半年的探索和实践，我们设计并实现了一套完整的AI Agent记忆系统，显著提升了对话的连贯性和个性化程度。本文将分享这套系统的设计思路、技术选型和落地经验。

一、记忆系统需求与架构

1. 三层记忆架构设计

基于人类记忆模型，我们设计了三层记忆架构：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class MemoryArchitecture:
    """AI Agent三层记忆架构"""
    
    def __init__(self):
        # 工作记忆：当前对话上下文 (4k tokens以内)
        self.working_memory = {
            "current_session": [],
            "context_window": 4096,
            "active_topics": []
        }
        
        # 短期记忆：近期会话历史 (24小时-7天)
        self.short_term_memory = {
            "recent_sessions": [],
            "user_preferences": {},
            "interaction_patterns": {}
        }
        
        # 长期记忆：持久化知识 (永久存储)
        self.long_term_memory = {
            "knowledge_base": None,
            "user_profile": {},
            "domain_knowledge": {}
        }

2. 技术选型

经过调研和测试，最终选择的技术栈：

组件 选择 理由
向量数据库 ChromaDB 轻量级，支持本地部署
内存缓存 Redis 高性能，支持复杂数据结构
向量化模型 sentence-transformers 多语言支持，效果稳定
框架集成 LangChain 生态完善，扩展性好

二、核心组件实现

1. 工作记忆管理

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
class WorkingMemoryManager:
    """工作记忆管理器 - 管理当前对话上下文"""
    
    def __init__(self, token_limit=4096):
        self.token_limit = token_limit
        self.current_context = []
        self.topic_tracker = {}
    
    def add_message(self, role: str, content: str, metadata: dict = None):
        """添加消息并维护上下文大小"""
        message = {
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat(),
            "metadata": metadata or {}
        }
        
        self.current_context.append(message)
        self._maintain_context_size()
        self._update_topics(content)
    
    def _maintain_context_size(self):
        """维护token限制"""
        while self._estimate_tokens() > self.token_limit:
            # 优先移除非重要消息
            for i, msg in enumerate(self.current_context):
                if not msg.get("metadata", {}).get("important", False):
                    self.current_context.pop(i)
                    break
            else:
                # 如果都是重要消息，移除最早的
                self.current_context.pop(0)
    
    def _estimate_tokens(self) -> int:
        """估算token数量 (1 token ≈ 4 chars)"""
        total_text = " ".join([msg["content"] for msg in self.current_context])
        return len(total_text) // 4
    
    def get_relevant_context(self, query: str, max_messages: int = 10) -> list:
        """获取与查询相关的上下文"""
        if not query:
            return self.current_context[-max_messages:]
        
        # 基于关键词计算相关性
        query_keywords = set(self._extract_keywords(query))
        scored_messages = []
        
        for msg in self.current_context:
            msg_keywords = set(self._extract_keywords(msg["content"]))
            relevance = len(query_keywords & msg_keywords) / max(len(query_keywords), 1)
            scored_messages.append((relevance, msg))
        
        # 按相关性排序并返回
        scored_messages.sort(key=lambda x: x[0], reverse=True)
        return [msg for _, msg in scored_messages[:max_messages]]
    
    def _extract_keywords(self, text: str) -> list:
        """简化的关键词提取"""
        import re
        words = re.findall(r'\b\w+\b', text.lower())
        stopwords = {"的", "是", "在", "有", "我", "你", "他"}
        return [w for w in words if len(w) > 1 and w not in stopwords]

2. 短期记忆实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
class ShortTermMemoryManager:
    """短期记忆管理器 - 使用Redis缓存近期数据"""
    
    def __init__(self, redis_client, ttl_hours=24):
        self.redis = redis_client
        self.ttl = ttl_hours * 3600
    
    def save_conversation_summary(self, user_id: str, session_id: str, summary: dict):
        """保存对话摘要"""
        key = f"conv_summary:{user_id}:{session_id}"
        self.redis.setex(key, self.ttl, json.dumps(summary, ensure_ascii=False))
    
    def get_recent_conversations(self, user_id: str, limit: int = 5) -> list:
        """获取最近对话摘要"""
        pattern = f"conv_summary:{user_id}:*"
        keys = self.redis.keys(pattern)
        
        conversations = []
        for key in keys[-limit:]:
            data = self.redis.get(key)
            if data:
                conversations.append(json.loads(data))
        
        return sorted(conversations, key=lambda x: x.get("timestamp", ""), reverse=True)
    
    def update_user_preferences(self, user_id: str, preferences: dict):
        """更新用户偏好"""
        key = f"user_prefs:{user_id}"
        existing = self.redis.get(key)
        
        if existing:
            current_prefs = json.loads(existing)
            current_prefs.update(preferences)
        else:
            current_prefs = preferences
        
        current_prefs["last_updated"] = datetime.now().isoformat()
        self.redis.setex(key, self.ttl * 7, json.dumps(current_prefs, ensure_ascii=False))

3. 长期记忆实现

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
class LongTermMemoryManager:
    """长期记忆管理器 - 使用向量数据库存储知识"""
    
    def __init__(self, chroma_client, embedding_model):
        self.chroma_client = chroma_client
        self.embedding_model = embedding_model
        
        # 创建不同类型的记忆集合
        self.collections = {
            "conversations": chroma_client.get_or_create_collection("conversations"),
            "knowledge": chroma_client.get_or_create_collection("knowledge"),
            "user_profiles": chroma_client.get_or_create_collection("user_profiles")
        }
    
    def add_knowledge(self, content: str, metadata: dict, collection_name: str = "knowledge"):
        """添加知识到长期记忆"""
        # 生成向量
        embedding = self.embedding_model.encode([content])[0].tolist()
        
        # 生成唯一ID
        knowledge_id = f"{collection_name}_{datetime.now().timestamp()}"
        
        # 存储到向量数据库
        self.collections[collection_name].add(
            documents=[content],
            embeddings=[embedding],
            metadatas=[{
                **metadata,
                "created_at": datetime.now().isoformat()
            }],
            ids=[knowledge_id]
        )
        
        return knowledge_id
    
    def search_relevant_knowledge(self, query: str, collection_name: str = "knowledge", 
                                top_k: int = 5) -> list:
        """搜索相关知识"""
        query_embedding = self.embedding_model.encode([query])[0].tolist()
        
        results = self.collections[collection_name].query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            include=["documents", "metadatas", "distances"]
        )
        
        # 格式化结果
        relevant_knowledge = []
        for i in range(len(results["documents"][0])):
            relevant_knowledge.append({
                "content": results["documents"][0][i],
                "metadata": results["metadatas"][0][i],
                "relevance_score": 1 - results["distances"][0][i]
            })
        
        return relevant_knowledge

三、系统集成与使用

1. 统一记忆接口

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
class UnifiedMemorySystem:
    """统一记忆系统接口"""
    
    def __init__(self, config: dict):
        self.working_memory = WorkingMemoryManager(config.get("context_limit", 4096))
        self.short_term_memory = ShortTermMemoryManager(
            redis_client, config.get("short_term_ttl", 24)
        )
        self.long_term_memory = LongTermMemoryManager(chroma_client, embedding_model)
    
    def process_user_message(self, user_id: str, message: str, session_id: str) -> dict:
        """处理用户消息并构建完整上下文"""
        
        # 添加到工作记忆
        self.working_memory.add_message("user", message)
        
        # 搜索相关知识
        relevant_knowledge = self.long_term_memory.search_relevant_knowledge(message)
        
        # 获取用户偏好和历史
        user_preferences = self.short_term_memory.get_user_preferences(user_id)
        recent_conversations = self.short_term_memory.get_recent_conversations(user_id)
        
        # 构建完整上下文
        context = {
            "current_conversation": self.working_memory.get_relevant_context(message),
            "relevant_knowledge": relevant_knowledge,
            "user_preferences": user_preferences,
            "recent_history": recent_conversations
        }
        
        return context
    
    def save_agent_response(self, user_id: str, response: str, session_id: str, context: dict):
        """保存Agent响应并更新记忆"""
        
        # 添加到工作记忆
        self.working_memory.add_message("assistant", response)
        
        # 提取并保存新知识
        knowledge_items = self._extract_knowledge(context["current_conversation"])
        for item in knowledge_items:
            self.long_term_memory.add_knowledge(item["content"], item["metadata"])
        
        # 更新用户偏好
        preferences = self._analyze_preferences(context["current_conversation"])
        if preferences:
            self.short_term_memory.update_user_preferences(user_id, preferences)

2. 实际应用示例

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
# 初始化记忆系统
config = {
    "context_limit": 4096,
    "short_term_ttl": 24,
    "chroma_path": "./memory_db"
}

memory_system = UnifiedMemorySystem(config)

def chat_with_memory(user_id: str, message: str, session_id: str):
    """带记忆的智能对话"""
    
    # 处理用户消息，获取完整上下文
    context = memory_system.process_user_message(user_id, message, session_id)
    
    # 构建包含记忆信息的prompt
    system_prompt = f"""
    你是一个拥有记忆的智能助手。请基于以下信息回答用户问题：
    
    相关知识：{context['relevant_knowledge'][:3]}  # 最相关的3条
    用户偏好：{context['user_preferences']}
    对话历史：{context['current_conversation'][-5:]}  # 最近5轮对话
    """
    
    # 调用LLM生成回复
    response = call_llm(system_prompt, message)
    
    # 保存回复并更新记忆
    memory_system.save_agent_response(user_id, response, session_id, context)
    
    return response

def call_llm(system_prompt: str, user_message: str) -> str:
    """调用大语言模型（示例）"""
    # 这里接入你的LLM服务
    # 可以是OpenAI API、本地模型等
    return "基于记忆系统生成的智能回复"

四、性能优化与监控

1. 缓存优化策略

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
class PerformanceOptimizer:
    """性能优化器"""
    
    def __init__(self):
        self.local_cache = {}
        self.metrics = {"cache_hits": 0, "cache_misses": 0}
    
    def cached_search(self, query: str, search_func, cache_ttl: int = 300):
        """带缓存的搜索"""
        cache_key = f"search:{hash(query)}"
        
        # 检查本地缓存
        if cache_key in self.local_cache:
            entry = self.local_cache[cache_key]
            if time.time() - entry["timestamp"] < cache_ttl:
                self.metrics["cache_hits"] += 1
                return entry["data"]
            else:
                del self.local_cache[cache_key]
        
        # 执行搜索
        self.metrics["cache_misses"] += 1
        results = search_func(query)
        
        # 更新缓存
        self.local_cache[cache_key] = {
            "data": results,
            "timestamp": time.time()
        }
        
        return results

2. 系统监控

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class MemorySystemMonitor:
    """记忆系统监控器"""
    
    def __init__(self, memory_system):
        self.memory_system = memory_system
    
    def get_system_stats(self) -> dict:
        """获取系统统计信息"""
        return {
            "working_memory_size": len(self.memory_system.working_memory.current_context),
            "redis_memory_usage": self._get_redis_memory_usage(),
            "vector_db_collections": len(self.memory_system.long_term_memory.collections),
            "cache_hit_rate": self._calculate_cache_hit_rate()
        }
    
    def cleanup_old_data(self, days_to_keep: int = 30):
        """清理过期数据"""
        cutoff_date = datetime.now() - timedelta(days=days_to_keep)
        # 实现数据清理逻辑
        pass

五、落地经验与最佳实践

1. 关键经验总结

技术选型经验

  • ChromaDB适合中小规模应用，大规模建议Pinecone/Weaviate
  • Redis缓存层效果显著，建议配置持久化
  • embedding模型选择要平衡效果和性能

架构设计经验

  • 三层记忆架构清晰分离不同记忆需求
  • 异步处理长期记忆更新很重要
  • 降级策略确保记忆故障不影响基本功能

性能优化经验

  • 本地缓存显著提升搜索性能
  • 向量维度选择要平衡存储和检索
  • 定期清理和压缩数据库

2. 常见问题解决

  1. 向量维度不匹配：更换embedding模型时需重新生成所有向量
  2. 内存泄漏：工作记忆需要正确的token限制和清理机制
  3. 搜索精度问题：调整top_k参数和相关性阈值

总结

AI Agent记忆系统的构建是一个系统工程，需要在功能完整性、性能效率和系统稳定性之间找到平衡。我们的三层记忆架构经过半年生产环境验证，能够有效提升对话质量和用户体验。

核心要点：

  • 分层设计：工作记忆、短期记忆、长期记忆各司其职
  • 技术选型：根据规模选择合适的存储方案
  • 性能优化：缓存策略和异步处理不可忽视
  • 监控维护：完善的监控和清理机制保证长期稳定

随着AI Agent应用的普及，记忆系统将成为提升用户体验的关键技术。希望我们的实践经验能为同行提供参考，共同推动AI Agent技术的发展。