"""The world-model agent loop: tool use over the workspace until a typed transition is submitted. The loop is a plain function-calling agent: the model receives the tool catalog, calls tools, sees their results as tool messages, and finishes by calling ``submit_transition``. The harness bounds the number of tool calls, forces the submission when the budget is spent, and keeps a trace of every call for the audit record. """ from __future__ import annotations import json from typing import Any from pydantic import ValidationError from trace2env.llm import AgentLLM from trace2env.models import SUBMIT_TOOL_NAME, TransitionSubmission from trace2env.workspace import WorkspaceTools, clip_value # Names a model may use for the final submission; the contract's tool is SUBMIT_TOOL_NAME, the rest are tolerated. SUBMIT_ALIASES = {SUBMIT_TOOL_NAME, "final", "submit", "submit_final", "final_submission", "transition_submission"} _PLACEHOLDERS = {"placeholder", "todo", "tbd", "n/a", "na", "...", "", "", "observation", "output", "the observation", "[observation]", "[output]", "lorem ipsum"} def _is_placeholder(observation: str) -> bool: """A submission whose observation is a stand-in token rather than environment text (seen from a live model).""" text = observation.strip().strip("`'\"").strip().lower() return text in _PLACEHOLDERS or text.replace("_", " ") in _PLACEHOLDERS class AgentProtocolError(RuntimeError): pass def _dumps(value: Any) -> str: return json.dumps(value, ensure_ascii=False, default=str) class AgentLoop: def __init__(self, llm: AgentLLM, tools: WorkspaceTools, *, system: str, max_tool_calls: int = 8, role: str = "runtime_agent_turn", max_invalid_submissions: int = 2, max_rounds: int | None = None, require_observation: bool = False): self.llm = llm self.tools = tools self.system = system self.max_tool_calls = max(0, max_tool_calls) self.role = role # Invalid submissions (a submit call whose arguments do not validate) are answered with the validation error # and do not spend the tool budget, so they need their own bound: more than this many is a protocol error. # A cap on model calls per step bounds every path (a model answered 171 rounds of empty submissions before). self.max_invalid_submissions = max(0, max_invalid_submissions) # Opt-in (the playground): a submission that applies no rule and carries no observation text is sent back # once more as invalid instead of being rendered from the JSON fallback. Off in the benchmark runs. self.require_observation = require_observation self.max_rounds = max_rounds if max_rounds is not None else 2 * self.max_tool_calls + 4 def run(self, user: str) -> tuple[TransitionSubmission, list[dict[str, Any]], dict[str, int]]: messages: list[dict[str, Any]] = [{"role": "system", "content": self.system}, {"role": "user", "content": user}] specs = self.tools.specs() trace: list[dict[str, Any]] = [] usage: dict[str, int] = {"prompt_tokens": 0, "completion_tokens": 0, "model_calls": 0} used = 0 nudges = 0 invalid = 0 force_nudged = False while True: if usage["model_calls"] >= self.max_rounds: raise AgentProtocolError(f"The world-model agent made {usage['model_calls']} calls without a valid submission") force = SUBMIT_TOOL_NAME if used >= self.max_tool_calls else None reply = self.llm.respond(messages=messages, tools=specs, role=self.role, force_tool=force) usage["model_calls"] += 1 for key in ("prompt_tokens", "completion_tokens"): usage[key] += int(reply.usage.get(key, 0) or 0) if not reply.tool_calls: # A bare answer is not a transition. Ask once for the typed submission, then give up. nudges += 1 if nudges > 1: raise AgentProtocolError("The world-model agent replied without submitting a transition") messages.append({"role": "assistant", "content": reply.content or ""}) messages.append({"role": "user", "content": f"Call {SUBMIT_TOOL_NAME} to finish; a plain reply is not a transition."}) continue messages.append( { "role": "assistant", "content": reply.content, "tool_calls": [ {"id": call.id, "type": "function", "function": {"name": call.name, "arguments": _dumps(call.arguments)}} for call in reply.tool_calls ], } ) for call in reply.tool_calls: if (call.name or "").strip().lower() in SUBMIT_ALIASES: try: submission = TransitionSubmission.model_validate(call.arguments) except ValidationError as exc: invalid += 1 trace.append({"tool": SUBMIT_TOOL_NAME, "arguments": clip_value(call.arguments, 200), "error": "invalid submission"}) if invalid > self.max_invalid_submissions: raise AgentProtocolError(f"The world-model agent submitted an invalid transition {invalid} times") from exc messages.append({"role": "tool", "tool_call_id": call.id, "content": _dumps({"error": "Invalid submission: " f"{exc.errors(include_input=False, include_url=False)[:3]}"})}) continue if self.require_observation and not submission.rule_ids and ( submission.observation is None or _is_placeholder(submission.observation)): invalid += 1 reason = "observation missing" if submission.observation is None else "observation is a placeholder" trace.append({"tool": SUBMIT_TOOL_NAME, "arguments": clip_value(call.arguments, 200), "error": f"invalid submission: {reason}"}) if invalid > self.max_invalid_submissions: raise AgentProtocolError(f"The world-model agent submitted an invalid transition {invalid} times") messages.append({"role": "tool", "tool_call_id": call.id, "content": _dumps({"error": "Invalid submission: 'observation' must be the exact text the " "environment returns when no rule renders the response (an " "empty string if it prints nothing), never a placeholder."})}) continue trace.append({"tool": SUBMIT_TOOL_NAME, "arguments": {"rule_ids": submission.rule_ids, "citations": submission.citations, "effects": len(submission.effects)}}) return submission, trace, usage if force is not None: if not force_nudged: # One reminder before failing: models that ignore a forced tool choice usually comply when told. force_nudged = True trace.append({"tool": call.name, "error": "refused: tool budget exhausted"}) messages.append({"role": "tool", "tool_call_id": call.id, "content": _dumps({"error": f"The tool budget is exhausted; call {SUBMIT_TOOL_NAME} now " "with the complete transition submission."})}) continue raise AgentProtocolError(f"Tool budget of {self.max_tool_calls} calls is exhausted; {call.name!r} was refused") used += 1 result = self.tools.call(call.name, call.arguments) trace.append(self.tools.calls[-1]) messages.append({"role": "tool", "tool_call_id": call.id, "content": _dumps(result)}) if used > self.max_tool_calls + 1: raise AgentProtocolError("The world-model agent exceeded its tool budget")