difficulty.py

#!/usr/bin/env python3
"""LLM difficulty estimator (Part 2). For each task, a cheap model reads the
prompt (+repo) and estimates difficulty signals as a DISTRIBUTION (per the
research: elicit ranges, not points; aggregate samples). These become ex-ante
features for the success + time/cost models. Resumable, concurrent, cached.

  python difficulty.py --labeled        # only the success-labeled tasks (cheap test)
  python difficulty.py --all            # every task
"""
import os, sys, json, re
import pandas as pd
from concurrent.futures import ThreadPoolExecutor, as_completed

HERE = os.path.dirname(os.path.abspath(__file__))
TASKS = os.path.join(HERE, "data", "tasks.csv")
OUT = os.path.join(HERE, "data", "difficulty.csv")
MODEL = "claude-haiku-4-5"
N_SAMPLES = 2  # self-consistency
FIELDS = ["human_minutes_p50", "human_minutes_p90", "n_files_est", "scope_1to5",
          "novelty_1to5", "ambiguity_0to1", "needs_tests_0to1", "is_continuation"]


def get_key():
    k = os.environ.get("ANTHROPIC_API_KEY")
    if k:
        return k
    sys.exit("set ANTHROPIC_API_KEY (only needed for the optional difficulty estimator)")


SYS = "You estimate how hard a coding task is for an AI coding agent. Reply with ONLY one JSON object, no prose."
TMPL = """Task given to an AI coding agent:
---
{prompt}
---
Repo: {repo}

Estimate the difficulty assuming a skilled engineer UNFAMILIAR with this repo.
If the task is a vague continuation ("yes", "continue", "what's next", "go on"), set is_continuation=1 and estimate conservatively from what little is given.
Output JSON with exactly these keys:
{{"human_minutes_p50": int, "human_minutes_p90": int, "n_files_est": int, "scope_1to5": int, "novelty_1to5": int, "ambiguity_0to1": float, "needs_tests_0to1": float, "is_continuation": 0_or_1}}"""


def parse(txt):
    m = re.search(r"\{.*\}", txt, re.S)
    if not m:
        return None
    try:
        o = json.loads(m.group(0))
        return {k: float(o.get(k)) for k in FIELDS if o.get(k) is not None}
    except Exception:
        return None


def estimate(client, prompt, repo):
    vals = {k: [] for k in FIELDS}
    for _ in range(N_SAMPLES):
        try:
            r = client.messages.create(
                model=MODEL, max_tokens=200, temperature=0.7, system=SYS,
                messages=[{"role": "user", "content": TMPL.format(prompt=str(prompt)[:2000], repo=repo or "unknown")}],
            )
            o = parse(r.content[0].text)
            if o:
                for k, v in o.items():
                    vals[k].append(v)
        except Exception:
            continue
    out = {}
    for k in FIELDS:
        out[k] = float(pd.Series(vals[k]).median()) if vals[k] else float("nan")
    return out


def main():
    import anthropic
    client = anthropic.Anthropic(api_key=get_key())
    d = pd.read_csv(TASKS)
    if "--labeled" in sys.argv:
        d = d[d.success_proxy.notna()].copy()
    done = set()
    if os.path.exists(OUT):
        prev = pd.read_csv(OUT)
        done = set(prev.promptId.astype(str))
        print(f"resuming: {len(done)} already done", file=sys.stderr)
    todo = d[~d.promptId.astype(str).isin(done)].reset_index(drop=True)
    print(f"estimating difficulty for {len(todo)} tasks ({MODEL}, {N_SAMPLES}x)...", file=sys.stderr)
    rows, n = [], 0
    with ThreadPoolExecutor(max_workers=12) as ex:
        futs = {ex.submit(estimate, client, r.prompt_text, r.repo): r.promptId for r in todo.itertuples()}
        for fut in as_completed(futs):
            pid = futs[fut]
            try:
                rows.append({"promptId": pid, **fut.result()})
            except Exception:
                rows.append({"promptId": pid, **{k: float("nan") for k in FIELDS}})
            n += 1
            if n % 50 == 0:
                print(f"  {n}/{len(todo)}", file=sys.stderr)
    new = pd.DataFrame(rows)
    if os.path.exists(OUT):
        new = pd.concat([pd.read_csv(OUT), new], ignore_index=True)
    new.to_csv(OUT, index=False)
    print(f"wrote {len(new)} difficulty rows -> {OUT}", file=sys.stderr)


if __name__ == "__main__":
    main()