Agent记忆系统设计：短期、长期与情景记忆的实现

┌───────────────────────────────────────────────┐
│              记忆系统架构                      │
├───────────────────────────────────────────────┤
│                                               │
│  ┌─────────────────────────────────────────┐ │
│  │  Layer 1: 短期记忆 (Working Memory)     │ │
│  │  - 当前对话上下文                       │ │
│  │  - 当前任务状态                         │ │
│  │  - 最近的工具调用结果                   │ │
│  │  容量: ~32K tokens | TTL: 会话结束      │ │
│  └─────────────────────────────────────────┘ │
│                    ↓ 溢出压缩                  │
│  ┌─────────────────────────────────────────┐ │
│  │  Layer 2: 情景记忆 (Episodic Memory)    │ │
│  │  - 具体事件记录（时间、地点、人物）      │ │
│  │  - 对话历史摘要                         │ │
│  │  - 任务执行日志                         │ │
│  │  容量: 无限 | TTL: 90天衰减             │ │
│  └─────────────────────────────────────────┘ │
│                    ↓ 知识提取                  │
│  ┌─────────────────────────────────────────┐ │
│  │  Layer 3: 长期记忆 (Long-term Memory)   │ │
│  │  - 用户偏好和习惯                       │ │
│  │  - 学到的知识和规则                     │ │
│  │  - 经验教训                             │ │
│  │  - 语义关联                             │ │
│  │  容量: 无限 | TTL: 永久                 │ │
│  └─────────────────────────────────────────┘ │
│                                               │
└───────────────────────────────────────────────┘

from dataclasses import dataclass, field
from datetime import datetime
from typing import Any

@dataclass
class ShortTermMemory:
    """短期记忆：管理当前会话上下文"""
    
    # 核心数据
    messages: list[dict] = field(default_factory=list)
    current_task: str | None = None
    task_state: dict = field(default_factory=dict)
    tool_results: list[dict] = field(default_factory=list)
    
    # 配置
    max_tokens: int = 32000
    current_tokens: int = 0
    
    def add_message(self, role: str, content: str):
        """添加消息到短期记忆"""
        message = {
            "role": role,
            "content": content,
            "timestamp": datetime.now().isoformat()
        }
        self.messages.append(message)
        self.current_tokens += self._estimate_tokens(content)
        
        # 超过容量时触发压缩
        if self.current_tokens > self.max_tokens * 0.8:
            self._compress()
    
    def add_tool_result(self, tool_name: str, result: Any):
        """添加工具调用结果"""
        entry = {
            "tool": tool_name,
            "result": result,
            "timestamp": datetime.now().isoformat()
        }
        self.tool_results.append(entry)
        
        # 工具结果也占用上下文
        result_str = str(result)
        if len(self.tool_results) > 10:
            # 保留最近10次工具调用
            self.tool_results = self.tool_results[-10:]
    
    def _compress(self):
        """压缩短期记忆：保留最近的 + 摘要旧的"""
        if len(self.messages) < 4:
            return
        
        # 将旧消息压缩为摘要
        old_messages = self.messages[:-4]  # 保留最近4条
        recent_messages = self.messages[-4:]
        
        summary = self._summarize(old_messages)
        
        self.messages = [
            {"role": "system", "content": f"[历史摘要] {summary}"},
            *recent_messages
        ]
        self.current_tokens = sum(
            self._estimate_tokens(m["content"]) for m in self.messages
        )
    
    def _summarize(self, messages: list[dict]) -> str:
        """使用LLM摘要历史消息"""
        # 实际实现中使用LLM
        return f"之前讨论了{len(messages)}条消息，主要关于..."
    
    def _estimate_tokens(self, text: str) -> int:
        """粗略估算token数"""
        return len(text) // 3  # 中文约3字符/token
    
    def get_context(self) -> str:
        """获取格式化的上下文"""
        context_parts = []
        
        if self.current_task:
            context_parts.append(f"当前任务: {self.current_task}")
        
        if self.task_state:
            context_parts.append(f"任务状态: {self.task_state}")
        
        if self.tool_results:
            recent_tools = self.tool_results[-3:]
            tool_summary = "; ".join(
                f"{t['tool']}→{str(t['result'])[:100]}" 
                for t in recent_tools
            )
            context_parts.append(f"最近工具调用: {tool_summary}")
        
        return "\n".join(context_parts)

from datetime import datetime
from typing import Optional
import json

@dataclass
class EpisodicEvent:
    """情景记忆中的单个事件"""
    event_id: str
    timestamp: datetime
    event_type: str          # conversation, task, decision, error
    description: str         # 事件描述
    participants: list[str]  # 参与者（Agent、用户）
    context: dict            # 事件上下文
    outcome: Optional[str]   # 事件结果
    emotional_valence: float # 情感价值（-1到1，正面/负面）
    
class EpisodicMemory:
    """情景记忆：存储和检索具体事件"""
    
    def __init__(self, storage_backend="file"):
        self.events: list[EpisodicEvent] = []
        self.storage_path = "memory/episodes/"
    
    async def record(self, event: EpisodicEvent):
        """记录新事件"""
        self.events.append(event)
        await self._persist(event)
        
        # 触发记忆巩固：从事件中提取知识
        if event.event_type == "decision":
            await self._consolidate_decision(event)
        elif event.event_type == "error":
            await self._consolidate_error(event)
    
    async def recall_by_time(
        self, start: datetime, end: datetime
    ) -> list[EpisodicEvent]:
        """按时间范围回忆"""
        return [
            e for e in self.events
            if start <= e.timestamp <= end
        ]
    
    async def recall_by_type(
        self, event_type: str, limit: int = 10
    ) -> list[EpisodicEvent]:
        """按事件类型回忆"""
        return [
            e for e in self.events
            if e.event_type == event_type
        ][-limit:]
    
    async def recall_similar(
        self, query: str, top_k: int = 5
    ) -> list[EpisodicEvent]:
        """语义相似性回忆"""
        # 使用向量检索
        query_embedding = await self._embed(query)
        event_embeddings = await self._get_event_embeddings()
        
        similarities = [
            (event, self._cosine_sim(query_embedding, emb))
            for event, emb in zip(self.events, event_embeddings)
        ]
        
        similarities.sort(key=lambda x: x[1], reverse=True)
        return [e for e, _ in similarities[:top_k]]
    
    async def _consolidate_decision(self, event: EpisodicEvent):
        """从决策事件中提取知识，存入长期记忆"""
        # 例如：用户选择了方案A而非方案B → 记住用户偏好
        pass
    
    async def _consolidate_error(self, event: EpisodicEvent):
        """从错误事件中提取教训"""
        # 例如：API调用失败 → 记住需要重试或使用备用方案
        pass

import math
from datetime import datetime, timedelta

class MemoryDecay:
    """记忆衰减策略"""
    
    @staticmethod
    def time_weighted_score(
        event_time: datetime,
        current_time: datetime,
        half_life_days: float = 30.0
    ) -> float:
        """
        时间加权得分：越久远的事件得分越低
        half_life_days: 半衰期，30天意味着30天前的事件得分减半
        """
        delta = (current_time - event_time).total_seconds() / 86400  # 转为天
        return math.exp(-0.693 * delta / half_life_days)
    
    @staticmethod
    def frequency_boost(
        event: EpisodicEvent,
        all_events: list[EpisodicEvent]
    ) -> float:
        """频繁出现的事件获得加分"""
        similar_count = sum(
            1 for e in all_events
            if e.event_type == event.event_type
            and e.description[:50] == event.description[:50]
        )
        return math.log(1 + similar_count)
    
    @staticmethod
    def emotional_boost(event: EpisodicEvent) -> float:
        """高情感价值的事件获得加分（重要决策、严重错误等）"""
        return 1.0 + abs(event.emotional_valence) * 0.5

from enum import Enum

class MemoryType(str, Enum):
    PREFERENCE = "preference"       # 用户偏好
    KNOWLEDGE = "knowledge"         # 学到的知识
    LESSON = "lesson"               # 经验教训
    RELATIONSHIP = "relationship"   # 人际关系
    PROCEDURE = "procedure"         # 操作流程

@dataclass
class LongTermMemory:
    """长期记忆条目"""
    memory_id: str
    memory_type: MemoryType
    content: str                    # 记忆内容
    confidence: float               # 置信度 0-1
    source: str                     # 来源（哪个事件提取的）
    created_at: datetime
    last_accessed: datetime
    access_count: int               # 被访问次数
    related_memories: list[str]     # 关联记忆ID

class LongTermMemoryStore:
    """长期记忆存储"""
    
    def __init__(self, vector_store, llm):
        self.vector_store = vector_store  # 向量数据库
        self.llm = llm
        self.memories: dict[str, LongTermMemory] = {}
    
    async def store(self, memory: LongTermMemory):
        """存储长期记忆"""
        self.memories[memory.memory_id] = memory
        
        # 向量化存储
        embedding = await self._embed(memory.content)
        await self.vector_store.upsert(
            id=memory.memory_id,
            embedding=embedding,
            metadata={
                "type": memory.memory_type.value,
                "content": memory.content,
                "confidence": memory.confidence,
                "created_at": memory.created_at.isoformat()
            }
        )
    
    async def retrieve(
        self,
        query: str,
        memory_type: MemoryType | None = None,
        top_k: int = 5
    ) -> list[LongTermMemory]:
        """检索相关长期记忆"""
        query_embedding = await self._embed(query)
        
        # 向量检索
        results = await self.vector_store.query(
            vector=query_embedding,
            filter={"type": memory_type.value} if memory_type else None,
            top_k=top_k * 2  # 多检索一些，后面做重排序
        )
        
        # 重排序：综合考虑相似度、置信度、时效性
        scored_results = []
        for r in results:
            memory = self.memories.get(r.id)
            if memory:
                score = (
                    0.5 * r.score +                          # 语义相似度
                    0.2 * memory.confidence +                 # 置信度
                    0.15 * self._recency_score(memory) +      # 时效性
                    0.15 * self._frequency_score(memory)      # 访问频率
                )
                scored_results.append((memory, score))
        
        scored_results.sort(key=lambda x: x[1], reverse=True)
        
        # 更新访问记录
        for memory, _ in scored_results[:top_k]:
            memory.last_accessed = datetime.now()
            memory.access_count += 1
        
        return [m for m, _ in scored_results[:top_k]]
    
    async def update(self, memory_id: str, new_content: str):
        """更新已有记忆"""
        if memory_id in self.memories:
            memory = self.memories[memory_id]
            memory.content = new_content
            memory.confidence = min(1.0, memory.confidence + 0.1)
            
            # 更新向量
            embedding = await self._embed(new_content)
            await self.vector_store.upsert(
                id=memory_id,
                embedding=embedding,
                metadata={"content": new_content}
            )
    
    async def forget(self, memory_id: str):
        """遗忘（删除）记忆"""
        if memory_id in self.memories:
            del self.memories[memory_id]
            await self.vector_store.delete(id=memory_id)

class MemoryConsolidation:
    """记忆巩固：从情景记忆中提取长期记忆"""
    
    async def consolidate(self, episodes: list[EpisodicEvent]) -> list[LongTermMemory]:
        """从一系列情景事件中提取长期记忆"""
        long_term_memories = []
        
        # 提取用户偏好
        preferences = await self._extract_preferences(episodes)
        long_term_memories.extend(preferences)
        
        # 提取经验教训
        lessons = await self._extract_lessons(episodes)
        long_term_memories.extend(lessons)
        
        # 提取操作流程
        procedures = await self._extract_procedures(episodes)
        long_term_memories.extend(procedures)
        
        return long_term_memories
    
    async def _extract_preferences(self, episodes: list[EpisodicEvent]) -> list[LongTermMemory]:
        """从事件中提取用户偏好"""
        # 使用LLM分析事件中的偏好信号
        prompt = f"""分析以下事件，提取用户偏好：
        事件: {[e.description for e in episodes]}
        
        提取用户偏好的格式：
        - 偏好内容
        - 置信度（0-1）
        - 依据事件
        """
        # 解析LLM输出...
        return []

class AgentMemorySystem:
    """Agent记忆系统：统一管理三层记忆"""
    
    def __init__(self):
        self.short_term = ShortTermMemory(max_tokens=32000)
        self.episodic = EpisodicMemory()
        self.long_term = LongTermMemoryStore(...)
    
    async def recall(self, query: str) -> dict:
        """统一记忆检索"""
        # 1. 短期记忆：直接可用
        short_term_context = self.short_term.get_context()
        
        # 2. 长期记忆：语义检索
        long_term_hits = await self.long_term.retrieve(query, top_k=5)
        
        # 3. 情景记忆：相似事件
        similar_episodes = await self.episodic.recall_similar(query, top_k=3)
        
        return {
            "short_term": short_term_context,
            "long_term": [m.content for m in long_term_hits],
            "episodic": [e.description for e in similar_episodes]
        }
    
    async def remember(self, event: EpisodicEvent):
        """记录新事件"""
        # 存入情景记忆
        await self.episodic.record(event)
        
        # 如果是重要事件，触发记忆巩固
        if event.emotional_valence > 0.5 or event.event_type == "decision":
            long_term = await self._consolidate([event])
            for memory in long_term:
                await self.long_term.store(memory)