Jyotisha/scripts/rectification/probe_question_contract.py

"""Shared probe → choice-card contract for Python event probes.

Must stay aligned with frontend/src/lib/rectification-agentic/v9/probe-question-contract.ts.
Canonical bytes live in contracts/probe-question-v1.json.
"""

from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any, Sequence, TypedDict

QUESTION_CONTRACT_VERSION = "probe-question-v1"
ANSWER_CLASSES = ("yes", "weak_yes", "no", "unsure")
EXISTENCE_STYLE_OPTIONS: tuple[dict[str, str], ...] = (
    {"label": "明确发生且时间吻合", "answer_class": "yes"},
    {"label": "发生过但程度较弱", "answer_class": "weak_yes"},
    {"label": "明确没有发生", "answer_class": "no"},
    {"label": "这段记不清楚", "answer_class": "unsure"},
)
QUALITY_STYLE_OPTIONS: tuple[dict[str, str], ...] = (
    {"label": "发挥明显失常或压力很大", "answer_class": "yes"},
    {"label": "有压力但不算明显失常", "answer_class": "weak_yes"},
    {"label": "发挥正常、没有明显失常", "answer_class": "no"},
    {"label": "这段记不清楚", "answer_class": "unsure"},
)
VARGA_NONE_STYLE_OPTION = {"label": "都不是这些特质", "answer_class": "no"}
UNSURE_STYLE_OPTION = {"label": "这段记不清楚", "answer_class": "unsure"}
FORBIDDEN_COPY_TOKENS = ("外貌", "体质", "胎记", "疤痕", "伤疤", "身高", "体型")
LABEL_MIN = 4
LABEL_MAX = 80
_FORBIDDEN = FORBIDDEN_COPY_TOKENS
_CLOCK = re.compile(r"(?:[01]?\d|2[0-3]):[0-5]\d")
_ANSWER_CLASS_SUFFIX = re.compile(r"·(?:yes|weak_yes|no|unsure)$")
_CONTRACT_PATH = Path(__file__).resolve().parents[2] / "contracts" / "probe-question-v1.json"


class StyleOptionsOk(TypedDict):
    ok: bool
    options: list[dict[str, str]]


class StyleOptionsErr(TypedDict):
    ok: bool
    reason: str


class ProbeRenderOk(TypedDict):
    ok: bool


class ProbeRenderErr(TypedDict):
    ok: bool
    reason: str


def probe_question_kind(value: Any) -> str:
    if value in {"varga_style", "event_quality"}:
        return str(value)
    return "existence"


def question_contract_version_is_compatible(value: Any) -> bool:
    if value is None:
        return True
    if isinstance(value, str):
        return value == QUESTION_CONTRACT_VERSION
    if not isinstance(value, dict):
        return False
    version = value.get("version") or value.get("question_contract_version")
    return version is None or version == QUESTION_CONTRACT_VERSION


def probe_question_contract_payload() -> dict[str, Any]:
    return {
        "version": QUESTION_CONTRACT_VERSION,
        "answer_classes": list(ANSWER_CLASSES),
        "label_min": LABEL_MIN,
        "label_max": LABEL_MAX,
        "forbidden_copy_tokens": list(FORBIDDEN_COPY_TOKENS),
        "existence_style_options": [dict(item) for item in EXISTENCE_STYLE_OPTIONS],
        "quality_style_options": [dict(item) for item in QUALITY_STYLE_OPTIONS],
        "varga_none_style_option": dict(VARGA_NONE_STYLE_OPTION),
        "unsure_style_option": dict(UNSURE_STYLE_OPTION),
    }


def canonical_probe_question_contract_json() -> str:
    return json.dumps(probe_question_contract_payload(), ensure_ascii=False, indent=2) + "\n"


def load_probe_question_contract_golden() -> str:
    return _CONTRACT_PATH.read_text(encoding="utf-8")


def clipped_probe_label(value: Any, minimum: int = LABEL_MIN, maximum: int = LABEL_MAX) -> str | None:
    if not isinstance(value, str):
        return None
    text = " ".join(value.split())
    if len(text) < minimum or len(text) > maximum:
        return None
    if any(token in text for token in _FORBIDDEN) or _CLOCK.search(text):
        return None
    return text


def _label_reject_reason(value: Any) -> str | None:
    if not isinstance(value, str):
        return "empty"
    text = " ".join(value.split())
    if len(text) < LABEL_MIN or len(text) > LABEL_MAX:
        return "label_length"
    if any(token in text for token in _FORBIDDEN) or _CLOCK.search(text):
        return "forbidden_copy"
    return None


def _incoming_option(row: Any) -> dict[str, str] | dict[str, str] | None:
    if not isinstance(row, dict):
        return None
    answer = row.get("answer_class") or row.get("answerClass")
    reject = _label_reject_reason(row.get("label"))
    if reject == "empty" or answer not in ANSWER_CLASSES:
        return None
    if reject:
        return {"reason": reject}
    label = clipped_probe_label(row.get("label"))
    if not label:
        return {"reason": "not_renderable"}
    payload = {"label": label, "answer_class": str(answer)}
    sign = row.get("sign")
    if isinstance(sign, str) and sign.strip():
        payload["sign"] = sign.strip()
    return payload


def complete_style_options(
    choice_kind: Any,
    style_options: Sequence[Any] | None = None,
) -> StyleOptionsOk | StyleOptionsErr:
    kind = probe_question_kind(choice_kind)
    incoming_reason: str | None = None
    incoming: list[dict[str, str]] = []
    for row in style_options or []:
        parsed = _incoming_option(row)
        if not parsed:
            continue
        if "reason" in parsed and "label" not in parsed:
            incoming_reason = incoming_reason or str(parsed["reason"])
            continue
        incoming.append(parsed)
    by_class: dict[str, dict[str, str]] = {}
    if kind == "varga_style":
        for option in incoming:
            by_class[option["answer_class"]] = option
        by_class.setdefault("unsure", dict(UNSURE_STYLE_OPTION))
        scoring = [item for item in ANSWER_CLASSES if item != "unsure" and item in by_class]
        if len(scoring) < 2:
            return {"ok": False, "reason": incoming_reason or "varga_insufficient_scoring"}
        by_class.setdefault("no", dict(VARGA_NONE_STYLE_OPTION))
        if "yes" not in by_class or "weak_yes" not in by_class:
            return {"ok": False, "reason": "varga_missing_weak_yes"}
    else:
        catalog = QUALITY_STYLE_OPTIONS if kind == "event_quality" else EXISTENCE_STYLE_OPTIONS
        for option in catalog:
            by_class[option["answer_class"]] = dict(option)
        for option in incoming:
            by_class[option["answer_class"]] = option
    ordered: list[dict[str, str]] = []
    seen: set[str] = set()
    for index, answer_class in enumerate(ANSWER_CLASSES):
        option = by_class.get(answer_class)
        if not option:
            return {"ok": False, "reason": "not_renderable"}
        label = option["label"]
        if label in seen and option.get("sign"):
            label = f"{label}（{option['sign']}）"
        if label in seen:
            label = f"{option['label']}·{index + 1}"
        seen.add(label)
        ordered.append({**option, "label": label})
    labels = {item["label"] for item in ordered}
    classes = {item["answer_class"] for item in ordered}
    if (
        len(ordered) != 4
        or len(labels) != 4
        or classes != set(ANSWER_CLASSES)
        or any(_ANSWER_CLASS_SUFFIX.search(item["label"]) for item in ordered)
    ):
        return {"ok": False, "reason": "not_renderable"}
    return {"ok": True, "options": ordered}


def completed_style_options(
    choice_kind: Any,
    style_options: Sequence[Any] | None = None,
) -> list[dict[str, str]] | None:
    result = complete_style_options(choice_kind, style_options)
    if result.get("ok"):
        return result.get("options")  # type: ignore[return-value]
    return None


def is_renderable_probe(probe: dict[str, Any]) -> ProbeRenderOk | ProbeRenderErr:
    gain = probe.get("information_gain")
    if not isinstance(gain, (int, float)) or gain <= 0:
        return {"ok": False, "reason": "zero_gain"}
    candidate_ids = probe.get("candidate_ids") or []
    outcomes = probe.get("expected_outcomes") or []
    if len(candidate_ids) < 2:
        return {"ok": False, "reason": "insufficient_candidates"}
    if len(outcomes) < 2:
        return {"ok": False, "reason": "insufficient_outcomes"}
    completed = complete_style_options(probe.get("choice_kind"), probe.get("style_options"))
    if not completed.get("ok"):
        return {"ok": False, "reason": str(completed.get("reason") or "not_renderable")}
    return {"ok": True}