benchmarks/validators.py

"""Deterministic validators for benchmark outputs.

The validators in this module deliberately do not call an LLM or a provider. They
only normalize an output when the evidence is unambiguous and otherwise return a
stable failure code that the retry layer can act on.
"""

from __future__ import annotations

from dataclasses import dataclass, field
import json
import re
from typing import Any, Iterable, Mapping, Optional


_MMLU_LETTERS = frozenset("ABCD")
_MMLU_EXPLICIT = re.compile(
    r"\b(?:answer|答案|risposta|final(?:\s+answer)?|choice|scelta)\b\s*"
    r"(?:is|=|:)\s*[*`_\[\(]*([A-D])[*`_\]\)]*",
    re.IGNORECASE,
)
_MMLU_MARKED = re.compile(r"(?:^|\n|\s)(?:\(?([A-D])\)?)[\.:\)](?:\s|$)", re.IGNORECASE)
_MMLU_ISOLATED = re.compile(r"(?<![A-Za-z])([A-D])(?![A-Za-z])", re.IGNORECASE)

_CODE_FENCE = re.compile(
    r"```\s*([A-Za-z0-9_+#.-]*)\s*\n?(.*?)```", re.IGNORECASE | re.DOTALL
)
_CODE_JSON_KEYS = ("code", "typescript", "source", "implementation")
_TS_DECLARATION = re.compile(
    r"\b(?:export\s+)?(?:async\s+)?(?:function|class|interface|type|const|let|var)\s+([A-Za-z_$][\w$]*)",
    re.MULTILINE,
)
_TS_IMPORT_EXPORT = re.compile(r"\b(?:import|export)\b")
_TS_SYNTAX_TOKENS = re.compile(r"[{}();]|=>|:\s*[A-Za-z_$][\w$<>,\[\]| ]*")
_PLACEHOLDER = re.compile(r"\b(?:TODO|TBD|your implementation|implement here)\b", re.IGNORECASE)
_REASONING_EXPLICIT = re.compile(
    r"(?:####|final\s+answer|answer|risposta|risultato|result|total|totale)\s*[:=]?\s*"
    r"(-?\d[\d,]*(?:\.\d+)?)",
    re.IGNORECASE,
)
_REASONING_BOLD = re.compile(r"\*\*\s*(-?\d[\d,]*(?:\.\d+)?)\s*\*\*")
_REASONING_LINE_NUMBER = re.compile(r"(?m)^\s*(-?\d[\d,]*(?:\.\d+)?)\s*$")
_REASONING_FAILURES = frozenset({"answer_missing", "wrong_numeric_answer", "calculation_conflict"})


@dataclass(frozen=True)
class ValidationResult:
    """Stable validator result consumed by shadow mode and retry logic."""

    valid: bool
    normalized: Optional[str] = None
    failure_code: Optional[str] = None
    evidence: dict[str, Any] = field(default_factory=dict)
    repair_hint: Optional[str] = None

    def as_dict(self) -> dict[str, Any]:
        return {
            "valid": self.valid,
            "normalized": self.normalized,
            "failure_code": self.failure_code,
            "evidence": self.evidence,
            "repair_hint": self.repair_hint,
        }


def _failure(code: str, *, evidence: Optional[dict[str, Any]] = None, hint: str = "") -> ValidationResult:
    return ValidationResult(
        valid=False,
        failure_code=code,
        evidence=evidence or {},
        repair_hint=hint or None,
    )


def _success(normalized: str, *, evidence: Optional[dict[str, Any]] = None) -> ValidationResult:
    return ValidationResult(valid=True, normalized=normalized, evidence=evidence or {})


def _clean_text(raw: Any) -> str:
    if raw is None:
        return ""
    if isinstance(raw, str):
        return raw.strip()
    return str(raw).strip()


def _mmlu_candidates(text: str) -> list[str]:
    """Return candidates in evidence order, preserving duplicates for ambiguity checks."""
    explicit = [m.group(1).upper() for m in _MMLU_EXPLICIT.finditer(text)]
    if explicit:
        return explicit
    marked = [m.group(1).upper() for m in _MMLU_MARKED.finditer(text)]
    if marked:
        return marked
    return [m.group(1).upper() for m in _MMLU_ISOLATED.finditer(text)]


def validate_mmlu_output(raw: Any, *, expected: Optional[str] = None) -> ValidationResult:
    """Validate a multiple-choice answer without guessing from explanation prose.

    Accepted outputs contain one unambiguous A/B/C/D choice. Explicit labels such
    as ``ANSWER: C`` have priority over marked choices and isolated letters. If
    multiple distinct candidates are present, the result is ambiguous and fails.
    ``expected`` is optional and is only used to expose correctness in evidence; it
    never changes the parsing result.
    """

    text = _clean_text(raw)
    if not text:
        return _failure(
            "answer_missing",
            hint="Return exactly one canonical choice using ANSWER: A, B, C, or D.",
        )

    candidates = _mmlu_candidates(text)
    distinct = sorted(set(candidates))
    evidence: dict[str, Any] = {
        "candidates": candidates,
        "distinct_candidates": distinct,
        "source_length": len(text),
    }
    if expected is not None:
        normalized_expected = str(expected).strip().upper()
        evidence["expected"] = normalized_expected
        if normalized_expected in _MMLU_LETTERS:
            evidence["correct"] = len(distinct) == 1 and distinct[0] == normalized_expected

    if not candidates:
        return _failure(
            "answer_missing",
            evidence=evidence,
            hint="Return exactly one canonical choice using ANSWER: A, B, C, or D.",
        )
    if len(distinct) != 1 or distinct[0] not in _MMLU_LETTERS:
        return _failure(
            "answer_ambiguous",
            evidence=evidence,
            hint="Remove competing choices and return one letter: A, B, C, or D.",
        )

    return _success(distinct[0], evidence=evidence)


def _extract_code(raw: Any) -> tuple[Optional[str], str, dict[str, Any]]:
    """Extract code from a TypeScript fence or a JSON envelope."""

    text = _clean_text(raw)
    if not text:
        return None, "none", {"source_length": 0}

    try:
        decoded = json.loads(text)
    except (TypeError, json.JSONDecodeError):
        decoded = None
    if isinstance(decoded, Mapping):
        for key in _CODE_JSON_KEYS:
            value = decoded.get(key)
            if isinstance(value, str) and value.strip():
                return value.strip(), f"json:{key}", {"source_length": len(text)}

    fences = _CODE_FENCE.findall(text)
    if fences:
        typed = [body.strip() for language, body in fences if language.lower() in {"ts", "typescript"}]
        if typed:
            return max(typed, key=len), "fence:typescript", {"fence_count": len(fences)}
        return None, "fence:wrong-language", {"languages": [language.lower() for language, _ in fences]}

    return None, "none", {"source_length": len(text)}


def _normalize_symbols(required_symbols: Iterable[str]) -> list[str]:
    return [symbol.strip() for symbol in required_symbols if str(symbol).strip()]


def _parse_numeric_token(value: str) -> int | float:
    normalized = value.replace(",", "").strip()
    number = float(normalized) if "." in normalized else int(normalized)
    return number


def _reasoning_candidates(text: str) -> tuple[list[int | float], str]:
    """Extract answer candidates conservatively, preferring explicit final markers."""
    explicit = [_parse_numeric_token(match.group(1)) for match in _REASONING_EXPLICIT.finditer(text)]
    if explicit:
        return explicit, "explicit"
    bold = [_parse_numeric_token(match.group(1)) for match in _REASONING_BOLD.finditer(text)]
    if bold:
        return bold, "bold"
    lines = [_parse_numeric_token(match.group(1)) for match in _REASONING_LINE_NUMBER.finditer(text)]
    if lines:
        return lines[-1:], "final_line"
    return [], "none"


def validate_reasoning_output(raw: Any, *, expected: Optional[int | float] = None) -> ValidationResult:
    """Validate a numeric reasoning answer without calling an LLM.

    Explicit final markers have priority over intermediate arithmetic. Multiple
    distinct explicit answers are classified as a conflict rather than guessed.
    """
    text = _clean_text(raw)
    if not text:
        return _failure(
            "answer_missing",
            hint="Show the calculation and finish with #### N, where N is the final integer.",
        )

    candidates, source = _reasoning_candidates(text)
    distinct = list(dict.fromkeys(candidates))
    evidence: dict[str, Any] = {
        "candidates": candidates,
        "distinct_candidates": distinct,
        "source": source,
        "source_length": len(text),
    }
    if expected is not None:
        try:
            normalized_expected = _parse_numeric_token(str(expected))
            evidence["expected"] = normalized_expected
        except ValueError:
            normalized_expected = expected

    if not candidates:
        return _failure(
            "answer_missing",
            evidence=evidence,
            hint="Show the calculation and finish with #### N, where N is the final integer.",
        )
    if len(distinct) > 1:
        return _failure(
            "calculation_conflict",
            evidence=evidence,
            hint="Recalculate the final value and provide exactly one final numeric answer.",
        )

    normalized = distinct[0]
    if expected is not None and normalized != normalized_expected:
        evidence["correct"] = False
        return _failure(
            "wrong_numeric_answer",
            evidence=evidence,
            hint="Recheck every arithmetic step and return the corrected final number.",
        )

    evidence["correct"] = True if expected is not None else None
    return _success(str(normalized), evidence=evidence)


def validate_reasoning_retry(
    goal: Any,
    raw: Any,
    *,
    expected: Optional[int | float] = None,
    is_last_attempt: bool,
) -> Optional[ValidationResult]:
    """Return a reasoning failure only when a non-final numeric retry is warranted."""
    if is_last_attempt or not any(marker in str(goal or "").lower() for marker in ("reasoning", "gsm8k", "risolvi il problema matematico")):
        return None
    result = validate_reasoning_output(raw, expected=expected)
    return result if result.failure_code in _REASONING_FAILURES else None


_CODING_RETRY_FAILURES = frozenset({
    "code_missing",
    "code_wrong_language",
    "code_empty",
    "code_placeholder",
    "required_symbol_missing",
    "code_syntax_suspect",
})


def is_typescript_goal(goal: Any) -> bool:
    lowered = str(goal or "").lower()
    return any(marker in lowered for marker in ("code_correct", "typescript", "```ts", "```typescript"))


def validate_coding_retry(goal: Any, raw: Any, *, is_last_attempt: bool) -> Optional[ValidationResult]:
    """Return the failed result only when a non-final TypeScript retry is warranted."""

    if is_last_attempt or not is_typescript_goal(goal):
        return None
    result = validate_coding_output(raw)
    return result if result.failure_code in _CODING_RETRY_FAILURES else None


def validate_coding_output(
    raw: Any,
    *,
    required_symbols: Iterable[str] = (),
    min_significant_lines: int = 1,
    reject_placeholders: bool = True,
) -> ValidationResult:
    """Validate extraction and minimum structural quality of TypeScript output.

    This is intentionally a contract validator, not a compiler. Syntax checks are
    conservative and deterministic; full compilation remains a separate isolated
    integration test because it depends on the repository's TypeScript toolchain.
    """

    code, source, extraction_evidence = _extract_code(raw)
    if code is None:
        failure = "code_wrong_language" if source == "fence:wrong-language" else "code_missing"
        return _failure(
            failure,
            evidence=extraction_evidence | {"extraction": source},
            hint="Return exactly one non-empty ```typescript code block.",
        )

    significant_lines = [line for line in code.splitlines() if line.strip() and not line.strip().startswith("//")]
    evidence: dict[str, Any] = extraction_evidence | {
        "extraction": source,
        "significant_lines": len(significant_lines),
        "has_import_or_export": bool(_TS_IMPORT_EXPORT.search(code)),
        "has_syntax_tokens": bool(_TS_SYNTAX_TOKENS.search(code)),
    }

    if not significant_lines or len(significant_lines) < max(1, min_significant_lines):
        return _failure(
            "code_empty",
            evidence=evidence,
            hint="Provide a complete non-empty TypeScript implementation.",
        )
    if reject_placeholders and _PLACEHOLDER.search(code):
        return _failure(
            "code_placeholder",
            evidence=evidence,
            hint="Replace TODO/TBD placeholders with executable TypeScript.",
        )

    declarations = {match.group(1) for match in _TS_DECLARATION.finditer(code)}
    required = _normalize_symbols(required_symbols)
    missing = [symbol for symbol in required if symbol not in declarations and not re.search(rf"\b{re.escape(symbol)}\b", code)]
    evidence["declarations"] = sorted(declarations)
    evidence["required_symbols"] = required
    evidence["missing_symbols"] = missing
    if missing:
        return _failure(
            "required_symbol_missing",
            evidence=evidence,
            hint=f"Implement and expose the required symbols: {', '.join(missing)}.",
        )

    if not _TS_SYNTAX_TOKENS.search(code):
        return _failure(
            "code_syntax_suspect",
            evidence=evidence,
            hint="Return syntactically structured TypeScript with declarations and delimiters.",
        )

    return _success(code, evidence=evidence)