src/judge.py

"""
Pluggable LLM judge with aggressive on-disk caching.

Design decisions worth stating, because they differ from the prior run:

1. ONE STORY PER CALL, absolute rubric. The prior run sent all 16 samples in a
   single call and asked for relative scores. That is cheaper but wrong here:
   - quality is compared against a fixed gate tau (5.0), so it must be
     absolutely calibrated, not calibrated relative to whatever else happened
     to be in the group;
   - eval compares quality ACROSS models, which relative scoring makes
     meaningless;
   - batched scoring has position bias.
   Per-story also makes the cache actually work: identical stories recur across
   arms and eval passes, and each is paid for once, ever.

2. NO group_diversity FIELD. A set-level scalar is constant within a GRPO group
   -> zero within-group std -> zero advantage. It cannot train anything. All
   diversity credit comes from embeddings (src/diversity.py), where it is
   per-sample by construction.

3. Judge novelty is an ABSOLUTE 'freshness vs. generic AI slop' axis, not
   'different from the others'. The latter is the embeddings' job. This keeps
   the two signals from being redundant-but-inconsistent.
"""

from __future__ import annotations

import asyncio
import hashlib
import json
import os
import sqlite3
import time
from dataclasses import dataclass, asdict
from pathlib import Path

RUBRIC_VERSION = "v4"

SYSTEM_RUBRIC = """You are a demanding fiction editor scoring ONE short story written to a prompt.
Score on absolute standards, not relative to anything you have seen before.

STEP 1 -- FINISHEDNESS. Read the LAST sentence first, before anything else.
A story is UNFINISHED if the final sentence stops mid-clause or mid-word, or if
the piece simply halts with the situation unresolved and no closing beat.
Beautiful prose does not make an unfinished story finished. If UNFINISHED:
set "ended":false, quality 1-3, novelty 0-3, and stop deliberating -- an
unfinished fragment is never novel, because you cannot know what it became.
Only if it is FINISHED, continue to step 2.

STEP 2 -- quality (0-10): coherence + prompt fulfilment + craft
 0-2  Incoherent, or ignores the prompt entirely.
 3-4  Followable but generic: stock premise, flat prose, perfunctory ending.
 5-6  Competent. Does what the prompt asks, reads cleanly, resolves. Forgettable.
 7-8  Strong. Controlled voice, purposeful structure, an ending that lands.
 9-10 Excellent. Precise images, earned emotional turn, nothing wasted.
If it never engages the prompt's actual premise, quality <= 4.

STEP 3 -- novelty (0-10): freshness of premise, voice and construction vs. generic AI fiction
 0-2  Pure slop signature: portentous one-line paragraphs, "It wasn't X, it was Y",
      italicised abstract nouns doing the emotional work, cosmic-melancholy fog.
 3-4  Familiar premise handled in the familiar way.
 5-6  One genuine choice (an unusual angle OR a distinct voice OR an odd structure).
 7-8  Several: takes a real swing on premise, form or register and controls it.
 9-10 Startling and still coherent. You could not have predicted this.
Novelty is NOT weirdness: incoherence, non-sequiturs and word salad score 0-2.
Judge freshness against fiction in general, NOT against other samples.

Return ONLY minified JSON, no markdown fence:
{"quality":<0-10>,"novelty":<0-10>,"ended":<true|false>,"note":"<=10 words"}"""


@dataclass
class JudgeScore:
    quality: float
    novelty: float
    ended: bool = True
    note: str = ""
    ok: bool = True          # False => call failed, scores are neutral fallbacks
    cached: bool = False

    def as_dict(self) -> dict:
        return asdict(self)


NEUTRAL = JudgeScore(quality=5.0, novelty=5.0, ended=True, note="judge_failed", ok=False)


# --------------------------------------------------------------- disk cache
class ScoreCache:
    """sqlite keyed by sha256(model|rubric|prompt|story). Survives restarts."""

    def __init__(self, path: str | Path):
        self.path = str(path)
        Path(self.path).parent.mkdir(parents=True, exist_ok=True)
        self._db = sqlite3.connect(self.path, check_same_thread=False)
        self._db.execute("PRAGMA journal_mode=WAL")
        self._db.execute(
            "CREATE TABLE IF NOT EXISTS scores ("
            " k TEXT PRIMARY KEY, quality REAL, novelty REAL,"
            " ended INT, note TEXT, ts REAL)"
        )
        self._db.commit()
        self.hits = 0
        self.misses = 0

    @staticmethod
    def key(model: str, prompt: str, story: str) -> str:
        h = hashlib.sha256()
        h.update(f"{model}\x00{RUBRIC_VERSION}\x00{prompt}\x00{story}".encode())
        return h.hexdigest()

    def get(self, k: str) -> JudgeScore | None:
        r = self._db.execute(
            "SELECT quality,novelty,ended,note FROM scores WHERE k=?", (k,)
        ).fetchone()
        if r is None:
            self.misses += 1
            return None
        self.hits += 1
        return JudgeScore(quality=r[0], novelty=r[1], ended=bool(r[2]),
                          note=r[3] or "", ok=True, cached=True)

    def put(self, k: str, s: JudgeScore) -> None:
        if not s.ok:
            return  # never cache a failure
        self._db.execute(
            "INSERT OR REPLACE INTO scores VALUES (?,?,?,?,?,?)",
            (k, s.quality, s.novelty, int(s.ended), s.note[:120], time.time()),
        )
        self._db.commit()

    def stats(self) -> dict:
        n = self._db.execute("SELECT COUNT(*) FROM scores").fetchone()[0]
        tot = self.hits + self.misses
        return {"rows": n, "hits": self.hits, "misses": self.misses,
                "hit_rate": (self.hits / tot) if tot else 0.0}


# ------------------------------------------------------------------- judges
def _parse(text: str) -> JudgeScore:
    t = (text or "").strip()
    if t.startswith("```"):
        t = t.split("\n", 1)[-1].rsplit("```", 1)[0].strip()
    i, j = t.find("{"), t.rfind("}")
    if i >= 0 and j > i:
        t = t[i:j + 1]
    d = json.loads(t)

    def num(v, lo=0.0, hi=10.0):
        return max(lo, min(hi, float(v)))

    return JudgeScore(
        quality=num(d.get("quality", 5)),
        novelty=num(d.get("novelty", d.get("novel", 5))),
        ended=bool(d.get("ended", True)),
        note=str(d.get("note", ""))[:120],
        ok=True,
    )


class OpenRouterJudge:
    """Async judge over OpenRouter. Default backend."""

    BASE = "https://openrouter.ai/api/v1/chat/completions"

    def __init__(
        self,
        model: str = "deepseek/deepseek-v4-flash-0731",
        api_key: str | None = None,
        cache_path: str | Path = "cache/judge.sqlite",
        concurrency: int = 12,
        max_retries: int = 4,
        temperature: float = 0.0,
        timeout: float = 120.0,
    ):
        self.model = model
        self.api_key = api_key or os.environ.get("OPENROUTER_API_KEY", "")
        if not self.api_key:
            raise RuntimeError("OPENROUTER_API_KEY not set")
        self.cache = ScoreCache(cache_path)
        # NOT created here. asyncio.Semaphore binds to the loop that first
        # awaits it, and score_many_sync() calls asyncio.run() -> a NEW loop
        # every training step. A semaphore built in __init__ survives step 1 and
        # then raises "bound to a different event loop" on step 2, which our
        # retry wrapper would convert into neutral 5.0 scores forever. Built
        # per-call in score_many() instead.
        self.concurrency = concurrency
        self.sem: asyncio.Semaphore | None = None
        self.max_retries = max_retries
        self.temperature = temperature
        self.timeout = timeout
        self.n_calls = 0
        self.tok_in = 0
        self.tok_out = 0
        self.n_failed = 0
        self.n_empty = 0
        self.n_truncated = 0
        self.last_error = ""

    def _user(self, prompt: str, story: str) -> str:
        return f"WRITING PROMPT:\n{prompt}\n\nSTORY:\n{story}\n\nJSON only:"

    async def _one(self, client, prompt: str, story: str) -> JudgeScore:
        payload = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": SYSTEM_RUBRIC},
                {"role": "user", "content": self._user(prompt, story)},
            ],
            "temperature": self.temperature,
            # deepseek-v4-flash-0731 is a REASONING model. Left alone it spends
            # its whole output budget on `reasoning` and returns content=null
            # with finish_reason="length" -- which our fallback silently turned
            # into a neutral 5.0, i.e. a judge that scored everything identically
            # while looking like it worked. Calibration caught it at 4/166 calls
            # succeeding. Disabling reasoning is also 25x cheaper and 5x faster
            # (30 output tokens vs 764).
            "reasoning": {"enabled": False},
            # Generous bound on a structured output that empirically needs ~30
            # tokens. Not a length cap on content we care about -- and if it is
            # ever hit, `n_truncated` below makes it loud instead of silent.
            "max_tokens": 300,
        }
        delay = 2.0
        for attempt in range(self.max_retries):
            try:
                async with self.sem:
                    r = await client.post(
                        self.BASE, json=payload, timeout=self.timeout,
                        headers={"Authorization": f"Bearer {self.api_key}"},
                    )
                if r.status_code in (429, 500, 502, 503, 529):
                    await asyncio.sleep(delay); delay *= 2
                    continue
                r.raise_for_status()
                body = r.json()
                self.n_calls += 1
                u = body.get("usage") or {}
                self.tok_in += int(u.get("prompt_tokens", 0) or 0)
                self.tok_out += int(u.get("completion_tokens", 0) or 0)
                ch = body["choices"][0]
                if ch.get("finish_reason") == "length":
                    self.n_truncated += 1
                content = ch["message"].get("content")
                if not content:
                    self.n_empty += 1
                    raise ValueError("empty content (reasoning ate the budget?)")
                return _parse(content)
            except Exception as e:
                self.last_error = f"{type(e).__name__}: {str(e)[:140]}"
                if attempt == self.max_retries - 1:
                    self.n_failed += 1
                    return NEUTRAL
                await asyncio.sleep(delay); delay *= 2
        self.n_failed += 1
        return NEUTRAL

    def health(self) -> dict:
        """Fail-loud accounting. A judge that returns neutral 5.0 for every call
        looks exactly like a judge that works, until you check this."""
        att = self.n_calls + self.n_failed
        return {"calls_ok": self.n_calls, "failed": self.n_failed,
                "empty_content": self.n_empty, "truncated": self.n_truncated,
                "fail_rate": (self.n_failed / att) if att else 0.0,
                "last_error": self.last_error}

    def assert_healthy(self, max_fail_rate: float = 0.05) -> None:
        h = self.health()
        if h["fail_rate"] > max_fail_rate:
            raise RuntimeError(
                f"JUDGE UNHEALTHY: {h}. Refusing to train on neutral fallbacks -- "
                f"a constant reward column produces zero GRPO advantage."
            )

    async def score_many(self, pairs: list[tuple[str, str]]) -> list[JudgeScore]:
        """pairs = [(prompt, story), ...] -> scores in the same order."""
        import httpx

        out: list[JudgeScore | None] = [None] * len(pairs)
        todo: list[int] = []
        for i, (p, s) in enumerate(pairs):
            hit = self.cache.get(self.cache.key(self.model, p, s))
            if hit is not None:
                out[i] = hit
            else:
                todo.append(i)

        if todo:
            self.sem = asyncio.Semaphore(self.concurrency)   # bind to THIS loop

            async def one_and_cache(i):
                """Persist each score the moment it lands, not after the whole
                batch. The pool stage issues ~16k calls in a single
                score_many(); caching only at the end meant a crash at 99%
                threw away every call. Incremental writes make any failure
                resumable for free -- a re-run replays as cache hits."""
                s = await self._one(client, pairs[i][0], pairs[i][1])
                if s.ok:
                    self.cache.put(self.cache.key(self.model, *pairs[i]), s)
                return i, s

            async with httpx.AsyncClient(http2=False) as client:
                for fut in asyncio.as_completed([one_and_cache(i) for i in todo]):
                    i, s = await fut
                    out[i] = s

        return [o if o is not None else NEUTRAL for o in out]

    def score_many_sync(self, pairs: list[tuple[str, str]]) -> list[JudgeScore]:
        """Blocking wrapper for use inside TRL reward functions."""
        try:
            asyncio.get_running_loop()
        except RuntimeError:
            return asyncio.run(self.score_many(pairs))
        # already inside a loop (rare in TRL): run in a private one
        import concurrent.futures
        with concurrent.futures.ThreadPoolExecutor(1) as ex:
            return ex.submit(asyncio.run, self.score_many(pairs)).result()

    def cost_estimate(self, price_in: float, price_out: float) -> dict:
        """price_* in $ per 1M tokens."""
        c = self.tok_in / 1e6 * price_in + self.tok_out / 1e6 * price_out
        return {"calls": self.n_calls, "tok_in": self.tok_in,
                "tok_out": self.tok_out, "usd": round(c, 4),
                **self.cache.stats()}


def build_judge(**kw) -> OpenRouterJudge:
    """Factory. Extend here if an Anthropic/OpenAI key shows up."""
    return OpenRouterJudge(**kw)