第十二章 Agent Memory系统¶

"""agent_memory/short_term.py - 短期记忆（工作记忆）"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from enum import Enum
from typing import Any

import tiktoken

# 同时继承str和Enum：枚举值既是Enum成员又是str，可直接用于字符串比较和JSON序列化
class MessageRole(str, Enum):
    SYSTEM = "system"
    USER = "user"
    ASSISTANT = "assistant"
    TOOL = "tool"

@dataclass
class ChatMessage:
    """增强型消息结构"""
    role: MessageRole
    content: str
    timestamp: float = field(default_factory=time.time)
    token_count: int = 0
    importance: float = 1.0       # 重要性评分 (0-1)
    metadata: dict[str, Any] = field(default_factory=dict)

    def to_openai_format(self) -> dict[str, str]:
        """转换为OpenAI API格式"""
        return {"role": self.role.value, "content": self.content}

class ShortTermMemory:
    """
    短期记忆管理器。

    实现三种截断策略:
    1. FIFO (先进先出) - 最简单
    2. 滑动窗口 + 摘要 - 保留首尾，压缩中间
    3. 重要性评分 - 优先丢弃低重要性消息
    """

    def __init__(
        self,
        max_tokens: int = 8000,
        strategy: str = "sliding_window",  # "fifo", "sliding_window", "importance"
        encoding_name: str = "cl100k_base",
    ) -> None:
        self.max_tokens = max_tokens
        self.strategy = strategy
        self._encoder = tiktoken.get_encoding(encoding_name)
        self._messages: list[ChatMessage] = []
        self._compressed_summary: str = ""

    def add(
        self,
        role: str | MessageRole,
        content: str,
        importance: float = 1.0,
        **metadata: Any,
    ) -> None:
        """添加消息到短期记忆"""
        if isinstance(role, str):
            role = MessageRole(role)

        token_count = len(self._encoder.encode(content))

        msg = ChatMessage(
            role=role,
            content=content,
            token_count=token_count,
            importance=importance,
            metadata=metadata,
        )
        self._messages.append(msg)
        self._apply_strategy()

    def get_messages(self) -> list[dict[str, str]]:
        """获取当前所有消息（OpenAI格式）"""
        result = []
        if self._compressed_summary:
            result.append({
                "role": "system",
                "content": f"[对话历史摘要]\n{self._compressed_summary}",
            })
        for msg in self._messages:
            result.append(msg.to_openai_format())
        return result

    @property
    def total_tokens(self) -> int:
        """当前总Token数"""
        total = sum(m.token_count + 4 for m in self._messages)  # +4为每条消息的格式开销Token
        if self._compressed_summary:
            total += len(self._encoder.encode(self._compressed_summary))
        return total

    def _apply_strategy(self) -> None:
        """根据策略执行截断"""
        if self.total_tokens <= self.max_tokens:
            return

        if self.strategy == "fifo":
            self._strategy_fifo()
        elif self.strategy == "sliding_window":
            self._strategy_sliding_window()
        elif self.strategy == "importance":
            self._strategy_importance()

    def _strategy_fifo(self) -> None:
        """先进先出：直接删除最旧的消息"""
        while self.total_tokens > self.max_tokens and len(self._messages) > 1:
            removed = self._messages.pop(0)

    def _strategy_sliding_window(self) -> None:
        """
        滑动窗口策略：
        - 保留第一条消息（通常是system prompt或首次提问）
        - 保留最近的N条消息
        - 中间部分标记为可压缩
        """
        while self.total_tokens > self.max_tokens and len(self._messages) > 2:
            # 保护第一条和最后两条
            self._messages.pop(1)  # 总是删除第二条（最旧的非首条）

    def _strategy_importance(self) -> None:
        """重要性策略：优先删除低重要性消息"""
        while self.total_tokens > self.max_tokens and len(self._messages) > 1:
            # 找到重要性最低的非最后一条消息
            min_importance = float("inf")
            min_idx = 0
            for i, msg in enumerate(self._messages[:-1]):
                if msg.importance < min_importance:
                    min_importance = msg.importance
                    min_idx = i
            self._messages.pop(min_idx)

    def set_summary(self, summary: str) -> None:
        """设置压缩摘要"""
        self._compressed_summary = summary

    def clear(self) -> None:
        """清空短期记忆"""
        self._messages.clear()
        self._compressed_summary = ""

    def get_stats(self) -> dict[str, Any]:
        """获取记忆统计信息"""
        return {
            "message_count": len(self._messages),
            "total_tokens": self.total_tokens,
            "max_tokens": self.max_tokens,
            "utilization": f"{self.total_tokens / self.max_tokens * 100:.1f}%",
            "has_summary": bool(self._compressed_summary),
            "strategy": self.strategy,
        }

"""agent_memory/summarizer.py - 对话摘要压缩器"""

from __future__ import annotations

from openai import OpenAI

class ConversationSummarizer:
    """使用LLM对对话历史进行摘要压缩"""

    SUMMARIZE_PROMPT = """请将以下对话历史压缩为一段简洁的摘要。

要求：
1. 保留用户的核心需求和关键上下文
2. 保留重要的数据、结论和决策
3. 丢弃寒暄、重复和不重要的细节
4. 控制在150字以内
5. 使用第三人称描述

对话历史:
{conversation}

摘要:"""

    def __init__(self, client: OpenAI, model: str = "gpt-4o-mini") -> None:
        self.client = client
        self.model = model

    def summarize(self, messages: list[dict[str, str]]) -> str:
        """
        将消息列表压缩为摘要。

        Args:
            messages: OpenAI格式的消息列表

        Returns:
            压缩后的摘要字符串
        """
        # 格式化对话
        conversation = "\n".join(
            f"{m['role'].upper()}: {m['content']}" for m in messages
        )

        prompt = self.SUMMARIZE_PROMPT.replace("{conversation}", conversation)

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3,  # 低温度保证摘要的准确性
            max_tokens=300,
        )

        return response.choices[0].message.content or ""

    def incremental_summarize(
        self,
        existing_summary: str,
        new_messages: list[dict[str, str]],
    ) -> str:
        """
        增量摘要：在现有摘要基础上融入新消息。
        比全量重新摘要更省Token。
        """
        new_conversation = "\n".join(
            f"{m['role'].upper()}: {m['content']}" for m in new_messages
        )

        prompt = f"""基于现有摘要和新的对话内容，生成更新后的摘要。

现有摘要:
{existing_summary}

新对话:
{new_conversation}

更新后的摘要（150字以内）:"""

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3,
            max_tokens=300,
        )

        return response.choices[0].message.content or ""

"""agent_memory/long_term.py - 长期记忆（向量存储）"""

from __future__ import annotations

import hashlib
import json
import time
from dataclasses import dataclass, field
from typing import Any

import chromadb
from chromadb.config import Settings
from openai import OpenAI

@dataclass
class MemoryItem:
    """长期记忆条目"""
    content: str                            # 记忆内容
    memory_type: str = "general"            # 类型：general, fact, preference, task_result
    importance: float = 0.5                 # 重要性评分 (0-1)
    created_at: float = field(default_factory=time.time)
    last_accessed: float = field(default_factory=time.time)
    access_count: int = 0                   # 被检索次数
    metadata: dict[str, Any] = field(default_factory=dict)

    @property
    def memory_id(self) -> str:
        """基于内容生成唯一ID"""
        # 注意：MD5用于生成非安全用途的内容ID是可以的，但如需更强的唯一性
        # 可改用 hashlib.sha256(...).hexdigest()[:16]
        return hashlib.md5(self.content.encode()).hexdigest()[:16]

class LongTermMemory:
    """
    长期记忆管理器，基于ChromaDB向量数据库。

    核心功能:
    - 记忆存储与向量索引
    - 语义检索（基于余弦相似度）
    - 重要性评分
    - 记忆衰减与遗忘
    """

    def __init__(
        self,
        collection_name: str = "agent_memory",
        persist_dir: str = "./memory_store",
        embedding_model: str = "text-embedding-3-small",
        openai_api_key: str | None = None,
        decay_rate: float = 0.995,       # 每次检索后的时间衰减率
    ) -> None:
        self.embedding_model = embedding_model
        self.decay_rate = decay_rate

        # 初始化OpenAI（用于embedding）
        self.openai_client = OpenAI(api_key=openai_api_key)

        # 初始化ChromaDB
        self.chroma_client = chromadb.PersistentClient(path=persist_dir)
        self.collection = self.chroma_client.get_or_create_collection(
            name=collection_name,
            metadata={"hnsw:space": "cosine"},  # 使用余弦相似度
        )

    def store(self, item: MemoryItem) -> str:
        """
        存储一条记忆。

        Args:
            item: 记忆条目

        Returns:
            记忆ID
        """
        memory_id = item.memory_id

        # 生成embedding
        embedding = self._get_embedding(item.content)

        # 元数据
        metadata = {
            "memory_type": item.memory_type,
            "importance": item.importance,
            "created_at": item.created_at,
            "last_accessed": item.last_accessed,
            "access_count": item.access_count,
            **{k: str(v) for k, v in item.metadata.items()},  # **解包字典推导式：将元数据值转str后合并到外层字典
        }

        # 存入ChromaDB
        self.collection.upsert(
            ids=[memory_id],
            embeddings=[embedding],
            documents=[item.content],
            metadatas=[metadata],
        )

        return memory_id

    def retrieve(
        self,
        query: str,
        top_k: int = 5,
        memory_type: str | None = None,
        min_importance: float = 0.0,
    ) -> list[dict[str, Any]]:
        """
        语义检索相关记忆。

        Args:
            query: 查询文本
            top_k: 返回数量
            memory_type: 过滤记忆类型
            min_importance: 最低重要性阈值

        Returns:
            相关记忆列表，按相关性排序
        """
        query_embedding = self._get_embedding(query)

        # 构建过滤条件
        where_filter: dict[str, Any] | None = None
        conditions = []
        if memory_type:
            conditions.append({"memory_type": {"$eq": memory_type}})
        if min_importance > 0:
            conditions.append({"importance": {"$gte": min_importance}})

        if len(conditions) == 1:
            where_filter = conditions[0]
        elif len(conditions) > 1:
            where_filter = {"$and": conditions}

        # 执行查询
        results = self.collection.query(
            query_embeddings=[query_embedding],
            n_results=top_k,
            where=where_filter,
            include=["documents", "metadatas", "distances"],
        )

        # 整理结果
        memories = []
        if results["documents"] and results["documents"][0]:
            for i, doc in enumerate(results["documents"][0]):
                meta = results["metadatas"][0][i] if results["metadatas"] else {}
                distance = results["distances"][0][i] if results["distances"] else 0

                # 余弦距离转相似度
                similarity = 1 - distance

                # 应用时间衰减
                created_at = float(meta.get("created_at", time.time()))
                age_days = (time.time() - created_at) / 86400
                decay_factor = self.decay_rate ** age_days
                final_score = similarity * decay_factor

                memories.append({
                    "content": doc,
                    "similarity": similarity,
                    "final_score": final_score,
                    "metadata": meta,
                })

                # 更新访问统计
                memory_id = results["ids"][0][i]
                self._update_access(memory_id, meta)

        # 按综合评分排序
        memories.sort(key=lambda x: x["final_score"], reverse=True)  # 原地降序排序，最相关记忆排在最前
        return memories

    def forget(self, min_importance: float = 0.1, max_age_days: float = 90) -> int:
        """
        遗忘机制：清除低重要性且长时间未访问的记忆。

        Args:
            min_importance: 低于此重要性的记忆会被清除
            max_age_days: 超过此天数未访问的记忆会被清除

        Returns:
            清除的记忆数量
        """
        cutoff_time = time.time() - (max_age_days * 86400)

        # 获取所有记忆
        all_records = self.collection.get(
            include=["metadatas"]
        )

        ids_to_delete = []
        if all_records["ids"]:
            for i, memory_id in enumerate(all_records["ids"]):
                meta = all_records["metadatas"][i] if all_records["metadatas"] else {}
                importance = float(meta.get("importance", 0.5))
                last_accessed = float(meta.get("last_accessed", time.time()))

                if importance < min_importance and last_accessed < cutoff_time:
                    ids_to_delete.append(memory_id)

        if ids_to_delete:
            self.collection.delete(ids=ids_to_delete)

        return len(ids_to_delete)

    def get_stats(self) -> dict[str, Any]:
        """获取长期记忆统计"""
        count = self.collection.count()
        return {
            "total_memories": count,
            "collection_name": self.collection.name,
            "decay_rate": self.decay_rate,
        }

    def _get_embedding(self, text: str) -> list[float]:
        """调用OpenAI Embedding API"""
        response = self.openai_client.embeddings.create(
            model=self.embedding_model,
            input=text,
        )
        return response.data[0].embedding

    def _update_access(self, memory_id: str, current_meta: dict) -> None:
        """更新记忆的访问时间和次数"""
        access_count = int(current_meta.get("access_count", 0)) + 1
        self.collection.update(
            ids=[memory_id],
            metadatas=[{
                **current_meta,  # **解包保留原有元数据，仅覆盖更新访问时间和次数
                "last_accessed": time.time(),
                "access_count": access_count,
            }],
        )

"""agent_memory/importance.py - 记忆重要性评估"""

from openai import OpenAI

class ImportanceScorer:
    """使用LLM评估记忆的重要性"""

    SCORING_PROMPT = """请评估以下信息对于AI助手的长期价值，给出0到1之间的重要性评分。

评分标准：
- 1.0: 极其重要（用户的核心偏好、关键决策、重要事实）
- 0.7: 比较重要（有用的上下文、中期相关的信息）
- 0.4: 一般重要（可能未来有用的细节）
- 0.1: 不太重要（临时性信息、寒暄、可轻易再获取的内容）

信息内容：
{content}

请只返回一个0到1之间的数字，不要有其他内容。"""

    def __init__(self, client: OpenAI, model: str = "gpt-4o-mini") -> None:
        self.client = client
        self.model = model

    def score(self, content: str) -> float:
        """
        评估一段内容的重要性。

        Returns:
            0-1之间的重要性评分
        """
        prompt = self.SCORING_PROMPT.replace("{content}", content)

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.1,
            max_tokens=10,
        )

        try:
            score = float(response.choices[0].message.content.strip())
            return max(0.0, min(1.0, score))
        except (ValueError, AttributeError):
            return 0.5  # 默认中等重要性

"""agent_memory/episodic.py - 情景记忆（经验回放）"""

from __future__ import annotations

import json
import time
from dataclasses import dataclass, field
from typing import Any

from .long_term import LongTermMemory, MemoryItem

@dataclass
class Episode:
    """一次完整的任务执行记录"""
    task: str                              # 任务描述
    steps: list[dict[str, Any]]            # 执行步骤
    result: str                            # 最终结果
    success: bool                          # 是否成功
    duration_seconds: float = 0.0          # 耗时
    reflection: str = ""                   # 反思总结
    timestamp: float = field(default_factory=time.time)

    def to_text(self) -> str:
        """将情景转为文本（用于存储和检索）"""
        status = "✅成功" if self.success else "❌失败"
        steps_summary = "\n".join(
            f"  Step {i+1}: {s.get('action', 'think')} → {s.get('result', '')[:100]}"
            for i, s in enumerate(self.steps)
        )
        return (
            f"任务: {self.task}\n"
            f"状态: {status}\n"
            f"步骤:\n{steps_summary}\n"
            f"结果: {self.result[:200]}\n"
            f"反思: {self.reflection}"
        )

class EpisodicMemory:
    """
    情景记忆管理器。

    功能:
    - 记录任务执行轨迹
    - 检索相似任务的历史经验
    - 提取成功/失败模式
    - 动态生成Few-shot示例
    """

    def __init__(self, long_term_memory: LongTermMemory) -> None:
        self.ltm = long_term_memory

    def save_episode(self, episode: Episode) -> str:
        """保存一次任务执行经验"""
        item = MemoryItem(
            content=episode.to_text(),
            memory_type="episode",
            importance=0.8 if episode.success else 0.6,
            metadata={
                "task": episode.task,
                "success": str(episode.success),
                "step_count": len(episode.steps),
                "duration": episode.duration_seconds,
            },
        )
        return self.ltm.store(item)

    def recall_similar(
        self,
        current_task: str,
        top_k: int = 3,
        success_only: bool = False,
    ) -> list[dict[str, Any]]:
        """
        检索与当前任务相似的历史经验。

        Args:
            current_task: 当前任务描述
            top_k: 返回数量
            success_only: 是否只返回成功的经验
        """
        results = self.ltm.retrieve(
            query=current_task,
            top_k=top_k * 2,  # 多检索一些，后面过滤
            memory_type="episode",
        )

        if success_only:
            results = [
                r for r in results
                if r["metadata"].get("success") == "True"
            ]

        return results[:top_k]

    def generate_few_shot_examples(
        self,
        current_task: str,
        max_examples: int = 2,
    ) -> str:
        """
        根据历史经验动态生成Few-shot示例。

        Returns:
            格式化的Few-shot示例文本
        """
        similar_episodes = self.recall_similar(
            current_task, top_k=max_examples, success_only=True
        )

        if not similar_episodes:
            return ""

        examples = ["以下是类似任务的成功经验供参考：\n"]
        for i, ep in enumerate(similar_episodes, 1):
            examples.append(f"--- 经验 {i} ---")
            examples.append(ep["content"])
            examples.append("")

        return "\n".join(examples)

class ReflectionEngine:
    """
    反思引擎：在任务完成后进行自我反思，
    提取可复用的经验和教训。
    """

    REFLECTION_PROMPT = """分析以下任务执行过程，提取关键经验：

任务: {task}
执行结果: {result}
是否成功: {success}
执行步骤:
{steps}

请简洁总结：
1. 哪些策略有效？
2. 哪些步骤可以改进？
3. 下次遇到类似任务的建议。

总结（100字以内）:"""

    def __init__(self, client: Any, model: str = "gpt-4o-mini") -> None:
        self.client = client
        self.model = model

    def reflect(self, episode: Episode) -> str:
        """对一次任务执行进行反思"""
        steps_text = "\n".join(
            f"Step {i+1}: {json.dumps(s, ensure_ascii=False)}"
            for i, s in enumerate(episode.steps)
        )

        prompt = self.REFLECTION_PROMPT.format(
            task=episode.task,
            result=episode.result,
            success="是" if episode.success else "否",
            steps=steps_text,
        )

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3,
            max_tokens=200,
        )

        reflection = response.choices[0].message.content or ""
        episode.reflection = reflection
        return reflection

"""agent_memory/semantic.py - 语义记忆（知识图谱）"""

from __future__ import annotations

import json
import re
from dataclasses import dataclass, field
from typing import Any

from openai import OpenAI

@dataclass
class Triple:
    """知识三元组"""
    head: str          # 头实体
    relation: str      # 关系
    tail: str          # 尾实体
    confidence: float = 1.0  # 置信度
    source: str = ""         # 信息来源

    def __str__(self) -> str:
        return f"({self.head}, {self.relation}, {self.tail})"

class KnowledgeGraph:
    """
    简化版知识图谱。

    使用邻接表存储三元组关系，
    支持实体查询和关系推理。
    """

    def __init__(self) -> None:
        # 邻接表: entity -> [(relation, target, confidence), ...]
        self._graph: dict[str, list[tuple[str, str, float]]] = {}
        # 全部三元组
        self._triples: list[Triple] = []

    def add_triple(self, triple: Triple) -> None:
        """添加一个知识三元组"""
        self._triples.append(triple)

        # 正向关系
        if triple.head not in self._graph:
            self._graph[triple.head] = []
        self._graph[triple.head].append(
            (triple.relation, triple.tail, triple.confidence)
        )

        # 反向索引（方便反向查询）
        if triple.tail not in self._graph:
            self._graph[triple.tail] = []
        self._graph[triple.tail].append(
            (f"被{triple.relation}", triple.head, triple.confidence)
        )

    def query_entity(self, entity: str) -> list[dict[str, Any]]:
        """查询实体的所有关系"""
        relations = self._graph.get(entity, [])
        return [
            {"relation": r, "target": t, "confidence": c}
            for r, t, c in relations
        ]

    def query_relation(
        self, head: str, relation: str
    ) -> list[str]:
        """查询特定关系的目标实体"""
        relations = self._graph.get(head, [])
        return [t for r, t, _ in relations if r == relation]  # 列表推导+元组解包：_忽略置信度，筛选匹配关系的目标实体

    def find_path(
        self, start: str, end: str, max_depth: int = 3
    ) -> list[list[str]] | None:
        """
        查找两个实体之间的关系路径（BFS）。
        用于多跳推理。
        """
        if start == end:
            return [[start]]

        from collections import deque
        queue: deque[list[str]] = deque([[start]])  # BFS队列，每个元素是一条路径（实体列表）
        visited = {start}
        paths = []

        while queue:
            path = queue.popleft()
            if len(path) > max_depth * 2:  # 路径包含实体和关系交替
                continue

            current = path[-1]
            for rel, target, _ in self._graph.get(current, []):
                if target == end:
                    paths.append(path + [rel, target])
                elif target not in visited:
                    visited.add(target)
                    queue.append(path + [rel, target])

        return paths if paths else None

    def to_text(self, entity: str | None = None) -> str:
        """将知识转为文本（用于注入LLM上下文）"""
        if entity:
            relations = self.query_entity(entity)
            lines = [f"关于 [{entity}] 的知识:"]
            for r in relations:
                lines.append(f"  - {r['relation']}: {r['target']}")
            return "\n".join(lines)
        else:
            lines = ["当前知识图谱:"]
            for triple in self._triples:
                lines.append(f"  - {triple}")
            return "\n".join(lines)

    @property
    def entity_count(self) -> int:
        return len(self._graph)

    @property
    def triple_count(self) -> int:
        return len(self._triples)

class KnowledgeExtractor:
    """使用LLM从文本中提取知识三元组"""

    EXTRACTION_PROMPT = """从以下文本中提取知识三元组（实体-关系-实体）。

文本:
{text}

请以JSON数组格式返回，每个元素包含 head, relation, tail 三个字段:
[
  {{"head": "实体A", "relation": "关系", "tail": "实体B"}},
  ...
]

注意：
- 只提取明确陈述的事实
- 关系用简洁的动词或动词短语表示
- 如果没有可提取的知识，返回空数组 []

只返回JSON数组:"""

    def __init__(self, client: OpenAI, model: str = "gpt-4o-mini") -> None:
        self.client = client
        self.model = model

    def extract(self, text: str) -> list[Triple]:
        """从文本中提取知识三元组"""
        prompt = self.EXTRACTION_PROMPT.replace("{text}", text)

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.1,
            max_tokens=500,
        )

        content = response.choices[0].message.content or "[]"
        content = content.strip().strip("```json").strip("```").strip()

        try:
            raw_triples = json.loads(content)
        except json.JSONDecodeError:
            return []

        triples = []
        for item in raw_triples:
            if all(k in item for k in ("head", "relation", "tail")):  # 校验三元组必须包含头实体、关系、尾实体
                triples.append(Triple(
                    head=item["head"],
                    relation=item["relation"],
                    tail=item["tail"],
                    source=text[:100],
                ))
        return triples

from openai import OpenAI

client = OpenAI(api_key="sk-xxx")

# 提取知识
extractor = KnowledgeExtractor(client)
triples = extractor.extract(
    "OpenAI成立于2015年，由Sam Altman担任CEO。"
    "GPT-4是OpenAI开发的大语言模型，于2023年发布。"
)

# 构建知识图谱
kg = KnowledgeGraph()
for t in triples:
    kg.add_triple(t)
    print(f"  提取: {t}")

# 查询
print(kg.to_text("OpenAI"))
# 输出:
# 关于 [OpenAI] 的知识:
#   - 成立于: 2015年
#   - CEO: Sam Altman
#   - 开发了: GPT-4

# 路径查找
paths = kg.find_path("Sam Altman", "GPT-4")
print(f"Sam Altman → GPT-4 路径: {paths}")

"""agent_memory/memgpt.py - MemGPT风格的记忆管理"""

from __future__ import annotations

from typing import Any
from dataclasses import dataclass, field

from .short_term import ShortTermMemory
from .long_term import LongTermMemory, MemoryItem

@dataclass
class CoreMemory:
    """核心记忆块：始终在上下文中的关键信息"""
    persona: str = "你是一个有帮助的AI助手。"    # Agent自我描述
    human: str = ""                               # 用户画像
    custom_blocks: dict[str, str] = field(default_factory=dict)  # 可扩展的自定义记忆块，每实例独立

    def to_text(self) -> str:
        text = f"[Persona]\n{self.persona}\n\n[Human]\n{self.human}"
        for name, content in self.custom_blocks.items():
            text += f"\n\n[{name}]\n{content}"
        return text

    @property
    def total_length(self) -> int:
        return len(self.to_text())

class MemGPTController:
    """
    MemGPT控制器：管理三级记忆层级。

    提供给Agent的记忆工具:
    - core_memory_append: 向核心记忆追加信息
    - core_memory_replace: 替换核心记忆中的内容
    - recall_memory_search: 搜索对话历史
    - archival_memory_insert: 存入长期知识库
    - archival_memory_search: 搜索长期知识库
    """

    def __init__(
        self,
        short_term: ShortTermMemory,
        long_term: LongTermMemory,
        core: CoreMemory | None = None,
    ) -> None:
        self.short_term = short_term
        self.long_term = long_term
        self.core = core or CoreMemory()

    def core_memory_append(self, section: str, content: str) -> str:
        """向核心记忆的指定部分追加信息"""
        if section == "persona":
            self.core.persona += f"\n{content}"
        elif section == "human":
            self.core.human += f"\n{content}"
        else:
            existing = self.core.custom_blocks.get(section, "")
            self.core.custom_blocks[section] = f"{existing}\n{content}".strip()
        return f"✅ 已追加到核心记忆 [{section}]"

    def core_memory_replace(
        self, section: str, old_content: str, new_content: str
    ) -> str:
        """替换核心记忆中的内容"""
        if section == "persona":
            self.core.persona = self.core.persona.replace(old_content, new_content)
        elif section == "human":
            self.core.human = self.core.human.replace(old_content, new_content)
        else:
            if section in self.core.custom_blocks:
                self.core.custom_blocks[section] = \
                    self.core.custom_blocks[section].replace(old_content, new_content)
        return f"✅ 核心记忆 [{section}] 已更新"

    def recall_memory_search(self, query: str, top_k: int = 5) -> str:
        """搜索对话历史"""
        messages = self.short_term.get_messages()
        # 简单的关键词匹配（生产中应使用向量检索）
        relevant = [
            m for m in messages
            if query.lower() in m.get("content", "").lower()
        ]
        if not relevant:
            return "未找到相关对话记录。"
        result = "\n".join(
            f"[{m['role']}] {m['content'][:200]}" for m in relevant[-top_k:]
        )
        return f"找到 {len(relevant)} 条相关记录:\n{result}"

    def archival_memory_insert(self, content: str) -> str:
        """存入长期知识库"""
        item = MemoryItem(content=content, memory_type="archival", importance=0.7)
        memory_id = self.long_term.store(item)
        return f"✅ 已存入长期知识库 (ID: {memory_id})"

    def archival_memory_search(self, query: str, top_k: int = 5) -> str:
        """搜索长期知识库"""
        results = self.long_term.retrieve(query, top_k=top_k, memory_type="archival")
        if not results:
            return "长期知识库中未找到相关内容。"
        lines = []
        for r in results:
            lines.append(f"[相似度: {r['similarity']:.2f}] {r['content'][:200]}")
        return "\n".join(lines)

    def get_memory_tools_schema(self) -> list[dict[str, Any]]:
        """获取所有记忆工具的Schema（用于Function Calling）"""
        return [
            {
                "type": "function",
                "function": {
                    "name": "core_memory_append",
                    "description": "向核心记忆追加重要信息（用户偏好、关键事实等）",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "section": {
                                "type": "string",
                                "enum": ["persona", "human"] +
                                        list(self.core.custom_blocks.keys()),
                                "description": "要追加到的记忆部分",
                            },
                            "content": {
                                "type": "string",
                                "description": "要追加的内容",
                            },
                        },
                        "required": ["section", "content"],
                    },
                },
            },
            {
                "type": "function",
                "function": {
                    "name": "archival_memory_search",
                    "description": "搜索长期知识库中的相关信息",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "query": {
                                "type": "string",
                                "description": "搜索查询",
                            },
                        },
                        "required": ["query"],
                    },
                },
            },
            {
                "type": "function",
                "function": {
                    "name": "archival_memory_insert",
                    "description": "将重要信息存入长期知识库",
                    "parameters": {
                        "type": "object",
                        "properties": {
                            "content": {
                                "type": "string",
                                "description": "要存储的内容",
                            },
                        },
                        "required": ["content"],
                    },
                },
            },
        ]

"""agent_memory/unified.py - 统一记忆接口"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

from .short_term import ShortTermMemory
from .long_term import LongTermMemory, MemoryItem
from .episodic import EpisodicMemory, Episode
from .semantic import KnowledgeGraph, Triple
from .importance import ImportanceScorer

@dataclass
class MemoryQueryResult:
    """记忆查询的统一结果"""
    source: str                # "short_term", "long_term", "episodic", "semantic"
    content: str
    relevance: float
    metadata: dict[str, Any]

class UnifiedMemory:
    """
    统一记忆系统：整合所有记忆模块，
    提供单一接口进行记忆存储和检索。
    """

    def __init__(
        self,
        short_term: ShortTermMemory,
        long_term: LongTermMemory,
        episodic: EpisodicMemory | None = None,
        knowledge_graph: KnowledgeGraph | None = None,
        importance_scorer: ImportanceScorer | None = None,
    ) -> None:
        self.short_term = short_term
        self.long_term = long_term
        self.episodic = episodic
        self.kg = knowledge_graph
        self.scorer = importance_scorer

    def remember(
        self,
        content: str,
        memory_type: str = "general",
        auto_score: bool = True,
    ) -> str:
        """
        存储一条记忆，自动评估重要性并分配到合适的存储层。

        Args:
            content: 要记忆的内容
            memory_type: 记忆类型
            auto_score: 是否自动评估重要性

        Returns:
            存储结果描述
        """
        importance = 0.5
        if auto_score and self.scorer:
            importance = self.scorer.score(content)

        # 根据重要性选择存储策略
        if importance >= 0.7:
            # 高重要性：同时存入长期记忆
            item = MemoryItem(
                content=content,
                memory_type=memory_type,
                importance=importance,
            )
            memory_id = self.long_term.store(item)
            return f"高重要性({importance:.1f})，已存入长期记忆 [{memory_id}]"
        else:
            # 低重要性：仅保留在短期记忆中
            return f"中低重要性({importance:.1f})，保留在短期记忆中"

    def recall(
        self,
        query: str,
        top_k: int = 5,
        sources: list[str] | None = None,
    ) -> list[MemoryQueryResult]:
        """
        综合检索所有记忆源。

        Args:
            query: 查询文本
            top_k: 每个源返回的最大数量
            sources: 指定检索的记忆源，None表示全部

        Returns:
            按相关性排序的记忆列表
        """
        sources = sources or ["long_term", "episodic", "semantic"]
        results: list[MemoryQueryResult] = []

        # 长期记忆检索
        if "long_term" in sources:
            lt_results = self.long_term.retrieve(query, top_k=top_k)
            for r in lt_results:
                results.append(MemoryQueryResult(
                    source="long_term",
                    content=r["content"],
                    relevance=r["final_score"],
                    metadata=r["metadata"],
                ))

        # 情景记忆检索
        if "episodic" in sources and self.episodic:
            ep_results = self.episodic.recall_similar(query, top_k=top_k)
            for r in ep_results:
                results.append(MemoryQueryResult(
                    source="episodic",
                    content=r["content"],
                    relevance=r["final_score"],
                    metadata=r["metadata"],
                ))

        # 知识图谱检索
        if "semantic" in sources and self.kg:
            # 简单的实体匹配
            kg_results = self.kg.query_entity(query)
            if kg_results:
                content = self.kg.to_text(query)
                results.append(MemoryQueryResult(
                    source="semantic",
                    content=content,
                    relevance=0.8,
                    metadata={"entity": query},
                ))

        # 按相关性排序
        results.sort(key=lambda x: x.relevance, reverse=True)  # 按相关性降序，截取前top_k条返回
        return results[:top_k]

    def build_context(self, query: str, max_tokens: int = 2000) -> str:
        """
        为LLM构建记忆上下文：
        将检索到的相关记忆格式化为可注入的文本块。
        """
        memories = self.recall(query, top_k=10)
        if not memories:
            return ""

        context_parts = ["[相关记忆]"]
        current_tokens = 10  # 预估标题token

        for mem in memories:
            # 粗略估算token（1个中文字≈2 tokens）
            estimated_tokens = len(mem.content) * 2
            if current_tokens + estimated_tokens > max_tokens:
                break
            source_label = {
                "long_term": "💾",
                "episodic": "📝",
                "semantic": "🔗",
            }.get(mem.source, "📌")
            context_parts.append(
                f"{source_label} [{mem.source}] (相关性:{mem.relevance:.2f})\n{mem.content}"
            )
            current_tokens += estimated_tokens

        return "\n\n".join(context_parts)

    def get_stats(self) -> dict[str, Any]:
        """获取全部记忆系统的统计信息"""
        stats: dict[str, Any] = {
            "short_term": self.short_term.get_stats(),
            "long_term": self.long_term.get_stats(),
        }
        if self.kg:
            stats["knowledge_graph"] = {
                "entities": self.kg.entity_count,
                "triples": self.kg.triple_count,
            }
        return stats