Agent Memory：四种记忆

"""agent_memory_travel.py — 四种记忆的旅游助手"""
from __future__ import annotations

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


# ── 极简运行时 ──────────────────────────────────────────

@dataclass(frozen=True)
class Message:
    role: str
    content: str

class MockLLM:
    def __init__(self, responses: list[str]):
        self._responses = list(responses)
        self._idx = 0
    def complete(self, messages: list[Message]) -> str:
        if self._idx >= len(self._responses):
            raise RuntimeError("MockLLM 回复用完了")
        out = self._responses[self._idx]
        self._idx += 1
        return out


# ── 记忆存储 ────────────────────────────────────────────

class InMemoryKV:
    """Episodic memory 的 KV 存储。"""
    def __init__(self):
        self._data: dict[str, Any] = {}
    def get(self, key: str, default=None) -> Any:
        return self._data.get(key, default)
    def set(self, key: str, value: Any) -> None:
        self._data[key] = value


@dataclass(frozen=True)
class Document:
    doc_id: str
    text: str

class SemanticStore:
    """Semantic memory：简单的词频检索。"""
    def __init__(self, docs: list[Document]):
        self._docs = docs
    def search(self, query: str, k: int = 3) -> list[Document]:
        terms = self._tokenize(query)
        if not terms:
            return []
        scored = []
        for doc in self._docs:
            doc_tokens = self._tokenize(doc.text)
            score = sum(1 for t in terms for d in doc_tokens if t == d)
            if score > 0:
                scored.append((doc, score))
        scored.sort(key=lambda x: x[1], reverse=True)
        return [doc for doc, _ in scored[:k]]
    @staticmethod
    def _tokenize(text: str) -> list[str]:
        tokens, buf = [], ""
        for ch in text.lower():
            if ch.isalnum() or '\u4e00' <= ch <= '\u9fff':
                buf += ch
            else:
                if buf:
                    tokens.append(buf)
                    buf = ""
        if buf:
            tokens.append(buf)
        return tokens


# ── Procedural Memory：工具和规则 ───────────────────────

TOOLS = {
    "get_weather": "查询城市天气。参数：city, date。",
    "search_places": "搜索景点。参数：query, k。",
}

RULES = [
    "订票前必须确认用户预算。",
    "轻松步行行程不超过 3 个景点。",
    "下雨时优先推荐室内景点。",
]


# ── 带四种记忆的 Agent ──────────────────────────────────

def travel_agent(
    model: MockLLM,
    user_query: str,
    *,
    episodic: InMemoryKV,
    semantic: SemanticStore,
) -> str:
    """
    一个带四种记忆的旅游 Agent：
    - Working: messages 列表
    - Episodic: 教训和用户偏好（KV）
    - Semantic: 文档检索（向量库）
    - Procedural: 工具定义和规则（代码）
    """

    # ── 1. Procedural Memory → system prompt ──
    tools_text = "\n".join(f"- {name}: {desc}" for name, desc in TOOLS.items())
    rules_text = "\n".join(f"- {r}" for r in RULES)
    system_prompt = (
        "你是杭州旅游助手。\n\n"
        f"可用工具：\n{tools_text}\n\n"
        f"必须遵守的规则：\n{rules_text}"
    )

    # ── 2. Episodic Memory → 读教训和偏好 ──
    lessons = episodic.get("lessons", [])
    preferences = episodic.get("user_preferences", {})

    episodic_context = ""
    if lessons:
        episodic_context += "历史教训：\n" + "\n".join(f"- {l}" for l in lessons) + "\n\n"
    if preferences:
        episodic_context += "用户偏好：\n" + "\n".join(
            f"- {k}: {v}" for k, v in preferences.items()
        ) + "\n\n"

    print(f"[Episodic] 教训 {len(lessons)} 条，偏好 {len(preferences)} 项")

    # ── 3. Semantic Memory → 检索相关文档 ──
    docs = semantic.search(user_query, k=3)
    semantic_context = ""
    if docs:
        semantic_context = "相关资料：\n" + "\n".join(
            f"- [{d.doc_id}] {d.text[:80]}" for d in docs
        ) + "\n\n"
    print(f"[Semantic] 检索到 {len(docs)} 条文档")

    # ── 4. Working Memory → 构建 messages ──
    messages: list[Message] = [
        Message(role="system", content=system_prompt),
        Message(
            role="user",
            content=(
                f"{episodic_context}"
                f"{semantic_context}"
                f"用户问题：{user_query}"
            ),
        ),
    ]
    print(f"[Working] messages 共 {len(messages)} 条")

    # ── 生成回答 ──
    answer = model.complete(messages)

    # ── 5. 写回 Episodic Memory ──
    # 保存这次交互的摘要
    history = episodic.get("interaction_history", [])
    history.append({"query": user_query, "answer_snippet": answer[:60]})
    episodic.set("interaction_history", history)
    print(f"[Episodic] 交互历史更新，共 {len(history)} 条")

    return answer


# ── 运行 ────────────────────────────────────────────────

def main() -> None:
    # 初始化 Episodic Memory（模拟之前积累的教训和偏好）
    episodic = InMemoryKV()
    episodic.set("lessons", [
        "轻松步行要求时，景点不超过 3 个。",
        "推荐室内景点时要注明开放时间。",
    ])
    episodic.set("user_preferences", {
        "步行偏好": "轻松，不赶路",
        "饮食": "喜欢茶和素食",
    })

    # 初始化 Semantic Memory
    semantic = SemanticStore([
        Document("tea_museum", "中国茶叶博物馆，9:00-17:00，免费，室内，适合雨天。"),
        Document("tea_village", "龙井村采茶体验，户外，步行友好。"),
        Document("veggie", "杭州素食推荐：灵隐寺素面、龙井虾仁（可做素版）。"),
        Document("west_lake", "西湖全天开放，免费，环湖步行约 3 小时。"),
    ])

    model = MockLLM([
        (
            "根据您的偏好（轻松步行、喜欢茶和素食），推荐：\n"
            "1. 上午：龙井村采茶体验（步行友好，约 2 小时）[tea_village]\n"
            "2. 中午：灵隐寺素面 [veggie]\n"
            "3. 下午：中国茶叶博物馆（室内，9:00-17:00，免费）[tea_museum]\n"
            "全程 3 个景点，不赶路。"
        ),
    ])

    answer = travel_agent(
        model,
        "杭州一日游，喜欢茶文化",
        episodic=episodic,
        semantic=semantic,
    )
    print(f"\n回答：\n{answer}")

    # 验证记忆状态
    print(f"\n记忆状态：")
    print(f"  教训：{episodic.get('lessons')}")
    print(f"  偏好：{episodic.get('user_preferences')}")
    print(f"  历史：{episodic.get('interaction_history')}")


if __name__ == "__main__":
    main()

[Episodic] 教训 2 条，偏好 2 项
[Semantic] 检索到 3 条文档
[Working] messages 共 2 条

回答：
根据您的偏好（轻松步行、喜欢茶和素食），推荐：
1. 上午：龙井村采茶体验（步行友好，约 2 小时）[tea_village]
2. 中午：灵隐寺素面 [veggie]
3. 下午：中国茶叶博物馆（室内，9:00-17:00，免费）[tea_museum]
全程 3 个景点，不赶路。

记忆状态：
  教训：['轻松步行要求时，景点不超过 3 个。', '推荐室内景点时要注明开放时间。']
  偏好：{'步行偏好': '轻松，不赶路', '饮食': '喜欢茶和素食'}
  历史：[{'query': '杭州一日游，喜欢茶文化', 'answer_snippet': '根据您的偏好（轻松步行、喜欢茶和素食），推荐：\n1. 上午：龙井村采茶体验（步行友好，约 2 小时）[tea_'}]