#!/usr/bin/env python3 """ agent-tool-use-loop.py - a provider-neutral agent tool-use loop with an explicit state machine, six stop conditions, and a token budget, plus a deterministic simulator that measures what the stop conditions are worth. What it does ------------ Runs N simulated agent episodes twice over the same random draws: once with only a step cap (the "step cap only" arm) and once with the full guard set (no-progress detection, a consecutive-tool-error cap, and a token budget). It prints the distribution of terminal states and the mean cost per episode for each arm, and optionally writes one CSV row per episode. The model here is a deterministic simulator, not a network call: it emits tool calls and final answers at fixed, seeded rates so the printed numbers are properties of the loop, not of any particular vendor's model. Its failure rates are inputs (see ModelBehaviour), not findings. Inputs ------ --episodes N episodes per arm (default 2000) --seed S base seed (default 11) --csv PATH write per-episode rows to PATH (optional) How to run ---------- python agent-tool-use-loop.py python agent-tool-use-loop.py --episodes 2000 --seed 11 --csv agent-tool-use-loop.csv Standard library only. Written and run on Python 3.11. """ from __future__ import annotations import argparse import csv import json import random import time from collections import Counter from dataclasses import dataclass, field from typing import Any, Callable # -------------------------------------------------------------------------- # Tools # -------------------------------------------------------------------------- class ToolError(Exception): """Raised by a tool handler; becomes a tool result with is_error set.""" @dataclass(frozen=True) class Tool: name: str description: str required: tuple[str, ...] handler: Callable[[dict[str, Any]], str] class Registry: """Holds tools and executes one call, never raising into the loop.""" def __init__(self, tools: list[Tool]) -> None: self.tools = {t.name: t for t in tools} def execute(self, name: str, args: dict[str, Any]) -> tuple[str, bool]: """Return (content, is_error). Every failure is a result, not an exception.""" tool = self.tools.get(name) if tool is None: known = ", ".join(sorted(self.tools)) return (f"Unknown tool {name!r}. Available tools: {known}.", True) missing = [k for k in tool.required if k not in args] if missing: return (f"Missing required argument(s) {missing} for {name}. " f"Required: {list(tool.required)}.", True) try: return (tool.handler(args), False) except ToolError as exc: return (f"{name} failed: {exc}", True) # A tiny orders database the simulated agent has to read through. ORDERS = { "A-1001": {"customer": "north wind", "status": "shipped", "total": 412.50}, "A-1002": {"customer": "north wind", "status": "refunded", "total": 88.00}, "A-1003": {"customer": "tarn & co", "status": "shipped", "total": 19.99}, } def _search_orders(args: dict[str, Any]) -> str: query = str(args["query"]).strip().lower() if not query: raise ToolError("query was empty; pass a customer name substring") hits = [oid for oid, row in ORDERS.items() if query in row["customer"]] if not hits: raise ToolError(f"no orders match {query!r}; try a shorter substring") return json.dumps({"order_ids": hits}) def _get_order(args: dict[str, Any]) -> str: oid = str(args["order_id"]) if oid not in ORDERS: raise ToolError(f"order {oid} not found") return json.dumps(ORDERS[oid]) TOOLS = Registry([ Tool("search_orders", "Find order ids by customer name substring.", ("query",), _search_orders), Tool("get_order", "Fetch one order by id.", ("order_id",), _get_order), ]) # -------------------------------------------------------------------------- # Budget and state # -------------------------------------------------------------------------- @dataclass(frozen=True) class Budget: max_steps: int = 12 max_tokens: int = 10_000 max_seconds: float = 30.0 max_consecutive_errors: int = 4 max_repeats: int = 3 # more than this many identical calls => no progress enforce_guards: bool = True # False = step cap and wall clock only @dataclass class State: """Everything the loop needs to decide whether to keep going.""" step: int = 0 prompt_tokens: int = 0 completion_tokens: int = 0 consecutive_errors: int = 0 call_counts: Counter = field(default_factory=Counter) transcript: list[dict[str, Any]] = field(default_factory=list) started: float = field(default_factory=time.monotonic) @property def tokens(self) -> int: return self.prompt_tokens + self.completion_tokens # Terminal states of the machine. ANSWERED = "answered" STEP_CAP = "step_cap" TOKEN_BUDGET = "token_budget" WALL_CLOCK = "wall_clock" ERROR_CAP = "error_cap" NO_PROGRESS = "no_progress" def check_stop(state: State, budget: Budget) -> str | None: """Return a terminal state name, or None to keep looping.""" if state.step >= budget.max_steps: return STEP_CAP if time.monotonic() - state.started > budget.max_seconds: return WALL_CLOCK if not budget.enforce_guards: return None if state.tokens >= budget.max_tokens: return TOKEN_BUDGET if state.consecutive_errors >= budget.max_consecutive_errors: return ERROR_CAP if state.call_counts and max(state.call_counts.values()) > budget.max_repeats: return NO_PROGRESS return None # -------------------------------------------------------------------------- # Token accounting (a stand-in for a real tokenizer) # -------------------------------------------------------------------------- SYSTEM_TOKENS = 420 # system prompt plus the two tool schemas CALL_TOKENS = 45 # one tool_use block RESULT_TOKENS = 120 # one tool_result block TEXT_TOKENS = 60 # a final assistant answer QUESTION_TOKENS = 30 def transcript_tokens(transcript: list[dict[str, Any]]) -> int: total = QUESTION_TOKENS for msg in transcript: if msg["kind"] == "calls": total += CALL_TOKENS * len(msg["calls"]) elif msg["kind"] == "results": total += RESULT_TOKENS * len(msg["results"]) else: total += TEXT_TOKENS return total # -------------------------------------------------------------------------- # The simulated model # -------------------------------------------------------------------------- @dataclass(frozen=True) class ModelBehaviour: """Failure rates are inputs, chosen to be pessimistic. Not measurements.""" p_repeat: float = 0.12 # re-issues its previous call verbatim p_bad_args: float = 0.10 # omits a required argument p_wander: float = 0.08 # calls a tool that cannot help p_stuck: float = 0.10 # episode latches into repeating one call forever plan_length: int = 3 # distinct successful calls needed before it answers PLAN = [ ("search_orders", {"query": "north"}), ("get_order", {"order_id": "A-1001"}), ("get_order", {"order_id": "A-1002"}), ] class SimulatedModel: """Emits tool calls or a final answer. Deterministic given the rng.""" def __init__(self, rng: random.Random, behaviour: ModelBehaviour) -> None: self.rng = rng self.b = behaviour self.stuck = rng.random() < behaviour.p_stuck self.done: set[str] = set() self.last: tuple[str, dict[str, Any]] | None = None def observe(self, name: str, args: dict[str, Any], is_error: bool) -> None: """Only a new, successful call counts as progress; a repeat does not.""" if not is_error: self.done.add(name + json.dumps(args, sort_keys=True)) def step(self) -> dict[str, Any]: if len(self.done) >= self.b.plan_length: return {"kind": "text", "text": "north wind has one shipped order, A-1001."} if self.stuck and self.last is not None: return {"kind": "calls", "calls": [self.last]} roll = self.rng.random() if self.last is not None and roll < self.b.p_repeat: call = self.last elif roll < self.b.p_repeat + self.b.p_bad_args: call = ("get_order", {}) # missing order_id elif roll < self.b.p_repeat + self.b.p_bad_args + self.b.p_wander: call = ("search_orders", {"query": "zzz"}) # no such customer else: call = PLAN[min(len(self.done), len(PLAN) - 1)] self.last = call return {"kind": "calls", "calls": [call]} # -------------------------------------------------------------------------- # The loop # -------------------------------------------------------------------------- @dataclass class Outcome: terminal: str steps: int tokens: int tool_calls: int tool_errors: int stuck: bool def run_episode(model: SimulatedModel, registry: Registry, budget: Budget) -> Outcome: state = State() tool_calls = tool_errors = 0 while True: stop = check_stop(state, budget) if stop is not None: return Outcome(stop, state.step, state.tokens, tool_calls, tool_errors, model.stuck) # 1. Request. The prompt is the whole transcript, every time. state.prompt_tokens += SYSTEM_TOKENS + transcript_tokens(state.transcript) state.step += 1 turn = model.step() # 2. Terminal turn: the model answered instead of calling a tool. if turn["kind"] == "text": state.completion_tokens += TEXT_TOKENS state.transcript.append(turn) return Outcome(ANSWERED, state.step, state.tokens, tool_calls, tool_errors, model.stuck) # 3. Execute every call in the turn, then return every result together. state.completion_tokens += CALL_TOKENS * len(turn["calls"]) state.transcript.append(turn) results = [] turn_had_error = False for name, args in turn["calls"]: tool_calls += 1 state.call_counts[(name, json.dumps(args, sort_keys=True))] += 1 content, is_error = registry.execute(name, args) if is_error: tool_errors += 1 turn_had_error = True model.observe(name, args, is_error) results.append({"tool": name, "content": content, "is_error": is_error}) state.transcript.append({"kind": "results", "results": results}) state.consecutive_errors = state.consecutive_errors + 1 if turn_had_error else 0 # -------------------------------------------------------------------------- # Experiment # -------------------------------------------------------------------------- ARMS = { "step cap only": Budget(enforce_guards=False), "full guards": Budget(enforce_guards=True), } def main() -> None: ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--episodes", type=int, default=2000) ap.add_argument("--seed", type=int, default=11) ap.add_argument("--csv", default=None) args = ap.parse_args() behaviour = ModelBehaviour() rows: list[dict[str, Any]] = [] summary: dict[str, tuple[Counter, float, float]] = {} killed_by_guard: Counter = Counter() for arm, budget in ARMS.items(): terminals: Counter = Counter() tokens = steps = 0 for ep in range(args.episodes): # Same seed per episode in both arms: the model's draws are paired. rng = random.Random(f"{args.seed}:{ep}") out = run_episode(SimulatedModel(rng, behaviour), TOOLS, budget) terminals[out.terminal] += 1 tokens += out.tokens steps += out.steps if budget.enforce_guards and out.terminal != ANSWERED: killed_by_guard["stuck" if out.stuck else "healthy"] += 1 rows.append({"arm": arm, "episode": ep, "terminal": out.terminal, "stuck": int(out.stuck), "steps": out.steps, "tokens": out.tokens, "tool_calls": out.tool_calls, "tool_errors": out.tool_errors}) summary[arm] = (terminals, tokens / args.episodes, steps / args.episodes) order = [ANSWERED, NO_PROGRESS, ERROR_CAP, TOKEN_BUDGET, STEP_CAP, WALL_CLOCK] width = max(len(s) for s in order) + 2 print(f"episodes per arm: {args.episodes} seed: {args.seed}") print(f"budget: {ARMS['full guards']}") print() print(f"{'terminal state':<{width}}" + "".join(f"{a:>16}" for a in ARMS)) for name in order: cells = "".join( f"{summary[a][0][name]:>10} {summary[a][0][name] / args.episodes:>5.1%}" for a in ARMS) print(f"{name:<{width}}{cells}") print() print(f"{'mean steps':<{width}}" + "".join(f"{summary[a][2]:>16.2f}" for a in ARMS)) print(f"{'mean tokens':<{width}}" + "".join(f"{summary[a][1]:>16.0f}" for a in ARMS)) base, guarded = summary["step cap only"][1], summary["full guards"][1] print(f"{'token change':<{width}}" + f"{'':>16}" + f"{(guarded - base) / base:>16.1%}") ans_b = summary["step cap only"][0][ANSWERED] / args.episodes ans_g = summary["full guards"][0][ANSWERED] / args.episodes print(f"{'answered change':<{width}}" + f"{'':>16}" + f"{ans_g - ans_b:>+16.2%}") killed = killed_by_guard["stuck"] + killed_by_guard["healthy"] print(f"\nfull-guard early exits: {killed} " f"({killed_by_guard['stuck']} genuinely stuck, " f"{killed_by_guard['healthy']} false positives = " f"{killed_by_guard['healthy'] / args.episodes:.1%} of all episodes)") # How tight should the no-progress detector be? Sweep it, holding the rest. print("\nno-progress detector, swept (full guards otherwise unchanged)") print(f"{'max_repeats':>12}{'answered':>12}{'no_progress':>13}{'mean tokens':>13}") for mr in (2, 3, 4, 5, 6): b = Budget(max_repeats=mr) t: Counter = Counter() tok = 0 for ep in range(args.episodes): rng = random.Random(f"{args.seed}:{ep}") out = run_episode(SimulatedModel(rng, behaviour), TOOLS, b) t[out.terminal] += 1 tok += out.tokens print(f"{mr:>12}{t[ANSWERED] / args.episodes:>11.1%}" f"{t[NO_PROGRESS] / args.episodes:>13.1%}{tok / args.episodes:>13.0f}") if args.csv: with open(args.csv, "w", newline="", encoding="utf-8") as fh: writer = csv.DictWriter(fh, fieldnames=list(rows[0])) writer.writeheader() writer.writerows(rows) print(f"\nwrote {len(rows)} rows to {args.csv}") if __name__ == "__main__": main()