第三章：protocol.py——消息、文本与 token 边界#

源码位置：https://github.com/KMnO4-zx/agentic-rl-lab/blob/main/05-retool/protocol.py

3.1 本章第一次进入 token 世界#

前两章的数据对象是：

JSONL → dict → MathExample

protocol.py 接过 MathExample.question，第一次把人类可读的对话转换成模型实际接收的 token ID：

question: str
    ↓ initial_messages()
messages: list[dict]
    ↓ tokenizer.apply_chat_template(..., tools=[CODE_TOOL])
prompt_tokens: list[int]
    ↓ sampler
completion_tokens: list[int]
    ↓ decode / parse_assistant()
answer 或 tool call
    ↓ 若调用工具
tool observation
    ↓ build_next_prompt()
下一轮连续 prompt_tokens

这一章的四个核心边界是：

1. 消息边界：system / user / assistant / tool；
2. 生成边界：已有 prompt 与新生成 completion；
3. 动作边界：模型生成的 assistant token；
4. 环境边界：系统追加的 tool observation token。

后续 PPO 是否能正确训练，取决于这些边界能否保持准确。

3.2 先认识协议的两个层次#

这里的“协议”同时包含语义规则和序列化规则。

第一层：语义协议#

它规定模型应该如何行动：

需要计算
    → 生成 code_interpreter 工具调用
    → 等待 tool 返回结果
    → 根据结果继续推理

得到最终答案
    → 最后一行输出 \boxed{...}

这部分主要由 SYSTEM_PROMPT 和 parse_assistant() 表达。

第二层：token 序列化协议#

它规定一组结构化消息如何变成 Qwen 真正接收的 token 序列：

messages + tools
    → Qwen chat template
    → role 标记、工具定义、内容、结束标记
    → token IDs

这部分主要由 _render_chat()、build_prompt() 和 build_next_prompt() 表达。

前者回答“模型的输出是什么意思”，后者回答“这些内容在 token 流里的准确位置是什么”。

3.3 CODE_TOOL：把 Python 沙箱声明成原生工具#

CODE_TOOL 是一个遵循 function tool schema 的普通 Python 字典：

CODE_TOOL = {
    "type": "function",
    "function": {
        "name": "code_interpreter",
        "description": "...",
        "parameters": {
            "type": "object",
            "properties": {
                "code": {
                    "type": "string",
                    "description": "The Python code to be executed.",
                }
            },
            "required": ["code"],
        },
    },
}

这段声明告诉模型：

• 工具名是 code_interpreter；
• 工具只有一个必需参数 code；
• code 必须是字符串；
• 结果通过 print() 输出；
• 每次执行相互独立，不保留变量或文件状态。

它只是一份工具说明书，本身不会执行 Python。真正执行代码的是下一章的 LocalPythonSandbox。

工具定义怎样进入模型输入#

CODE_TOOL 没有被塞进 messages，而是通过：

tokenizer.apply_chat_template(
    messages,
    tools=[CODE_TOOL],
    ...,
)

交给 Qwen 的 chat template。模板会把工具说明渲染成模型熟悉的原生工具上下文。因此初始 messages 虽然只有 system 和 user 两条，实际 prompt token 中还会包含工具定义及模板控制标记。

3.4 SYSTEM_PROMPT：行动规则而不是执行代码#

system prompt 规定了几组关键行为：

工具使用时机#

计算、符号操作或枚举有帮助时，可以调用 code_interpreter。

每轮最多调用一次#

模型一次 assistant turn 只能发一个工具调用，然后必须等待执行结果。这正好对应 rollout 状态机：

assistant turn
    → 至多一个 tool call
    → tool observation
    → 新的 assistant turn

工具无状态#

每次执行都是新进程，所以下一段代码不能依赖上一段代码里定义的变量。模型必须重新定义所需内容。

最终答案格式#

最终回答最后一行必须是：

\boxed{<your final answer>}

不能同轮既调用工具又交最终答案#

工具调用意味着“我还需要环境反馈”；最终答案意味着“轨迹可以结束”。把两者放在同一轮会造成状态含义冲突。

需要区分：system prompt 是给模型的行为指导，不是百分之百可靠的程序校验。模型仍可能生成错误格式，所以后面还有 parse_assistant() 和 reward.py 两层检查。

3.5 initial_messages()：初始对话不含参考答案#

函数：

def initial_messages(question: str) -> list[dict[str, Any]]:
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": question},
    ]

假设：

example.question = "If 2x + 1 = 7, what is x?"
example.answer = "3"

得到的初始消息只有：

[
    {"role": "system", "content": SYSTEM_PROMPT},
    {"role": "user", "content": "If 2x + 1 = 7, what is x?"},
]

example.answer 不会传入这个函数。因此信息边界是：

question → 模型可见
answer   → 模型不可见，只供 reward 使用

3.6 _render_chat()：从 messages 到一维 token IDs#

核心调用：

rendered = tokenizer.apply_chat_template(
    messages,
    tools=[CODE_TOOL],
    tokenize=True,
    add_generation_prompt=add_generation_prompt,
    enable_thinking=False,
)

tokenize=True#

要求模板直接返回 token ID，而不是只返回渲染后的文本。

tools=[CODE_TOOL]#

把工具定义交给模型原生模板渲染。

enable_thinking=False#

关闭模板提供的显式 thinking 模式。它不等于禁止模型写普通的分步推理文本，只是不启用对应的模板级 thinking 结构。

add_generation_prompt#

这是最关键的边界开关：

• True：在末尾加入“接下来轮到 assistant 生成”的模板前缀；
• False：把现有消息渲染成已经结束的完整历史，不额外开启新一轮生成。

概念上可以理解为：

add_generation_prompt=False
    system + user

add_generation_prompt=True
    system + user + <assistant 开始位置>

实际控制 token 的具体形式由 Qwen tokenizer 的 chat template 决定，不应在业务代码里手写这些特殊标记。

3.7 为什么还要统一 tokenizer 的返回类型#

不同 tokenizer 或调用配置可能返回：

• Python list[int]；
• 二维的 list[list[int]]；
• NumPy/PyTorch 风格、带 .tolist() 的数组或张量；
• 含有 input_ids 的 Mapping/BatchEncoding。

源码依次做标准化：

if isinstance(rendered, Mapping):
    rendered = rendered["input_ids"]
if hasattr(rendered, "tolist"):
    rendered = rendered.tolist()
if rendered and isinstance(rendered[0], list):
    rendered = rendered[0]
return [int(token) for token in rendered]

最终函数建立了稳定的返回契约：

list[int]

这是一条单对话的一维 token 序列，而不是带 batch 维的二维张量。

3.8 build_prompt()：确定模型从哪里开始生成#

函数非常短：

def build_prompt(tokenizer, messages):
    return _render_chat(
        tokenizer,
        messages,
        add_generation_prompt=True,
    )

假设 messages 是 system + user，概念结果可以写成：

P₀ = [工具说明 + system 消息 + user 消息 + assistant 开始标记]

P₀ 是首轮送给 sampler 的 prompt_tokens。采样器只生成它后面的内容：

P₀ | C₁
   ↑
   生成从这里开始

其中：

• P₀：环境提供的 prompt；
• C₁：模型真实采样出的第一轮 completion token。

这条竖线就是最初的 prompt/action 边界。

3.9 ParsedAssistant：把一轮输出归为三类#

定义：

@dataclass(frozen=True)
class ParsedAssistant:
    kind: str
    content: str
    code: str | None = None

kind 只有三种约定值：

字段含义：

• content：工具调用前的普通推理，或完整答案/非法输出；
• code：仅 kind="tool" 时保存提取出的 Python 代码。

3.10 正则表达式如何识别工具调用#

正则的目标格式是：

<tool_call>
<function=code_interpreter>
<parameter=code>
Python 代码
</parameter>
</function>
</tool_call>

关键部分：

(.*?)

它捕获 code 标签内的内容；re.DOTALL 让 . 也能匹配换行，因此 Python 代码可以是多行。

标签间使用 \s*，允许存在空格和换行；但工具名、参数名和标签拼写必须符合预期。

.*? 是非贪婪匹配，会尽量在第一个合法闭合标签处停止。解析器只负责提取字符串，不检查 Python 语法；语法是否正确由沙箱执行结果决定。

3.11 parse_assistant() 的判定顺序#

情况一：完全没有匹配到合法工具调用#

if not matches:
    kind = "invalid" if "<tool_call>" in text else "answer"

普通答案：

Solving gives x=3.
\boxed{3}

返回：

ParsedAssistant(kind="answer", content="...", code=None)

注意：parse_assistant() 不检查是否真的存在 \boxed{}。只要没有工具调用意图，它就先归类为 answer；最终格式和正确性由 reward.py 检查。

若文本含字面量 <tool_call>，却没有完整匹配协议，例如缺少闭合标签，就返回 invalid。

情况二：匹配到不止一个工具调用#

if len(matches) != 1:
    return ParsedAssistant(kind="invalid", ...)

这对应 system prompt 的“一次 assistant turn 最多调用一次”。

情况三：工具调用后还有非空文本#

text[matches[0].end():].strip()

如果闭合 </tool_call> 后仍有内容，就判为 invalid。例如：

<tool_call>...</tool_call>
\boxed{3}

不能在同一轮先调用工具又给最终答案，因为系统尚未返回 observation。

情况四：代码为空#

code = matches[0].group(1).strip()
if not code:
    return ParsedAssistant(kind="invalid", ...)

空工具调用没有执行意义，因此非法。

情况五：恰好一个合法工具调用#

模型可以在工具标签前写普通推理：

I will verify the arithmetic with Python.
<tool_call>
<function=code_interpreter>
<parameter=code>
print((7 - 1) / 2)
</parameter>
</function>
</tool_call>

解析结果为：

ParsedAssistant(
    kind="tool",
    content="I will verify the arithmetic with Python.",
    code="print((7 - 1) / 2)",
)

整个 assistant 原文仍保存在轨迹的 AssistantTurn.text 中；ParsedAssistant 是为了决定下一次状态转移而产生的解释结果。

3.12 解析规则速查表#

这里的“answer”表示“状态机把它当作最终回答”，不等于“答案一定正确”。

3.13 tool_message()：把执行结果包装成环境消息#

沙箱执行完成后，stdout、stderr 或超时信息会被包装为：

{
    "role": "tool",
    "tool_call_id": call_id,
    "name": "code_interpreter",
    "content": content,
}

例如：

tool_message("code-0-3-1", "3.0\n")

得到：

{
    "role": "tool",
    "tool_call_id": "code-0-3-1",
    "name": "code_interpreter",
    "content": "3.0\n",
}

字段职责：

• role="tool"：这是环境反馈，不是模型输出；
• tool_call_id：关联哪一次调用；
• name：指出工具名称；
• content：实际 observation。

call_id 由 rollout 根据问题、组内轨迹和调用次数生成。protocol.py 只负责包装，不负责创建唯一编号。

3.14 为什么不能把 assistant 文本重新 token 化#

这是整章最关键的理论点。

采样器返回：

completion_tokens = C
old_logprobs      = L

其中 L[i] 是模型采样 C[i] 时旧策略给出的 logprob。PPO 训练必须保证：

target token C[i] ↔ old logprob L[i]

一种看似自然但危险的做法是：

真实 completion tokens
    ↓ decode
assistant 文本
    ↓ 加进 messages，再渲染并 encode 整段历史
新 token 序列

问题是一般不能保证：

encode(decode(C)) == C

可能改变 token 的因素包括：

• tokenizer 对空格和换行的归一化；
• 特殊 token 在 decode/encode 时被移除或重新解释；
• chat template 对 assistant 内容执行 strip；
• </think> 等内容触发模板把文本重构为 reasoning/content 两段；
• 同一可见字符串可能存在不止一种 token 切分方式。

一旦重编码后的 token 变成 C'，训练数据却仍携带为 C 采样得到的旧 logprob L，位置就错位了：

错误情况：target C'[i] ↔ old logprob of C[i]

PPO 的概率比率和 loss 随之失去含义。因此必须保留 sampler 返回的真实 token，直接在后面增量拼接。

这就是注释中的：

token-in token-out

3.15 build_next_prompt() 的五个输入#

函数签名：

def build_next_prompt(
    tokenizer,
    messages_before_assistant,
    previous_prompt_tokens,
    completion_tokens,
    next_tool_message,
) -> list[int]:

五个输入分别是：

函数最终返回下一轮可以直接送给 sampler 的完整 prompt token。

3.16 目标公式：只在真实 token 后面追加增量#

把几个 token 片段记为：

P = previous_prompt_tokens
C = completion_tokens
E = assistant_closing_tokens
O = observation_tokens

目标是构造：

P_next = P + C + E_missing + O

其中：

• P 原样保留；
• C 原样保留；
• E_missing 只补 sampler 没有返回的 assistant 结束 token；
• O 是 tool observation 及下一轮 assistant 生成前缀的增量 token。

函数不会用重新渲染后的历史替换 P + C。

3.17 为什么使用占位 assistant 内容 "x"#

为了得到 E 和 O，代码仍然需要向 chat template 询问：“一条 assistant 消息结束后，以及加入 tool 消息后，模板会增加哪些 token？”

但它不能把真实 assistant 文本放回模板，因为模板可能 strip 或重构真实文本。于是使用稳定占位内容：

placeholder_message = {"role": "assistant", "content": "x"}

关键假设是：

assistant 后面的结构性结束片段和 tool 消息增量，由消息角色和 tool 内容决定，不需要知道真实 assistant 内容。

占位符的作用只是帮助定位边界，最终输出中不会出现 x。

3.18 第一步：定位 assistant 结束 token#

先得到当前历史的标准生成 prompt：

canonical_prompt = build_prompt(tokenizer, messages_before_assistant)

记作：

P_can

然后渲染“历史 + assistant x”，但不开启下一轮生成：

canonical_assistant_end = _render_chat(
    tokenizer,
    [*messages_before_assistant, placeholder_message],
    add_generation_prompt=False,
)

占位文本本身单独编码为：

placeholder_tokens = _encoded_text_tokens(tokenizer, "x")

概念上模板结果应该满足：

canonical_assistant_end
= canonical_prompt + tokens("x") + assistant_closing_tokens

源码先验证这个前缀关系：

canonical_action = [*canonical_prompt, *placeholder_tokens]
if canonical_assistant_end[:len(canonical_action)] != canonical_action:
    raise ValueError(...)

只有验证成立，才能安全切出：

assistant_closing_tokens = canonical_assistant_end[len(canonical_action):]

如果未来 tokenizer/chat template 改版，不再满足这一结构，代码会明确报错，而不是悄悄生成错位训练数据。

3.19 第二步：定位 tool observation 增量#

接着渲染：

原历史 + assistant "x" + tool message + 新 assistant 生成位置

代码：

canonical_next_prompt = build_prompt(
    tokenizer,
    [*messages_with_assistant, next_tool_message],
)

它应该满足：

canonical_next_prompt
= canonical_assistant_end + observation_tokens

源码同样先验证前缀：

if canonical_next_prompt[:len(canonical_assistant_end)] \
        != canonical_assistant_end:
    raise ValueError(...)

再切出：

observation_tokens = canonical_next_prompt[len(canonical_assistant_end):]

observation_tokens 这个名字是简写#

它不一定只包含 stdout 的正文 token。它是模板在 assistant 结束后新增的整个片段，通常包括：

tool role/结构标记
+ tool content
+ tool 结束标记
+ 下一轮 assistant 生成前缀

这些 token 都不是模型在上一轮主动采样的动作，因此后续训练时属于 observation/context 区域。

3.20 _suffix_prefix_overlap()：避免重复结束标记#

不同 sampler 对停止 token 的返回方式可能不同：

• completion 不包含 assistant 结束 token；
• completion 包含一部分结束 token；
• completion 已包含完整结束 token。

所以不能无条件再追加完整 assistant_closing_tokens。

辅助函数寻找：

completion_tokens 的后缀
与
assistant_closing_tokens 的前缀

之间最长的完全相同部分。

举一个仅用于说明算法的虚构 token 例子：

completion_tokens = [11, 12, 90]
assistant_closing_tokens = [90, 91]

最长重叠是 [90]，长度为 1，因此只补：

assistant_closing_tokens[1:] == [91]

若 completion 已经以 [90, 91] 结尾，重叠长度为 2，不再补任何结束 token；若没有重叠，则补完整结束片段。

函数从最长可能长度向 1 递减，找到第一个匹配就返回；完全没有匹配则返回 0。

3.21 最终拼接：真实历史不被替换#

最后返回：

return [
    *previous_prompt_tokens,
    *completion_tokens,
    *assistant_closing_tokens[overlap:],
    *observation_tokens,
]

注意 canonical_prompt 和占位符只用于推导增量，不会进入最终返回值。

最终序列使用的是：

真实 previous prompt
+ 真实 sampled completion
+ 必要的结构性补全
+ 新 tool observation

这保证下一轮模型看到的历史与上一轮真实采样 token 完全连续。

3.22 用一条两轮轨迹追踪 token 边界#

假设首轮：

P₀ = 初始 system/user/tool-definition prompt tokens
C₁ = 模型生成的推理 + tool call tokens
E₁ = 尚缺少的 assistant closing tokens
O₁ = tool observation + 下一轮 assistant 前缀 tokens

下一轮 prompt 是：

P₁ = P₀ | C₁ | E₁ + O₁

模型再生成最终回答：

P₁ | C₂

展开后整条序列为：

[P₀] [C₁] [E₁ + O₁] [C₂]
 context action observation action

从训练 mask 的角度预览：

P₀          → advantage 0
C₁          → trajectory advantage
E₁ + O₁     → advantage 0
C₂          → trajectory advantage

严格来说，如果 sampler 自己已经返回了某些 assistant closing token，那么它们位于 C₁ 中，属于真实采样动作；只有系统补上的缺失 closing token 才进入后续 prompt 增量并被 mask 为 0。

这正是第七章 build_datum() 能识别 observation 区间的基础：每轮新的 prompt 必须是前面真实 token 轨迹的前缀扩展。

3.23 两套历史：可读 messages 与真实 token stream#

rollout 同时维护两套表示：

结构化消息历史#

messages: list[dict]

用途：

• 记录 system/user/assistant/tool 的语义结构；
• 便于解析、日志和调试；
• 构造 canonical 模板增量。

连续 token 历史#

next_prompt_tokens: list[int]

用途：

• 下一轮直接送入 sampler；
• 保留真实采样 token；
• 保证 completion 与旧 logprob 对齐；
• 后续构造 PPO Datum。

二者描述同一段对话，但不能互相随意替代：

messages       → 语义记录，可能经过 strip/模板规范化
token stream   → 训练事实，必须保持逐 token 原样

首轮可以从 messages 完整渲染；一旦出现真实 assistant completion，后续优先使用增量构造的 token stream。

3.24 _encoded_text_tokens() 为什么只服务于占位符#

这个函数调用：

tokenizer.encode(text, add_special_tokens=False)

并做与 _render_chat() 类似的列表标准化。

它编码的是受代码控制的稳定占位内容 "x"，用于计算模板边界。它没有拿真实 assistant 文本重建历史，因此不违背 token-in/token-out 原则。

add_special_tokens=False 很重要：占位文本前后的 role/结束特殊 token 应由 chat template 提供，不能由普通文本编码再次插入。

3.25 stop_sequences()：一轮 assistant 在哪里停止#

函数：

def stop_sequences(tokenizer):
    eos_token = getattr(tokenizer, "eos_token", None)
    return [eos_token] if eos_token else []

它把 tokenizer 的 EOS 字符串交给采样参数：

SamplingParams(stop=stop_sequences(tokenizer), ...)

作用是让 sampler 在模型产生一轮结束信号时停止继续生成。

如果 tokenizer 没有 eos_token，返回空列表，不设置额外停止字符串。代码使用 getattr(..., None)，避免访问不存在的属性时报错。

这里没有把 </tool_call> 设置成 stop sequence，因此模型理论上可能在工具标签后继续输出文字；这种情况会被 parse_assistant() 判为 invalid。正常行为依靠 system prompt 和模型原生工具格式共同约束。

3.26 从协议结果到 rollout 状态转移#

protocol.py 本身不改变 Trajectory.done，它只提供判断和序列构造工具。真正状态转移发生在 rollout.py：

assistant completion
    ↓ parse_assistant()
    ├─ kind="tool"
    │    → 创建 PendingExecution
    │    → sandbox 执行
    │    → tool_message()
    │    → build_next_prompt()
    │    → 下一轮采样
    │
    ├─ kind="answer"
    │    → final_text = assistant text
    │    → done = True
    │
    └─ kind="invalid"
         → final_text = assistant text
         → done = True
         → 通常 reward = -1

所以 ParsedAssistant.kind 是协议层交给状态机的控制信号。

3.27 重点与疑难点#

重点 1：工具定义是 prompt 的一部分#

messages 中看不到 CODE_TOOL，但 apply_chat_template(..., tools=[CODE_TOOL]) 会把工具说明渲染进实际 prompt。计算上下文长度时不能只数 system 和 question 文本。

重点 2：answer 只是一种解析类别#

parse_assistant() 不负责判断 boxed 格式或数学正确性。“解析为 answer”只表示轨迹不再等待工具，随后应该结束并交给 reward。

重点 3：旧 logprob 绑定真实采样 token#

PPO 使用 sampler 返回的逐 token 旧 logprob。任何 decode 后重编码都可能让 token 与 logprob 失配，因此真实 completion token 是不能替换的训练事实。

重点 4：observation 留在上下文，但不是模型动作#

工具结果必须进入下一轮 prompt，否则模型无法利用计算结果；但这些 token 是环境提供的，后续 advantage 必须为 0。

疑难点 1：占位符不是伪造训练文本#

"x" 只用于向 chat template 查询结构性边界。最终 prompt 使用真实的 previous_prompt_tokens + completion_tokens，不会包含占位符。

疑难点 2：observation_tokens 比 stdout 更宽#

它是加入 tool message 后的完整模板增量，可能包含 role 标记、消息闭合和下一轮 assistant 前缀。训练 mask 应覆盖整个增量，而不只是工具正文。

疑难点 3：解析器对格式严格，对答案内容宽松#

工具标签必须严格符合协议；普通无工具文本则一律先视为 answer。这样职责清晰：协议层判断动作类型，reward 层判断最终结果。

疑难点 4：模板兼容性通过前缀断言保护#

代码假设 chat template 具有“历史前缀不被改写”的性质。两个 ValueError 检查用于在 tokenizer/template 变化时尽早失败，防止静默产生错位 token。

疑难点 5：正则不是完整 XML/HTML 解析器#

它只识别项目约定的 Qwen 工具格式。若代码字符串本身包含和协议完全相同的闭合标签，可能提前结束匹配；常规数学 Python 代码基本不会遇到这种输入。

3.28 本章对象与 token 账本#

3.29 本章小结#

用一句话概括 protocol.py：

它把数学题和工具说明渲染成 Qwen 原生 prompt，严格区分最终回答、合法工具调用与非法输出，并在工具返回后保留真实采样 token、只追加结构性 closing 和 observation 增量，从而守住 PPO 所需的 token/logprob 对齐。

读完本章应能回答：

1. CODE_TOOL 与 LocalPythonSandbox 的职责有什么不同？
2. 为什么初始 messages 中没有工具字典，但实际 prompt 中有工具说明？
3. add_generation_prompt=True 确定了什么边界？
4. parse_assistant() 为什么会把没有 boxed 的普通文本也归为 answer？
5. 为什么不能对真实 assistant 文本重新执行 tokenize？
6. build_next_prompt() 中占位符 "x" 的作用是什么？
7. observation_tokens 为什么不仅仅是 stdout 的 token？
8. _suffix_prefix_overlap() 防止了什么问题？
9. 在 [P₀][C₁][O₁][C₂] 中，哪些 token 是模型动作，哪些必须 mask？