agent/agent_iconclass.py

#!/usr/bin/env python3
# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "smolagents[openai]",
#     "openai",
#     "datasets",
#     "huggingface_hub",
#     "iconclass",
#     "sentence-transformers",
#     "numpy",
#     "pillow",
# ]
# ///
"""Interactive Iconclass tool-agent (NEXT_SESSION §2) — the real multi-step agent.

GATE_PRECISION.md showed the anchored-fusion gate is stuck at H-F1 48.5 because a
*single-shot* discriminator can't separate the right leaf from its siblings. The
agent DE-RISK was GREEN: holistic reasoning recovered 27/51 (53%) of the GT codes
the gate missed. This script builds the REAL agent that adds the two untested wins:

  (a) LIVE FT-retriever search  — the agent issues its own search_iconclass queries
      (the de-risk only reasoned over the cached K=24 pool; ~48% of GT isn't in it).
  (b) a VERIFICATION step       — the agent must split its proposals into ACCEPTED
      and REJECTED (with a reason each), which both fixes the precision tax AND
      yields model-grade hard negatives (the rejected proposals).

final_codes = anchor (the precise FT-4B) ∪ agent-accepted extras  (anchored fusion).

Trace capture is first-class: one JSONL row per image (incl the rejected proposals)
→ hf://buckets/davanstrien/iconclass-agent-traces, for eventual SFT distillation.

Eval on the clean 788 ruler over the SAME ~40-image de-risk split (25 HARD where the
gate missed a recoverable in-pool GT + 15 RANDOM control); reports anchor / gate /
agent H-F1 and "GT the gate missed that the agent recovered", like agent_derisk.py.

USAGE
    # probe one HARD image verbosely (pick agent type: toolcalling vs code)
    uv run agent_iconclass.py --endpoint-url https://XXXX.endpoints.huggingface.cloud --probe
    # full 40-image validation
    uv run agent_iconclass.py --endpoint-url https://XXXX.endpoints.huggingface.cloud \
        --n-hard 25 --n-random 15 --out agent_traces.jsonl --results agent_iconclass_results.json
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import threading
import time
import uuid
from concurrent.futures import ThreadPoolExecutor, as_completed

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

import agent_tools  # noqa: E402  (our improved tools)
from describe_retrieve import (  # noqa: E402
    DATASET,
    SPLIT,
    VLM_EXTRA_BODY,
    _calculate_hierarchical_f1,
    _code_description,
    _extract_json,
    first_image,
)

DEFAULT_MODEL_ID = "Qwen/Qwen3.6-35B-A3B-FP8"

# Generic-pose pollution down-weight (verbatim from agent_derisk.py / fuse_rank.py)
POSE_PREFIXES = (
    "31A231",
    "33A81",
    "31A2313",
    "31A235",
    "31A2351",
    "31A2623",
    "31A2711",
    "31A27113",
    "31B621",
    "31A233",
    "43C214",
)
POSE_PENALTY = 0.3


def base(c):
    return c.split("(")[0].strip() if isinstance(c, str) else ""


def is_pose(c):
    b = base(c)
    return any(b == p or b.startswith(p) for p in POSE_PREFIXES)


def adj(c, s):
    return s * POSE_PENALTY if is_pose(c) else s


# ===================================================================== data
def load_inputs(pool_path, ft_path, grade_path):
    """Mirror agent_derisk.py: FT pool + 4B anchor + cached 35B judge grades."""
    d = json.load(open(pool_path))
    ft = json.load(open(ft_path))
    ftc = {
        int(r["index"]): list(r.get("predicted_codes", []) or []) for r in ft["results"]
    }
    gc = json.load(open(grade_path))
    recs = {int(r["index"]): r for r in d["records"] if not r.get("error")}
    return recs, ftc, gc


def missed(recs, idx, ftc):
    """Retrieved candidates NOT already in the anchor, with pose-adjusted sim."""
    r = recs[idx]
    ab = set(ftc.get(idx, []))
    return [
        (c, adj(c, r["retrieved_scores"].get(c, 0.5)))
        for c in r["retrieved_codes"]
        if c not in ab
    ]


def gate_pred(recs, idx, ftc, gc, n=3):
    """The deployed gate: anchor + top-N missed codes with judge grade >= 4."""
    ms = [(c, s) for c, s in missed(recs, idx, ftc) if gc.get(f"{idx}|{c}", 0) >= 4]
    ms.sort(key=lambda x: (gc.get(f"{idx}|{x[0]}", 0) / 5) * x[1], reverse=True)
    return ftc.get(idx, []) + [c for c, _ in ms[:n]]


def recoverable_missed_gt(recs, idx, ftc):
    """GT codes present (by base) in the missed-candidate pool — the reachable ones."""
    gb = {base(g) for g in recs[idx]["gt_codes"]}
    return [c for c, _ in missed(recs, idx, ftc) if base(c) in gb]


def stride(lst, k):
    if len(lst) <= k:
        return lst
    step = len(lst) / k
    return [lst[int(i * step)] for i in range(k)]


def select_split(recs, ftc, gc, n_hard, n_random):
    """HARD = gate missed a recoverable GT; RANDOM = control. Deterministic."""
    hard, rand_pool = [], []
    for idx in sorted(recs):
        gpred = set(gate_pred(recs, idx, ftc, gc))
        rec_gt = recoverable_missed_gt(recs, idx, ftc)
        if rec_gt and any(c not in gpred for c in rec_gt):
            hard.append(idx)
        rand_pool.append(idx)
    hard_sel = stride(hard, n_hard)
    rand_sel = [i for i in stride(rand_pool, n_random * 2) if i not in set(hard_sel)][
        :n_random
    ]
    return hard, sorted(set(hard_sel)), sorted(set(rand_sel) - set(hard_sel))


# ===================================================================== image
def prep_image(example, max_size=768):
    from PIL import Image

    img = first_image(example)
    if img is None:
        return None
    if max(img.size) > max_size:
        r = max_size / max(img.size)
        img = img.resize((int(img.size[0] * r), int(img.size[1] * r)), Image.LANCZOS)
    return img.convert("RGB") if img.mode != "RGB" else img


# ===================================================================== prompt
TASK_TEMPLATE = """You are an expert Iconclass cataloguer examining ONE artwork (the image shown to you).

A specialist model has ALREADY assigned these codes — treat them as correct and depicted (do NOT repeat them):
{anchor}

Some additional candidate codes were retrieved by similarity. You may accept, reject, or ignore them, and you should ALSO search for more codes yourself:
{pool}

YOUR JOB: find the ADDITIONAL Iconclass codes that are GENUINELY depicted in THIS artwork, and REJECT plausible-but-absent ones.

METHOD:
1. Look at the image. Identify the actual scene/story, the figures, objects, attributes, gestures.
2. Use search_iconclass(query) to find codes for what you see — especially SPECIFIC leaves (the exact saint, the exact pose, the exact object), not just general categories.
3. Use lookup_iconclass_code(code) to read a candidate's definition and its SIBLINGS, and compare against the image to pick the RIGHT leaf (e.g. head-left vs head-right).
4. Be STRICT. Only commit codes truly depicted in THIS image. Reject the wrong saint / wrong story, generic poses that are not central, and anything you cannot verify against the image.

Keep tool use focused (a handful of searches/lookups). When done, call final_answer with a JSON object EXACTLY of this shape:
{{"accepted": ["code", ...], "rejected": [{{"code": "code", "reason": "one short reason it is NOT depicted"}}, ...]}}
- accepted: the ADDITIONAL depicted codes (exclude the already-assigned ones). [] if none.
- rejected: candidates you considered but rejected, each with a one-line reason."""


def build_task(anchor_codes, pool):
    anchor_txt = (
        "\n".join(f"  - {c} = {_code_description(c) or c}" for c in anchor_codes)
        or "  (none)"
    )
    if pool:
        pool_txt = "\n".join(
            f"  {i + 1}. {p['code']} = {p['definition']} (sim {p['sim']}, judge {p['judge_grade']})"
            for i, p in enumerate(pool)
        )
    else:
        pool_txt = "  (none retrieved — search for codes yourself)"
    return TASK_TEMPLATE.format(anchor=anchor_txt, pool=pool_txt)


# ===================================================================== agent
def make_model(args):
    from smolagents import OpenAIModel

    return OpenAIModel(
        model_id=args.model_id,
        api_base=args.endpoint_url.rstrip("/") + "/v1",
        api_key=args.api_key,
        temperature=0.0,
        max_tokens=args.max_tokens,
        extra_body=VLM_EXTRA_BODY,
    )


def make_agent(model, agent_type, max_steps, verbosity=0):
    from smolagents import CodeAgent, ToolCallingAgent, tool

    tools = [
        tool(agent_tools.search_iconclass),
        tool(agent_tools.lookup_iconclass_code),
        tool(agent_tools.get_parent_code),
        tool(agent_tools.define_codes),
        tool(agent_tools.validate_iconclass_code),
    ]
    cls = CodeAgent if agent_type == "code" else ToolCallingAgent
    kwargs = dict(
        tools=tools, model=model, max_steps=max_steps, verbosity_level=verbosity
    )
    if agent_type == "code":
        kwargs["additional_authorized_imports"] = ["json"]
    return cls(**kwargs)


def _coerce_obj(obj):
    """Coerce a final-answer payload to a dict/list. Handles JSON strings AND
    Python-repr strings (single quotes) the model sometimes emits."""
    if isinstance(obj, (dict, list)):
        return obj
    if not isinstance(obj, str):
        return {}
    j = _extract_json(obj)
    if isinstance(j, (dict, list)):
        return j
    import ast
    import re as _re

    m = _re.search(r"[{\[].*[}\]]", obj, _re.S)
    if m:
        for parser in (json.loads, ast.literal_eval):
            try:
                return parser(m.group(0))
            except Exception:
                continue
    return {}


def parse_final(result):
    """Robustly pull (accepted, rejected[{code,reason}]) from the final answer."""
    obj = _coerce_obj(result)
    # the model occasionally nests the dict as a string inside {"answer": "..."}
    if isinstance(obj, dict) and "accepted" not in obj and "rejected" not in obj:
        for v in obj.values():
            cv = _coerce_obj(v)
            if isinstance(cv, dict) and ("accepted" in cv or "rejected" in cv):
                obj = cv
                break
    accepted, rejected = [], []
    if isinstance(obj, dict):
        accepted = [
            c for c in (obj.get("accepted") or []) if isinstance(c, str) and c.strip()
        ]
        for r in obj.get("rejected") or []:
            if isinstance(r, dict) and r.get("code"):
                rejected.append(
                    {"code": r["code"], "reason": str(r.get("reason", ""))[:200]}
                )
            elif isinstance(r, str) and r.strip():
                rejected.append({"code": r, "reason": ""})
    elif isinstance(obj, list):
        accepted = [c for c in obj if isinstance(c, str) and c.strip()]
    return accepted, rejected


def extract_tool_calls(agent):
    """Compact trajectory from agent memory — tool calls + observations, NO images."""
    out = []
    try:
        from smolagents import ActionStep
    except Exception:
        ActionStep = None
    for step in getattr(agent.memory, "steps", []):
        if ActionStep is not None and not isinstance(step, ActionStep):
            continue
        tcs = getattr(step, "tool_calls", None) or []
        calls = []
        for tc in tcs:
            calls.append(
                {
                    "tool": getattr(tc, "name", None),
                    "args": getattr(tc, "arguments", None),
                }
            )
        obs = getattr(step, "observations", None)
        mo = getattr(step, "model_output", None)
        out.append(
            {
                "step": getattr(step, "step_number", None),
                "model_output": (mo[:600] if isinstance(mo, str) else None),
                "tool_calls": calls,
                "observations": (obs[:800] if isinstance(obs, str) else None),
                "error": str(getattr(step, "error", None))
                if getattr(step, "error", None)
                else None,
            }
        )
    return out


def run_one(args, model, idx, example, anchor, pool, gt):
    """Run the agent on one image; return the full trace row."""
    img = prep_image(example)
    row = {
        "index": idx,
        "run_id": args.run_id,
        "ts": args.ts,
        "model": args.model_id,
        "gt_codes": gt,
        "anchor_codes": anchor,
        "candidate_pool": pool,
        "tool_calls": [],
        "proposed": [],
        "accepted": [],
        "rejected": [],
        "final_codes": list(anchor),
        "raw_response": "",
        "eval": {},
    }
    if img is None:
        row["error"] = "no image"
        return row

    agent = make_agent(model, args.agent_type, args.max_steps, args.verbosity)
    task = build_task(anchor, pool)
    try:
        result = agent.run(task, images=[img])
    except Exception as e:
        row["error"] = f"agent: {str(e)[:200]}"
        row["tool_calls"] = extract_tool_calls(agent)
        return row

    accepted, rejected = parse_final(result)
    # keep only valid, non-anchor codes for the committed extras
    anchor_set = set(anchor)
    valid_accepted = [
        c
        for c in accepted
        if c not in anchor_set and agent_tools.validate_iconclass_code(c)
    ]
    proposed = list(dict.fromkeys(accepted + [r["code"] for r in rejected]))
    # Defensive cap: a degenerate enumeration (runaway code-family walk — e.g.
    # 100+ codes from one fake subtree, seen when the server's generation_config
    # overrode temp=0 -> 1.0) is not real labeling. 15 is well above any sane
    # iconclass count (anchor<=6 + a few extras); clear + flag for diagnosis.
    DEGEN_CAP = 15
    degenerate = len(valid_accepted) > DEGEN_CAP
    if degenerate:
        print(
            f"    [idx={idx}] DEGENERATE: {len(valid_accepted)} accepted (> {DEGEN_CAP}) "
            f"— clearing (likely temp-override code-family enumeration).",
            flush=True,
        )
        valid_accepted = []
    final_codes = list(anchor) + [c for c in valid_accepted if c not in anchor_set]

    row["tool_calls"] = extract_tool_calls(agent)
    row["proposed"] = proposed
    row["accepted"] = valid_accepted
    row["rejected"] = rejected
    row["final_codes"] = final_codes
    row["degenerate"] = degenerate
    row["raw_response"] = (str(result) if not isinstance(result, str) else result)[
        :4000
    ]
    return row


# ===================================================================== eval
def evaluate_row(row, recs, ftc, gc):
    idx = row["index"]
    gt = row["gt_codes"]
    anchor = row["anchor_codes"]
    agent_pred = row["final_codes"]
    gpred = gate_pred(recs, idx, ftc, gc)
    rec_gt = recoverable_missed_gt(recs, idx, ftc)
    gate_set = set(gpred)
    gt_missed_by_gate = [c for c in rec_gt if c not in gate_set]
    agent_set = set(agent_pred)
    gt_recovered_by_agent = [c for c in gt_missed_by_gate if c in agent_set]
    row["eval"] = {
        "anchor_hf1": _calculate_hierarchical_f1(anchor, gt),
        "gate_hf1": _calculate_hierarchical_f1(gpred, gt),
        "agent_hf1": _calculate_hierarchical_f1(agent_pred, gt),
        "oracle_hf1": _calculate_hierarchical_f1(anchor + rec_gt, gt),
        "n_gt_missed_by_gate": len(gt_missed_by_gate),
        "n_gt_recovered_by_agent": len(gt_recovered_by_agent),
        "recovered_missed_gt": gt_recovered_by_agent,
        "n_agent_extras": len(row["accepted"]),
        "n_rejected": len(row["rejected"]),
    }
    return row


def summarize(rows, hard_set):
    def avg(rs, k):
        return sum(r["eval"][k] for r in rs) / len(rs) if rs else 0.0

    for label, subset in (
        ("HARD", [r for r in rows if r["index"] in hard_set]),
        ("RANDOM", [r for r in rows if r["index"] not in hard_set]),
        ("ALL", rows),
    ):
        subset = [r for r in subset if r.get("eval")]
        if not subset:
            continue
        print(f"\n=== {label} (n={len(subset)}) ===")
        print(f"  anchor-only H-F1 : {avg(subset, 'anchor_hf1'):.1%}")
        print(
            f"  GATE H-F1        : {avg(subset, 'gate_hf1'):.1%}   (deployed; 48.5 on full 788)"
        )
        print(
            f"  AGENT H-F1       : {avg(subset, 'agent_hf1'):.1%}   <-- agent vs gate"
        )
        print(
            f"  oracle H-F1      : {avg(subset, 'oracle_hf1'):.1%}   (anchor + all recoverable GT)"
        )
        tot_missed = sum(r["eval"]["n_gt_missed_by_gate"] for r in subset)
        tot_rec = sum(r["eval"]["n_gt_recovered_by_agent"] for r in subset)
        print(
            f"  GT the gate MISSED that the AGENT RECOVERED: {tot_rec}/{tot_missed}  (de-risk: 27/51)"
        )
        print(
            f"  avg extras committed: {avg(subset, 'n_agent_extras'):.1f}   avg rejected: {avg(subset, 'n_rejected'):.1f}"
        )


# ===================================================================== lazy stream
class _LazyImages:
    """Stream an image dataset once, exposing examples by DENSE index.

    Matches infer_anchor.py's dense indexing (only image-bearing rows consume a
    dense index). Decodes images on demand under a lock so the agent starts on
    image 0 immediately and memory stays bounded by the in-flight window (~the
    executor's concurrency), NOT N — fixing the upfront-materialization cliff
    that wedged large ``--limit`` runs (decoding thousands of full-res PILs
    before any agent work). Relies on the ThreadPoolExecutor dispatching tasks
    in submission (index) order so the lookahead stays ~concurrency.
    """

    def __init__(self, dataset_name, split_name):
        from datasets import load_dataset

        self._it = iter(load_dataset(dataset_name, split=split_name, streaming=True))
        self._lock = threading.Lock()
        self._buf = {}  # dense_idx -> example (pending consumers)
        self._next_dense = 0
        self._exhausted = False

    def get(self, dense_idx):
        """Return the example for dense_idx (popping it), or None if past end."""
        with self._lock:
            if dense_idx in self._buf:
                return self._buf.pop(dense_idx)
            while not self._exhausted:
                try:
                    ex = next(self._it)
                except StopIteration:
                    self._exhausted = True
                    break
                if first_image(ex) is None:
                    continue  # no image: doesn't consume a dense index (matches infer_anchor)
                if self._next_dense == dense_idx:
                    self._next_dense += 1
                    return ex  # hot path: requested next-in-line, don't buffer
                self._buf[self._next_dense] = ex
                self._next_dense += 1
            return self._buf.pop(dense_idx, None)


# ===================================================================== main
def main():
    sys.stdout.reconfigure(line_buffering=True)
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument(
        "--endpoint-url", required=True, help="base URL of the deployed endpoint"
    )
    ap.add_argument("--model-id", default=DEFAULT_MODEL_ID)
    ap.add_argument("--api-key", default=None, help="defaults to the cached HF token")
    ap.add_argument(
        "--agent-type", default="toolcalling", choices=["toolcalling", "code"]
    )
    ap.add_argument("--pool", default="dr_test788_bgebase.json")
    ap.add_argument("--ft", default="sftbrill_test788.json")
    ap.add_argument("--grade-cache", default="grade_cache.json")
    ap.add_argument("--n-hard", type=int, default=25)
    ap.add_argument("--n-random", type=int, default=15)
    ap.add_argument(
        "--pool-hint",
        type=int,
        default=12,
        help="# missed candidates shown to the agent",
    )
    ap.add_argument("--max-steps", type=int, default=10)
    ap.add_argument("--max-tokens", type=int, default=3000)
    ap.add_argument("--concurrency", type=int, default=4)
    ap.add_argument("--verbosity", type=int, default=0)
    ap.add_argument(
        "--probe", action="store_true", help="run ONE hard image verbosely, then exit"
    )
    ap.add_argument("--out", default="agent_traces.jsonl")
    ap.add_argument("--results", default="agent_iconclass_results.json")
    ap.add_argument(
        "--resume",
        action="store_true",
        help="If --out exists, skip indices already present and append (don't wipe). "
        "Crash insurance for long unlabeled runs; ignored if the file is absent.",
    )
    # ---- Exp B: unlabeled agent-data-build mode (no GT, no cached pool) ----
    ap.add_argument(
        "--unlabeled",
        action="store_true",
        help="Unlabeled mode: 9B anchor from --ft, LIVE search_iconclass (no cached "
        "pool), no GT eval. For Exp B ORPO pair generation on e.g. european_art.",
    )
    ap.add_argument(
        "--dataset",
        default=None,
        help="Override the image dataset (unlabeled mode). Default: biglam/european_art.",
    )
    ap.add_argument(
        "--split", default=None, help="Override the split (unlabeled mode)."
    )
    ap.add_argument(
        "--limit",
        type=int,
        default=20,
        help="Max images to run in unlabeled mode (sel = first --limit anchor rows).",
    )
    args = ap.parse_args()

    if not args.api_key:
        from huggingface_hub import get_token

        args.api_key = get_token()
    args.run_id = uuid.uuid4().hex[:12]
    args.ts = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    if args.probe:
        args.verbosity = max(args.verbosity, 2)

    # ---- mode branch: test-wired (GT + cached pool + 4B preds) vs unlabeled ----
    if args.unlabeled:
        # Exp B data-build: 9B anchor from --ft, LIVE search (no cached pool), no GT.
        ftc = {
            int(r["index"]): list(r.get("predicted_codes", []) or [])
            for r in json.load(open(args.ft))["results"]
        }
        gc = {}
        recs = None  # unlabeled: no GT records
        sel = sorted(ftc.keys())[: args.limit]
        hard_set = set()
        dataset_name = args.dataset or "biglam/european_art"
        split_name = args.split or "train"
        print(
            f"UNLABELED mode: anchor={args.ft} ({len(ftc)} rows) | "
            f"{dataset_name}[{split_name}] | running {len(sel)} | "
            f"run_id={args.run_id} agent_type={args.agent_type} endpoint={args.endpoint_url}"
        )

        agent_tools._get_index()

        # Lazy streaming: decode images on demand by dense index (matching
        # infer_anchor.py — only image-bearing rows consume a dense index), so the
        # agent starts on image 0 immediately and memory is bounded by the in-flight
        # window, not N. Fixes the upfront-materialization cliff at large --limit.
        ds = _LazyImages(dataset_name, split_name)
        print(f"lazy-streaming {dataset_name}[{split_name}] (decode-on-demand)")

        def pool_for(idx):
            return []  # live search_iconclass only (no cached retrieval pool)

    else:
        recs, ftc, gc = load_inputs(args.pool, args.ft, args.grade_cache)
        hard_all, hard_sel, rand_sel = select_split(
            recs, ftc, gc, args.n_hard, args.n_random
        )
        if args.probe:
            sel = hard_sel[:1]
        else:
            sel = sorted(set(hard_sel) | set(rand_sel))
        hard_set = set(hard_sel)
        print(
            f"hard available={len(hard_all)}  selected: hard={len(hard_sel)} random={len(rand_sel)} total={len(sel)}"
        )
        print(
            f"run_id={args.run_id}  agent_type={args.agent_type}  endpoint={args.endpoint_url}"
        )

        # pre-warm the FT index once (avoid a thread race on first search)
        agent_tools._get_index()

        # attach images + per-image pool hint
        from datasets import load_dataset

        print(f"loading {DATASET}[{SPLIT}] …")
        ds = load_dataset(DATASET, split=SPLIT)

        def pool_for(idx):
            ms = sorted(missed(recs, idx, ftc), key=lambda x: x[1], reverse=True)[
                : args.pool_hint
            ]
            return [
                {
                    "code": c,
                    "definition": _code_description(c) or c,
                    "sim": round(s, 3),
                    "judge_grade": gc.get(f"{idx}|{c}", 0),
                }
                for c, s in ms
            ]

    model = make_model(args)
    lock = threading.Lock()
    rows = []
    done = set()
    if args.resume and os.path.exists(args.out) and os.path.getsize(args.out) > 0:
        with open(args.out) as _f:
            for _l in _f:
                try:
                    done.add(int(json.loads(_l)["index"]))
                except Exception:
                    continue
        before = len(sel)
        sel = [i for i in sel if i not in done]
        print(f"RESUME: {len(done)} indices already in {args.out}; {before} -> {len(sel)} to run")
    else:
        open(args.out, "w").close()  # fresh local JSONL for this run

    def task_fn(idx):
        ex = ds.get(idx) if isinstance(ds, _LazyImages) else ds[idx]
        if ex is None:
            row = {
                "index": idx,
                "run_id": args.run_id,
                "ts": args.ts,
                "model": args.model_id,
                "error": "image not in streamed subset",
                "anchor_codes": ftc.get(idx, []),
                "accepted": [],
                "rejected": [],
                "final_codes": list(ftc.get(idx, [])),
            }
            with lock:
                with open(args.out, "a") as f:
                    f.write(json.dumps(row, default=str) + "\n")
            return row
        anchor = ftc.get(idx, [])
        gt = recs[idx]["gt_codes"] if recs is not None else []
        pool = pool_for(idx)
        row = run_one(args, model, idx, ex, anchor, pool, gt)
        if recs is not None:
            evaluate_row(row, recs, ftc, gc)
        with lock:
            with open(args.out, "a") as f:
                f.write(json.dumps(row, default=str) + "\n")
        return row

    print(f"running agent on {len(sel)} images (concurrency={args.concurrency})…")
    t0 = time.time()
    if args.concurrency <= 1 or args.probe:
        for i, idx in enumerate(sel):
            rows.append(task_fn(idx))
            print(f"  [{i + 1}/{len(sel)}] idx={idx} done")
    else:
        with ThreadPoolExecutor(max_workers=args.concurrency) as ex_pool:
            futs = {ex_pool.submit(task_fn, idx): idx for idx in sel}
            for i, fut in enumerate(as_completed(futs)):
                rows.append(fut.result())
                print(f"  [{i + 1}/{len(sel)}] idx={futs[fut]} done")
    print(f"agent runs done in {time.time() - t0:.0f}s")

    n_err = sum(1 for r in rows if r.get("error"))
    if n_err:
        print(f"  WARNING: {n_err}/{len(rows)} rows had errors")

    if args.probe:
        r = rows[0]
        print("\n--- PROBE TRACE ---")
        print(
            json.dumps(
                {k: r[k] for k in ("accepted", "rejected", "final_codes")}, indent=2
            )
        )
        print("tool_calls:")
        for tc in r["tool_calls"]:
            print(f"  step {tc['step']}: {[c['tool'] for c in tc['tool_calls']]}")
        print(f"eval: {r.get('eval')}")
        print(f"\nLocal trace → {args.out}")
        return

    if args.unlabeled:
        # Exp B pair-format summary (no GT metrics).
        n_ok = [r for r in rows if not r.get("error")]
        tot_acc = sum(len(r.get("accepted", [])) for r in n_ok)
        tot_rej = sum(len(r.get("rejected", [])) for r in n_ok)
        tot_anc = sum(len(r.get("anchor_codes", [])) for r in n_ok)
        print("\n=== UNLABELED run summary ===")
        print(f"  rows: {len(n_ok)}/{len(rows)} ok ({n_err} errors)")
        print(
            f"  avg anchor codes/img : {tot_anc / max(len(n_ok), 1):.1f}   "
            f"avg accepted/img: {tot_acc / max(len(n_ok), 1):.1f}   "
            f"avg rejected/img: {tot_rej / max(len(n_ok), 1):.1f}"
        )
        print(
            "  ORPO pairs (per image): chosen = anchor ∪ accepted,  rejected = anchor"
        )
    else:
        summarize(rows, hard_set)
    json.dump(
        {
            "run_id": args.run_id,
            "ts": args.ts,
            "model": args.model_id,
            "agent_type": args.agent_type,
            "unlabeled": bool(args.unlabeled),
            "n": len(rows),
            "n_err": n_err,
            "rows": rows,
        },
        open(args.results, "w"),
        indent=2,
        default=str,
    )
    print(f"\nSaved → {args.results}   traces → {args.out}")


if __name__ == "__main__":
    main()