diff --git a/bench/a320_bench/cli.py b/bench/a320_bench/cli.py new file mode 100644 index 0000000..6085519 --- /dev/null +++ b/bench/a320_bench/cli.py @@ -0,0 +1,113 @@ +"""``a320-bench``: run recorded benchmark episodes from the command line. + + a320-bench run --scenario scenarios/elec/apu_gen_fault.yaml \ + --model anthropic/claude-opus-4-8 --runs 3 --out runs/ + +Each run gets a fresh Sim, a fresh benchmark-profile MCP server and its own +JSONL trajectory under ``//``. The command needs the +``[providers]`` extra (litellm); everything else in the package runs without +it. +""" + +import argparse +import asyncio +import json +import sys +from typing import Any + +from a320_bench.episode import run_episode +from a320_bench.scenario import ScenarioError, load_scenario + + +def _positive_int(text: str) -> int: + value = int(text) + if value < 1: + raise argparse.ArgumentTypeError(f"must be >= 1, got {value}") + return value + + +def _sampling_dict(text: str) -> "dict[str, Any]": + """Parse --sampling: must be a JSON object (litellm.completion kwargs).""" + try: + value = json.loads(text) + except json.JSONDecodeError as exc: + raise argparse.ArgumentTypeError(f"not valid JSON: {exc}") from exc + if not isinstance(value, dict): + raise argparse.ArgumentTypeError( + f"must be a JSON object, got {type(value).__name__}" + ) + return value + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="a320-bench", + description="Phase 5 benchmark harness: recorded agent episodes over MCP.", + ) + sub = parser.add_subparsers(dest="command", required=True) + + run = sub.add_parser("run", help="run one scenario against a real model") + run.add_argument("--scenario", required=True, help="path to a scenario YAML") + run.add_argument( + "--model", + required=True, + help="litellm model id, e.g. anthropic/claude-opus-4-8 or gpt-...", + ) + run.add_argument( + "--runs", type=_positive_int, default=1, help="episodes to run (default 1)" + ) + run.add_argument("--out", default="runs", help="output directory (default runs/)") + run.add_argument( + "--sampling", + type=_sampling_dict, + default=None, + help='JSON dict passed to litellm.completion verbatim, e.g. \'{"temperature": 0}\'', + ) + return parser + + +def main(argv: "list[str] | None" = None) -> int: + args = build_parser().parse_args(argv) + + # Imported here, not at module top: the CLI is the only piece that needs + # litellm, and the error message tells the user exactly what to install. + try: + from a320_bench.providers.litellm_adapter import LiteLLMAdapter + except ImportError as exc: + print(f"a320-bench: {exc}", file=sys.stderr) + return 2 + + try: + scenario = load_scenario(args.scenario) + except ScenarioError as exc: + print(f"a320-bench: {exc}", file=sys.stderr) + return 2 + + infra_failures = 0 + for i in range(args.runs): + adapter = LiteLLMAdapter(args.model, sampling=args.sampling) + result = asyncio.run(run_episode(scenario, adapter, args.out)) + if not result.valid: + verdict = "INVALID" + elif result.reason == "provider_error": + verdict = "ERROR" + else: + verdict = "PASS" if result.all_passed else "FAIL" + print( + f"[{i + 1}/{args.runs}] {scenario.id} {verdict} " + f"reason={result.reason} tool_calls={result.tool_calls_used} " + f"sim_t={result.sim_time_end:.1f}s -> {result.trajectory_path}", + file=sys.stderr, + ) + if verdict in ("INVALID", "ERROR"): + infra_failures += 1 + + # Infrastructure problems deserve a red exit code: an invalid scenario + # (the world never manifested the failure) or a provider error (bad key, + # network down) — a paid batch of N broken runs must not end green. An + # agent that failed the procedure is a *result*, not an error. + return 1 if infra_failures else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/bench/a320_bench/episode.py b/bench/a320_bench/episode.py index cf4df49..a9ab6ac 100644 --- a/bench/a320_bench/episode.py +++ b/bench/a320_bench/episode.py @@ -286,6 +286,7 @@ def over_budget() -> "str | None": step=step, text=turn.text, stop_reason=turn.stop_reason, + provider_raw=turn.raw, # finish reason / token usage, adapter-defined ) if not turn.tool_calls: diff --git a/bench/a320_bench/providers/litellm_adapter.py b/bench/a320_bench/providers/litellm_adapter.py new file mode 100644 index 0000000..a62d7cc --- /dev/null +++ b/bench/a320_bench/providers/litellm_adapter.py @@ -0,0 +1,152 @@ +"""Real LLM providers through litellm's OpenAI-format completion API. + +Why litellm (decision in docs/decisiones.md): one adapter covers every +provider the baselines need, at the price of a translation layer — which is +why the version is pinned **exactly** in ``bench/pyproject.toml`` and recorded +in every trajectory's meta. Verified against litellm 1.93.0: +``completion(model, messages, ..., tools, tool_choice)`` and messages whose +``tool_calls[].function`` carry ``name`` + ``arguments`` (a JSON string). + +This module is NOT imported by ``a320_bench.providers`` eagerly: CI runs +without the ``[providers]`` extra, so litellm must stay an opt-in import. +""" + +import json +from importlib import metadata +from typing import Any + +from a320_bench.providers.base import ProviderAdapter, ToolCall, ToolResult, Turn + +try: + import litellm +except ImportError as exc: # pragma: no cover - environment guard + raise ImportError( + "litellm is not installed. The real-provider adapter needs the " + "[providers] extra: pip install -e 'bench/[providers]'" + ) from exc + + +def _mcp_tools_to_openai(tools: list[dict[str, Any]]) -> list[dict[str, Any]]: + """MCP tool schemas map 1:1 onto OpenAI function-calling tools.""" + return [ + { + "type": "function", + "function": { + "name": t["name"], + "description": t["description"], + "parameters": t["inputSchema"], + }, + } + for t in tools + ] + + +class LiteLLMAdapter(ProviderAdapter): + """One conversation with one model through ``litellm.completion``. + + Blocking calls on purpose: the episode is single-threaded around an + `unsendable` Sim and there is nothing to serve while the model thinks. + `sampling` is passed through to completion verbatim and recorded in + `info` — the harness does not choose sampling defaults, the experiment + config does. + """ + + def __init__(self, model: str, *, sampling: "dict[str, Any] | None" = None): + self.model = model + self._sampling = dict(sampling or {}) + self._messages: list[dict[str, Any]] = [] + self._tools: list[dict[str, Any]] = [] + self.info: dict[str, Any] = { + "provider": "litellm", + "model": model, + "sampling": self._sampling, + "litellm_version": metadata.version("litellm"), + } + + def start(self, *, instructions: str, tools: list[dict[str, Any]], user_message: str) -> Turn: + self._tools = _mcp_tools_to_openai(tools) + self._messages = [ + {"role": "system", "content": instructions}, + {"role": "user", "content": user_message}, + ] + return self._complete() + + def next(self, results: list[ToolResult], *, nudge: "str | None" = None) -> Turn: + for result in results: + self._messages.append( + { + "role": "tool", + "tool_call_id": result.call.id, + "content": result.content if not result.is_error + else f"ERROR: {result.content}", + } + ) + if nudge is not None: + self._messages.append({"role": "user", "content": nudge}) + return self._complete() + + def _complete(self) -> Turn: + response = litellm.completion( + model=self.model, + messages=self._messages, + tools=self._tools, + **self._sampling, + ) + choice = response.choices[0] + message = choice.message + + # The assistant message goes back into history in provider format so + # the next completion sees its own tool calls. `tool_calls` is omitted + # (not None) when there are none: the OpenAI passthrough sends the + # message dict verbatim (litellm 1.93.0 + # llms/openai/chat/gpt_transformation.py:451-455), and OpenAI rejects + # an explicit null; the Anthropic path reads it with `.get(...)` + # (prompt_templates/factory.py:2539) so absent and None are equivalent. + assistant_message: dict[str, Any] = { + "role": "assistant", + "content": message.content, + } + if message.tool_calls: + assistant_message["tool_calls"] = [ + { + "id": tc.id, + "type": "function", + "function": { + "name": tc.function.name, + "arguments": tc.function.arguments, + }, + } + for tc in message.tool_calls + ] + self._messages.append(assistant_message) + + calls = [] + malformed: dict[str, str] = {} + for tc in message.tool_calls or []: + try: + args = json.loads(tc.function.arguments) if tc.function.arguments else {} + except json.JSONDecodeError: + args = None + if not isinstance(args, dict): + # Malformed (or non-object) arguments become {}: for tools with + # required params the schema rejects the call and the recorded + # is_error stays the agent's mistake. The original payload is + # preserved in `raw` so the trajectory keeps the evidence for + # the scorer. + malformed[tc.id] = tc.function.arguments + args = {} + calls.append(ToolCall(id=tc.id, name=tc.function.name, args=args)) + + usage = getattr(response, "usage", None) + raw: dict[str, Any] = { + "finish_reason": choice.finish_reason, + "usage": usage.model_dump() if hasattr(usage, "model_dump") else None, + } + if malformed: + raw["malformed_tool_arguments"] = malformed + return Turn( + text=message.content or "", + tool_calls=tuple(calls), + stop_reason=choice.finish_reason or "", + raw=raw, + ) diff --git a/bench/pyproject.toml b/bench/pyproject.toml index 57fcf16..4428abc 100644 --- a/bench/pyproject.toml +++ b/bench/pyproject.toml @@ -21,9 +21,14 @@ dependencies = [ ] [project.optional-dependencies] -# Real LLM providers (slice D). CI installs without this extra: every test -# runs against the ScriptedAdapter, no network, no keys. -providers = [] +# Real LLM providers. CI installs without this extra: every test runs against +# the ScriptedAdapter, no network, no keys. Pinned EXACTLY: litellm is the +# benchmark's provider translation layer and its version is part of a run's +# identity (recorded in every trajectory's meta). +providers = ["litellm==1.93.0"] + +[project.scripts] +a320-bench = "a320_bench.cli:main" [tool.setuptools.packages.find] include = ["a320_bench*"] diff --git a/bench/tests/test_litellm_adapter.py b/bench/tests/test_litellm_adapter.py new file mode 100644 index 0000000..8c18f66 --- /dev/null +++ b/bench/tests/test_litellm_adapter.py @@ -0,0 +1,255 @@ +"""LiteLLMAdapter mapping tests with a mocked litellm.completion (#71). + +No network, no keys: what is under test is the translation — MCP tool schemas +to OpenAI-format tools, provider tool calls to the runner's ToolCall, tool +results and the nudge back into the message history. Skipped entirely when +litellm is not installed (CI runs without the [providers] extra). +""" + +from types import SimpleNamespace + +import pytest + +litellm = pytest.importorskip("litellm", reason="needs the [providers] extra") + +from a320_bench.providers.base import ToolCall, ToolResult # noqa: E402 +from a320_bench.providers.litellm_adapter import LiteLLMAdapter # noqa: E402 + +MCP_TOOLS = [ + { + "name": "advance", + "description": "Advance simulated time.", + "inputSchema": {"type": "object", "properties": {"seconds": {"type": "number"}}}, + } +] + + +def _response(*, content=None, tool_calls=None, finish_reason="tool_use"): + return SimpleNamespace( + choices=[ + SimpleNamespace( + message=SimpleNamespace(content=content, tool_calls=tool_calls), + finish_reason=finish_reason, + ) + ], + usage=None, + ) + + +def _tool_call(id_, name, arguments): + return SimpleNamespace(id=id_, function=SimpleNamespace(name=name, arguments=arguments)) + + +@pytest.fixture +def captured(monkeypatch): + """Mock litellm.completion, capturing every kwargs it was called with.""" + calls = [] + responses = [] + + def fake_completion(**kwargs): + calls.append(kwargs) + return responses.pop(0) + + monkeypatch.setattr(litellm, "completion", fake_completion) + return calls, responses + + +def test_start_maps_schemas_and_messages(captured): + calls, responses = captured + responses.append( + _response(tool_calls=[_tool_call("c1", "advance", '{"seconds": 5}')]) + ) + + adapter = LiteLLMAdapter("some/model") + turn = adapter.start( + instructions="SYSTEM TEXT", tools=MCP_TOOLS, user_message="TASK" + ) + + kwargs = calls[0] + assert kwargs["model"] == "some/model" + assert kwargs["messages"][0] == {"role": "system", "content": "SYSTEM TEXT"} + assert kwargs["messages"][1] == {"role": "user", "content": "TASK"} + tool = kwargs["tools"][0] + assert tool["type"] == "function" + assert tool["function"]["name"] == "advance" + assert tool["function"]["parameters"] == MCP_TOOLS[0]["inputSchema"] + + assert turn.tool_calls == (ToolCall(id="c1", name="advance", args={"seconds": 5}),) + assert turn.stop_reason == "tool_use" + + +def test_next_feeds_results_and_nudge_into_history(captured): + calls, responses = captured + responses.append(_response(tool_calls=[_tool_call("c1", "advance", "{}")])) + responses.append(_response(content="done", tool_calls=None, finish_reason="stop")) + + adapter = LiteLLMAdapter("some/model") + turn1 = adapter.start(instructions="S", tools=MCP_TOOLS, user_message="U") + adapter.next( + [ToolResult(call=turn1.tool_calls[0], content="t=5.0s", is_error=False)], + nudge="NUDGE TEXT", + ) + + history = calls[1]["messages"] + # assistant turn with its tool call went back in provider format + assistant = history[2] + assert assistant["role"] == "assistant" + assert assistant["tool_calls"][0]["function"]["name"] == "advance" + # tool result tied to the call id + assert history[3] == {"role": "tool", "tool_call_id": "c1", "content": "t=5.0s"} + # the nudge is a user message, exactly as recorded in the trajectory + assert history[4] == {"role": "user", "content": "NUDGE TEXT"} + + +def test_error_results_are_marked_for_the_model(captured): + calls, responses = captured + responses.append(_response(tool_calls=[_tool_call("c1", "advance", "{}")])) + responses.append(_response(content="ok", finish_reason="stop")) + + adapter = LiteLLMAdapter("some/model") + turn1 = adapter.start(instructions="S", tools=MCP_TOOLS, user_message="U") + adapter.next( + [ToolResult(call=turn1.tool_calls[0], content="seconds must be positive", is_error=True)] + ) + + tool_msg = calls[1]["messages"][3] + assert tool_msg["content"] == "ERROR: seconds must be positive" + + +def test_malformed_arguments_become_empty_args(captured): + """Bad JSON from the model turns into {} — the tool schema rejects it and + the refusal is recorded as the agent's error, not a harness crash. The + original payload survives in raw so the trajectory keeps the evidence.""" + calls, responses = captured + responses.append( + _response(tool_calls=[_tool_call("c1", "advance", '{"seconds": NOT JSON')]) + ) + + adapter = LiteLLMAdapter("some/model") + turn = adapter.start(instructions="S", tools=MCP_TOOLS, user_message="U") + assert turn.tool_calls[0].args == {} + assert turn.raw["malformed_tool_arguments"] == {"c1": '{"seconds": NOT JSON'} + + +def test_non_object_arguments_are_treated_as_malformed(captured): + """Valid JSON that is not an object (e.g. a bare list) cannot be tool args.""" + calls, responses = captured + responses.append(_response(tool_calls=[_tool_call("c1", "advance", "[1, 2]")])) + + adapter = LiteLLMAdapter("some/model") + turn = adapter.start(instructions="S", tools=MCP_TOOLS, user_message="U") + assert turn.tool_calls[0].args == {} + assert turn.raw["malformed_tool_arguments"] == {"c1": "[1, 2]"} + + +def test_two_tool_calls_in_one_turn_keep_order_and_history(captured): + """A parallel-call turn: both calls surface in order, and the *next* + request's history carries the assistant turn (content and both calls) + followed by one tool message per result, tied by id.""" + calls, responses = captured + responses.append( + _response( + content="doing both", + tool_calls=[ + _tool_call("c1", "advance", '{"seconds": 5}'), + _tool_call("c2", "read_ecam", "{}"), + ], + ) + ) + responses.append(_response(content="done", finish_reason="stop")) + + adapter = LiteLLMAdapter("some/model") + turn = adapter.start(instructions="S", tools=MCP_TOOLS, user_message="U") + assert [c.id for c in turn.tool_calls] == ["c1", "c2"] + assert turn.text == "doing both" + + adapter.next( + [ + ToolResult(call=turn.tool_calls[0], content="t=5.0s", is_error=False), + ToolResult(call=turn.tool_calls[1], content="[]", is_error=False), + ] + ) + history = calls[1]["messages"] + assistant = history[2] + assert assistant["content"] == "doing both" + assert [tc["id"] for tc in assistant["tool_calls"]] == ["c1", "c2"] + assert history[3] == {"role": "tool", "tool_call_id": "c1", "content": "t=5.0s"} + assert history[4] == {"role": "tool", "tool_call_id": "c2", "content": "[]"} + + +def test_turn_without_tool_calls_omits_the_key_in_history(captured): + """OpenAI rejects an explicit tool_calls null on assistant messages, so a + no-call turn must go back into history without the key at all.""" + calls, responses = captured + responses.append(_response(content="thinking out loud", finish_reason="stop")) + responses.append(_response(content="done", finish_reason="stop")) + + adapter = LiteLLMAdapter("some/model") + adapter.start(instructions="S", tools=MCP_TOOLS, user_message="U") + adapter.next([], nudge="NUDGE") + + assistant = calls[1]["messages"][2] + assert assistant == {"role": "assistant", "content": "thinking out loud"} + assert "tool_calls" not in assistant + + +def test_sampling_is_forwarded_to_completion(captured): + calls, responses = captured + responses.append(_response(content="ok", finish_reason="stop")) + + adapter = LiteLLMAdapter("some/model", sampling={"temperature": 0, "seed": 7}) + adapter.start(instructions="S", tools=MCP_TOOLS, user_message="U") + + assert calls[0]["temperature"] == 0 + assert calls[0]["seed"] == 7 + + +def test_usage_lands_in_raw(captured): + calls, responses = captured + + class FakeUsage: + def model_dump(self): + return {"prompt_tokens": 11, "completion_tokens": 3} + + response = _response(content="ok", finish_reason="stop") + response.usage = FakeUsage() + responses.append(response) + + adapter = LiteLLMAdapter("some/model") + turn = adapter.start(instructions="S", tools=MCP_TOOLS, user_message="U") + assert turn.raw["finish_reason"] == "stop" + assert turn.raw["usage"] == {"prompt_tokens": 11, "completion_tokens": 3} + + +def test_info_records_the_translation_layer_version(captured): + adapter = LiteLLMAdapter("some/model", sampling={"temperature": 0}) + assert adapter.info["provider"] == "litellm" + assert adapter.info["sampling"] == {"temperature": 0} + assert adapter.info["litellm_version"], "the pinned version must be recorded" + + +def test_cli_parser_shape(): + from a320_bench.cli import build_parser + + args = build_parser().parse_args( + ["run", "--scenario", "s.yaml", "--model", "m", "--runs", "3", "--sampling", '{"temperature": 0}'] + ) + assert args.command == "run" + assert args.runs == 3 + assert args.sampling == {"temperature": 0} + + +@pytest.mark.parametrize( + "argv", + [ + ["run", "--scenario", "s.yaml", "--model", "m", "--sampling", "{not json"], + ["run", "--scenario", "s.yaml", "--model", "m", "--sampling", "[1, 2]"], + ["run", "--scenario", "s.yaml", "--model", "m", "--runs", "0"], + ], +) +def test_cli_rejects_bad_arguments(argv): + from a320_bench.cli import build_parser + + with pytest.raises(SystemExit) as excinfo: + build_parser().parse_args(argv) + assert excinfo.value.code == 2