diff --git a/AGENTS.md b/AGENTS.md index 10e5fa7..395807c 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -23,6 +23,18 @@ 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 `` 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`. @@ -30,8 +42,8 @@ Chat templates must accept a `system` role. Tool results are sent as role `tool` | 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) diff --git a/modulo/agent.py b/modulo/agent.py index 7f4a9f4..2956c10 100644 --- a/modulo/agent.py +++ b/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 + 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"\s*$", text[last:start]) + after = re.match(r"\s*", 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"", "", stripped) + if not keep: + stripped = re.sub(r"", "", 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, diff --git a/modulo/chat.py b/modulo/chat.py index 0be84bc..a667472 100644 --- a/modulo/chat.py +++ b/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 (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): diff --git a/tests/test_agent_logic.py b/tests/test_agent_logic.py index e9b3bf1..ea5ef65 100644 --- a/tests/test_agent_logic.py +++ b/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. {"name": "search", "arguments": {"q": "decay"}} ' + 'And {"name": "drop", "arguments": {"concept": "a"}} 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 '{"name": "search", "arguments": {"q": "decay"}}' in out.prose + assert '"drop"' not in out.prose and out.prose.count("") == 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})] diff --git a/tests/test_chat.py b/tests/test_chat.py index d8f45e0..dae3849 100644 --- a/tests/test_chat.py +++ b/tests/test_chat.py @@ -81,3 +81,17 @@ def test_probe_runs_once_per_tokenizer(): chat_prompt(tok, CONVERSATION) chat_prompt(tok, CONVERSATION) assert tok.calls == 3 # one probe, two renders + + +def test_harness_tools_reach_the_template(): + seen = [] + + class Recording(StubTokenizer): + def apply_chat_template(self, messages, **kw): + seen.append(kw.get("tools")) + return super().apply_chat_template(messages, **kw) + tok = Recording() + tools = [{"type": "function", "function": {"name": "search"}}] + chat_prompt(tok, CONVERSATION, tools=tools) + chat_prompt(tok, CONVERSATION) + assert seen[-2] == tools and seen[-1] is None # no tools: the argument isn't passed at all