json_guard.py

"""Validate model output and guarantee the UI never crashes.

If the model returns malformed JSON, or names a fault that is not in the
rule-derived candidate set, we fall back to the top deterministic candidate.
The deterministic layer is the floor; the model can refine but not fabricate.
"""
from __future__ import annotations

import json
import re
from dataclasses import dataclass
from typing import List

from fault_rules import Candidate

VALID_URGENCY = {"CRITICAL", "HIGH", "MEDIUM", "LOW", "UNKNOWN"}


@dataclass
class DiagnosisResult:
    fault: str
    urgency: str
    checks: List[str]
    safety: str
    confidence: int
    grounded: bool  # True if the model's fault matched a candidate

    def to_dict(self) -> dict:
        return self.__dict__


def _extract_json(text: str) -> dict | None:
    if not text:
        return None
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if not match:
        return None
    try:
        return json.loads(match.group())
    except Exception:
        return None


_FALLBACK_CANDIDATE = Candidate(
    name="Inconclusive", urgency="LOW", weight=0.0,
    evidence="No deterministic candidate was available.",
)

# Hard caps so a runaway model response can never break the UI layout.
MAX_FAULT_LEN = 80
MAX_TEXT_LEN = 240


def _clip(text: str, limit: int) -> str:
    text = " ".join(str(text).split())  # collapse whitespace/newlines
    return text if len(text) <= limit else text[: limit - 1].rstrip() + "…"


def validate(raw_text: str, candidates: List[Candidate]) -> DiagnosisResult:
    # rank_candidates guarantees >=1, but never trust the caller blindly.
    top = candidates[0] if candidates else _FALLBACK_CANDIDATE
    candidates = candidates or [top]
    parsed = _extract_json(raw_text)

    if not isinstance(parsed, dict):
        return DiagnosisResult(
            fault=_clip(top.name, MAX_FAULT_LEN), urgency=top.urgency,
            checks=["Re-record a longer, clearer sample.",
                    "Compare against the appliance's normal sound.",
                    "If unsure, consult a technician."],
            safety="None", confidence=int(top.weight * 100), grounded=True,
        )

    fault = str(parsed.get("fault", top.name)).strip() or top.name
    urgency = str(parsed.get("urgency", top.urgency)).strip().upper()
    if urgency not in VALID_URGENCY:
        urgency = top.urgency

    checks = parsed.get("checks", [])
    if not isinstance(checks, list) or not checks:
        checks = ["Inspect the most likely component first.",
                  "Listen again after checking.",
                  "Escalate to a technician if it persists."]
    checks = [_clip(c, MAX_TEXT_LEN) for c in checks if str(c).strip()][:3]
    if not checks:
        checks = ["Inspect the most likely component first."]

    safety = _clip(parsed.get("safety", "None"), MAX_TEXT_LEN) or "None"

    try:
        confidence = int(float(parsed.get("confidence", top.weight * 100)))
    except (TypeError, ValueError):
        confidence = int(top.weight * 100)
    confidence = max(0, min(100, confidence))

    candidate_names = {c.name.lower() for c in candidates}
    grounded = fault.lower() in candidate_names or top.name == "Inconclusive"
    if not grounded:
        # Model named something outside the evidence — snap back to the floor.
        fault = top.name
        urgency = top.urgency
        confidence = min(confidence, int(top.weight * 100))
        grounded = True

    # An Inconclusive verdict should never carry HIGH/CRITICAL urgency.
    if fault.lower() == "inconclusive":
        urgency = "LOW"

    return DiagnosisResult(
        fault=_clip(fault, MAX_FAULT_LEN), urgency=urgency, checks=checks,
        safety=safety, confidence=confidence, grounded=grounded,
    )