esh-pfi-infrastructure/tools/judge-bench/run.py

"""Run the judge benchmark against two gateway models and score them.

Fairness rules, stated so they can be argued with:
  - IDENTICAL prompts to both models. Selene has an Atla-native template;
    using it for selene and a generic one for gen would confound model
    quality with prompt fit. Consumers reach both through the same gateway
    with the same prompt, so the neutral prompt IS the production condition.
    (If selene loses badly, re-test with its native template before
    concluding anything — that is the honest follow-up, not this run.)
  - Same sampling params, same parser, same retry policy.
  - temperature 0: a judge should be as deterministic as the stack allows.
    Residual variation across repeats is then a real property, not noise we
    dialled in.
  - Repeats measure stability, not accuracy: a judge that flips its verdict
    on identical input is unusable regardless of how often it is right.
"""
import json
import pathlib
import re
import statistics
import sys
import time
import urllib.error
import urllib.request

import argparse

ap = argparse.ArgumentParser(description="Judge benchmark: compare two gateway models on LLM-as-judge work.")
ap.add_argument("key", help="LiteLLM virtual key")
ap.add_argument("--models", nargs=2, metavar=("A", "B"), required=True,
                help="the two model names to compare. REQUIRED — there is no sensible\n"
                     "default pair, and a stale default would silently benchmark a\n"
                     "retired seat (the original default, selene-1-mini-8b, now 400s).")
ap.add_argument("--repeats", type=int, default=3,
                help="runs per item; measures stability, not accuracy (default: 3)")
ap.add_argument("--limit", type=int, default=0,
                help="cap items per mode, for smoke tests (default: 0 = all)")
ap.add_argument("--gateway", default="http://10.250.50.70:4000/v1/chat/completions")
args = ap.parse_args()

GATEWAY = args.gateway
KEY = args.key
MODELS = args.models
REPEATS = args.repeats

sys.path.insert(0, str(pathlib.Path(__file__).resolve().parent))
from items import PAIRWISE, ABSOLUTE  # noqa: E402

if args.limit:
    PAIRWISE = PAIRWISE[:args.limit]
    ABSOLUTE = ABSOLUTE[:args.limit]

PAIR_PROMPT = """You are an impartial judge evaluating two AI responses.

Question:
{q}

Response A:
{a}

Response B:
{b}

Decide which response is better. Consider factual accuracy, whether the
response follows the instructions given, relevance, and completeness.

Reply with ONLY a JSON object, no other text:
{{"verdict": "A" or "B" or "tie", "reason": "<one short sentence>"}}"""

ABS_PROMPT = """You are an impartial judge scoring a single AI response.

Question:
{q}

Response:
{r}

Score the response from 1 to 5:
1 = wrong, evasive, or ignores the question
2 = mostly unhelpful or contains a clear error
3 = partially correct but incomplete or thin
4 = correct and helpful with minor gaps
5 = correct, complete, and follows all instructions

Reply with ONLY a JSON object, no other text:
{{"score": <integer 1-5>, "reason": "<one short sentence>"}}"""


def call(model, prompt, timeout=120):
    body = json.dumps({
        "model": model,
        "messages": [{"role": "user", "content": prompt}],
        "temperature": 0,
        "max_tokens": 200,
    }).encode()
    req = urllib.request.Request(GATEWAY, data=body, method="POST")
    req.add_header("Authorization", "Bearer " + KEY)
    req.add_header("Content-Type", "application/json")
    t0 = time.time()
    try:
        with urllib.request.urlopen(req, timeout=timeout) as r:
            d = json.loads(r.read().decode())
        return d["choices"][0]["message"]["content"], time.time() - t0, None
    except urllib.error.HTTPError as e:
        return None, time.time() - t0, "HTTP %s: %s" % (e.code, e.read().decode()[:120])
    except Exception as e:
        return None, time.time() - t0, "%s: %s" % (type(e).__name__, e)


def parse(text, key):
    """Extract the field. Tolerant of fenced/prefixed output on purpose:
    format strictness is measured separately as `clean`, so a model is not
    scored wrong merely for wrapping valid JSON in a code fence."""
    if text is None:
        return None, False
    clean = text.strip().startswith("{") and text.strip().endswith("}")
    m = re.search(r'\{.*\}', text, re.S)
    if m:
        try:
            v = json.loads(m.group(0)).get(key)
            if v is not None:
                return v, clean
        except Exception:
            pass
    if key == "verdict":
        m = re.search(r'\b(tie|A|B)\b', text)
        return (m.group(1) if m else None), False
    m = re.search(r'\b([1-5])\b', text)
    return (int(m.group(1)) if m else None), False


results = {m: {"pair": [], "abs": [], "lat": [], "errors": 0} for m in MODELS}

for model in MODELS:
    print("\n=== %s ===" % model, flush=True)
    for it in PAIRWISE:
        verdicts, clean_n = [], 0
        for _ in range(REPEATS):
            txt, dt, err = call(model, PAIR_PROMPT.format(q=it["q"], a=it["a"], b=it["b"]))
            results[model]["lat"].append(dt)
            if err:
                results[model]["errors"] += 1
                verdicts.append(None)
                continue
            v, clean = parse(txt, "verdict")
            v = v.strip() if isinstance(v, str) else v
            if isinstance(v, str) and v.lower() == "tie":
                v = "tie"
            verdicts.append(v)
            clean_n += int(clean)
        modal = max(set(verdicts), key=verdicts.count) if verdicts else None
        results[model]["pair"].append(dict(
            id=it["id"], truth=it["truth"], verdicts=verdicts, modal=modal,
            correct=(modal == it["truth"]),
            stable=(len(set(verdicts)) == 1), clean=clean_n))
        print("  %-16s truth=%-4s got=%-14s %s%s" % (
            it["id"], it["truth"], str(verdicts),
            "OK " if modal == it["truth"] else "MISS",
            "" if len(set(verdicts)) == 1 else "  <UNSTABLE>"), flush=True)

    for it in ABSOLUTE:
        scores, clean_n = [], 0
        for _ in range(REPEATS):
            txt, dt, err = call(model, ABS_PROMPT.format(q=it["q"], r=it["r"]))
            results[model]["lat"].append(dt)
            if err:
                results[model]["errors"] += 1
                scores.append(None)
                continue
            s, clean = parse(txt, "score")
            try:
                s = int(s)
            except Exception:
                s = None
            scores.append(s)
            clean_n += int(clean)
        got = [s for s in scores if s is not None]
        med = statistics.median(got) if got else None
        lo, hi = it["band"]
        results[model]["abs"].append(dict(
            id=it["id"], band=it["band"], scores=scores, median=med,
            in_band=(med is not None and lo <= med <= hi),
            stable=(len(set(scores)) == 1), clean=clean_n))
        print("  %-16s band=%-6s got=%-16s %s%s" % (
            it["id"], "%d-%d" % it["band"], str(scores),
            "OK " if (med is not None and lo <= med <= hi) else "MISS",
            "" if len(set(scores)) == 1 else "  <UNSTABLE>"), flush=True)

out = str(pathlib.Path(__file__).resolve().parent / "raw.json")
with open(out, "w") as f:
    json.dump(results, f, indent=1)
print("\nraw -> %s" % out)