Agent 记忆系统设计：从短期上下文到长期知识

class WorkingMemory:
    def __init__(self, max_tokens=8000):
        self.max_tokens = max_tokens
        self.messages = []
    
    def add(self, message):
        self.messages.append(message)
        self._compress_if_needed()
    
    def _compress_if_needed(self):
        while self.token_count() > self.max_tokens:
            if len(self.messages) <= 2:
                break
            
            # 取出最早的几条消息做摘要
            old_messages = self.messages[:3]
            summary = self.summarize(old_messages)
            
            self.messages = [
                {"role": "system", "content": f"之前的对话摘要：{summary}"}
            ] + self.messages[3:]

class AttentionManager:
    """管理上下文中各部分的"重要性""""
    
    def __init__(self):
        self.attention_scores = {}
    
    def update(self, message_id, content, user_query):
        """计算每条消息对当前问题的相关性"""
        self.attention_scores[message_id] = {
            "recency": self.recency_score(message_id),
            "relevance": self.relevance_score(content, user_query),
            "importance": self.importance_score(content),
        }
    
    def get_context(self, max_tokens=4000):
        """按注意力分数排序，选择最重要的上下文"""
        ranked = sorted(
            self.attention_scores.items(),
            key=lambda x: x[1]["recency"] * 0.2 + 
                         x[1]["relevance"] * 0.5 + 
                         x[1]["importance"] * 0.3,
            reverse=True
        )
        
        context = []
        token_count = 0
        for msg_id, _ in ranked:
            msg = self.get_message(msg_id)
            token_count += count_tokens(msg)
            if token_count > max_tokens:
                break
            context.append(msg)
        
        return context

from dataclasses import dataclass
from datetime import datetime

@dataclass
class EpisodicMemory:
    id: str
    session_id: str
    timestamp: datetime
    user_input: str
    agent_response: str
    tools_used: list[str]
    outcome: str  # success / failure / partial
    embedding: list[float]  # 用于检索
    metadata: dict  # 用户信息、环境等

class EpisodicMemoryStore:
    async def recall(self, query, user_id, limit=5):
        """回忆与当前情况相关的历史对话"""
        
        # 时间衰减权重：越近的记忆权重越高
        recency_weight = 0.3
        
        # 语义相似度
        query_emb = self.embedder.embed(query)
        
        results = await self.db.search(
            user_id=user_id,
            embedding=query_emb,
            limit=limit * 3,  # 多取一些做重排
        )
        
        # 综合排序
        for r in results:
            days_ago = (datetime.now() - r.timestamp).days
            r.score = r.similarity * (1 - recency_weight) + \
                      math.exp(-days_ago / 30) * recency_weight
        
        results.sort(key=lambda x: x.score, reverse=True)
        return results[:limit]

class MemoryConsolidation:
    """将情景记忆转化为语义记忆（类似人类睡眠时的记忆巩固）"""
    
    async def consolidate(self, user_id):
        # 1. 提取最近7天的对话
        recent = await self.episodic_store.get_recent(user_id, days=7)
        
        # 2. 用 LLM 提取可复用的知识
        knowledge = await self.llm.extract(
            f"从以下对话中提取可复用的知识和模式：\n{recent}"
        )
        
        # 3. 存入语义记忆
        for item in knowledge:
            await self.semantic_store.add(
                content=item.content,
                embedding=self.embedder.embed(item.content),
                source="consolidated",
                confidence=item.confidence,
            )
        
        # 4. 标记已巩固的情景记忆
        await self.episodic_store.mark_consolidated(recent)

class SemanticMemory:
    """三层知识存储"""
    
    def __init__(self):
        # 第一层：快速事实（Key-Value）
        self.facts = RedisStore()  # "用户公司" → "ABC科技"
        
        # 第二层：文档知识（向量检索）
        self.documents = VectorStore()  # API文档、产品手册
        
        # 第三层：深层知识（知识图谱）
        self.graph = GraphStore()  # 实体关系

class KnowledgeUpdater:
    async def learn(self, new_info, source="conversation"):
        """学习新知识"""
        
        # 1. 检查是否与已有知识冲突
        conflicts = await self.find_conflicts(new_info)
        
        if conflicts:
            # 2. 冲突解决：用更新的信息覆盖
            for conflict in conflicts:
                if conflict.confidence < new_info.confidence:
                    await self.update(conflict.id, new_info)
                else:
                    # 不确定，存为"待确认"
                    await self.flag_for_review(new_info, conflict)
        else:
            # 3. 新知识，直接存入
            await self.semantic_store.add(new_info)

class ReflectionSystem:
    """Agent 从失败中学习"""
    
    async def reflect(self, task, attempt, outcome):
        if outcome.success:
            return  # 成功不需要反思
        
        # 分析失败原因
        analysis = await self.llm.analyze(
            task=task,
            actions=attempt.actions,
            error=outcome.error,
            prompt="分析为什么这次任务失败了，如何避免"
        )
        
        # 提取规则
        rules = self.extract_rules(analysis)
        
        # 存入反思记忆
        for rule in rules:
            await self.reflection_store.add({
                "situation": rule.situation,
                "mistake": rule.mistake,
                "better_action": rule.better_action,
                "confidence": rule.confidence,
            })
    
    async def consult(self, task):
        """执行任务前，查阅反思记忆"""
        relevant_mistakes = await self.reflection_store.search(task)
        
        if relevant_mistakes:
            return f"""
            注意：根据历史经验，类似任务有以下陷阱：
            {format_mistakes(relevant_mistakes)}
            请避免这些错误。
            """
        return None

class AgentMemory:
    """完整的 Agent 记忆系统"""
    
    def __init__(self):
        self.working = WorkingMemory(max_tokens=8000)
        self.episodic = EpisodicMemoryStore()
        self.semantic = SemanticMemory()
        self.reflection = ReflectionSystem()
    
    async def prepare_context(self, query, user_id):
        """为 LLM 调用准备完整上下文"""
        context_parts = []
        
        # 1. 工作记忆（当前对话）
        context_parts.append(self.working.get_recent(n=6))
        
        # 2. 相关情景记忆
        episodes = await self.episodic.recall(query, user_id, limit=3)
        if episodes:
            context_parts.append(f"历史相关对话：{format_episodes(episodes)}")
        
        # 3. 语义知识
        knowledge = await self.semantic.search(query, limit=5)
        if knowledge:
            context_parts.append(f"相关知识：{format_knowledge(knowledge)}")
        
        # 4. 反思提醒
        warning = await self.reflection.consult(query)
        if warning:
            context_parts.append(warning)
        
        return "\n\n".join(context_parts)

class MemoryPrivacy:
    """记忆系统中的隐私保护"""
    
    PII_PATTERNS = [
        (r'\d{11}', 'PHONE'),      # 手机号
        (r'\d{18}', 'ID_CARD'),    # 身份证
        (r'[\w.]+@[\w.]+', 'EMAIL'), # 邮箱
    ]
    
    def sanitize(self, memory):
        """存储前脱敏"""
        for pattern, pii_type in self.PII_PATTERNS:
            memory = re.sub(pattern, f'[{pii_type}]', memory)
        return memory
    
    def user_forget(self, user_id):
        """用户要求删除所有记忆"""
        await self.episodic.delete_by_user(user_id)
        await self.semantic.delete_by_user(user_id)
        await self.reflection.delete_by_user(user_id)
        # 生成删除确认报告
        return {"deleted": True, "user_id": user_id}