第六章：记忆系统——从会话到知识图谱

# coder.md 提示词中的记忆写入模式
"""
## Session Memory Write
After completing your work, write a session summary to:
  {spec_dir}/session_insights/session_{n}.json

Include:
- What you built and key decisions made
- Any gotchas or surprises encountered
- Patterns you discovered in the codebase
- What you'd tell yourself if you had to redo this task
"""

class GraphitiMemory:
    """
    持久记忆的高层接口。

    所有操作都是 async 并有降级处理：
    如果 Graphiti 未启用或不可用，操作静默跳过。
    """

    def __init__(
        self,
        spec_dir: Path,
        project_dir: Path,
        group_id_mode: str = GroupIdMode.PROJECT,  # 默认：项目级共享
    ):
        self.group_id_mode = group_id_mode
        # ...

    @property
    def group_id(self) -> str:
        """
        记忆命名空间：决定记忆的作用域。

        SPEC 模式：每个任务独立记忆（适合隔离实验）
        PROJECT 模式：所有任务共享项目记忆（跨任务学习）
        """
        if self.group_id_mode == GroupIdMode.PROJECT:
            # 用项目路径 hash 确保唯一性
            path_hash = hashlib.md5(
                str(self.project_dir.resolve()).encode(),
                usedforsecurity=False
            ).hexdigest()[:8]
            return f"project_{self.project_dir.name}_{path_hash}"
        else:
            return self.spec_dir.name  # "001-add-auth"

# integrations/graphiti/queries_pkg/schema.py

EPISODE_TYPE_SESSION_INSIGHT = "session_insight"
# 用途：智能体会话的总体洞察
# 示例："在此任务中，发现数据库连接需要指定 charset=utf8mb4"

EPISODE_TYPE_CODEBASE_DISCOVERY = "codebase_discovery"
# 用途：代码库结构和模式的发现
# 示例："auth/ 模块使用 JWT，token 存储在 Redis 中，TTL 24小时"

EPISODE_TYPE_PATTERN = "pattern"
# 用途：项目中反复出现的代码模式
# 示例："所有 API handler 都需要 @require_auth 装饰器"

EPISODE_TYPE_GOTCHA = "gotcha"
# 用途：陷阱、已知问题、不直观的行为
# 示例："切勿在 async 函数中直接调用 db.query()，必须用 await db.query()"

EPISODE_TYPE_TASK_OUTCOME = "task_outcome"
# 用途：任务完成结果（成功/失败、方法、耗时）

EPISODE_TYPE_QA_RESULT = "qa_result"
# 用途：QA 发现的问题和修复方式

EPISODE_TYPE_HISTORICAL_CONTEXT = "historical_context"
# 用途：历史决策和架构演化的背景

async def save_session_insights(
    self,
    session_num: int,
    insights: dict,
) -> bool:
    """
    将会话洞察保存为 Graphiti Episode。
    """
    if not await self._ensure_initialized():
        return False  # Graphiti 未启用，静默跳过

    try:
        result = await self._queries.add_session_insight(session_num, insights)

        if result and self.state:
            self.state.last_session = session_num
            self.state.episode_count += 1
            self.state.save(self.spec_dir)  # 持久化状态

        return result
    except Exception as e:
        logger.warning(f"Failed to save session insights: {e}")
        # 记忆写入失败不应阻断工作流
        return False

# 保存代码库发现
await memory.save_codebase_discoveries({
    "auth_module": "JWT-based, Redis token storage, 24h TTL",
    "database_pattern": "Connection pool size=10, charset must be utf8mb4",
})

# 保存 Gotcha（陷阱）
await memory.save_gotcha(
    "database_async",
    "Never call db.query() directly in async functions, use 'await db.query()'"
)

# 保存任务结果
await memory.save_task_outcome(
    task_id="001-add-auth",
    success=True,
    approach="JWT with Redis session store",
    duration_minutes=45,
)

# 保存 QA 结果
await memory.save_qa_result(
    issue="Missing input validation in POST /users",
    fix="Added Pydantic model for request validation",
    severity="HIGH",
)

async def get_relevant_context(
    self,
    query: str,
    max_results: int = MAX_CONTEXT_RESULTS,  # 默认最多 10 条
) -> list[dict]:
    """
    语义搜索：找出与查询最相关的历史记忆。

    使用向量相似度搜索，而不是关键词匹配。
    返回按相关度排序的 Episode 列表。
    """
    if not await self._ensure_initialized():
        return []

    try:
        return await self._search.semantic_search(query, max_results)
    except Exception as e:
        logger.warning(f"Memory search failed: {e}")
        return []  # 搜索失败返回空，不阻断工作

# integrations/graphiti/providers_pkg/factory.py（概念化）

EMBEDDER_PROVIDERS = {
    "openai":      OpenAIEmbedder,       # text-embedding-3-small/large
    "voyage":      VoyageEmbedder,       # voyage-code-2（代码优化）
    "google":      GoogleEmbedder,       # text-embedding-004
    "azure_openai":AzureOpenAIEmbedder,  # 企业 Azure 部署
    "ollama":      OllamaEmbedder,       # 本地模型（隐私保护）
    "openrouter":  OpenRouterEmbedder,   # 路由器（多模型切换）
}

def create_embedder(config: GraphitiConfig) -> BaseEmbedder:
    provider_name = config.embedder_provider
    if provider_name not in EMBEDDER_PROVIDERS:
        raise ProviderError(f"Unknown embedder: {provider_name}")
    return EMBEDDER_PROVIDERS[provider_name](config)

if self.state and self.state.has_provider_changed(self.config):
    migration_info = self.state.get_migration_info(self.config)
    logger.warning(
        f"⚠️  Embedding provider changed: "
        f"{migration_info['old_provider']} → {migration_info['new_provider']}"
    )
    logger.warning(
        "   Run: python integrations/graphiti/migrate_embeddings.py"
    )

# SPEC 模式：每个任务看到独立的记忆空间
memory_001 = GraphitiMemory(spec_dir="specs/001-auth", group_id_mode=GroupIdMode.SPEC)
memory_002 = GraphitiMemory(spec_dir="specs/002-payment", group_id_mode=GroupIdMode.SPEC)
# 搜索 memory_001 时，只能找到任务 001 的记忆
# 搜索 memory_002 时，只能找到任务 002 的记忆

# PROJECT 模式：所有任务共享项目记忆
memory_001 = GraphitiMemory(spec_dir="specs/001-auth", group_id_mode=GroupIdMode.PROJECT)
memory_002 = GraphitiMemory(spec_dir="specs/002-payment", group_id_mode=GroupIdMode.PROJECT)
# 任务 001 写入的 Gotcha，任务 002 可以检索到
# 任务 002 发现的 Pattern，任务 003 可以受益于

# 创建嵌入式数据库（无需服务器）
from integrations.graphiti.queries_pkg.kuzu_driver_patched import (
    create_patched_kuzu_driver,
)

db_path = config.get_db_path()  # ~/.auto-claude/graphiti/project_xxx.db
driver = create_patched_kuzu_driver(db=str(db_path))

graphiti = Graphiti(
    graph_driver=driver,     # 嵌入式，文件存储
    llm_client=llm_client,
    embedder=embedder,
)

# lab_06_memory.py

import json
import math
from dataclasses import dataclass, field, asdict
from datetime import datetime
from pathlib import Path
from typing import Optional
from collections import Counter

class EpisodeType:
    GOTCHA = "gotcha"
    PATTERN = "pattern"
    CODEBASE_DISCOVERY = "codebase_discovery"
    TASK_OUTCOME = "task_outcome"

@dataclass
class Episode:
    """一条知识记忆。"""
    id: str
    episode_type: str
    content: str
    metadata: dict = field(default_factory=dict)
    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
    tags: list[str] = field(default_factory=list)

class KnowledgeStore:
    """
    轻量级知识存储，使用 JSON 文件持久化。
    """

    def __init__(self, storage_dir: Path):
        self.storage_dir = storage_dir
        self.storage_dir.mkdir(parents=True, exist_ok=True)
        self._episodes: list[Episode] = []
        self._load()

    def add_episode(
        self,
        episode_type: str,
        content: str,
        metadata: dict = None,
        tags: list[str] = None,
    ) -> Episode:
        """
        TODO Part A: 添加一条知识记忆

        要求：
        1. 生成唯一 ID（类型 + 序号，如 "gotcha_001"）
        2. 创建 Episode 对象
        3. 添加到 self._episodes
        4. 持久化到磁盘（调用 self._save()）
        5. 返回 Episode
        """
        pass

    def get_by_type(self, episode_type: str) -> list[Episode]:
        """按类型筛选知识。"""
        return [e for e in self._episodes if e.episode_type == episode_type]

    def _save(self):
        """持久化到 JSON 文件。"""
        data = [asdict(e) for e in self._episodes]
        store_file = self.storage_dir / "knowledge.json"
        with open(store_file, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)

    def _load(self):
        """从文件恢复。"""
        store_file = self.storage_dir / "knowledge.json"
        if store_file.exists():
            with open(store_file, encoding="utf-8") as f:
                data = json.load(f)
            self._episodes = [Episode(**d) for d in data]

class TFIDFSearcher:
    """
    使用 TF-IDF 实现简化的语义检索。

    TF-IDF 原理：
    - TF (词频)：词在当前文档中出现的频率
    - IDF (逆文档频率)：词在所有文档中的稀缺程度
    - TF-IDF = TF × IDF：重要词得分高，常见词得分低
    """

    def search(
        self,
        query: str,
        episodes: list[Episode],
        max_results: int = 5,
    ) -> list[tuple[Episode, float]]:
        """
        TODO Part B: 实现 TF-IDF 搜索

        步骤：
        1. 对 query 和所有 episode.content 进行分词（简单按空格/标点分割）
        2. 计算 IDF：对每个词，统计出现在多少个文档中
           IDF(word) = log(总文档数 / 包含该词的文档数)
        3. 对每个 episode，计算与 query 的相似度：
           - 找出 query 和 episode 共有的词
           - 相似度 = sum(TF(word, episode) * IDF(word))
        4. 返回相似度最高的 max_results 个 (episode, score) 对

        提示：
        - 使用 Counter 统计词频
        - 忽略长度 <= 2 的词（停用词过滤）
        - 将词转为小写进行匹配
        """
        if not episodes:
            return []

        # TODO: 实现 TF-IDF 检索逻辑
        pass

    def _tokenize(self, text: str) -> list[str]:
        """简单分词：小写 + 按非字母数字字符分割。"""
        import re
        words = re.split(r'[^a-zA-Z0-9\u4e00-\u9fff]+', text.lower())
        return [w for w in words if len(w) > 2]

class MemoryInjector:
    """
    将检索到的记忆注入智能体提示词。
    控制注入量，避免上下文溢出。
    """

    MAX_INJECTION_CHARS = 2000  # 最多注入 2000 字符

    def build_memory_context(
        self,
        query: str,
        store: KnowledgeStore,
        searcher: TFIDFSearcher,
        max_results: int = 5,
    ) -> str:
        """
        TODO Part C: 构建记忆注入文本

        要求：
        1. 用 searcher 检索与 query 相关的 episodes
        2. 按类型分组（先展示 GOTCHA，再展示 PATTERN，再展示其他）
        3. 格式化为可读的提示词片段
        4. 如果总字符数超过 MAX_INJECTION_CHARS，截断到限制内
        5. 返回格式化后的字符串

        格式参考：
        ## Relevant Historical Context

        ### Known Gotchas
        - [gotcha_001] Always use utc=True with JWT library

        ### Patterns
        - [pattern_001] Rate limiting via Redis, 100 req/min
        """
        all_episodes = store._episodes
        if not all_episodes:
            return ""

        # TODO: 实现记忆注入逻辑
        pass


def run_demo():
    """演示完整的记忆积累和检索流程。"""
    import tempfile

    with tempfile.TemporaryDirectory() as tmp:
        store = KnowledgeStore(Path(tmp))
        searcher = TFIDFSearcher()
        injector = MemoryInjector()

        # 模拟任务 001：认证系统开发的发现
        store.add_episode(
            EpisodeType.GOTCHA,
            "JWT library requires utc=True parameter, otherwise timezone bugs occur in production",
            tags=["jwt", "authentication", "timezone"],
        )
        store.add_episode(
            EpisodeType.PATTERN,
            "Rate limiting is implemented via Redis with 100 requests per minute limit",
            tags=["redis", "rate-limit", "api"],
        )
        store.add_episode(
            EpisodeType.CODEBASE_DISCOVERY,
            "Database connections must use charset=utf8mb4 to support emoji characters",
            tags=["database", "mysql", "encoding"],
        )
        store.add_episode(
            EpisodeType.TASK_OUTCOME,
            "Authentication system implemented using JWT + Redis session store, took 45 minutes",
            metadata={"task_id": "001-auth", "success": True},
        )

        # 模拟任务 002：搜索与 API 开发相关的历史记忆
        query = "building REST API with authentication and database"
        context = injector.build_memory_context(query, store, searcher)

        print("=== 注入的历史记忆 ===")
        print(context)
        print(f"\n字符数: {len(context)}")

        # 验证
        assert "JWT" in context or "jwt" in context.lower(), "JWT gotcha 应该被检索到"
        assert len(context) <= MemoryInjector.MAX_INJECTION_CHARS, "不应超过字符限制"
        print("\n✓ 所有验证通过")


if __name__ == "__main__":
    run_demo()