harness.py

#!/usr/bin/env python3
"""ComposeBench substrate harness.

The benchmark holds tools, model, task text, turn cap, and grader fixed while
changing only the action surface:

  repl     - one persistent Python code block per turn.
  batchfc  - one dependent JSON DAG per turn.
  fc       - one JSON tool call per turn with $rN registers.
  cli      - one shell-style tool call per turn with $rN registers.

The harness writes one JSONL row per (task, substrate, rep). It does not include
the original paper result logs; use it to rerun them.
"""
import argparse
import ast
import contextlib
import io
import json
import pathlib
import re
import sys
import types

ROOT = pathlib.Path(__file__).resolve().parent
sys.path.insert(0, str(ROOT))

from llm_client import USAGE, chat as llm_chat, has_credentials  # noqa: E402
import gen_tasks as G  # noqa: E402


def _sandbox_tool(fn):
    def _make(real):
        def _w(*a, **k):
            return real(*a, **k)
        return _w

    w = _make(fn)
    return types.FunctionType(
        w.__code__, {"__builtins__": {}}, fn.__name__, w.__defaults__, w.__closure__
    )


TOOL_FNS = {name: _sandbox_tool(fn) for name, (_sig, _desc, fn) in G.TOOLS.items()}

SAFE_BUILTINS = {
    "abs": abs,
    "all": all,
    "any": any,
    "bool": bool,
    "dict": dict,
    "enumerate": enumerate,
    "float": float,
    "int": int,
    "len": len,
    "list": list,
    "max": max,
    "min": min,
    "print": print,
    "range": range,
    "repr": repr,
    "round": round,
    "set": set,
    "sorted": sorted,
    "str": str,
    "sum": sum,
    "tuple": tuple,
    "zip": zip,
}
FORBIDDEN_NAMES = {
    "__builtins__",
    "__import__",
    "breakpoint",
    "compile",
    "dir",
    "eval",
    "exec",
    "getattr",
    "globals",
    "help",
    "input",
    "locals",
    "open",
    "setattr",
    "vars",
}
FORBIDDEN_ATTRS = {
    "__class__",
    "__closure__",
    "__code__",
    "__dict__",
    "__func__",
    "__globals__",
    "__mro__",
    "__self__",
    "__subclasses__",
}


class UnsafeCode(Exception):
    pass


def _validate_code_ast(tree):
    for node in ast.walk(tree):
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            raise UnsafeCode("imports are disabled")
        if isinstance(node, ast.Name) and node.id in FORBIDDEN_NAMES:
            raise UnsafeCode(f"use of {node.id} is disabled")
        if isinstance(node, ast.Attribute) and (
            node.attr in FORBIDDEN_ATTRS or node.attr.startswith("__")
        ):
            raise UnsafeCode(f"attribute {node.attr} is disabled")
        if isinstance(node, ast.Constant) and isinstance(node.value, str):
            markers = ("__globals__", "__closure__", "__subclasses__")
            if any(marker in node.value for marker in markers):
                raise UnsafeCode("dunder reflection strings are disabled")


def _safe_namespace():
    ns = {"__builtins__": SAFE_BUILTINS}
    ns.update(TOOL_FNS)
    return ns


def _run_python_action(code, ns):
    tree = ast.parse(code)
    _validate_code_ast(tree)
    last_expr = None
    if tree.body and isinstance(tree.body[-1], ast.Expr):
        last_expr = ast.Expression(tree.body[-1].value)
        _validate_code_ast(last_expr)
        tree.body = tree.body[:-1]
    buf = io.StringIO()
    with contextlib.redirect_stdout(buf):
        if tree.body:
            exec(compile(tree, "<composebench-repl>", "exec"), ns)
        if last_expr is not None:
            val = eval(compile(last_expr, "<composebench-repl>", "eval"), ns)
            if val is not None:
                print(repr(val), file=buf)
    return buf.getvalue().strip() or "(no output)"


def chat(model, messages):
    return llm_chat(model, messages)


def tool_catalog_text():
    lines = ["Available tools:"]
    for _name, (sig, desc, _fn) in G.TOOLS.items():
        lines.append(f"  - {sig}  # {desc}")
    return "\n".join(lines)


def _final_answer(text):
    if "Answer:" not in text:
        return None
    seg = text.rsplit("Answer:", 1)[1].strip()
    seg = seg.splitlines()[0].strip() if seg else ""
    if not seg or seg.startswith("<"):
        return None
    return seg


def _check(answer_text, gold):
    m = re.findall(r"-?\d+", str(answer_text).replace(",", ""))
    try:
        return bool(m) and int(m[-1]) == int(gold)
    except Exception:
        return False


def _call_tool(name, args):
    if name not in TOOL_FNS:
        raise KeyError(f"unknown tool {name!r}")
    return TOOL_FNS[name](*args)


def _resolve_registers(value, reg):
    if isinstance(value, str) and re.fullmatch(r"\$r\d+", value):
        return reg[value[1:]]
    if isinstance(value, list):
        return [_resolve_registers(x, reg) for x in value]
    return value


def _resolve_args(args, reg):
    return [_resolve_registers(a, reg) for a in args]


def _json_after(label, text):
    idx = text.find(label)
    if idx < 0:
        return None
    tail = text[idx + len(label):].strip()
    start = tail.find("{")
    if start < 0:
        return None
    try:
        obj, _end = json.JSONDecoder().raw_decode(tail[start:])
        return obj
    except json.JSONDecodeError:
        return None


def _obs(text, noise):
    if not noise:
        return str(text)
    return (
        "status=ok; telemetry={latency_ms:17, shard:'composebench'}; "
        f"payload={text}; note=ignore non-payload fields"
    )


FC_SYS = (
    "Solve the task by calling tools, one per turn, as JSON.\n"
    "{cat}\n\n"
    "Each action must be exactly one JSON object on a line beginning with Action:\n"
    "  Action: {{\"tool\":\"<name>\",\"args\":[<arg>, ...]}}\n"
    "Results are stored as r0, r1, ... . To pass a previous result, use the "
    "string \"$r0\". Do not nest tool calls. When done, output:\n"
    "  Answer: <final integer>\n"
    "Task: {task}"
)


def run_fc(task, model, n_limit, noise=False):
    reg = {}
    msgs = [{"role": "user", "content": FC_SYS.format(
        cat=tool_catalog_text(), task=task["instruction"])}]
    turns = 0
    while turns < n_limit:
        turns += 1
        gen = chat(model, msgs)
        msgs.append({"role": "assistant", "content": gen})
        fin = _final_answer(gen)
        call = _json_after("Action:", gen)
        if fin is not None and call is None:
            return _check(fin, task["answer"]), turns, ""
        if call is None:
            msgs.append({"role": "user", "content": "Output one Action JSON or Answer: <integer>."})
            continue
        try:
            args = _resolve_args(call.get("args", []), reg)
            res = _call_tool(call["tool"], args)
            rid = f"r{len(reg)}"
            reg[rid] = res
            obs = f"{rid} = {json.dumps(res)}"
        except Exception as e:
            obs = f"Error: {e}"
        msgs.append({"role": "user", "content": _obs(obs, noise)[:1800]})
    return False, turns, "turn_limit"


CLI_SYS = (
    "Solve the task by running one command per turn. Commands call tools:\n"
    "  <tool> <arg1> <arg2> ...\n"
    "{cat}\n\n"
    "Results are stored as r0, r1, ... . Reference a prior result with $r0. "
    "List arguments must be JSON-like with no spaces, e.g. [1,2,3] or "
    "[$r0,$r1]. Output exactly one command line beginning with '$ '. When done, "
    "output:\n  Answer: <final integer>\n"
    "Task: {task}"
)


def _parse_cli(line, reg):
    toks = line.strip().split(None, 1)
    if not toks:
        raise ValueError("empty command")
    name = toks[0]
    rest = toks[1] if len(toks) > 1 else ""

    def parse_token(tok):
        if re.fullmatch(r"\$r\d+", tok):
            return reg[tok[1:]]
        if tok.startswith("[") and tok.endswith("]"):
            try:
                loaded = json.loads(tok)
            except json.JSONDecodeError:
                inner = tok[1:-1].strip()
                loaded = [x.strip() for x in inner.split(",") if x.strip()]
            return _resolve_registers(loaded, reg)
        try:
            return int(tok)
        except ValueError:
            return tok.strip("\"'")

    return name, [parse_token(tok) for tok in re.findall(r"\[[^\]]*\]|\S+", rest)]


def run_cli(task, model, n_limit, noise=False):
    reg = {}
    msgs = [{"role": "user", "content": CLI_SYS.format(
        cat=tool_catalog_text(), task=task["instruction"])}]
    turns = 0
    while turns < n_limit:
        turns += 1
        gen = chat(model, msgs)
        msgs.append({"role": "assistant", "content": gen})
        fin = _final_answer(gen)
        m = re.search(r"^\$\s*(.+)$", gen, re.MULTILINE)
        if fin is not None and m is None:
            return _check(fin, task["answer"]), turns, ""
        if m is None:
            msgs.append({"role": "user", "content": "Output one '$ <command>' line or Answer: <integer>."})
            continue
        try:
            name, args = _parse_cli(m.group(1), reg)
            res = _call_tool(name, args)
            rid = f"r{len(reg)}"
            reg[rid] = res
            obs = f"{rid} = {json.dumps(res)}"
        except Exception as e:
            obs = f"Error: {e}"
        msgs.append({"role": "user", "content": _obs(obs, noise)[:1800]})
    return False, turns, "turn_limit"


REPL_SYS = (
    "Solve the task by writing Python that calls the provided tools. The tools "
    "are already defined as Python functions in your namespace.\n"
    "{cat}\n\n"
    "Each turn, write one code block:\n"
    "  Action:\n  ```python\n  <code>\n  ```\n"
    "You may call multiple tools, bind variables, loop, and compose values "
    "within one action. print() what you need to observe. The namespace persists "
    "across turns. When done, output:\n  Answer: <final integer>\n"
    "Task: {task}"
)


def run_repl(task, model, n_limit, noise=False):
    ns = _safe_namespace()
    msgs = [{"role": "user", "content": REPL_SYS.format(
        cat=tool_catalog_text(), task=task["instruction"])}]
    turns = 0
    while turns < n_limit:
        turns += 1
        gen = chat(model, msgs)
        msgs.append({"role": "assistant", "content": gen})
        m = re.search(r"```(?:python)?\s*(.*?)```", gen, re.DOTALL)
        if not m:
            fin = _final_answer(gen)
            if fin is not None:
                return _check(fin, task["answer"]), turns, ""
            msgs.append({"role": "user", "content": "Output an Action code block or Answer: <integer>."})
            continue
        try:
            out = _run_python_action(m.group(1), ns)
        except Exception as e:
            out = f"Error: {e}"
        msgs.append({"role": "user", "content": ("Output:\n" + _obs(out, noise))[:1800]})
    return False, turns, "turn_limit"


BATCHFC_SYS = (
    "Solve the task by submitting one executable function-call DAG per turn.\n"
    "{cat}\n\n"
    "Each action must be one JSON object on a line beginning with Action:\n"
    "  Action: {{\"steps\":[{{\"id\":\"a\",\"tool\":\"read_sensor\",\"args\":[3]}},"
    "{{\"id\":\"b\",\"tool\":\"calibrate\",\"args\":[\"$a\",2]}}],"
    "\"return\":\"$b\"}}\n"
    "Within the same DAG, refer to earlier step ids as \"$a\", \"$b\", etc. "
    "You may also refer to previous-turn result registers as \"$r0\". The DAG is "
    "executed in order and returns one value. When done, output:\n"
    "  Answer: <final integer>\n"
    "Task: {task}"
)


def _resolve_dag_value(value, local, reg):
    if isinstance(value, str) and value.startswith("$"):
        key = value[1:]
        if key in local:
            return local[key]
        if key in reg:
            return reg[key]
        raise KeyError(f"unknown reference {value}")
    if isinstance(value, list):
        return [_resolve_dag_value(x, local, reg) for x in value]
    return value


def _run_dag(steps, return_ref, reg=None):
    reg = reg or {}
    local = {}
    for step in steps:
        sid = str(step["id"])
        if not re.fullmatch(r"[A-Za-z][A-Za-z0-9_]*", sid):
            raise ValueError(f"invalid step id {sid!r}")
        args = [_resolve_dag_value(a, local, reg) for a in step.get("args", [])]
        local[sid] = _call_tool(step["tool"], args)
    return _resolve_dag_value(return_ref, local, reg)


def run_batchfc(task, model, n_limit, noise=False):
    reg = {}
    msgs = [{"role": "user", "content": BATCHFC_SYS.format(
        cat=tool_catalog_text(), task=task["instruction"])}]
    turns = 0
    while turns < n_limit:
        turns += 1
        gen = chat(model, msgs)
        msgs.append({"role": "assistant", "content": gen})
        fin = _final_answer(gen)
        dag = _json_after("Action:", gen)
        if fin is not None and dag is None:
            return _check(fin, task["answer"]), turns, ""
        if dag is None:
            msgs.append({"role": "user", "content": "Output one Action DAG JSON or Answer: <integer>."})
            continue
        try:
            res = _run_dag(dag.get("steps", []), dag.get("return"), reg)
            rid = f"r{len(reg)}"
            reg[rid] = res
            obs = f"{rid} = {json.dumps(res)}"
        except Exception as e:
            obs = f"Error: {e}"
        msgs.append({"role": "user", "content": _obs(obs, noise)[:1800]})
    return False, turns, "turn_limit"


RUNNERS = {
    "repl": run_repl,
    "batchfc": run_batchfc,
    "fc": run_fc,
    "cli": run_cli,
}


def load_tasks(path):
    p = pathlib.Path(path)
    if not p.exists():
        data = G.build()
        p.write_text(json.dumps(data, indent=2) + "\n")
    return json.loads(p.read_text())["tasks"]


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", default="claude-sonnet-4-6")
    ap.add_argument("--substrates", default="repl,batchfc,fc,cli")
    ap.add_argument("--reps", type=int, default=1)
    ap.add_argument("--n_turns_limit", type=int, default=20)
    ap.add_argument("--tasks_file", default=str(ROOT / "tasks.json"))
    ap.add_argument("--out", default=str(ROOT / "results.jsonl"))
    ap.add_argument("--noise", action="store_true")
    ap.add_argument("--kinds", default="chain,fanin")
    ap.add_argument("--min_depth", type=int, default=1)
    ap.add_argument("--max_depth", type=int, default=None)
    ap.add_argument("--widths", default="")
    args = ap.parse_args()

    if not has_credentials(args.model):
        raise SystemExit(
            f"No API credentials configured for {args.model!r}; set LLM_YAML to a private config or use provider env vars."
        )

    substrates = [s.strip() for s in args.substrates.split(",") if s.strip()]
    unknown = [s for s in substrates if s not in RUNNERS]
    if unknown:
        raise SystemExit(f"unknown substrates: {unknown}; known={sorted(RUNNERS)}")

    kinds = {k.strip() for k in args.kinds.split(",") if k.strip()}
    widths = {int(w.strip()) for w in args.widths.split(",") if w.strip()}
    tasks = [
        t for t in load_tasks(args.tasks_file)
        if (
            t["kind"] in kinds
            and (t.get("depth") is None or t["depth"] >= args.min_depth)
            and (args.max_depth is None or t.get("depth") is None or t["depth"] <= args.max_depth)
            and (not widths or t.get("width") in widths)
        )
    ]
    out = pathlib.Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    seen = set()
    if out.exists():
        for line in out.read_text().splitlines():
            if not line.strip():
                continue
            row = json.loads(line)
            seen.add((row["task_id"], row["substrate"], row["rep"]))

    with out.open("a") as f:
        for rep in range(args.reps):
            for task in tasks:
                for substrate in substrates:
                    key = (task["task_id"], substrate, rep)
                    if key in seen:
                        continue
                    before = dict(USAGE)
                    error = ""
                    try:
                        ok, turns, error = RUNNERS[substrate](
                            task, args.model, args.n_turns_limit, args.noise
                        )
                    except Exception as e:
                        ok, turns, error = False, 0, f"runner_error: {e}"
                    after = dict(USAGE)
                    row = {
                        "task_id": task["task_id"],
                        "kind": task["kind"],
                        "depth": task.get("depth"),
                        "width": task.get("width"),
                        "substrate": substrate,
                        "rep": rep,
                        "model": args.model,
                        "is_correct": bool(ok),
                        "n_turns": turns,
                        "turn_limit": args.n_turns_limit,
                        "noise": bool(args.noise),
                        "input_tokens": after["in"] - before["in"],
                        "output_tokens": after["out"] - before["out"],
                        "api_calls": after["calls"] - before["calls"],
                        "error": error,
                    }
                    f.write(json.dumps(row, sort_keys=True) + "\n")
                    f.flush()
                    print(json.dumps(row, sort_keys=True))


if __name__ == "__main__":
    main()