From 2e0eca8dac75b244818c2df2ca488348d8d852ba Mon Sep 17 00:00:00 2001 From: Jen D Date: Sun, 1 Nov 2026 22:06:26 +0800 Subject: [PATCH] Send tool results as user messages when the template can't show them OLMo-2's chat template silently drops role "tool" messages, and Mistral-7B-Instruct-v0.3's rejects them (it wants tool-call IDs and strict user/assistant alternation). Either way the model never learned whether its steer calls worked. chat.supports_tool_role() probes the template once per tokenizer. Where tool messages aren't shown, adapt_messages() sends them as user messages starting with "[Tool result]" and merges consecutive user messages. Session logs record tool_role. Checked on a GPU: OLMo-2 now sees its tool results and stacks a second vector on the first; Mistral no longer errors. Co-Authored-By: Claude Opus 5.5 --- AGENTS.md | 11 +++++- agent/loop.py | 3 ++ modulo/chat.py | 58 ++++++++++++++++++++++++++++++-- tests/test_chat.py | 83 ++++++++++++++++++++++++++++++++++++++++++++++ 4 files changed, 151 insertions(+), 4 deletions(-) create mode 100644 tests/test_chat.py diff --git a/AGENTS.md b/AGENTS.md index de8b6c0..8e88d1d 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -25,7 +25,15 @@ The agent harness calls `agent.step(messages)` each turn instead of `model.gener `agent.apply()` / `agent.drop()` called between steps behave like the model's own tool calls: an applied vector gets one full-strength steered turn, and a drop blocks re-applying for `COOLDOWN_TURNS` turns. -Chat templates must accept `system` and `tool` roles (Qwen, Llama-3.x, recent Mistral do). Gemma-2 rejects both, so it isn't supported yet. +Chat templates must accept a `system` role. Tool results are sent as role `tool`. If a model's template can't show that role, they're sent instead as user messages starting with `[Tool result]`, with consecutive user messages merged. Modulo checks this once per tokenizer (`modulo.chat.supports_tool_role`), and session logs record it as `tool_role`. + +| model | role `tool` | notes | +|---|---|---| +| Qwen2.5 | shown | | +| Llama-3.1 | shown | | +| OLMo-2 | fallback | its template silently drops `tool` messages | +| Mistral-7B-Instruct-v0.3 | fallback | its template requires tool-call IDs and strictly alternating turns, and puts the system prompt in the latest user message. Its tokenizer needs `sentencepiece` and `protobuf`. | +| Gemma-2 | — | rejects the `system` role, so it isn't supported yet | ### 2. As a standalone demo @@ -282,6 +290,7 @@ Host the vector on HuggingFace and add an entry to `community/vectors.yaml` (or | `run.sh` | Start a standalone demo session | | `modulo/__init__.py` | Public API exports | | `modulo/agent.py` | `SteeringAgent` class (core) | +| `modulo/chat.py` | Prompt rendering: chat template, thinking mode off, tool-role fallback, tokenizing without a second BOS | | `modulo/registry.py` | `VectorRegistry`: your own vectors plus the community index for a model | | `modulo/extraction.py` | Vector extraction (manual + auto), per-vector question sets | | `modulo/reports.py` | Report format, validation, `reports.jsonl` read/write | diff --git a/agent/loop.py b/agent/loop.py index fa6cf71..6e490bf 100644 --- a/agent/loop.py +++ b/agent/loop.py @@ -36,6 +36,7 @@ import config # sets HF_HOME before transformers is imported import modulo.agent as ma from modulo import SteeringAgent, __version__ +from modulo.chat import supports_tool_role from modulo.reports import file_sha256 from load_model import load_model_and_tokenizer @@ -80,6 +81,8 @@ def session_metadata(agent: SteeringAgent, args, rng_seed: int, gen_params: dict "modulo_version": __version__, "model": agent.model_name, "dtype": config.DTYPE, + # False: the chat template can't show role "tool", so tool results went in as labelled user messages + "tool_role": supports_tool_role(agent.tokenizer), "torch": torch.__version__, "transformers": transformers.__version__, "control": agent.control, diff --git a/modulo/chat.py b/modulo/chat.py index 7d23782..0be84bc 100644 --- a/modulo/chat.py +++ b/modulo/chat.py @@ -1,16 +1,68 @@ """Chat-template rendering shared by the agent, extraction and verification.""" +# Tool results are sent as role "tool". Templates that can't show that role get +# them as a user message starting with this label instead. +TOOL_RESULT_LABEL = "[Tool result]" +_TOOL_PROBE = "MODULO_TOOL_ROLE_PROBE" + + +def _render(tokenizer, messages: list[dict]) -> str: + return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, + enable_thinking=False) + + +def supports_tool_role(tokenizer) -> bool: + """Whether the chat template shows a role-"tool" message to the model (cached on the tokenizer). + + Some don't: OLMo-2's template silently drops tool messages, and + Mistral-7B-Instruct-v0.3's raises unless they carry tool-call IDs. + """ + cached = getattr(tokenizer, "_modulo_tool_role", None) + if cached is None: + probe = [{"role": "user", "content": "a"}, {"role": "assistant", "content": "b"}, + {"role": "tool", "content": _TOOL_PROBE}] + try: + cached = _TOOL_PROBE in _render(tokenizer, probe) + except Exception: + cached = False + try: + tokenizer._modulo_tool_role = cached + except AttributeError: + pass + return cached + + +def adapt_messages(tokenizer, messages: list[dict]) -> list[dict]: + """Messages the template can render: unchanged, unless it can't show tool messages. + + Then each tool message becomes a user message starting with TOOL_RESULT_LABEL, + and consecutive user messages are merged (some templates, like Mistral's, + require user and assistant turns to alternate). + """ + if supports_tool_role(tokenizer): + return messages + out: list[dict] = [] + for m in messages: + if m["role"] == "tool": + m = {"role": "user", "content": f"{TOOL_RESULT_LABEL}\n{m['content']}"} + if out and m["role"] == "user" and out[-1]["role"] == "user": + out[-1] = {"role": "user", "content": f"{out[-1]['content']}\n\n{m['content']}"} + else: + out.append(dict(m)) + return out + + def chat_prompt(tokenizer, messages: list[dict]) -> str: """The prompt text for messages, ready for the assistant's turn. Thinking mode is turned off: Qwen3-style templates think by default, which spends the turn's token budget inside (so steered answers, tool calls and extraction activations never come from a real reply). Templates - without the variable ignore it. + without the variable ignore it. Tool messages are adapted for templates + that can't show them (see adapt_messages). """ - return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, - enable_thinking=False) + return _render(tokenizer, adapt_messages(tokenizer, messages)) def encode(tokenizer, text: str, device=None): diff --git a/tests/test_chat.py b/tests/test_chat.py new file mode 100644 index 0000000..d8f45e0 --- /dev/null +++ b/tests/test_chat.py @@ -0,0 +1,83 @@ +"""Chat rendering for templates that can't show tool messages (CPU only, stub tokenizers).""" +import pytest + +from modulo.chat import TOOL_RESULT_LABEL, adapt_messages, chat_prompt, supports_tool_role + + +class StubTokenizer: + """Renders 'role: content' lines, like a chat template would.""" + + def __init__(self): + self.calls = 0 + + def check(self, messages): + pass + + def apply_chat_template(self, messages, tokenize=False, add_generation_prompt=True, **kw): + self.calls += 1 + self.check(messages) + return "".join(f"{m['role']}: {m['content']}\n" for m in self.shown(messages)) + "assistant:" + + def shown(self, messages): + return messages + + +class DropsTool(StubTokenizer): + """Like OLMo-2: tool messages silently vanish.""" + + def shown(self, messages): + return [m for m in messages if m["role"] != "tool"] + + +class StrictTemplate(StubTokenizer): + """Like Mistral-7B-Instruct-v0.3: rejects tool messages without IDs and requires alternation.""" + + def check(self, messages): + roles = [m["role"] for m in messages if m["role"] != "system"] + if "tool" in roles: + raise ValueError("Tool call IDs should be alphanumeric strings with length 9!") + if any(a == b for a, b in zip(roles, roles[1:])): + raise ValueError("conversation roles must alternate user/assistant") + + +CONVERSATION = [ + {"role": "system", "content": "SYS"}, + {"role": "user", "content": "hello"}, + {"role": "assistant", "content": '{"name": "steer", "arguments": {"concept": "humor", "strength": 6}}'}, + {"role": "tool", "content": '{"name": "steer", "ok": true}'}, + {"role": "tool", "content": "[status] Turn 2 of 8."}, +] + + +def test_template_with_tool_role_is_left_alone(): + tok = StubTokenizer() + assert supports_tool_role(tok) + assert adapt_messages(tok, CONVERSATION) is CONVERSATION + assert "tool: [status] Turn 2 of 8." in chat_prompt(tok, CONVERSATION) + + +@pytest.mark.parametrize("tok", [DropsTool(), StrictTemplate()], ids=["drops-tool", "strict"]) +def test_tool_results_reach_templates_without_tool_role(tok): + assert not supports_tool_role(tok) + text = chat_prompt(tok, CONVERSATION) + assert '{"name": "steer", "ok": true}' in text + assert "[status] Turn 2 of 8." in text + assert TOOL_RESULT_LABEL in text + + +def test_consecutive_tool_and_user_messages_are_merged(): + tok = StrictTemplate() + msgs = adapt_messages(tok, CONVERSATION + [{"role": "user", "content": "and now?"}]) + roles = [m["role"] for m in msgs] + assert roles == ["system", "user", "assistant", "user"] + last = msgs[-1]["content"] + assert last.startswith(TOOL_RESULT_LABEL) + assert last.count(TOOL_RESULT_LABEL) == 2 and last.endswith("and now?") + assert CONVERSATION[3]["role"] == "tool" # the caller's messages are not modified + + +def test_probe_runs_once_per_tokenizer(): + tok = DropsTool() + chat_prompt(tok, CONVERSATION) + chat_prompt(tok, CONVERSATION) + assert tok.calls == 3 # one probe, two renders