Unit Testing AI Agents (Deterministic, Cheap, and Actually Useful)

from dataclasses import dataclass
from typing import Any, Dict, List, Protocol


@dataclass(frozen=True)
class Budget:
  max_steps: int = 10
  max_tool_calls: int = 10


class LLM(Protocol):
  def next_action(self, state: Dict[str, Any]) -> Dict[str, Any]: ...


class Tools(Protocol):
  def call(self, name: str, args: Dict[str, Any]) -> Dict[str, Any]: ...


def run_agent(task: str, *, llm: LLM, tools: Tools, budget: Budget) -> Dict[str, Any]:
  trace: List[Dict[str, Any]] = []
  tool_calls = 0
  state: Dict[str, Any] = {"task": task, "notes": []}

  for step in range(budget.max_steps):
      action = llm.next_action(state)
      trace.append({"step": step, "action": action})

      if action.get("type") == "finish":
          return {"output": action.get("answer", ""), "trace": trace, "stop_reason": "finish"}

      if action.get("type") != "tool":
          return {"output": "", "trace": trace, "stop_reason": "invalid_action"}

      tool_calls += 1
      if tool_calls > budget.max_tool_calls:
          return {"output": "", "trace": trace, "stop_reason": "max_tool_calls"}

      obs = tools.call(action["tool"], action.get("args", {}))
      trace.append({"step": step, "observation": obs, "tool": action["tool"]})
      state["notes"].append(obs)

  return {"output": "", "trace": trace, "stop_reason": "max_steps"}


# --- unit test (pytest style) ---
class FakeLLM:
  def __init__(self):
      self.n = 0

  def next_action(self, state):
      self.n += 1
      if self.n == 1:
          return {"type": "tool", "tool": "http.get", "args": {"url": "https://example.com"}}
      return {"type": "finish", "answer": "ok"}


class FakeTools:
  def __init__(self):
      self.calls = []

  def call(self, name, args):
      self.calls.append((name, args))
      return {"ok": True, "status": 200, "body": "hello"}


def test_unit_loop_contract():
  out = run_agent(
      "fetch once and finish",
      llm=FakeLLM(),
      tools=FakeTools(),
      budget=Budget(max_steps=5, max_tool_calls=3),
  )

  assert out["stop_reason"] == "finish"
  assert len(out["trace"]) >= 2

export function runAgent(task, { llm, tools, budget }) {
const trace = [];
let toolCalls = 0;
const state = { task, notes: [] };

for (let step = 0; step < budget.maxSteps; step++) {
  const action = llm.nextAction(state);
  trace.push({ step, action });

  if (action?.type === "finish") {
    return { output: action.answer ?? "", trace, stop_reason: "finish" };
  }

  if (action?.type !== "tool") {
    return { output: "", trace, stop_reason: "invalid_action" };
  }

  toolCalls += 1;
  if (toolCalls > budget.maxToolCalls) {
    return { output: "", trace, stop_reason: "max_tool_calls" };
  }

  const obs = tools.call(action.tool, action.args || {});
  trace.push({ step, tool: action.tool, observation: obs });
  state.notes.push(obs);
}

return { output: "", trace, stop_reason: "max_steps" };
}

// --- unit test (jest style) ---
test("unit loop contract", () => {
const llm = {
  n: 0,
  nextAction() {
    this.n += 1;
    if (this.n === 1) return { type: "tool", tool: "http.get", args: { url: "https://example.com" } };
    return { type: "finish", answer: "ok" };
  },
};
const tools = { calls: [], call(name, args) { this.calls.push([name, args]); return { ok: true, status: 200 }; } };

const out = runAgent("fetch once and finish", {
  llm,
  tools,
  budget: { maxSteps: 5, maxToolCalls: 3 },
});

expect(out.stop_reason).toBe("finish");
expect(tools.calls.length).toBe(1);
});