《AI Agent 开发入门》8.1 对话缓冲与摘要记忆

# minagent/memory.py
from __future__ import annotations
from typing import Protocol

from .types import Message


class Memory(Protocol):
    def add(self, message: Message) -> None: ...
    def get_context(self) -> list[Message]: ...

# minagent/memory.py（续）
from .budget import approx_tokens


class BufferMemory:
    def __init__(self, max_messages: int | None = None,
                 max_tokens: int | None = None) -> None:
        self.max_messages = max_messages
        self.max_tokens = max_tokens
        self._items: list[Message] = []

    def add(self, message: Message) -> None:
        self._items.append(message)
        self._trim()

    def _trim(self) -> None:
        if self.max_messages is not None:
            self._items = self._items[-self.max_messages:]
        if self.max_tokens is not None:
            while self._items and self._count() > self.max_tokens:
                self._items.pop(0)

    def _count(self) -> int:
        return sum(approx_tokens(m.content) for m in self._items)

    def get_context(self) -> list[Message]:
        return list(self._items)

# minagent/memory.py（续）
from .types import Model


class SummaryMemory:
    def __init__(self, model: Model, max_messages: int = 4,
                 summarize_every: int = 3) -> None:
        self.model = model
        self.max_messages = max_messages
        self.summarize_every = summarize_every
        self.summary = ""
        self._items: list[Message] = []

    def add(self, message: Message) -> None:
        self._items.append(message)
        if len(self._items) > self.max_messages + self.summarize_every:
            self._compact()

    def _compact(self) -> None:
        old = self._items[:-self.max_messages]
        self._items = self._items[-self.max_messages:]
        rendered = "\n".join(f"{m.role}: {m.content}" for m in old)
        prompt = f"把下面这段对话压缩成一句话摘要：\n{rendered}"
        resp = self.model.complete([Message(role="user", content=prompt)], [])
        merged = resp.message.content
        self.summary = f"{self.summary} {merged}".strip() if self.summary else merged

    def get_context(self) -> list[Message]:
        ctx: list[Message] = []
        if self.summary:
            ctx.append(Message(role="system", content=f"对话摘要：{self.summary}"))
        ctx.extend(self._items)
        return ctx

# demo_memory.py
convo = [
    Message(role="user", content="我在做一个 Python 的 Agent 项目。"),
    Message(role="assistant", content="好的，你希望它解决什么问题？"),
    Message(role="user", content="主要用来做本地知识库问答。"),
    Message(role="assistant", content="了解，需要接入向量检索吗？"),
    Message(role="user", content="先不用，我想先把记忆模块跑通。"),
    Message(role="assistant", content="明白，先做短期记忆。"),
    Message(role="user", content="对，而且我偏好简洁的代码示例。"),
    Message(role="assistant", content="收到，后面都用最小可运行示例。"),
]

buf = BufferMemory(max_messages=6)
for m in convo:
    buf.add(m)
print("BufferMemory(6) 注入条数:", len(buf.get_context()))

tokbuf = BufferMemory(max_tokens=40)
for m in convo:
    tokbuf.add(m)
print("BufferMemory(max_tokens=40) 注入条数:", len(tokbuf.get_context()))

fake = FakeModel([
    "用户在做 Python 本地知识库问答 Agent，先跑通记忆模块，偏好简洁示例。",
])
summ = SummaryMemory(fake, max_messages=4, summarize_every=3)
for m in convo:
    summ.add(m)
ctx = summ.get_context()
print("SummaryMemory 注入条数:", len(ctx))
print("摘要:", ctx[0].content)
print("保留的最近消息:", [m.content for m in ctx[1:]])

BufferMemory(6) 注入条数: 6
BufferMemory(max_tokens=40) 注入条数: 3
SummaryMemory 注入条数: 5
摘要: 对话摘要：用户在做 Python 本地知识库问答 Agent，先跑通记忆模块，偏好简洁示例。
保留的最近消息: ['先不用，我想先把记忆模块跑通。', '明白，先做短期记忆。', '对，而且我偏好简洁的代码示例。', '收到，后面都用最小可运行示例。']