90 lines
4.3 KiB
Python
90 lines
4.3 KiB
Python
"""CLI.
|
|
python3 -m harness.cli run T01 --agent oracle|null|llm [--model NAME] [--run N] [--no-teardown]
|
|
python3 -m harness.cli teardown runs/<run_dir>
|
|
python3 -m harness.cli cleanup-list Z900001_ runs/cleanup_probe
|
|
"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
|
|
from .agents import LlmAgent, NullAgent, OracleAgent
|
|
from .adt_client import load_env
|
|
from .runner import Runner, delete_uris
|
|
|
|
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
|
|
def next_run_no():
|
|
runs = os.path.join(ROOT, "runs")
|
|
nums = [int(d[:3]) for d in os.listdir(runs) if d[:3].isdigit()]
|
|
return max(nums, default=0) + 1
|
|
|
|
|
|
def main():
|
|
load_env(os.path.join(ROOT, ".env"))
|
|
ap = argparse.ArgumentParser()
|
|
sub = ap.add_subparsers(dest="cmd", required=True)
|
|
r = sub.add_parser("run")
|
|
r.add_argument("task")
|
|
r.add_argument("--agent", default="oracle", choices=["oracle", "null", "llm"])
|
|
r.add_argument("--model")
|
|
r.add_argument("--base-url")
|
|
r.add_argument("--run", type=int)
|
|
r.add_argument("--no-teardown", action="store_true")
|
|
r.add_argument("--max-minutes", type=float, help="optional time limit for the llm agent (default: none)")
|
|
t = sub.add_parser("teardown")
|
|
t.add_argument("run_dir")
|
|
sub.add_parser("teardown-all", help="delete objects of every run that has delete_uris.json and no teardown.json")
|
|
rs = sub.add_parser("rescore", help="score an existing run again (objects must still exist)")
|
|
rs.add_argument("run_dir")
|
|
c = sub.add_parser("cleanup-list", help="write delete_uris.json for all objects with a prefix")
|
|
c.add_argument("prefix")
|
|
c.add_argument("out_dir")
|
|
a = ap.parse_args()
|
|
|
|
if a.cmd == "run":
|
|
agent = {"oracle": OracleAgent, "null": NullAgent}.get(a.agent)
|
|
agent = agent() if agent else LlmAgent(a.model, a.base_url,
|
|
max_seconds=a.max_minutes * 60 if a.max_minutes else None)
|
|
runner = Runner(os.path.join(ROOT, "tasks"), os.path.join(ROOT, "runs"))
|
|
rep, run_dir = runner.run(a.task, agent, a.run or next_run_no(), teardown=not a.no_teardown)
|
|
print(json.dumps({k: rep[k] for k in ("task", "agent", "prefix", "gates", "score",
|
|
"tool_calls", "seconds", "teardown")}, indent=1))
|
|
print("run dir:", run_dir)
|
|
elif a.cmd == "cleanup-list":
|
|
from .mcp_client import McpClient
|
|
from .runner import DELETE_ORDER, Runner as R
|
|
os.makedirs(a.out_dir, exist_ok=True)
|
|
with McpClient() as m:
|
|
objs = R("", "")._objects_with_prefix(m, a.prefix.upper())
|
|
objs.sort(key=lambda o: DELETE_ORDER.index(o["objectType"]) if o["objectType"] in DELETE_ORDER else 99)
|
|
json.dump([o["uri"] for o in objs], open(os.path.join(a.out_dir, "delete_uris.json"), "w"), indent=1)
|
|
print(len(objs), "objects ->", os.path.join(a.out_dir, "delete_uris.json"))
|
|
elif a.cmd == "rescore":
|
|
rd = os.path.abspath(a.run_dir)
|
|
name = os.path.basename(rd)
|
|
run_no, task_id, agent_name = name.split("_", 2)
|
|
runner = Runner(os.path.join(ROOT, "tasks"), os.path.join(ROOT, "runs"))
|
|
agent = NullAgent()
|
|
agent.name = agent_name
|
|
rep, _ = runner.run(task_id, agent, int(run_no), teardown=True, rescore_dir=rd)
|
|
print(json.dumps({k: rep.get(k) for k in ("task", "agent", "gates", "score", "tool_calls",
|
|
"activations", "max_fail_streak", "agent_seconds")}, indent=1))
|
|
elif a.cmd == "teardown-all":
|
|
runs = os.path.join(ROOT, "runs")
|
|
for d in sorted(os.listdir(runs)):
|
|
rd = os.path.join(runs, d)
|
|
if os.path.exists(os.path.join(rd, "delete_uris.json")) and not os.path.exists(os.path.join(rd, "teardown.json")):
|
|
res = delete_uris(json.load(open(os.path.join(rd, "delete_uris.json"))))
|
|
json.dump(res, open(os.path.join(rd, "teardown.json"), "w"), indent=1)
|
|
print(d, sum(v["deleted"] for v in res.values()), "/", len(res), "deleted")
|
|
else:
|
|
uris = json.load(open(os.path.join(a.run_dir, "delete_uris.json")))
|
|
res = delete_uris(uris)
|
|
json.dump(res, open(os.path.join(a.run_dir, "teardown.json"), "w"), indent=1)
|
|
print(json.dumps(res, indent=1))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|