scripts/explain_prompt_experiment.py

"""Step 11 reasoning-prompt experiment harness — fast iteration on the LLM polish prompt WITHOUT a full run.

Reuses the REAL production prompt-building/parsing/acceptance code from ``rank.explain.tiers`` (imported, not
copied) so a result here reflects actual pipeline behavior, not a diverged approximation. Candidate data comes
from the last full run's ``data/reports/explanations.json`` (facts + pre-polish ``reasoning``) plus
``data/candidates.jsonl`` (evidence, reconstructed the same way ``engine._evidence_snippet`` does) — no
retrieval funnel, no preprocessing, so one iteration is a few seconds instead of the ~2 minutes a full
``rank.py --out ... --limit 0`` run costs.

Run:
    uv run python scripts/explain_prompt_experiment.py                       # single-note mode, current _SYSTEM
    uv run python scripts/explain_prompt_experiment.py --mode batch          # batch mode, current _BATCH_SYSTEM
    uv run python scripts/explain_prompt_experiment.py --variant my_try -v   # a variant defined in VARIANTS below
    uv run python scripts/explain_prompt_experiment.py --n 20 --batch-size 5

To try a prompt tweak: add a named entry to ``VARIANTS`` below (don't edit ``tiers.py`` until a variant proves
out), then run with ``--variant <name>``. Requires an existing ``data/reports/explanations.json`` — regenerate
one first (``rank.py --out submission.csv --limit 0`` or a Streamlit run) if it's missing or stale.
"""

from __future__ import annotations

import argparse
import re
import sys
import time
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "src"))

from common.runtime import configure_compute  # noqa: E402

configure_compute()

import orjson  # noqa: E402

from config import load_settings  # noqa: E402
from jd import extract_spec  # noqa: E402
from models.llm import index_gen_client  # noqa: E402
from rank.explain.tiers import (  # noqa: E402
    _BATCH_SIZE,
    _BATCH_SYSTEM,
    _RULES,
    _SYSTEM,
    _accept_note,
    _build_batch_prompt,
    _build_prompt,
    _clean_blurb,
    _deformulate,
    _extract_notes_array,
    _vary_opener,
)

REPO_ROOT = Path(__file__).resolve().parents[1]
EXPLANATIONS_PATH = REPO_ROOT / "data" / "reports" / "explanations.json"
CANDIDATES_PATH = REPO_ROOT / "data" / "candidates.jsonl"

# Named prompt variants to try head-to-head against the shipped default (_SYSTEM / _BATCH_SYSTEM, imported from
# tiers.py — never copied here, so "default" always reflects production exactly). Add an entry to try a tweak
# without touching tiers.py; promote it there (into _RULES) only once it wins on a real n=25+ comparison, not
# just n=10 (a "tight" 55-70-word variant looked better at n=10 — 80% concrete vs 70% — but reversed at n=25 —
# 62% vs 67% — pure sample noise; always confirm at n=25 before trusting a result).
_SYSTEM_NUMBER = (
    "Write a 3 to 4 sentence recruiter note (third person, about 70 to 85 words) on why a candidate fits a role, "
    + _RULES
    + "If the candidate's own words contain ANY number or metric (scale, percentage, latency, team size, time "
    "span), sentence 1 MUST include it — a concrete number beats a vague achievement claim every time. If no "
    "number is present, describe the concrete outcome without inventing one.\n"
    "Output ONLY the note, no preamble, labels, quotes, or markdown."
)
_SYSTEM_TIGHT = (
    "Write a 3 sentence recruiter note (third person, about 55 to 70 words) on why a candidate fits a role, "
    + _RULES
    + "Be economical — no filler words, no throat-clearing, every clause earns its place. "
    "Output ONLY the note, no preamble, labels, quotes, or markdown."
)
VARIANTS: dict[str, dict[str, str]] = {
    "default": {"single": _SYSTEM, "batch": _BATCH_SYSTEM},
    "number": {"single": _SYSTEM_NUMBER, "batch": _BATCH_SYSTEM},
    "tight": {"single": _SYSTEM_TIGHT, "batch": _BATCH_SYSTEM},
}


def _load_candidates_index(candidate_ids: set[str]) -> dict[str, dict]:
    """One pass over candidates.jsonl, keeping only the requested ids (cheap substring pre-check before the
    JSON parse, so scanning the ~465 MB file for a handful of ids still takes only a couple of seconds)."""
    found: dict[str, dict] = {}
    with CANDIDATES_PATH.open(encoding="utf-8") as handle:
        for line in handle:
            if len(found) == len(candidate_ids):
                break
            if not any(cid in line for cid in candidate_ids):
                continue
            record = orjson.loads(line)
            cid = record.get("candidate_id")
            if cid in candidate_ids:
                found[cid] = record
    return found


def _evidence_for(record: dict, summary_limit: int = 500, experience_limit: int = 900) -> str:
    """Mirrors ``engine._evidence_snippet`` — kept in sync by hand (small, stable function; if that one changes
    its truncation scheme, update this to match so experiment results stay representative)."""
    profile = record.get("profile") or {}
    summary = str(profile.get("summary") or "").strip()[:summary_limit]
    roles = record.get("career_history") or []
    experience = "; ".join(
        f"{r.get('title', '')} @ {r.get('company', '')} ({r.get('industry', '')}) - {r.get('description', '')}"
        for r in roles[:8]
    )[:experience_limit]
    return f"{summary} {experience}".strip()


def _load_rows(n: int) -> list[dict]:
    if not EXPLANATIONS_PATH.exists():
        raise SystemExit(
            f"{EXPLANATIONS_PATH} not found — run `uv run python rank.py --out submission.csv --limit 0` "
            "(or a Streamlit run) at least once first."
        )
    data = orjson.loads(EXPLANATIONS_PATH.read_bytes())[:n]
    ids = {row["candidate_id"] for row in data}
    records = _load_candidates_index(ids)
    for row in data:
        row["_evidence"] = _evidence_for(records.get(row["candidate_id"], {}))
    return data


def _requirements(settings) -> str:
    spec = extract_spec(settings)
    titles = ["role: " + ", ".join(spec.job_titles[:2])] if spec.job_titles else []
    return "; ".join([*spec.must_have_skills[:6], *titles])


def _report(candidate_id: str, rank: int, cleaned: str, nlg: str, evidence: str, requirements: str, verbose: bool) -> bool:
    ok = _accept_note(cleaned, nlg, evidence, requirements)
    tag = "PASS" if ok else "FAIL"
    sentences = len(re.findall(r"[.!?](?:\s|$)", cleaned))
    print(f"{candidate_id} (rank {rank}) [{tag}] {len(cleaned)} chars, {sentences} sentences")
    print(f"  {cleaned[:220]}{'...' if len(cleaned) > 220 else ''}")
    if verbose and not ok:
        print(f"  (nlg source: {nlg[:150]})")
    return ok


_GAP_WORDS = ("though", "however", "gap", "lacks", "missing", "without", "albeit", "unfortunately", "limited")


def run_single(rows: list[dict], client, requirements: str, system_prompt: str, verbose: bool) -> list[tuple[bool, str]]:
    results: list[tuple[bool, str]] = []
    for row in rows:
        nlg, evidence = row["reasoning"], row["_evidence"]
        raw = client.generate(
            _build_prompt(requirements, nlg, evidence), system=system_prompt, temperature=0.0, max_tokens=240,
        )
        if verbose:
            print(f"--- raw ({row['candidate_id']}) ---\n{raw}\n")
        cleaned = _vary_opener(_deformulate(_clean_blurb(raw, nlg)))
        ok = _report(row["candidate_id"], row["rank"], cleaned, nlg, evidence, requirements, verbose)
        results.append((ok, cleaned))
    return results


def run_batch(
    rows: list[dict], client, requirements: str, system_prompt: str, batch_size: int, verbose: bool
) -> list[tuple[bool, str]]:
    results: list[tuple[bool, str]] = []
    for start in range(0, len(rows), batch_size):
        batch = rows[start : start + batch_size]
        items = [(row["reasoning"], row["_evidence"]) for row in batch]
        raw = client.generate(
            _build_batch_prompt(requirements, items), system=system_prompt, temperature=0.0,
            max_tokens=240 * len(batch) + 60,
        )
        if verbose:
            print(f"--- raw (batch {start // batch_size + 1}, {len(batch)} candidates) ---\n{raw}\n")
        notes = _extract_notes_array(raw, len(batch))
        if notes is None:
            print(f"batch {start // batch_size + 1}: PARSE FAILED (expected {len(batch)} notes) — all keep NLG")
            results.extend((False, row["reasoning"]) for row in batch)
            continue
        for row, raw_note in zip(batch, notes):
            nlg, evidence = row["reasoning"], row["_evidence"]
            cleaned = _vary_opener(_deformulate(_clean_blurb(raw_note, nlg)))
            ok = _report(row["candidate_id"], row["rank"], cleaned, nlg, evidence, requirements, verbose)
            results.append((ok, cleaned))
    return results


def _score(name: str, results: list[tuple[bool, str]], elapsed: float) -> dict:
    n = len(results)
    accepted = [text for ok, text in results if ok]
    lengths = [len(t) for t in accepted]
    words = [len(t.split()) for t in accepted]
    concrete = sum(1 for t in accepted if any(ch.isdigit() for ch in t))
    gappy = sum(1 for t in accepted if any(w in t.lower() for w in _GAP_WORDS))
    return {
        "variant": name, "accepted": len(accepted), "n": n, "seconds": round(elapsed, 1),
        "avg_chars": round(sum(lengths) / len(lengths), 0) if lengths else 0,
        "avg_words": round(sum(words) / len(words), 1) if words else 0,
        "concrete_pct": round(100 * concrete / len(accepted), 0) if accepted else 0,
        "gap_mentions": gappy,
    }


def _print_scoreboard(scores: list[dict]) -> None:
    print("\n=== Scoreboard ===")
    header = f"{'variant':<10} {'accept':>8} {'avg_chars':>10} {'avg_words':>10} {'concrete%':>10} {'gap_mentions':>13} {'seconds':>8}"
    print(header)
    for s in scores:
        ratio = f"{s['accepted']}/{s['n']}"
        print(
            f"{s['variant']:<10} {ratio:>8} {s['avg_chars']:>10} {s['avg_words']:>10} "
            f"{s['concrete_pct']:>10} {s['gap_mentions']:>13} {s['seconds']:>8}"
        )
    print(
        "\naccept = passed the real _accept_note gate (grounded + length + sentences + third-person) | "
        "concrete% = share of accepted notes containing at least one digit | "
        "gap_mentions = accepted notes containing a gap/caveat word (should be 0 for a strengths-only tier)"
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="Fast iteration harness for the Step 11 reasoning prompt.")
    parser.add_argument("--mode", choices=["single", "batch"], default="single")
    parser.add_argument(
        "--variant", default="default", choices=[*sorted(VARIANTS), "all"],
        help="named prompt to test, or 'all' to score every variant head-to-head",
    )
    parser.add_argument("--n", type=int, default=10, help="how many real (already-ranked) candidates to test")
    parser.add_argument("--batch-size", type=int, default=_BATCH_SIZE)
    parser.add_argument("-v", "--verbose", action="store_true", help="also print raw LLM output")
    args = parser.parse_args()

    settings = load_settings()
    rows = _load_rows(args.n)
    requirements = _requirements(settings)
    client = index_gen_client(settings)
    client.ensure()

    variant_names = sorted(VARIANTS) if args.variant == "all" else [args.variant]
    scores = []
    for name in variant_names:
        system_prompt = VARIANTS[name][args.mode]
        print(f"\n########## variant: {name} ##########")
        started = time.perf_counter()
        if args.mode == "batch":
            results = run_batch(rows, client, requirements, system_prompt, args.batch_size, args.verbose)
        else:
            results = run_single(rows, client, requirements, system_prompt, args.verbose)
        elapsed = time.perf_counter() - started
        accepted = sum(1 for ok, _ in results if ok)
        print(f"{accepted}/{len(rows)} accepted in {elapsed:.1f}s  (mode={args.mode}, variant={name})")
        scores.append(_score(name, results, elapsed))

    if len(variant_names) > 1:
        _print_scoreboard(scores)


if __name__ == "__main__":
    main()