状态与记忆

演示 Agent 中常见的记忆与状态管理策略： - BufferMemory：保留最近 N 轮对话（滑动窗口）。 - SummaryMemory：将历史压缩为摘要 + 最近若干轮。 - StateStore：键值状态存储，支持 TTL 过期。 全部基于内存实现，便于离线测试。 对应教程：第5卷《Agent 开发》第4章 状态与记忆。

文件路径: examples/volume-05-agent\state_memory.py

---

源代码

"""状态与记忆示例。

演示 Agent 中常见的记忆与状态管理策略：
    - ``BufferMemory``：保留最近 N 轮对话（滑动窗口）。
    - ``SummaryMemory``：将历史压缩为摘要 + 最近若干轮。
    - ``StateStore``：键值状态存储，支持 TTL 过期。

全部基于内存实现，便于离线测试。

对应教程：第5卷《Agent 开发》第4章 状态与记忆。
"""

from __future__ import annotations

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


@dataclass
class BufferMemory:
    """滑动窗口式对话记忆。

    Attributes:
        window_size: 保留的最近消息条数。
    """

    window_size: int = 10
    _messages: list[dict[str, str]] = field(default_factory=list)

    def add(self, role: str, content: str) -> None:
        """添加一条消息。

        Args:
            role: 角色。
            content: 内容。
        """
        self._messages.append({"role": role, "content": content})
        if len(self._messages) > self.window_size:
            # 保留最后 window_size 条
            self._messages = self._messages[-self.window_size :]

    def messages(self) -> list[dict[str, str]]:
        """返回当前记忆中的消息列表。"""
        return list(self._messages)

    def clear(self) -> None:
        """清空记忆。"""
        self._messages.clear()

    def token_estimate(self, chars_per_token: int = 4) -> int:
        """估算当前记忆占用 token 数。

        Args:
            chars_per_token: 每 token 字符数。

        Returns:
            估算 token 数。
        """
        total_chars = sum(len(m["content"]) for m in self._messages)
        return total_chars // chars_per_token


@dataclass
class SummaryMemory:
    """摘要式记忆：保留一个滚动摘要 + 最近若干条原始消息。

    Attributes:
        recent_size: 保留的最近原始消息条数。
    """

    recent_size: int = 4
    _summary: str = ""
    _recent: list[dict[str, str]] = field(default_factory=list)
    _summarizer_calls: int = 0

    def add(self, role: str, content: str, summarizer=None) -> None:
        """添加消息，并在最近消息超量时触发摘要压缩。

        Args:
            role: 角色。
            content: 内容。
            summarizer: 可选的摘要函数 ``f(old_summary, dropped_messages) -> new_summary``。
                为空时使用简单的拼接。
        """
        self._recent.append({"role": role, "content": content})
        if len(self._recent) > self.recent_size * 2:
            # 把较早的一半折叠进摘要
            to_drop = self._recent[: self.recent_size]
            self._recent = self._recent[self.recent_size :]
            old_summary = self._summary
            if summarizer is not None:
                self._summary = summarizer(old_summary, to_drop)
                self._summarizer_calls += 1
            else:
                dropped_text = " ".join(m["content"] for m in to_drop)
                self._summary = (old_summary + " " + dropped_text).strip()

    @property
    def summary(self) -> str:
        """当前摘要。"""
        return self._summary

    @property
    def summarizer_calls(self) -> int:
        """摘要函数被调用的次数（便于测试）。"""
        return self._summarizer_calls

    def messages(self) -> list[dict[str, str]]:
        """返回对外可见的消息列表（摘要 + 最近消息）。"""
        result: list[dict[str, str]] = []
        if self._summary:
            result.append({"role": "system", "content": f"对话摘要：{self._summary}"})
        result.extend(self._recent)
        return result


@dataclass
class StateStore:
    """带 TTL 的键值状态存储。

    Attributes:
        default_ttl: 默认存活秒数；``None`` 表示永不过期。
    """

    default_ttl: float | None = None
    _data: dict[str, tuple[Any, float | None]] = field(default_factory=dict)

    def set(self, key: str, value, ttl: float | None = "default") -> None:  # type: ignore[assignment]
        """设置键值。

        Args:
            key: 键。
            value: 值。
            ttl: 存活秒数；``"default"`` 使用 ``default_ttl``，``None`` 表示永不过期。
        """
        if ttl == "default":
            ttl = self.default_ttl
        expire_at = (time.monotonic() + ttl) if ttl is not None else None
        self._data[key] = (value, expire_at)

    def get(self, key: str, default=None):
        """获取键值，过期则删除并返回 default。"""
        if key not in self._data:
            return default
        value, expire_at = self._data[key]
        if expire_at is not None and time.monotonic() >= expire_at:
            del self._data[key]
            return default
        return value

    def delete(self, key: str) -> bool:
        """删除键。

        Returns:
            是否删除成功（键是否存在）。
        """
        return self._data.pop(key, None) is not None

    def cleanup_expired(self) -> int:
        """清理所有过期键。

        Returns:
            清理的键数量。
        """
        now = time.monotonic()
        expired = [k for k, (_, t) in self._data.items() if t is not None and now >= t]
        for k in expired:
            del self._data[k]
        return len(expired)


def main() -> None:
    """演示三种记忆/状态组件。"""
    print("== BufferMemory ==")
    buf = BufferMemory(window_size=3)
    for i in range(5):
        buf.add("user", f"msg{i}")
    print(buf.messages())

    print("\n== SummaryMemory ==")
    sm = SummaryMemory(recent_size=2)
    for i in range(6):
        sm.add("user", f"片段{i}")
    print("摘要:", sm.summary)
    print("消息:", sm.messages())

    print("\n== StateStore ==")
    store = StateStore(default_ttl=100)
    store.set("k1", "v1")
    store.set("k2", "v2", ttl=0.0)  # 立即过期
    print("k1:", store.get("k1"))
    print("k2:", store.get("k2", "已过期"))


if __name__ == "__main__":
    main()