Part E: LLM-as-judge

"""Module 10: LLM-as-judge prompts, a deterministic stand-in judge, bias tests, calibration, and cost.

    PYTHONPATH=. python examples/m10_judge.py demo                 # the three judge prompts and their size
    PYTHONPATH=. python examples/m10_judge.py calibrate            # stand-in judge vs human labels
    PYTHONPATH=. python examples/m10_judge.py bias                 # position and verbosity harness
    PYTHONPATH=. python examples/m10_judge.py cost                 # judge cost from real token counts
    LLM_PROVIDER=groq PYTHONPATH=. python examples/m10_judge.py calibrate --live
    LLM_PROVIDER=groq PYTHONPATH=. python examples/m10_judge.py bias --live

Without --live nothing here is a language model. The stand-in judge is a rule
(key facts, policy regexes, padding). The "biased" pairwise judge is a
ScriptedLLM that SIMULATES position and length bias so we can check that the
harness detects bias; it says nothing about how real models behave.
"""
from __future__ import annotations

import json
import re
import sys
from collections import Counter
from collections.abc import Callable

from examples.m10_evals import KEY_FACTS, ROADMAP_DATE, UNAUTHORIZED_ACTION, make_baseline
from examples.m10_human import DRAFTS, LABELS, RUBRIC, agreement_report, load_jsonl, print_agreement
from supportdesk.data import Ticket, get_article
from supportdesk.kb_search import tokenize
from supportdesk.llm import Usage
from supportdesk.pricing import PRICES, cost_usd
from supportdesk.stand_in import ScriptedLLM
from supportdesk.tokens import count_messages

# Judge prompts (pointwise, pairwise, reference-based) ------------------------------------

def pointwise_messages(ticket: str, article: str, reply: str) -> list[dict]:
    return [
        {"role": "system", "content": "You are a strict reviewer of support replies.\n" + RUBRIC +
         '\nThink about each claim first. Return only JSON: {"reasoning": "<two sentences>", "score": 1|2|3}'},
        {"role": "user", "content": f"Help-center article:\n{article}\n\nTicket:\n{ticket}\n\nDraft reply:\n{reply}"},
    ]


def pairwise_messages(ticket: str, article: str, reply_a: str, reply_b: str) -> list[dict]:
    return [
        {"role": "system", "content": "You compare two draft replies to the same support ticket.\n" + RUBRIC +
         "\nPrefer the reply an agent could send with fewer edits. Length is not a virtue. The order of "
         'the replies is random. Return only JSON: {"reasoning": "<two sentences>", "winner": "A"|"B"|"tie"}'},
        {"role": "user", "content": f"Help-center article:\n{article}\n\nTicket:\n{ticket}\n\n"
                                    f"Reply A:\n{reply_a}\n\nReply B:\n{reply_b}"},
    ]


def reference_messages(ticket: str, reference: str, reply: str) -> list[dict]:
    return [
        {"role": "system", "content": "You check a draft support reply against a reference answer written by a "
         "senior agent. The draft may use different words. Score 3 if it conveys every fact in the reference and "
         "nothing contradicting it, 2 if it misses a fact but contradicts nothing, 1 if it contradicts the reference "
         'or adds a claim the reference does not support. Return only JSON: {"reasoning": "<two sentences>", "score": 1|2|3}'},
        {"role": "user", "content": f"Ticket:\n{ticket}\n\nReference answer:\n{reference}\n\nDraft reply:\n{reply}"},
    ]


def parse_judge(text: str) -> dict:
    """Pull the first JSON object out of a judge reply; judges sometimes wrap it in prose."""
    match = re.search(r"\{.*\}", text or "", re.S)
    if not match:
        raise ValueError(f"no JSON in judge output: {text[:80]!r}")
    return json.loads(match.group(0))


def llm_pointwise_judge(chat_fn: Callable | None = None, **kw) -> Callable[[dict], int]:
    """A pointwise judge on a real model through supportdesk.llm.chat."""
    if chat_fn is None:
        from supportdesk.llm import chat as chat_fn

    def judge(d: dict) -> int:
        messages = pointwise_messages(d["ticket"], get_article(d["gold_article"]).body, d["reply"])
        return int(parse_judge(chat_fn(messages, response_format={"type": "json_object"}, **kw).text)["score"])
    return judge


# A deterministic stand-in judge ------------------------------------------------------------

def heuristic_score(ticket: str, reply: str, gold_article: str) -> int:
    """A rule-based, reference-based stand-in judge (not a model).

    1 if it breaks policy or states no key fact of the right article; 2 if it is long
    or padded (under half its sentences share a word with the ticket); else 3.
    """
    if ROADMAP_DATE.search(reply) or UNAUTHORIZED_ACTION.search(reply):
        return 1
    if not re.search(KEY_FACTS[gold_article], reply, re.I):
        return 1
    ticket_words = {w for w in tokenize(ticket) if len(w) > 3}
    sentences = [s for s in re.split(r"(?<=[.!?])\s+", reply) if s.strip()]
    on_topic = sum(bool(ticket_words & set(tokenize(s))) for s in sentences) / len(sentences)
    if len(reply.split()) > 80 or on_topic < 0.5:
        return 2
    return 3


def heuristic_judge(d: dict) -> int:
    return heuristic_score(d["ticket"], d["reply"], d["gold_article"])


def calibrate(judge: Callable[[dict], int], name: str) -> dict:
    """Agreement between a judge and the adjudicated human labels, plus the 'send it' decision."""
    drafts = {d["draft_id"]: d for d in load_jsonl(DRAFTS)}
    labels = load_jsonl(LABELS)
    human = [x["adjudicated"] for x in labels]
    judged = [judge(drafts[x["draft_id"]]) for x in labels]
    report = agreement_report(human, judged, "human (adjudicated)", name)
    print_agreement(report)
    tp = sum(h == 3 and j == 3 for h, j in zip(human, judged))
    said_send = sum(j == 3 for j in judged)
    should_send = sum(h == 3 for h in human)
    print(f"'send as is' (score 3): judge said {said_send}, humans said {should_send}, both {tp}  "
          f"precision {tp / max(said_send, 1):.2f} recall {tp / max(should_send, 1):.2f}")
    misses = [(x["draft_id"], h, j) for x, h, j in zip(labels, human, judged) if h != j]
    print("disagreements (draft, human, judge):", misses)
    return report


# Position and verbosity bias harness --------------------------------------------------------

PairJudge = Callable[[str, str, str, str], str]   # (ticket, article, reply_a, reply_b) -> "A" | "B" | "tie"


def heuristic_pair_judge(ticket: str, article_id: str, a: str, b: str) -> str:
    sa, sb = heuristic_score(ticket, a, article_id), heuristic_score(ticket, b, article_id)
    return "A" if sa > sb else "B" if sb > sa else "tie"


def simulated_biased_judge() -> PairJudge:
    """SIMULATED bias for testing the harness: prefers a reply 60% longer, otherwise says 'A'."""
    def responder(messages, kwargs):
        text = messages[-1]["content"]
        a = text.split("Reply A:\n", 1)[1].split("\n\nReply B:\n", 1)[0]
        b = text.split("\n\nReply B:\n", 1)[1]
        winner = "B" if len(b) > 1.6 * len(a) else "A"
        return json.dumps({"reasoning": "simulated", "winner": winner})
    fake = ScriptedLLM(responder=responder)
    return llm_pair_judge(fake)


def llm_pair_judge(chat_fn: Callable | None = None, **kw) -> PairJudge:
    if chat_fn is None:
        from supportdesk.llm import chat as chat_fn

    def judge(ticket: str, article_id: str, a: str, b: str) -> str:
        messages = pairwise_messages(ticket, get_article(article_id).body, a, b)
        return parse_judge(chat_fn(messages, response_format={"type": "json_object"}, **kw).text)["winner"]
    return judge


def swap_test(judge: PairJudge, pairs: list[tuple[str, str, str, str]]) -> dict:
    """Judge every pair in both orders. A verdict counts only if it survives the swap."""
    tally = Counter()
    for ticket, article_id, x, y in pairs:
        first = judge(ticket, article_id, x, y)               # x shown as A
        second = judge(ticket, article_id, y, x)              # x shown as B
        winner1 = {"A": "x", "B": "y"}.get(first, "tie")
        winner2 = {"A": "y", "B": "x"}.get(second, "tie")
        if winner1 == winner2:
            tally["consistent"] += 1
            tally[f"consistent_{winner1}"] += 1
        else:
            tally["inconsistent"] += 1
            tally["first_position_won_both"] += first == "A" and second == "A"
            tally["second_position_won_both"] += first == "B" and second == "B"
    tally["pairs"] = len(pairs)
    return dict(tally)


def pad(reply: str) -> str:
    """Make a reply longer without adding information (a verbosity-bias probe)."""
    return (reply + " We truly appreciate your patience and understanding. Our team is always committed to "
            "providing you with the best possible experience, and we are constantly working to improve. "
            "Please do not hesitate to reach out if there is anything else at all we can help with.")


def verbosity_test(judge: PairJudge, pairs: list[tuple[str, str, str, str]]) -> dict:
    """Where the judge prefers x over y, pad y. A fair judge keeps preferring x."""
    kept = flipped = 0
    for ticket, article_id, x, y in pairs:
        if judge(ticket, article_id, x, y) == "A" and judge(ticket, article_id, y, x) == "B":
            padded = pad(y)
            if judge(ticket, article_id, x, padded) == "A" and judge(ticket, article_id, padded, x) == "B":
                kept += 1
            else:
                flipped += 1
    return {"clear_preferences": kept + flipped, "kept_after_padding": kept, "flipped_or_tied": flipped}


def build_pairs() -> list[tuple[str, str, str, str]]:
    """Same-ticket pairs: each non-baseline draft against the v2 baseline's reply to that ticket."""
    baseline = make_baseline(2)
    pairs = []
    for d in load_jsonl(DRAFTS):
        if d["source"] == "baseline_v2":
            continue
        ticket_id, rest = d["ticket_id"], d["ticket"].split("\n\n", 1)
        t = Ticket(ticket_id, rest[0].removeprefix("Subject: "), rest[1], "team", "en", "dev", {})
        pairs.append((d["ticket"], d["gold_article"], d["reply"], baseline(t).draft.reply))
    return pairs


# Judge cost ------------------------------------------------------------------------------

def judge_cost(output_tokens: int = 150, reasoning_tokens: int = 400) -> None:
    """Real input-token counts for this module's judge prompts; output and reasoning tokens are ASSUMPTIONS."""
    drafts = load_jsonl(DRAFTS)
    point = [count_messages(pointwise_messages(d["ticket"], get_article(d["gold_article"]).body, d["reply"]))
             for d in drafts]
    pair = [count_messages(pairwise_messages(t, get_article(a).body, x, y)) for t, a, x, y in build_pairs()]
    mean_point, mean_pair = sum(point) / len(point), sum(pair) / len(pair)
    print(f"input tokens per call: pointwise mean {mean_point:.0f} (n={len(point)}), pairwise mean {mean_pair:.0f} (n={len(pair)})")
    print(f"assumed per call: {output_tokens} output tokens, plus {reasoning_tokens} reasoning tokens on reasoning models")
    print(f"{'model':<26}{'pointwise':>11}{'pairwise x2':>13}{'CI/month':>10}   (USD per 91-case suite; month = pairwise x2)")
    for model in ["openai/gpt-oss-120b", "openai/gpt-oss-20b", "llama-3.3-70b-versatile", "gemini-3.5-flash",
                  "gemini-3.5-flash-lite"]:
        extra = reasoning_tokens if model.startswith(("openai/gpt-oss", "gemini")) else 0
        one_point = cost_usd(Usage(input_tokens=round(mean_point), output_tokens=output_tokens + extra), model)
        one_pair = cost_usd(Usage(input_tokens=round(mean_pair), output_tokens=output_tokens + extra), model)
        suite_point, suite_pair = one_point * 91, one_pair * 91 * 2
        monthly = suite_pair * 20 * 22       # 20 CI runs a day, 22 working days
        print(f"{model:<26}{suite_point:>11.4f}{suite_pair:>13.4f}{monthly:>10.2f}")
    print("Prices from supportdesk/pricing.py (checked 21 Sep 2026; verify before relying on them).")


def demo(live: bool) -> None:
    """Build all three judge prompts for one draft; with --live, send them through supportdesk.llm.chat."""
    drafts = {d["draft_id"]: d for d in load_jsonl(DRAFTS)}
    bad, good = drafts["D-06"], drafts["D-03"]     # a policy violation, and a hand-written 'send' reply
    article = get_article(bad["gold_article"]).body
    prompts = {"pointwise": pointwise_messages(bad["ticket"], article, bad["reply"]),
               "pairwise": pairwise_messages(bad["ticket"], article, bad["reply"],
                                             "After 5 failed attempts an account is locked for 15 minutes. Support "
                                             "cannot unlock it sooner, so please try again then or use Forgot password."),
               "reference": reference_messages(good["ticket"], good["reply"],
                                               "You renewed 5 days ago, so you qualify for a full refund.")}
    for name, messages in prompts.items():
        print(f"{name:<10} {count_messages(messages):>4} input tokens   user message starts: "
              f"{messages[1]['content'][:60]!r}")
        if live:
            from supportdesk.llm import chat
            result = chat(messages, response_format={"type": "json_object"})
            print(f"           judge said: {parse_judge(result.text)}  ({result.usage.input_tokens} in, "
                  f"{result.usage.output_tokens} out, {result.latency_ms:.0f} ms)")


def main() -> None:
    cmd = sys.argv[1] if len(sys.argv) > 1 else "calibrate"
    live = "--live" in sys.argv
    if cmd == "demo":
        demo(live)
    elif cmd == "calibrate":
        calibrate(llm_pointwise_judge() if live else heuristic_judge, "LLM judge" if live else "stand-in judge")
    elif cmd == "bias":
        pairs = build_pairs()
        judges = {"LLM judge": llm_pair_judge()} if live else {
            "stand-in rule judge": heuristic_pair_judge, "SIMULATED biased judge": simulated_biased_judge()}
        for name, judge in judges.items():
            print(f"{name}: swap test {swap_test(judge, pairs)}")
            print(f"{name}: verbosity test {verbosity_test(judge, pairs)}")
    elif cmd == "cost":
        judge_cost()


if __name__ == "__main__":
    main()

python examples/m10_judge.py demo

python examples/m10_judge.py bias

python examples/m10_judge.py calibrate

LLM_PROVIDER=groq python examples/m10_judge.py calibrate --live

python examples/m10_judge.py cost