Send tool results as user messages when the template can't show them

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 |

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,

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 <think> (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):

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