yau-plant-assistant/api/contracts.py

"""Response contracts — one per question class, validated in Python.

The three lines this system does not cross are enforced HERE, not in a prompt.
A prompt is advisory and models drift; this module is a gate every response
passes through before it can reach an operator.

The flow, implemented in agent.py:

    generate -> validate -> (on failure) regenerate once -> validate
             -> (on failure) raise ContractViolation and return an error

A response that fails its contract is NEVER returned, not even partially, not
even with a warning attached. There is no debug flag that disables this.
"""

from __future__ import annotations

import re
from datetime import date
from enum import Enum
from typing import Any, Literal

from pydantic import BaseModel, Field, model_validator


class QuestionClass(str, Enum):
    """Assigned by the classifier before any generation happens.

    When uncertain, choose the MORE RESTRICTIVE class. The ordering used by
    classifier.py when confidence is split:

        UNCLEAR < HISTORICAL < REFERENCE < ADVISORY < PROCEDURAL

    Procedural beats Reference. Advisory beats Historical. Partly-advisory is
    advisory.
    """

    HISTORICAL = "historical"
    REFERENCE = "reference"
    PROCEDURAL = "procedural"
    ADVISORY = "advisory"
    UNCLEAR = "unclear"


RESTRICTIVENESS: dict[QuestionClass, int] = {
    QuestionClass.UNCLEAR: 0,
    QuestionClass.HISTORICAL: 1,
    QuestionClass.REFERENCE: 2,
    QuestionClass.ADVISORY: 3,
    QuestionClass.PROCEDURAL: 4,
}


class ContractViolation(Exception):
    """A generated response broke its class contract.

    Carries the offending output so guardrails.py can log the whole thing to
    Langfuse. It must not be included in the operator-facing error.
    """

    def __init__(self, rule: str, detail: str, offending_output: str = "") -> None:
        super().__init__(f"{rule}: {detail}")
        self.rule = rule
        self.detail = detail
        self.offending_output = offending_output


# ---------------------------------------------------------------------------
# Shared pieces
# ---------------------------------------------------------------------------


class Citation(BaseModel):
    """A pointer to a controlled document. Never a paraphrase of one."""

    doc_number: str
    title: str
    revision: str
    effective_date: date | None = None
    page: int | None = None
    section_title: str | None = None
    source_file: str
    superseded: bool = False

    @model_validator(mode="after")
    def reject_superseded(self) -> "Citation":
        # Citing a withdrawn revision is worse than finding nothing. Retrieval
        # filters these out; this is the second line, in case a caller passes
        # include_superseded and forgets what that means.
        if self.superseded:
            raise ContractViolation(
                "superseded_citation",
                f"{self.doc_number} rev {self.revision} is superseded",
            )
        return self


class TimeWindow(BaseModel):
    """Every data-dependent answer states the window it used, in site local."""

    start: str
    end: str
    timezone: str
    description: str = Field(
        description="Plain words, e.g. 'rolling 7 days to 2026-08-20 09:00 AEST'"
    )


class BaseAnswer(BaseModel):
    question: str
    question_class: QuestionClass
    answer: str
    used_fixture_data: bool = Field(
        default=False,
        description=(
            "TRUE when any row behind this answer came from db/002_fixtures.sql. "
            "The UI shows a banner. Never suppress it to make a demo cleaner."
        ),
    )
    stub_mode: bool = Field(
        default=False,
        description=(
            "TRUE when NO_LLM_STUB produced this answer: no model was called, "
            "the class was chosen by hand and retrieval was lexical. Carried in "
            "the contract rather than added by the UI, for the same reason as "
            "used_fixture_data - an answer nobody generated must not be "
            "indistinguishable from one that was. See api/stub.py."
        ),
    )


# ---------------------------------------------------------------------------
# Historical — Cube, optionally with document context
# ---------------------------------------------------------------------------


class HistoricalAnswer(BaseAnswer):
    question_class: Literal[QuestionClass.HISTORICAL] = QuestionClass.HISTORICAL
    query: dict[str, Any] = Field(description="The Cube query actually executed.")
    row_count: int
    rows: list[dict[str, Any]]
    time_window: TimeWindow
    citations: list[Citation] = Field(default_factory=list)

    @model_validator(mode="after")
    def check(self) -> "HistoricalAnswer":
        if self.row_count != len(self.rows):
            raise ContractViolation(
                "row_count_mismatch",
                f"row_count={self.row_count} but {len(self.rows)} rows attached",
                self.answer,
            )
        if self.row_count == 0 and not _says_no_records(self.answer):
            # Zero rows means "no records found", never an invented figure.
            raise ContractViolation(
                "zero_rows_not_declared",
                "query returned no rows but the answer does not say so",
                self.answer,
            )
        if self.row_count == 0 and _contains_quantity(
            self.answer, ignoring=[self.time_window.description,
                                   self.time_window.start, self.time_window.end]
        ):
            raise ContractViolation(
                "figure_without_data",
                "answer states a quantity but the query returned no rows",
                self.answer,
            )
        return self


# ---------------------------------------------------------------------------
# Reference — retrieval plus tag metadata
# ---------------------------------------------------------------------------


class ReferenceAnswer(BaseAnswer):
    question_class: Literal[QuestionClass.REFERENCE] = QuestionClass.REFERENCE
    citations: list[Citation]
    tags_referenced: list[str] = Field(default_factory=list)

    @model_validator(mode="after")
    def check(self) -> "ReferenceAnswer":
        if not self.citations and not _says_not_found(self.answer):
            raise ContractViolation(
                "uncited_reference",
                "no citations and the answer does not say nothing was found",
                self.answer,
            )
        return self


# ---------------------------------------------------------------------------
# Procedural — the class that must never produce instructions
# ---------------------------------------------------------------------------


class ProcedureIdentity(BaseModel):
    doc_number: str
    title: str
    revision: str
    effective_date: date | None
    authorising_role: str | None = None
    controlled_copy_location: str


class ProceduralAnswer(BaseAnswer):
    """Locate and cite. Never paraphrase, never reconstruct, never instruct.

    An interlock exists because someone assessed a hazard. A bypass procedure
    reassembled from retrieved fragments is a safety document nobody approved.
    """

    question_class: Literal[QuestionClass.PROCEDURAL] = QuestionClass.PROCEDURAL
    procedure: ProcedureIdentity | None
    prerequisites_verbatim: list[str] = Field(
        default_factory=list,
        description="Quoted exactly from the controlled document. Not summarised.",
    )
    steps_provided: Literal[False] = False
    citations: list[Citation]
    scope_banner: str = Field(
        default=(
            "This assistant has identified the controlled procedure. It has not "
            "reproduced or summarised the steps. Work from the controlled copy."
        )
    )

    @model_validator(mode="after")
    def check(self) -> "ProceduralAnswer":
        if self.procedure is None and not _says_not_found(self.answer):
            raise ContractViolation(
                "no_procedure_no_refusal",
                "no procedure identified and the answer does not say so",
                self.answer,
            )
        offending = _find_instruction_language(self.answer)
        if offending:
            raise ContractViolation(
                "synthesised_steps",
                f"instruction language in a procedural answer: {offending!r}",
                self.answer,
            )
        if _looks_like_a_step_list(self.answer):
            raise ContractViolation(
                "step_sequence_emitted",
                "answer contains an enumerated action sequence",
                self.answer,
            )
        if not self.scope_banner.strip():
            raise ContractViolation(
                "missing_scope_banner",
                "procedural answers must carry the scope banner",
                self.answer,
            )
        return self


# ---------------------------------------------------------------------------
# Advisory — evidence and a deferral, never a number presented as the answer
# ---------------------------------------------------------------------------


class Evidence(BaseModel):
    """What was actually observed. Facts with provenance, not conclusions."""

    description: str
    metric: str
    value: float | str | None = None
    unit: str | None = None
    sample_size: int | None = Field(
        default=None,
        description="How many operations/rows this is drawn from. Always shown.",
    )
    time_window: TimeWindow | None = None
    source: Literal["cube", "document"] = "cube"


class DocumentedLimit(BaseModel):
    """A limit somebody approved, with the document that approved it."""

    description: str
    value: float | str
    unit: str | None = None
    citation: Citation


class AdvisoryAnswer(BaseAnswer):
    """'Best' depends on equipment condition and concurrent operations this
    system cannot see. A number presented as an answer gets typed into a
    control system by someone who trusts it.
    """

    question_class: Literal[QuestionClass.ADVISORY] = QuestionClass.ADVISORY
    evidence: list[Evidence]
    documented_limits: list[DocumentedLimit] = Field(default_factory=list)
    recommendation_given: Literal[False] = False
    deferral: str = Field(
        description="Explicit statement of who decides and why not this system."
    )
    citations: list[Citation] = Field(default_factory=list)
    scope_banner: str = Field(
        default=(
            "This assistant has presented what the plant history shows and what "
            "the controlled documents state. It has not recommended a setpoint "
            "or operating parameter. That decision needs a competent person with "
            "sight of current equipment condition and concurrent operations."
        )
    )

    @model_validator(mode="after")
    def check(self) -> "AdvisoryAnswer":
        if not self.evidence and not _says_not_found(self.answer):
            raise ContractViolation(
                "advisory_without_evidence",
                "no evidence and no statement that none was found",
                self.answer,
            )
        if not self.deferral.strip():
            raise ContractViolation(
                "missing_deferral", "advisory answers must defer explicitly", self.answer
            )
        offending = _find_recommendation_language(self.answer)
        if offending:
            raise ContractViolation(
                "recommendation_given",
                f"recommendation language in an advisory answer: {offending!r}",
                self.answer,
            )
        if not self.scope_banner.strip():
            raise ContractViolation(
                "missing_scope_banner",
                "advisory answers must carry the scope banner",
                self.answer,
            )
        return self


# ---------------------------------------------------------------------------
# Unclear — ask, do not guess
# ---------------------------------------------------------------------------


class UnclearAnswer(BaseAnswer):
    question_class: Literal[QuestionClass.UNCLEAR] = QuestionClass.UNCLEAR
    clarifying_question: str
    candidate_interpretations: list[str] = Field(default_factory=list)

    @model_validator(mode="after")
    def check(self) -> "UnclearAnswer":
        if not self.clarifying_question.strip():
            raise ContractViolation(
                "no_clarifying_question",
                "unclear class must ask something specific",
                self.answer,
            )
        return self


Answer = (
    HistoricalAnswer | ReferenceAnswer | ProceduralAnswer | AdvisoryAnswer | UnclearAnswer
)

CONTRACT_FOR: dict[QuestionClass, type[BaseAnswer]] = {
    QuestionClass.HISTORICAL: HistoricalAnswer,
    QuestionClass.REFERENCE: ReferenceAnswer,
    QuestionClass.PROCEDURAL: ProceduralAnswer,
    QuestionClass.ADVISORY: AdvisoryAnswer,
    QuestionClass.UNCLEAR: UnclearAnswer,
}


# ---------------------------------------------------------------------------
# Detectors.
#
# These are blunt on purpose. A false positive costs a regeneration; a false
# negative puts a synthesised bypass procedure in front of an operator. When
# tightening one of these, add the case to eval/testset.jsonl first.
# ---------------------------------------------------------------------------

_INSTRUCTION_PATTERNS = [
    r"\byou (?:should|must|need to|can|will) (?:then )?(?:navigate|press|set|turn|open|close|isolate|bypass|lift|remove|switch|select|enter|write|start|stop|reset)\b",
    r"\bto (?:lift|bypass|remove|defeat|override|disable) the interlock,?\s",
    r"\bfirst,?\s+(?:navigate|press|set|turn|open|close|isolate|select|go to|log in)\b",
    r"\bthen,?\s+(?:navigate|press|set|turn|open|close|isolate|select|go to)\b",
    r"\bnext,?\s+(?:navigate|press|set|turn|open|close|isolate|select)\b",
    r"\bfollow these steps\b",
    r"\bhere(?:'s| is) how to\b",
    r"\bthe procedure is as follows\b",
    r"\bstep 1\b",
]

_RECOMMENDATION_PATTERNS = [
    r"\b(?:i|we) (?:recommend|suggest|advise)\b",
    r"\b(?:the )?(?:recommended|suggested|optimal|ideal|best) (?:flowrate|flow rate|rate|setpoint|set point|speed|level|value|setting)\b",
    r"\byou should (?:set|use|run|target|aim for|operate at)\b",
    r"\bset (?:it|the setpoint|the level|the rate|the flow) to\b",
    r"\bthe best (?:way|option|choice) (?:is|would be) to\b",
    r"\baim for (?:a |an )?\d",
    r"\btarget (?:a |an )?\d+(?:\.\d+)?\s*(?:%|m3/h|l/s|kpa|hz|mm)\b",
]

_NO_RECORDS_PATTERNS = [
    r"\bno (?:records|rows|data|results|matching records)\b",
    r"\bnothing (?:was )?(?:found|returned|recorded)\b",
    r"\bdid not return any\b",
    r"\bthere (?:are|were) no\b",
    r"\bno (?:such )?(?:alarms?|events?|operations?|occurrences?)\b",
]

_NOT_FOUND_PATTERNS = _NO_RECORDS_PATTERNS + [
    r"\bcould not (?:find|locate|identify)\b",
    r"\bno (?:controlled )?(?:procedure|document|documents?)\b",
    r"\bnot (?:available|held|in the document set)\b",
    r"\bi (?:do not|don't) have\b",
]

# A number that reads as a quantity: 6, 6.2, 6 times, 42%. Deliberately excludes
# dates and tag numbers, which are identity, not measurement.
_QUANTITY = re.compile(
    r"(?<![\w./-])\d+(?:\.\d+)?\s*(?:times|occurrences|%|m3/h|L/s|kPa|Hz|mm|hours|h|s)?(?![\w./-])",
    re.IGNORECASE,
)
_STEP_LIST = re.compile(r"^\s*(?:\d+[.)]|step\s+\d+\b)", re.IGNORECASE | re.MULTILINE)


def _first_match(text: str, patterns: list[str]) -> str | None:
    for pattern in patterns:
        found = re.search(pattern, text, re.IGNORECASE)
        if found:
            return found.group(0)
    return None


def _find_instruction_language(text: str) -> str | None:
    return _first_match(text, _INSTRUCTION_PATTERNS)


def _find_recommendation_language(text: str) -> str | None:
    return _first_match(text, _RECOMMENDATION_PATTERNS)


def _says_no_records(text: str) -> bool:
    return _first_match(text, _NO_RECORDS_PATTERNS) is not None


def _says_not_found(text: str) -> bool:
    return _first_match(text, _NOT_FOUND_PATTERNS) is not None


# The time window is full of numbers - dates, times, "rolling 7 days" - and
# stating it is exactly what a good zero-row answer does. Strip it before
# looking for fabricated figures, or the honest answer trips the check.
_WINDOW_NOISE = re.compile(
    r"rolling\s+\d+\s+days?|\d{4}-\d{2}-\d{2}(?:[T ]\d{2}:\d{2}(?::\d{2})?)?",
    re.IGNORECASE,
)


def _contains_quantity(text: str, ignoring: list[str] | None = None) -> bool:
    """Any bare quantity in a zero-row answer is a fabricated figure.

    Tolerates a literal zero - "the query returned 0 rows" is honest - and
    tolerates the time window, which the answer is required to state.
    """
    for phrase in ignoring or []:
        if phrase:
            text = text.replace(phrase, " ")
    text = _WINDOW_NOISE.sub(" ", text)
    for match in _QUANTITY.finditer(text):
        token = match.group(0).strip()
        if token.split()[0] not in {"0", "0.0", "zero"}:
            return True
    return False


def _looks_like_a_step_list(text: str) -> bool:
    """Two or more enumerated lines is an action sequence, whoever numbered it.

    A single numbered line is usually a clause reference like '4.2' inside a
    quoted prerequisite, which is legitimate.
    """
    return len(_STEP_LIST.findall(text)) >= 2


def validate_answer(payload: dict[str, Any], klass: QuestionClass) -> BaseAnswer:
    """Validate a generated payload against its class contract.

    Raises ContractViolation (never returns a partially-valid object).
    """
    model = CONTRACT_FOR[klass]
    try:
        return model.model_validate(payload)
    except ContractViolation:
        raise
    except Exception as exc:  # pydantic ValidationError and anything else
        raise ContractViolation(
            "schema_invalid",
            f"{klass.value} payload did not match its contract: {exc}",
            str(payload.get("answer", "")),
        ) from exc