构建自己的 AI Agent · 第三篇：对话管理 — 上下文窗口的数学与工程

import tiktoken

class TokenCounter:
    def __init__(self, model: str = "gpt-4o"):
        try:
            self.encoder = tiktoken.encoding_for_model(model)
        except KeyError:
            self.encoder = tiktoken.get_encoding("cl100k_base")

    def count(self, text: str) -> int:
        return len(self.encoder.encode(text))

    def count_messages(self, messages: list[dict]) -> int:
        """精确计算 messages 数组的总 token 数"""
        total = 0
        for msg in messages:
            total += 4  # 每条消息的固定开销（role + 分隔符）
            for key, value in msg.items():
                if value is None:
                    continue
                if isinstance(value, str):
                    total += self.count(value)
                elif isinstance(value, list):
                    # tool_calls 数组
                    total += self.count(json.dumps(value))
        total += 2  # assistant 回复的 priming
        return total

class SlidingWindow:
    def __init__(self, system_prompt: str, max_messages: int = 20):
        self.system_prompt = system_prompt
        self.max_messages = max_messages

    def trim(self, messages: list[dict]) -> list[dict]:
        """保留 system prompt + 最近 max_messages 条消息"""
        system_msgs = [m for m in messages if m["role"] == "system"]
        non_system = [m for m in messages if m["role"] != "system"]

        if len(non_system) <= self.max_messages:
            return messages

        return system_msgs + non_system[-self.max_messages:]

class WeightedSlidingWindow:
    WEIGHTS = {
        "system": float("inf"),
        "tool": 0.8,
        "user": 0.5,
        "assistant": 0.3
    }

    def trim(self, messages: list[dict], max_total_weight: float = 10.0) -> list[dict]:
        """按重要性权重保留消息，直到总权重达到上限"""
        kept = []
        total_weight = 0.0

        # 反向遍历（优先保留最近的消息）
        for msg in reversed(messages):
            w = self.WEIGHTS.get(msg["role"], 0.3)
            if total_weight + w > max_total_weight:
                break
            kept.append(msg)
            total_weight += w

        kept.reverse()
        return kept

class ConversationManager:
    def __init__(self, client, system_prompt: str,
                 max_tokens: int = 100000, summary_threshold: int = 80000):
        self.client = client
        self.system_prompt = system_prompt
        self.max_tokens = max_tokens        # 上下文窗口上限
        self.summary_threshold = summary_threshold  # 触发摘要的阈值
        self.counter = TokenCounter()
        self.summaries: list[str] = []       # 累积摘要列表
        self.recent_messages: list[dict] = [] # 近期消息（不被压缩的）

    def add_message(self, role: str, content: str) -> None:
        self.recent_messages.append({"role": role, "content": content})
        self._maybe_summarize()

    def _maybe_summarize(self) -> None:
        """如果总 token 数超过阈值，触发摘要"""
        total = (
            self.counter.count(self.system_prompt) +
            self.counter.count_messages(self._build_context())
        )

        if total > self.summary_threshold:
            # 将前 80% 的近期消息压缩
            split = len(self.recent_messages) * 4 // 5
            to_summarize = self.recent_messages[:split]
            self.recent_messages = self.recent_messages[split:]

            summary = self._summarize(to_summarize)
            self.summaries.append(summary)

    def _summarize(self, messages: list[dict]) -> str:
        """调用 LLM 将消息压缩为摘要"""
        text = "\n".join(
            f"[{m['role']}]: {m['content']}" for m in messages
        )
        response = self.client.chat.completions.create(
            model="gpt-4o-mini",  # 用便宜模型做摘要
            messages=[{
                "role": "system",
                "content": (
                    "将以下对话历史压缩为一段简洁的摘要（不超过 500 字）。"
                    "保留：关键决策、重要数据、用户偏好、未完成的任务。"
                    "忽略：冗余表述、已完成的中间步骤、礼貌用语。"
                )
            }, {
                "role": "user",
                "content": text
            }],
            temperature=0.1
        )
        return response.choices[0].message.content

    def _build_context(self) -> list[dict]:
        """构建发送给 LLM 的完整上下文"""
        messages = [{"role": "system", "content": self.system_prompt}]

        # 注入历史摘要
        if self.summaries:
            summary_text = "## 历史对话摘要\n\n" + "\n---\n".join(self.summaries)
            messages.append({
                "role": "system",
                "content": summary_text
            })

        # 追加近期消息
        messages.extend(self.recent_messages)
        return messages

    def get_context(self) -> list[dict]:
        return self._build_context()

class HierarchicalMemory:
    def __init__(self):
        self.ram: list[dict] = []       # 主上下文（有限容量）
        self.disk: list[str] = []       # 外部存储（事件摘要）

    def add_to_ram(self, msg: dict) -> None:
        self.ram.append(msg)
        if len(self.ram) > 30:  # RAM 满了
            self._swap_to_disk()

    def _swap_to_disk(self) -> None:
        """将 RAM 中最旧的 10 条消息移到 disk（压缩为摘要）"""
        old = self.ram[:10]
        self.ram = self.ram[10:]
        summary = self._summarize_events(old)
        self.disk.append(summary)

    def recall(self, query: str) -> list[str]:
        """从 disk 中召回与 query 相关的记忆"""
        # 简化版：返回最近的摘要（第 4 篇将实现语义检索）
        return self.disk[-3:]

class ConversationalAgent:
    def __init__(self, system_prompt: str, registry: ToolRegistry,
                 max_iter: int = 10, max_tokens: int = 100000):
        self.client = OpenAI()
        self.registry = registry
        self.max_iter = max_iter
        self.conv = ConversationManager(
            client=self.client,
            system_prompt=system_prompt,
            max_tokens=max_tokens
        )

    def run(self, user_query: str) -> str:
        self.conv.add_message("user", user_query)
        schemas = self.registry.get_schemas()

        for _ in range(self.max_iter):
            messages = self.conv.get_context()  # ← 自动管理上下文
            response = self.client.chat.completions.create(
                model="gpt-4o",
                messages=messages,
                tools=schemas if schemas else None
            )
            msg = response.choices[0].message

            if msg.tool_calls:
                for tc in msg.tool_calls:
                    fn_args = json.loads(tc.function.arguments)
                    result = self.registry.execute(tc.function.name, fn_args)

                    self.conv.add_message("assistant",
                        json.dumps({"tool_call": tc.function.name,
                                    "args": fn_args}))
                    self.conv.add_message("tool", result)
            else:
                self.conv.add_message("assistant", msg.content or "")
                return msg.content or ""

        return "达到最大迭代次数。"