engine/judge.py

"""LLM-judge correctness gate for the router.

`router._verify()` only proves the produced code RUNS (clean exit / tests it wrote
itself), not that it's actually CORRECT — so a small model can ship a clean-but-wrong
solution and the router accepts it instead of escalating (exactly how the bench's
roman_to_int slipped through: ran fine, wrong output).

This judge asks a more capable model whether the solution truly satisfies the task; a
concrete "no" is turned into an escalation by the router. Mirrors
smolcode-cli/src/judge.rs (JSON-only reply, temperature 0, lenient parse), but the
verdict drives ESCALATION rather than stop/continue.

Conservative by design: only a clear defect escalates. On judge error / timeout /
unparseable reply we ACCEPT — the judge is a net to catch obvious wrongness, not a
hard gate, and we don't want to over-escalate (and lose the small-model win).
"""
from __future__ import annotations

import json
import os
import re

import liteforge as lf

_SYSTEM = (
    "You are a strict senior code reviewer. You are given a coding TASK and the FILES "
    "an agent produced. The code already runs without crashing — your job is to judge "
    "whether it is actually CORRECT and COMPLETE for the task: check the exact "
    "requirements, edge cases, and obvious logic bugs.\n"
    "Reply with ONLY a JSON object: {\"correct\": true|false, \"reason\": \"<one short sentence>\"}.\n"
    "Set \"correct\": false if you find ANY bug, wrong/missing edge case, or unmet "
    "requirement. Ignore style. Do not write code."
)


def judge_enabled() -> bool:
    """Judge is on by default; SMALLCODE_JUDGE=0 disables it."""
    return os.environ.get("SMALLCODE_JUDGE", "1").lower() not in ("0", "false", "no", "")


def _files_block(files: dict[str, str], cap: int = 6000) -> str:
    blob = "\n\n".join(f"### {path}\n{content}" for path, content in files.items())
    return blob[:cap]


def _parse(text: str) -> bool | None:
    """True (correct), False (defect found), or None (couldn't tell)."""
    m = re.search(r"\{.*\}", text, re.DOTALL)
    if m:
        try:
            obj = json.loads(m.group(0))
            if isinstance(obj.get("correct"), bool):
                return obj["correct"]
        except Exception:
            pass
    low = text.lower()
    if "correct\": false" in low or "correct: false" in low or "incorrect" in low:
        return False
    if "correct\": true" in low or "correct: true" in low:
        return True
    return None


async def judge_correct(preset, judge_model: str, task: str,
                        files: dict[str, str], final: str) -> bool:
    """Return True if the solution likely satisfies the task, False on a clear defect.

    Accepts (True) on empty files, judge error, or unparseable reply.
    """
    if not files:
        return True
    user = (
        f"TASK:\n{task}\n\nFILES:\n{_files_block(files)}\n\n"
        f"AGENT'S FINAL CLAIM:\n{(final or '')[:500]}\n\n"
        "Is the solution correct and complete for the task? Reply with JSON only."
    )
    try:
        client = lf.AsyncForgeClient(
            base_url=preset.base_url, api_key=preset.api_key, default_model=judge_model,
        )
        resp = await client.complete(
            messages=[{"role": "system", "content": _SYSTEM},
                      {"role": "user", "content": user}],
            model=judge_model, temperature=0.0,
        )
        content = resp["choices"][0]["message"].get("content", "") or ""
    except Exception:
        return True  # judge unavailable -> don't block the accept
    verdict = _parse(content)
    return True if verdict is None else verdict