Work alongside a harness's own tools

AGENTS.md

@@ -23,15 +23,27 @@ Install it into the harness's environment with `pip install -e /path/to/modulo`.
 
 The agent harness calls `agent.step(messages)` each turn instead of `model.generate()`, then appends the assistant output and `result.tool_messages()` (role `tool`) so the model sees whether its calls worked. If there are tool messages, give the model the next turn. If `messages` starts with a system message, it is merged into the steering system prompt. Steering is transparent — the harness doesn't need to know about hooks or vectors.
 
+**If the harness has tools of its own**, pass their schemas (or functions, as for `apply_chat_template`) with `agent.step(messages, tools=[...])`. They're given to the chat template, and calls to them come back in `result.other_tool_calls`, not run. They're left in `result.prose` (with any `<tool_call>` tags) and in `raw_output`. Modulo's own calls are run and removed from `prose`, and `result.tool_messages()` holds only their results. Run the harness's calls yourself and append their results after Modulo's:
+
+```python
+result = agent.step(messages, tools=my_tools)
+messages += [{"role": "assistant", "content": result.raw_output}]
+messages += result.tool_messages()
+for call in result.other_tool_calls:     # ToolCall(name, arguments, raw)
+    messages.append({"role": "tool", "content": run_my_tool(call.name, call.arguments)})
+```
+
+A harness that describes its tools in its own system prompt instead can pass `SteeringAgent(..., pass_through_unknown_tools=True)`, which treats every call that isn't Modulo's as the harness's. Without either, a call to a tool Modulo doesn't know gets an "Unknown tool" result, which suits the standalone demo, where the model has no other tools.
+
 `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 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 |
+| Llama-3.1 | shown | with `tools=`, its template asks for `{"name": ..., "parameters": {...}}`; Modulo reads `parameters` as `arguments` |
+| OLMo-2 | fallback | its template silently drops `tool` messages, and ignores `tools=`: describe a harness's tools in its system prompt and use `pass_through_unknown_tools=True` |
 | 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 |
 
@@ -197,7 +209,7 @@ vec = extract_vector(model, tokenizer,
 
 ### SteeringAgent methods
 
-- `agent.step(messages, max_new_tokens=300, temperature=0.7, top_p=0.9) -> StepResult`
+- `agent.step(messages, max_new_tokens=300, temperature=0.7, top_p=0.9, tools=None) -> StepResult`
 - `agent.build_system_prompt() -> str` — system prompt describing the steer tool
 - `agent.status() -> str` — active vector status
 - `agent.apply(name, strength) -> (bool, str)` — manually apply a vector
@@ -207,7 +219,8 @@ vec = extract_vector(model, tokenizer,
 
 - `prose` — model output with tool calls stripped
 - `raw_output` — full model output
-- `tool_calls` — list of `ToolCall(name, arguments, raw)`
+- `tool_calls` — list of `ToolCall(name, arguments, raw)` that Modulo handled (its own tools, and unknown names)
+- `other_tool_calls` — calls to the harness's own tools (`tools=` or `pass_through_unknown_tools`), not run
 - `tool_results` — list of `ToolResult(name, arguments, ok, message)`
 - `degenerate_output` / `degenerate_reason` — the discarded steered output and why (`"repetition"` or `"nll"`), if the degeneration detector fired
 - `unsteered_nll` — mean per-token NLL of the steered output under the unsteered model (steered turns only)

modulo/agent.py

@@ -106,7 +106,7 @@ class ToolResult:
 class StepResult:
     prose: str           # model output with tool calls stripped
     raw_output: str      # full model output (the regenerated one if degenerate)
-    tool_calls: list     # list[ToolCall]
+    tool_calls: list     # list[ToolCall]: the calls Modulo handled (its own tools, and unknown names)
     steering_before: dict  # {name: strength} applied during this turn's generation
     steering_after: dict   # after tool calls + decay
     tool_results: list = field(default_factory=list)  # list[ToolResult]
@@ -118,6 +118,9 @@ class StepResult:
     # model; None on unsteered turns. High values mean unusual text, but see
     # AGENTS.md: it doesn't separate word salad from heavy role-play.
     unsteered_nll: Optional[float] = None
+    # Calls to the harness's own tools (step(tools=...), or any non-Modulo tool with
+    # pass_through_unknown_tools): not run, left in prose, for the harness to run
+    other_tool_calls: list = field(default_factory=list)
 
     def tool_messages(self) -> list[dict]:
         """Tool results as chat messages (role 'tool'), to append after the assistant turn."""
@@ -140,7 +143,8 @@ def tool_messages(self) -> list[dict]:
 
 
 def _find_tool_calls(text: str) -> list[tuple[int, int, ToolCall]]:
-    """Find JSON objects of the form {"name": str, "arguments": {...}} anywhere in text.
+    """Find JSON objects of the form {"name": str, "arguments": {...}} (or "parameters")
+    anywhere in text.
 
     Uses a real JSON decoder rather than a regex so argument values can contain
     braces, quotes, etc. Returns (start, end, ToolCall) spans in order.
@@ -153,9 +157,13 @@ def _find_tool_calls(text: str) -> list[tuple[int, int, ToolCall]]:
         except json.JSONDecodeError:
             i = text.find("{", i + 1)
             continue
-        if (isinstance(obj, dict) and isinstance(obj.get("name"), str)
-                and isinstance(obj.get("arguments"), dict)):
-            found.append((i, end, ToolCall(name=obj["name"], arguments=obj["arguments"], raw=text[i:end])))
+        # Llama-3.1's template asks for "parameters" instead of "arguments". An object
+        # with a "description" is an echoed tool definition, not a call.
+        args = obj.get("arguments") if isinstance(obj, dict) else None
+        if args is None and isinstance(obj, dict) and "description" not in obj:
+            args = obj.get("parameters")
+        if isinstance(obj, dict) and isinstance(obj.get("name"), str) and isinstance(args, dict):
+            found.append((i, end, ToolCall(name=obj["name"], arguments=args, raw=text[i:end])))
             i = text.find("{", end)
         else:
             i = text.find("{", i + 1)
@@ -181,18 +189,43 @@ def malformed_tool_calls(text: str) -> list[tuple[str, str]]:
     return found
 
 
+MODULO_TOOLS = ("steer", "drop", "report", "create")
+
+
+def tool_names(tools: Optional[list]) -> set[str]:
+    """Names in a list of tool schemas ({"type": "function", "function": {"name": ...}},
+    {"name": ...}) or Python functions, as accepted by apply_chat_template."""
+    names = set()
+    for t in tools or []:
+        if callable(t):
+            names.add(t.__name__)
+        elif isinstance(t, dict):
+            name = (t.get("function") or {}).get("name") or t.get("name")
+            if isinstance(name, str):
+                names.add(name)
+    return names
+
+
 def parse_tool_calls(text: str) -> list[ToolCall]:
     return [tc for _, _, tc in _find_tool_calls(text)]
 
 
-def strip_tool_calls(text: str) -> str:
+def strip_tool_calls(text: str, keep: Optional[set] = None) -> str:
+    """Text without its tool calls. Calls named in `keep` are left in place, with any
+    <tool_call> tags around them (they're the harness's own); the others are removed
+    together with their tags."""
     out, last = [], 0
-    for start, end, _ in _find_tool_calls(text):
-        out.append(text[last:start])
-        last = end
+    for start, end, tc in _find_tool_calls(text):
+        if keep and tc.name in keep:
+            continue
+        before = re.search(r"<tool_call>\s*$", text[last:start])
+        after = re.match(r"\s*</tool_call>", text[end:])
+        out.append(text[last:start - (len(before.group(0)) if before else 0)])
+        last = end + (len(after.group(0)) if after else 0)
     out.append(text[last:])
     stripped = "".join(out)
-    stripped = re.sub(r"</?tool_call>", "", stripped)
+    if not keep:
+        stripped = re.sub(r"</?tool_call>", "", stripped)
     return stripped.strip()
 
 
@@ -388,7 +421,8 @@ def __init__(self, model, tokenizer, registry: Optional[VectorRegistry] = None,
                  model_name: Optional[str] = None, auto_extract: bool = True,
                  vectors_dir: Optional[Path] = None, include_community: bool = False,
                  control: str = "real", control_seed: int = 0, disclose_controls: bool = False,
-                 nll_threshold: Optional[float] = None, menu_seed: Optional[int] = None):
+                 nll_threshold: Optional[float] = None, menu_seed: Optional[int] = None,
+                 pass_through_unknown_tools: bool = False):
         """
         Args:
             model: HF causal LM (must have model.model.layers)
@@ -413,12 +447,18 @@ def __init__(self, model, tokenizer, registry: Optional[VectorRegistry] = None,
                          shuffled once per session so no vector is always first (models
                          tend to pick the first one). None: a random order, kept in
                          self.menu_seed so a session can be reproduced.
+            pass_through_unknown_tools: treat every call to a tool that isn't Modulo's as
+                         the harness's own (StepResult.other_tool_calls) instead of
+                         answering "Unknown tool". For harnesses that describe their tools
+                         in their own system prompt rather than passing step(tools=...).
         """
         if control not in CONTROLS:
             raise ValueError(f"control must be one of {CONTROLS}, got {control!r}")
         self.control = control
         self.control_seed = control_seed
         self.disclose_controls = disclose_controls
+        self.pass_through_unknown_tools = pass_through_unknown_tools
+        self._step_tools: Optional[list] = None  # the harness's tools for the current step
         self.menu_seed = menu_seed if menu_seed is not None else random.SystemRandom().randrange(2**31)
         self._menu_order: Optional[list[str]] = None  # fixed at the first prompt of the session
         self.nll_threshold = nll_threshold
@@ -688,8 +728,9 @@ def _build_steering_context(self, prompt_len: int):
 
     def _execute(self, tc: ToolCall) -> ToolResult:
         args = tc.arguments
-        if tc.name not in ("steer", "drop", "report", "create"):
-            return ToolResult(tc.name, args, False, f"Unknown tool '{tc.name}'. Available: steer, drop, report, create.")
+        if tc.name not in MODULO_TOOLS:
+            available = ", ".join([*MODULO_TOOLS, *sorted(tool_names(self._step_tools))])
+            return ToolResult(tc.name, args, False, f"Unknown tool '{tc.name}'. Available: {available}.")
         concept = args.get("concept")
         if not isinstance(concept, str):
             return ToolResult(tc.name, args, False, "'concept' must be a string.")
@@ -769,19 +810,23 @@ def _full_messages(self, messages: list[dict]) -> list[dict]:
         return [{"role": "system", "content": sys_prompt}] + messages
 
     def _encode(self, messages: list[dict]):
-        text = chat_prompt(self.tokenizer, self._full_messages(messages))
+        text = chat_prompt(self.tokenizer, self._full_messages(messages), tools=self._step_tools)
         return encode(self.tokenizer, text, self.model.device)
 
     def step(self, messages: list[dict], max_new_tokens: int = 300,
-             temperature: float = 0.7, top_p: float = 0.9) -> StepResult:
+             temperature: float = 0.7, top_p: float = 0.9, tools: Optional[list] = None) -> StepResult:
         """Run one turn: build prompt, generate (steered), parse tool calls, apply them, decay.
 
         Args:
             messages: conversation history (list of {role, content} dicts)
             max_new_tokens, temperature, top_p: generation params
+            tools: the harness's own tools (JSON schemas or functions, as for
+                apply_chat_template), passed to the chat template. Calls to them are
+                returned in StepResult.other_tool_calls, not run.
 
         Returns: StepResult with prose, raw output, tool calls/results, and steering state
         """
+        self._step_tools = tools
         full = self._full_messages(messages)
         system_prompt = full[0]["content"]
         inputs = self._encode(messages)
@@ -822,8 +867,12 @@ def step(self, messages: list[dict], max_new_tokens: int = 300,
             with self._build_steering_context(retry.input_ids.shape[1]):
                 raw_output = self._generate(retry, max_new_tokens, temperature, top_p)
 
-        tool_calls = parse_tool_calls(raw_output)
-        prose = strip_tool_calls(raw_output)
+        harness = tool_names(tools)
+        tool_calls, other_tool_calls = [], []
+        for tc in parse_tool_calls(raw_output):
+            mine = tc.name in MODULO_TOOLS or (tc.name not in harness and not self.pass_through_unknown_tools)
+            (tool_calls if mine else other_tool_calls).append(tc)
+        prose = strip_tool_calls(raw_output, keep={tc.name for tc in other_tool_calls})
         tool_results = [self._execute(tc) for tc in tool_calls]
         tool_results += [ToolResult(name, {}, False,
                                     f"This {name} call isn't valid JSON, so it wasn't run: {snippet!r}. Write it "
@@ -834,7 +883,7 @@ def step(self, messages: list[dict], max_new_tokens: int = 300,
         steering_after = {n: av.current_strength for n, av in self.active.items()}
 
         return StepResult(
-            prose=prose, raw_output=raw_output, tool_calls=tool_calls,
+            prose=prose, raw_output=raw_output, tool_calls=tool_calls, other_tool_calls=other_tool_calls,
             steering_before=steering_before, steering_after=steering_after,
             tool_results=tool_results, degenerate_output=degenerate_output,
             degenerate_reason=degenerate_reason, repeated_output=repeated_output,

modulo/chat.py

@@ -1,4 +1,5 @@
 """Chat-template rendering shared by the agent, extraction and verification."""
+from typing import Optional
 
 
 # Tool results are sent as role "tool". Templates that can't show that role get
@@ -7,9 +8,10 @@
 _TOOL_PROBE = "MODULO_TOOL_ROLE_PROBE"
 
 
-def _render(tokenizer, messages: list[dict]) -> str:
+def _render(tokenizer, messages: list[dict], tools: Optional[list] = None) -> str:
+    extra = {"tools": tools} if tools else {}
     return tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
-                                         enable_thinking=False)
+                                         enable_thinking=False, **extra)
 
 
 def supports_tool_role(tokenizer) -> bool:
@@ -53,16 +55,17 @@ def adapt_messages(tokenizer, messages: list[dict]) -> list[dict]:
     return out
 
 
-def chat_prompt(tokenizer, messages: list[dict]) -> str:
+def chat_prompt(tokenizer, messages: list[dict], tools: Optional[list] = None) -> 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. Tool messages are adapted for templates
-    that can't show them (see adapt_messages).
+    that can't show them (see adapt_messages). `tools` (a harness's own tool
+    schemas) is passed to the template, for templates that list tools.
     """
-    return _render(tokenizer, adapt_messages(tokenizer, messages))
+    return _render(tokenizer, adapt_messages(tokenizer, messages), tools)
 
 
 def encode(tokenizer, text: str, device=None):

tests/test_agent_logic.py

@@ -720,3 +720,63 @@ def test_vectors_created_mid_session_join_the_end_of_the_menu(creator):
     before = _menu(creator)
     creator._execute(ToolCall("create", {"concept": "wonder", "more": "x", "less": "y"}, ""))
     assert _menu(creator) == before + ["wonder"]
+
+
+# --- harnesses with tools of their own ---
+
+SEARCH = {"type": "function", "function": {"name": "search", "parameters": {"type": "object"}}}
+MIXED = ('Let me look. <tool_call>{"name": "search", "arguments": {"q": "decay"}}</tool_call> '
+         'And <tool_call>{"name": "drop", "arguments": {"concept": "a"}}</tool_call> done.')
+
+
+def _step_with(agent, monkeypatch, output, **kw):
+    class FakeInputs:
+        input_ids = torch.zeros(1, 1)
+    monkeypatch.setattr(agent, "_encode", lambda msgs: FakeInputs())
+    monkeypatch.setattr(agent, "_build_steering_context", lambda prompt_len: nullcontext())
+    stub_generation(agent, monkeypatch, lambda *args: output)
+    return agent.step([{"role": "user", "content": "hi"}], **kw)
+
+
+def test_harness_tool_calls_are_returned_not_run(agent, monkeypatch):
+    agent.apply("a", 5.0)
+    out = _step_with(agent, monkeypatch, MIXED, tools=[SEARCH])
+    assert [tc.name for tc in out.other_tool_calls] == ["search"]
+    assert out.other_tool_calls[0].arguments == {"q": "decay"}
+    assert [r.name for r in out.tool_results] == ["drop"] and out.tool_results[0].ok
+    assert [m for m in out.tool_messages() if "search" in m["content"]] == []  # the harness answers its own
+    # the harness's call stays in prose, tags and all; Modulo's is removed with its tags
+    assert '<tool_call>{"name": "search", "arguments": {"q": "decay"}}</tool_call>' in out.prose
+    assert '"drop"' not in out.prose and out.prose.count("<tool_call>") == 1
+
+
+def test_without_tools_unknown_calls_get_an_error(agent, monkeypatch):
+    out = _step_with(agent, monkeypatch, MIXED)
+    assert out.other_tool_calls == []
+    r = next(r for r in out.tool_results if r.name == "search")
+    assert not r.ok and "Unknown tool 'search'. Available: steer, drop, report, create." in r.message
+    out = _step_with(agent, monkeypatch, '{"name": "lookup", "arguments": {}}', tools=[SEARCH])
+    assert "Available: steer, drop, report, create, search." in out.tool_results[0].message
+
+
+def test_pass_through_unknown_tools(monkeypatch, tmp_path):
+    a = SteeringAgent(None, None, registry=StubRegistry(vectors_dir=tmp_path), pass_through_unknown_tools=True)
+    out = _step_with(a, monkeypatch, MIXED)
+    assert [tc.name for tc in out.other_tool_calls] == ["search"] and [r.name for r in out.tool_results] == ["drop"]
+
+
+def test_tool_names_from_schemas_and_functions():
+    from modulo.agent import tool_names
+
+    def get_weather(city: str):
+        """Weather."""
+    assert tool_names([SEARCH, {"name": "calc"}, get_weather, "junk"]) == {"search", "calc", "get_weather"}
+    assert tool_names(None) == set()
+
+
+def test_llama_style_parameters_are_read_as_arguments():
+    calls = parse_tool_calls('{"name": "search", "parameters": {"q": "x"}} '
+                             '{"name": "steer", "parameters": {"concept": "a", "strength": 3}} '
+                             '{"name": "search", "description": "Search", "parameters": {"type": "object"}}')
+    assert [(c.name, c.arguments) for c in calls] == [("search", {"q": "x"}),
+                                                      ("steer", {"concept": "a", "strength": 3})]