AI Agent 记忆系统设计：从工作记忆、情景记忆到分层 Graph RAG

                    ┌─────────────────────────┐
                    │      Core Memory        │
                    │   (人设、用户画像、硬红线)   │  <-- 永久驻留 Prompt (数千 Token)
                    └────────────┬────────────┘
                                 │
                    ┌────────────▼────────────┐
                    │     Working Memory      │
                    │ (当前步骤工具入参、短期Scratchpad)│  <-- 节点内即时计算，跳变后清理
                    └────────────┬────────────┘
                                 │ 归纳整合 (Consolidation)
                    ┌────────────▼────────────┐
                    │     Episodic Memory     │
                    │  (当前会话事件流、历史规划快照)  │  <-- 动态摘要滑窗，Checkpointer 落盘
                    └────────────┬────────────┘
                                 │ 提纯沉淀 (Deep Archival)
                    ┌────────────▼────────────┐
                    │ Long-term Semantic/     │
                    │ Graph Memory (跨会话图谱)│  <-- 向量库 / 时序图数据库 (按需Top-K召回)
                    └─────────────────────────┘

[2026-03-01 记录]
(User) --[owns_tech_stack {valid_from: "2026-03-01", valid_to: "2026-03-15"}]--> (Python)

[2026-03-15 变更事件]
(User) --[migrated_to {timestamp: "2026-03-15"}]--> (Rust)
(User) --[owns_tech_stack {valid_from: "2026-03-15", valid_to: infinity}]--> (Rust)

"""
production_agent_memory.py
生产级 AI Agent 分层记忆系统（集成多维加权检索、时间遗忘曲线与自动化反思沉淀）
"""

import math
import time
from datetime import datetime
from typing import Any, Dict, List, Optional
from pydantic import BaseModel, Field


# ==========================================================
# 1. 记忆单元数据结构
# ==========================================================

class MemoryItem(BaseModel):
    id: str
    content: str
    created_at: float = Field(default_factory=time.time)
    last_accessed_at: float = Field(default_factory=time.time)
    importance: float = Field(ge=1.0, le=10.0, description="重要度等级 1.0-10.0")
    access_count: int = 0
    metadata: Dict[str, Any] = Field(default_factory=dict)

    def calculate_recency_score(self, current_time: float, decay_lambda: float = 0.005) -> float:
        """基于指数衰减曲线计算时间新鲜度 [0.0, 1.0]"""
        hours_passed = (current_time - self.last_accessed_at) / 3600.0
        return math.exp(-decay_lambda * hours_passed)


# ==========================================================
# 2. 生产级分层记忆管理器
# ==========================================================

class ProductionMemoryManager:
    def __init__(
        self,
        decay_lambda: float = 0.005,
        weight_semantic: float = 0.5,
        weight_recency: float = 0.3,
        weight_importance: float = 0.2,
    ):
        self.decay_lambda = decay_lambda
        self.w_sem = weight_semantic
        self.w_rec = weight_recency
        self.w_imp = weight_importance
        
        # 内存存储结构
        self.core_memory: Dict[str, str] = {}
        self.episodic_memory: List[MemoryItem] = []
        self.reflection_threshold = 25.0
        self.accumulated_importance = 0.0

    def set_core_memory(self, key: str, value: str):
        """设置不可被稀释的核心记忆（人设、核心业务边界）"""
        self.core_memory[key] = value

    def add_episodic_memory(self, content: str, importance: float, metadata: Optional[Dict[str, Any]] = None):
        """记录一条情景记忆"""
        item_id = f"mem_{len(self.episodic_memory) + 1}_{int(time.time())}"
        item = MemoryItem(
            id=item_id,
            content=content,
            importance=importance,
            metadata=metadata or {}
        )
        self.episodic_memory.append(item)
        self.accumulated_importance += importance

        # 检查是否需要触发后台反思整合
        if self.accumulated_importance >= self.reflection_threshold:
            self._trigger_consolidation_reflection()

    def _mock_embedding_similarity(self, query: str, content: str) -> float:
        """模拟向量相似度计算（实际生产中替换为 dense vector cosine similarity）"""
        q_words = set(query.lower().split())
        c_words = set(content.lower().split())
        overlap = len(q_words & c_words)
        return min(1.0, overlap / (len(q_words) + 1e-5))

    def retrieve(self, query: str, top_k: int = 3) -> List[MemoryItem]:
        """多维综合评分检索：综合考虑语义相关度、时间衰减与固有重要度"""
        now = time.time()
        scored_items: List[tuple[float, MemoryItem]] = []

        for item in self.episodic_memory:
            # 1. 语义分数
            s_sem = self._mock_embedding_similarity(query, item.content)
            # 2. 新鲜度分数（指数衰减）
            s_rec = item.calculate_recency_score(now, self.decay_lambda)
            # 3. 归一化重要度分数 [0.1, 1.0]
            s_imp = item.importance / 10.0

            # 最终复合得分
            final_score = (self.w_sem * s_sem) + (self.w_rec * s_rec) + (self.w_imp * s_imp)
            scored_items.append((final_score, item))

        # 按总得分降序排列
        scored_items.sort(key=lambda x: x[0], reverse=True)
        results = [item for _, item in scored_items[:top_k]]

        # 更新被命中的记忆的最后访问时间与频次
        for item in results:
            item.last_accessed_at = now
            item.access_count += 1

        return results

    def _trigger_consolidation_reflection(self):
        """后台记忆沉淀：将低阶事件序列抽象提纯为高阶规则，并执行弱记忆修剪"""
        print("\n⚡ [后台任务启动] 累积重要度达到阈值，正在执行记忆整合与反思提纯...")
        
        recent_events = self.episodic_memory[-5:]
        consolidated_insight = (
            f"归纳洞察：分析了最近 {len(recent_events)} 项操作记录，"
            f"提炼用户核心目标正从基础查询转向复杂分析。"
        )
        # 沉淀高阶规则（赋予最高等级 9.5 重要度）
        reflection_item = MemoryItem(
            id=f"insight_{int(time.time())}",
            content=consolidated_insight,
            importance=9.5,
            metadata={"source": "consolidation_reflection"}
        )
        self.episodic_memory.append(reflection_item)
        self.accumulated_importance = 0.0
        print(f"✓ 已沉淀高阶记忆规则：'{consolidated_insight}'")

    def build_prompt_context(self, current_task: str) -> str:
        """构建注入到当前 LLM 调用的精炼上下文"""
        relevant_memories = self.retrieve(current_task, top_k=2)
        
        prompt_parts = ["=== Core Memory (核心不可变约束) ==="]
        for k, v in self.core_memory.items():
            prompt_parts.append(f"- {k}: {v}")

        prompt_parts.append("\n=== Relevant Recalled Memories (相关上下文回忆) ===")
        for mem in relevant_memories:
            dt = datetime.fromtimestamp(mem.created_at).strftime("%Y-%m-%d %H:%M:%S")
            prompt_parts.append(f"[{dt}] (重要度: {mem.importance}) {mem.content}")

        return "\n".join(prompt_parts)


# ==========================================================
# 3. 运行验证与实战模拟
# ==========================================================

if __name__ == "__main__":
    print("--- 步骤 1: 初始化智能体记忆管理器 ---")
    mem_sys = ProductionMemoryManager()
    mem_sys.set_core_memory("UserRole", "资深大数据架构师")
    mem_sys.set_core_memory("SystemConstraint", "所有数据库操作必须包含只读限制，严禁 DROP 操作")

    print("\n--- 步骤 2: 模拟连续写入情景交互 ---")
    mem_sys.add_episodic_memory("用户要求查询 2026 年第一季度的日志归档", importance=6.0)
    mem_sys.add_episodic_memory("执行 SQL 查询超时，建议切换只读副本集群", importance=7.5)
    mem_sys.add_episodic_memory("用户确认切换至 Read-Replica-02 成功读取数据", importance=8.0)
    mem_sys.add_episodic_memory("系统发现某次慢查询主要是由于缺乏时间复合索引导致", importance=8.5)

    print("\n--- 步骤 3: 模拟当前新任务检索与 Prompt 上下文组装 ---")
    task = "查询 生产数据库 集群状态与慢日志"
    injected_context = mem_sys.build_prompt_context(task)
    
    print("\n[最终注入 LLM 的紧凑记忆 Prompt 视图]:")
    print(injected_context)

--- 步骤 1: 初始化智能体记忆管理器 ---

--- 步骤 2: 模拟连续写入情景交互 ---

⚡ [后台任务启动] 累积重要度达到阈值，正在执行记忆整合与反思提纯...
✓ 已沉淀高阶记忆规则：'归纳洞察：分析了最近 4 项操作记录，提炼用户核心目标正从基础查询转向复杂分析。'

--- 步骤 3: 模拟当前新任务检索与 Prompt 上下文组装 ---

[最终注入 LLM 的紧凑记忆 Prompt 视图]:
=== Core Memory (核心不可变约束) ===
- UserRole: 资深大数据架构师
- SystemConstraint: 所有数据库操作必须包含只读限制，严禁 DROP 操作

=== Relevant Recalled Memories (相关上下文回忆) ===
[2026-09-21 11:30:00] (重要度: 9.5) 归纳洞察：分析了最近 4 项操作记录，提炼用户核心目标正从基础查询转向复杂分析。
[2026-09-21 11:30:00] (重要度: 8.5) 系统发现某次慢查询主要是由于缺乏时间复合索引导致