Judge the Output: Pick the Better Prompt with Numbers You Can Defend

"""Judge two prompts with numbers: a test set, rule checks, an LLM judge you can trust, and a decision."""
import json
import random
import re
from concurrent.futures import ThreadPoolExecutor

from openai import OpenAI

CLIENT = OpenAI()
WRITER = "meta/llama-3.1-8b-instruct"    # writes the summaries
JUDGE = "meta/llama-3.3-70b-instruct"    # grades them: a stronger model than the one being graded


def load_jsonl(path):
    return [json.loads(line) for line in open(path, encoding="utf-8") if line.strip()]


# ---------- Step 1: a test set that covers the cases that matter ----------
def build_testset(reviews, per_group=2, seed=0):
    """Up to per_group reviews from every (sentiment, mentions_price) group, chosen with random.Random(seed),
    plus the shortest and the longest review. Return the chosen reviews sorted by id, no duplicates."""
    # TODO (Step 1):
    #   rng = random.Random(seed)
    #   group the reviews by (r["sentiment"], r["mentions_price"]) into a dict of lists
    #   for each group (in sorted(groups, key=str) order), sorted by id: rng.sample(members, min(per_group, len(members)))
    #   also add the shortest and the longest review (by len(text))
    #   return the chosen reviews sorted by id, each once (a dict keyed by id avoids duplicates)
    raise NotImplementedError("Step 1: write build_testset()")


# ---------- Step 2: checks that need no model ----------
NUMBER = re.compile(r"\d+(?:\.\d+)?")


def rule_checks(summary, review_text):
    """{"one_sentence", "max_20_words", "numbers_grounded"} as booleans."""
    raise NotImplementedError("rule_checks() arrives in Step 2")


def passes_rules(summary, review_text):
    return all(rule_checks(summary, review_text).values())


# ---------- Step 3: an LLM judge, checked against people ----------
JUDGE_PROMPT = """You check summaries of customer reviews for faithfulness.
A summary is faithful if every claim in it is supported by the review. It is unfaithful if it adds a detail
the review does not contain (a name, a number, a gender, an outcome) or reverses what the review says.
Leaving details out is fine. Reply with JSON only: {"faithful": true or false, "reason": "<one short sentence>"}"""


def judge_faithful(review_text, summary, client=None):
    """True if the judge model finds the summary faithful to the review."""
    client = client or CLIENT
    raise NotImplementedError("judge_faithful() arrives in Step 3")


def agreement(labels, reviews_by_id, client=None):
    """Share of human-labelled examples where the judge agrees with the person (0..1), and the disagreements."""
    raise NotImplementedError("agreement() arrives in Step 3")


# ---------- Step 4: compare two outputs without favouring a position ----------
PAIR_PROMPT = """Two summaries of the same customer review follow, labelled A and B. Which is more useful to a
shop manager: accurate, short, and keeps the customer's main point? Reply with one letter only: A or B."""


def prefer(review_text, first, second, client=None):
    """Ask the judge once; return "A" if it prefers `first`, "B" if it prefers `second`."""
    client = client or CLIENT
    r = client.chat.completions.create(
        model=JUDGE, temperature=0, max_tokens=3,
        messages=[{"role": "system", "content": PAIR_PROMPT},
                  {"role": "user", "content": f"Review:\n{review_text}\n\nA: {first}\n\nB: {second}"}],
    )
    return "B" if r.choices[0].message.content.strip().upper().startswith("B") else "A"


def pairwise(review_text, s1, s2, client=None):
    """"1" if s1 wins in BOTH orders, "2" if s2 wins in both, otherwise "tie"."""
    raise NotImplementedError("pairwise() arrives in Step 4")


def summarise(prompt, review_text, client=None):
    client = client or CLIENT
    r = client.chat.completions.create(model=WRITER, temperature=0, max_tokens=120,
                                       messages=[{"role": "system", "content": prompt},
                                                 {"role": "user", "content": review_text}])
    return r.choices[0].message.content.strip()


# ---------- Step 5: decide with an interval, not a hunch ----------
def bootstrap_diff(a, b, n=2000, seed=0):
    """Paired bootstrap of mean(b) - mean(a): resample case indices with replacement n times
    (the SAME indices for a and b), and return the 2.5th and 97.5th percentiles, rounded to 3 places."""
    raise NotImplementedError("bootstrap_diff() arrives in Step 5")


def decide(a, b, n=2000, seed=0):
    """"B" if the whole interval is above 0, "A" if it is below 0, otherwise "tie"."""
    raise NotImplementedError("decide() arrives in Step 5")