Agent-System-Dev/.task/genui-u1/run_trial.py

#!/usr/bin/env python3
"""U1 试验运行器（原型）：B=Markdown 基线，G=UINode 生成（含一次修复），G-hs=轻量档子集。

用法：python3 run_trial.py --arm G [--limit N] [--only rc-01,rc-02]
密钥经环境 MINIMAX_API_KEY 注入；预算护栏 >¥200 或单臂调用超限时中止。
"""
import argparse
import json
import os
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from harness import (BASELINE_PROMPT, CATALOG_PROMPT, build_user_message,
                     call_llm, extract_json, validate_spec)

HERE = os.path.dirname(os.path.abspath(__file__))
RESULTS = os.path.join(HERE, "results")
# E5 任务卡记录的 MiniMax 档位价（$/M tokens：入/出），汇率按 7.2 折 CNY；保守不计缓存折扣。
PRICE = {"MiniMax-M3": (0.3, 1.2), "MiniMax-M2.7-highspeed": (0.6, 2.4)}
BUDGET_CNY, ARM_CALL_CAPS = 200.0, {"B": 60, "G": 60, "G-hs": 20}


def spent_cny():
    total = 0.0
    for arm in os.listdir(RESULTS) if os.path.isdir(RESULTS) else []:
        d = os.path.join(RESULTS, arm)
        if not os.path.isdir(d):
            continue
        for fn in os.listdir(d):
            if not fn.endswith(".json"):
                continue
            try:
                rec = json.load(open(os.path.join(d, fn)))
            except Exception:
                continue
            for c in rec.get("calls", []):
                u = c.get("usage", {})
                pin, pout = PRICE.get(rec.get("model", ""), (0, 0))
                total += (u.get("prompt_tokens", 0) / 1e6 * pin + u.get("completion_tokens", 0) / 1e6 * pout) * 7.2
    return total


def call_count(arm):
    d = os.path.join(RESULTS, arm)
    if not os.path.isdir(d):
        return 0
    return sum(json.load(open(os.path.join(d, fn))).get("callCount", 1)
               for fn in os.listdir(d) if fn.endswith(".json"))


def run_task(arm, model, task):
    user = build_user_message(task)
    sids = {m["sourceId"] for m in task["materials"]}
    rec = {"task": task["id"], "family": task["family"], "arm": arm, "model": model,
           "goalKind": task["expect"]["goalKind"], "calls": [], "callCount": 0}

    def call(system, user, max_tokens):
        r = call_llm(model, system, user, max_tokens)
        rec["calls"].append({"usage": r["usage"], "latency_s": r["latency_s"]})
        rec["callCount"] += 1
        return r

    if arm == "B":
        r = call(BASELINE_PROMPT, user, 2500)
        rec["markdown"] = r["content"]
        rec["ok"] = bool(r["content"].strip())
        return rec

    # G / G-hs：生成 + 至多一次修复。max_tokens 需容纳内联 <think> + 完整 JSON（6000 会在
    # M3 上被思考段挤占导致截断，见 results/ 首轮 rc-03..08 失败样本）。
    r1 = call(CATALOG_PROMPT, user, 14000)
    rec["raw_first"] = r1["content"]
    errors, stats, spec = [], None, None
    try:
        spec = extract_json(r1["content"])
        errors, stats = validate_spec(spec, sids)
    except Exception as e:
        errors = [f"JSON 解析失败: {e}"]
    rec["valid_first"] = not errors
    rec["errors_first"] = errors[:15]
    rec["stats_first"] = {"nodes": stats["nodes"], "maxDepth": stats["maxDepth"],
                          "types": stats["types"], "refsUsed": sorted(stats["refsUsed"]),
                          "unknownRefs": sorted(stats["unknownRefs"])} if stats else None
    if errors:
        fix_user = user + "\n\n你上一次的输出被宿主拒绝，错误如下：\n" + \
            "\n".join(f"- {e}" for e in errors[:15]) + \
            "\n\n请修正后重新输出完整 JSON 对象（仍只输出 JSON）。"
        r2 = call(CATALOG_PROMPT, fix_user, 14000)
        rec["raw_repair"] = r2["content"]
        try:
            spec = extract_json(r2["content"])
            errors, stats = validate_spec(spec, sids)
        except Exception as e:
            errors = [f"JSON 解析失败: {e}"]
        rec["repaired"] = True
    else:
        rec["repaired"] = False
    rec["valid_final"] = not errors
    rec["errors_final"] = errors[:15]
    rec["spec"] = spec if not errors else None
    if stats and not errors:
        rec["stats_final"] = {"nodes": stats["nodes"], "maxDepth": stats["maxDepth"],
                              "types": stats["types"], "refsUsed": sorted(stats["refsUsed"]),
                              "idsProvided": len(stats["idsSeen"])}
    return rec


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--arm", required=True, choices=list(ARM_CALL_CAPS))
    ap.add_argument("--limit", type=int, default=0)
    ap.add_argument("--only", default="")
    args = ap.parse_args()

    tasks = json.load(open(os.path.join(HERE, "fixtures.json")))["tasks"]
    if args.only:
        keep = set(args.only.split(","))
        tasks = [t for t in tasks if t["id"] in keep]
    if args.arm == "G-hs":
        tasks = [t for t in tasks if t["family"] != "research-compare" or t.get("sharedGroup")][:10]
    if args.limit:
        tasks = tasks[:args.limit]

    model = "MiniMax-M2.7-highspeed" if args.arm == "G-hs" else "MiniMax-M3"
    outdir = os.path.join(RESULTS, args.arm)
    os.makedirs(outdir, exist_ok=True)

    for t in tasks:
        out = os.path.join(outdir, t["id"] + ".json")
        if os.path.exists(out):
            old = json.load(open(out))
            if old.get("arm") == args.arm and old.get("model") == model and (
                    old.get("ok") or old.get("valid_final")):
                print(f"[skip] {t['id']} (已完成)")
                continue
        if call_count(args.arm) >= ARM_CALL_CAPS[args.arm]:
            print(f"[stop] 臂 {args.arm} 达调用上限 {ARM_CALL_CAPS[args.arm]}")
            break
        spent, est = spent_cny(), 0.35  # 单次最坏估价（¥）先行校验
        if spent + est > BUDGET_CNY:
            print(f"[stop] 预算护栏：已花 ¥{spent:.2f} + 单次估价 ¥{est:.2f} 将超 ¥{BUDGET_CNY}")
            break
        t0 = time.time()
        rec = run_task(args.arm, model, t)
        rec["wall_s"] = round(time.time() - t0, 1)
        json.dump(rec, open(out, "w"), ensure_ascii=False, indent=1)
        flag = "OK " if (rec.get("ok") or rec.get("valid_final")) else "ERR"
        print(f"[{flag}] {t['id']} calls={rec['callCount']} wall={rec['wall_s']}s "
              f"spent≈¥{spent_cny():.2f}")
    print(f"[done] arm={args.arm} calls={call_count(args.arm)}/{ARM_CALL_CAPS[args.arm]} "
          f"spent≈¥{spent_cny():.2f}")


if __name__ == "__main__":
    main()