ts-as-modality/src/tsmm/eval/qa_judge.py

"""QA judging (T5.3): rule track (JSON/number tolerance) + LLM-judge track (0-5).

- Rule track: objective, for anomaly/describe/forecast tasks. JSON region IoU,
  numerical relative-tolerance match. Zero cost, fully reproducible.
- LLM-judge track: subjective, for root_cause/compare/event. Scores 0-5 via an
  LLM (GPT-4o in production). Because this host has no network, the default
  ``backend="stub"`` produces a deterministic heuristic score (exact/substring
  match → high; dissimilar → low) so the pipeline is runnable end-to-end offline;
  swap to ``backend="openai"`` (or pass a callable) for real judging.

Design ref: §5.3 (双轨) and §6.2 (GPT-4o judge, eval-only).
"""
from __future__ import annotations

import json
import re
from typing import Any, Callable, Dict, List, Optional, Tuple

from .parse_answer import _extract_balanced_json, parse_anomaly_segments


_JSON_RE = re.compile(r"\{[^{}]*\}", re.DOTALL)


def _extract_json(pred: str) -> Optional[dict]:
    if pred is None:
        return None
    # prefer balanced (nested-capable) extraction first
    for blob in _extract_balanced_json(pred):
        try:
            obj = json.loads(blob)
            if isinstance(obj, dict):
                return obj
        except json.JSONDecodeError:
            continue
    # fallback to flat blobs
    for blob in _JSON_RE.findall(pred):
        try:
            return json.loads(blob)
        except json.JSONDecodeError:
            continue
    return None


# ─── rule-track per-category scorers ──────────────────────────────────────

def rule_score_describe(pred: str, ref: dict, rtol: float = 0.05) -> Dict[str, Any]:
    """Compare numerical statistics with relative tolerance.

    ``ref`` is a dict of {stat_name: value}. A stat scores 1 if |pred-ref|/|ref|
    <= rtol (or both zero). Overall score = fraction of matched stats.
    """
    obj = _extract_json(pred)
    if obj is None:
        return {"score": 0.0, "parsed": False}
    if not ref:
        return {"score": 1.0 if not obj else 0.0, "parsed": True}
    hits = 0
    total = 0
    for k, rv in ref.items():
        total += 1
        pv = obj.get(k)
        if pv is None:
            continue
        try:
            pv = float(pv); rv = float(rv)
        except (TypeError, ValueError):
            continue
        if rv == 0 and pv == 0:
            hits += 1
        elif rv != 0 and abs(pv - rv) / abs(rv) <= rtol:
            hits += 1
    return {"score": hits / total, "parsed": True}


def rule_score_anomaly(pred: str, ref_segs: List[Tuple[int, int]], T: int) -> Dict[str, Any]:
    """Anomaly region IoU as the score."""
    pred_segs, ok = parse_anomaly_segments(pred)
    if not ok and not pred_segs:
        return {"score": 0.0, "parsed": False}
    # IoU over point sets
    def to_set(segs):
        s = set()
        for a, b in segs:
            a = max(0, a); b = min(T - 1, b)
            s.update(range(a, b + 1))
        return s
    ps, rs = to_set(pred_segs), to_set(ref_segs)
    if not ps and not rs:
        return {"score": 1.0, "parsed": True}
    inter = len(ps & rs)
    union = len(ps | rs)
    return {"score": inter / union if union else 0.0, "parsed": True}


def rule_score_forecast(pred: str, ref: List[float], rtol: float = 0.1) -> Dict[str, Any]:
    obj = _extract_json(pred)
    if obj is None:
        return {"score": 0.0, "parsed": False}
    vals = obj.get("forecast")
    if not isinstance(vals, list):
        return {"score": 0.0, "parsed": True}
    try:
        preds = [float(v) for v in vals]
    except (TypeError, ValueError):
        return {"score": 0.0, "parsed": True}
    n = min(len(preds), len(ref))
    if n == 0:
        return {"score": 0.0, "parsed": True}
    hits = 0
    for p, r in zip(preds[:n], ref[:n]):
        if r == 0 and p == 0:
            hits += 1
        elif r != 0 and abs(p - r) / abs(r) <= rtol:
            hits += 1
    return {"score": hits / len(ref) if ref else 0.0, "parsed": True}


# ─── rule-track dispatcher ──────────────────────────────────────────────────

class RuleJudge:
    """Dispatch to the right rule scorer by category."""

    def judge(self, category: str, pred: str, **ref) -> Dict[str, Any]:
        if category == "anomaly":
            out = rule_score_anomaly(pred, ref.get("ref_segs", []), ref.get("T", 1))
        elif category == "describe":
            out = rule_score_describe(pred, ref.get("ref_answer", {}))
        elif category == "forecast":
            out = rule_score_forecast(pred, ref.get("ref_forecast", []))
        else:
            # root_cause / compare / event → not rule-checkable
            return {"track": "skip", "score": None, "parsed": False}
        return {"track": "rule", **out}


# ─── LLM-judge track (0-5) ──────────────────────────────────────────────────

class LLMJudge:
    """Score free-text answers 0-5.

    ``backend``:
      - ``"stub"`` (default): deterministic offline heuristic (no network).
      - ``"openai"``: real GPT-4o judge (requires OPENAI_API_KEY + network).
      - a callable ``(question, ref, pred) -> {"score": float, "reason": str}``.
    """

    def __init__(self, backend: str = "stub", model: str = "gpt-4o",
                 api_key: Optional[str] = None, base_url: Optional[str] = None) -> None:
        self.backend = backend
        self.model = model
        self.api_key = api_key
        self.base_url = base_url

    def score(self, question: str, ref_answer: str, pred_answer: str) -> Dict[str, Any]:
        if callable(self.backend):
            return self.backend(question, ref_answer, pred_answer)
        if self.backend == "stub":
            return self._stub_score(question, ref_answer, pred_answer)
        if self.backend == "openai":
            return self._openai_score(question, ref_answer, pred_answer)
        raise ValueError(f"unknown backend: {self.backend}")

    def _stub_score(self, question, ref, pred) -> Dict[str, Any]:
        # Offline heuristic: exact/substring/token-overlap similarity → 0-5.
        ref = (ref or "").strip().lower()
        pred = (pred or "").strip().lower()
        if not ref and not pred:
            return {"score": 5.0, "reason": "stub: both empty", "backend": "stub"}
        if ref == pred:
            return {"score": 5.0, "reason": "stub: exact match", "backend": "stub"}
        rt, pt = set(ref.split()), set(pred.split())
        if not rt:
            return {"score": 1.0, "reason": "stub: empty ref", "backend": "stub"}
        overlap = len(rt & pt) / len(rt)
        score = round(5.0 * overlap, 1)
        return {"score": score, "reason": f"stub: token-overlap={overlap:.2f}", "backend": "stub"}

    def _openai_score(self, question, ref, pred) -> Dict[str, Any]:  # pragma: no cover
        import os
        from openai import OpenAI
        key = self.api_key or os.environ.get("OPENAI_API_KEY")
        client = OpenAI(api_key=key, base_url=self.base_url)
        prompt = (
            "You are a strict grader. Score the model's answer 0-5.\n"
            f"Question: {question}\nReference answer: {ref}\nModel answer: {pred}\n"
            'Reply JSON: {"score": <0-5>, "reason": "..."}'
        )
        resp = client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.0,
        )
        text = resp.choices[0].message.content
        try:
            obj = json.loads(text)
            return {"score": float(obj["score"]), "reason": obj.get("reason", ""),
                    "backend": "openai"}
        except (json.JSONDecodeError, KeyError):
            return {"score": 0.0, "reason": f"parse-fail: {text[:100]}",
                    "backend": "openai"}