yau-plant-assistant/api/agent.py

"""LangGraph agent — one branch per question class.

The graph is deliberately shallow. There is no general-purpose tool loop and no
"let the model decide what to do", because the class decides the tool path and
the class is assigned before generation starts. A loop that could route a
procedural question through the historical branch is the failure mode this
whole design exists to prevent.

    classify ─┬─ historical ─┐
              ├─ reference   ├─ generate → validate → (retry once) → answer
              ├─ procedural  │
              ├─ advisory   ─┘
              └─ unclear ────── clarify (no generation, no tools)

Each branch gathers its own evidence and hands a payload to the contract. The
prose model never sees a tool it was not given for its class.
"""

from __future__ import annotations

import json
import logging
from typing import Any, TypedDict

from langgraph.graph import END, StateGraph
from openai import AzureOpenAI

import classifier
import tools.equipment as equipment_tool
import tools.metrics as metrics
import tools.retrieval as retrieval
import stub
from config import settings
from contracts import (ContractViolation, QuestionClass, documented_limits,
                       procedure_identity)
from guardrails import enforce_contract

log = logging.getLogger("agent")


class State(TypedDict, total=False):
    question: str
    classification: classifier.Classification
    evidence: dict[str, Any]
    payload: dict[str, Any]
    answer: Any
    trace: Any
    forced_class: str | None


def _client() -> AzureOpenAI:
    cfg = settings()
    return AzureOpenAI(
        azure_endpoint=cfg.azure_openai_endpoint,
        api_key=cfg.azure_openai_api_key,
        api_version=cfg.azure_openai_api_version,
    )


def _embed(text: str, *, client: AzureOpenAI) -> list[float]:
    cfg = settings()
    return client.embeddings.create(model=cfg.embed_deployment, input=text).data[0].embedding


# ---------------------------------------------------------------------------
# System prompts.
#
# Byte-identical between calls so prompt caching applies - never interpolate
# the question, the evidence or the date into these. Everything variable goes
# in the user message.
#
# These prompts describe the job. They are NOT the safety mechanism; the
# contracts are. If a rule matters, it is in contracts.py as well as here, and
# if the two ever disagree, contracts.py wins by construction.
# ---------------------------------------------------------------------------

_BASE = """\
You write answers for operators at the Waterloo Road Pump Station, a three-pump
wastewater station. Be brief and concrete. Use the equipment names an operator
uses. Always state units, and state the time window any figure came from.
Never state a figure that is not in the evidence you were given.
Return JSON matching the schema you are given, and nothing else.
"""

HISTORICAL_PROMPT = _BASE + """\
You are answering a question about what happened. You have query results.
If the result set is empty, say plainly that no records were found for the
window - do not estimate, do not reason about what the number probably was.
"""

REFERENCE_PROMPT = _BASE + """\
You are answering a question about what something is or means. You have
document extracts and tag metadata. Cite every claim: document number,
revision, effective date. If nothing relevant was retrieved, say so.
"""

PROCEDURAL_PROMPT = _BASE + """\
You are IDENTIFYING a controlled procedure, not explaining it.

Say which procedure governs the action and that the operator should work from
the controlled copy. Do NOT state the document number, revision, effective
date, title or authorising role in your JSON fields: every one of those is
attached from the confirmed document header, because naming the wrong revision
of a procedure is the failure this system exists to prevent.

Quote prerequisites word for word into prerequisites_verbatim. Write no steps,
no paraphrase of steps, no summary of what the procedure involves, and no
advice about what to do first. An interlock exists because someone assessed a
hazard; a bypass procedure you reconstructed is a safety document nobody
approved.

You will not be shown the step sections at all - retrieval withholds them.

Two different situations, which must not be confused:

  * Retrieval returned NOTHING. Say so plainly and stop.
  * Retrieval returned a document that is marked draft, demo, uncontrolled,
    superseded or not for plant use. IDENTIFY IT ANYWAY and quote its
    prerequisites, and state that marking plainly in your answer. The marking
    is information the operator needs, not a reason to withhold what was found
    - "nothing was retrieved" when something was is a worse answer than either.
"""

ADVISORY_PROMPT = _BASE + """\
You are presenting EVIDENCE, not a recommendation.

Report what has actually been done: the rates used, how many operations that is
drawn from, the outcomes, when alarms occurred, and the documented limits. For
each limit, set source_file to the source_file of the extract you read it from,
copied exactly. A limit whose source_file does not match a retrieved extract is
discarded, because a limit attributed to the wrong document is worse than one
not stated. Then defer explicitly to a competent person.

Do not state a recommended value. Do not say what is best, optimal or safest.
Do not offer a range as a disguised recommendation. "Best" depends on equipment
condition and concurrent operations you cannot see, and any number you present
as an answer will be typed into a control system by someone who trusts it.
"""

PROMPTS = {
    QuestionClass.HISTORICAL: HISTORICAL_PROMPT,
    QuestionClass.REFERENCE: REFERENCE_PROMPT,
    QuestionClass.PROCEDURAL: PROCEDURAL_PROMPT,
    QuestionClass.ADVISORY: ADVISORY_PROMPT,
}


# ---------------------------------------------------------------------------
# Evidence gathering — one node per class. No node can reach another's tools.
# ---------------------------------------------------------------------------


def _resolve_entities(state: State) -> tuple[str | None, list[str]]:
    """Resolve the first equipment entity the classifier found.

    Returns (equipment_id, ambiguities). An ambiguous term is surfaced, not
    guessed at - agent behaviour on ambiguity is to say what it matched.
    """
    named = (state["classification"].entities or {}).get("equipment") or []
    for term in named:
        matches = equipment_tool.resolve(str(term))
        equipment = [m for m in matches if m.kind == "equipment"]
        if len(equipment) == 1:
            return equipment[0].canonical_id, []
        if len(equipment) > 1:
            return None, [m.canonical_id for m in equipment]
    return None, []


def gather_historical(state: State) -> State:
    trace = state.get("trace")
    equipment_id, _ = _resolve_entities(state)
    result = metrics.run(metrics.alarm_detail(equipment_id=equipment_id, days=7), trace=trace)
    _, _, window_description = metrics.rolling_window(7)
    state["evidence"] = {
        "query": result.query,
        "rows": result.rows,
        "row_count": result.row_count,
        "time_window": {**result.time_window, "description": window_description},
        "used_fixture_data": result.used_fixture_data,
    }
    return state


def gather_reference(state: State) -> State:
    equipment_id, _ = _resolve_entities(state)
    if stub.enabled():
        chunks = retrieval.rerank(
            retrieval.lexical_search(
                state["question"], top_k=8, equipment_id=equipment_id
            ),
            state["question"],
        )
    else:
        embedding = _embed(state["question"], client=_client())
        chunks = retrieval.rerank(
            retrieval.search(embedding, top_k=8, equipment_id=equipment_id),
            state["question"],
        )
    tag_rows = equipment_tool.tags_for_equipment(equipment_id) if equipment_id else []
    state["evidence"] = {
        "chunks": retrieval.as_dicts(chunks),
        "citations": [c.citation() for c in chunks],
        "tags": tag_rows,
        "used_fixture_data": False,
    }
    return state


def gather_procedural(state: State) -> State:
    """Procedures only, live revisions only. Nothing else is in scope here."""
    equipment_id, _ = _resolve_entities(state)
    if stub.enabled():
        chunks = retrieval.find_procedure_lexical(
            state["question"], equipment_id=equipment_id
        )
    else:
        embedding = _embed(state["question"], client=_client())
        chunks = retrieval.find_procedure(
            embedding, state["question"], equipment_id=equipment_id
        )
    state["evidence"] = {
        "chunks": retrieval.as_dicts(chunks),
        "citations": [c.citation() for c in chunks],
        "used_fixture_data": False,
    }
    return state


def gather_advisory(state: State) -> State:
    """Both paths: what was done (Cube) and what is allowed (documents)."""
    trace = state.get("trace")
    # Was 30 days. The historian keeps seven — every WRPS history group is
    # LIFE_TIME "1 weeks" — so a 30-day window returned the same evidence a
    # 7-day one does while telling the reader it covered a month. The sample
    # size behind an advisory answer is part of the evidence, and overstating
    # it by four times is the kind of error nobody would catch downstream.
    result = metrics.run(
        metrics.pump_down_evidence(days=metrics.HISTORY_RETENTION_DAYS), trace=trace
    )
    _, _, window_description = metrics.rolling_window(metrics.HISTORY_RETENTION_DAYS)
    if stub.enabled():
        chunks = retrieval.rerank(
            retrieval.lexical_search(state["question"], top_k=8, doc_type="design"),
            state["question"],
        )
    else:
        embedding = _embed(state["question"], client=_client())
        chunks = retrieval.rerank(
            retrieval.search(embedding, top_k=8, doc_type="design"), state["question"]
        )
    state["evidence"] = {
        "query": result.query,
        "rows": result.rows,
        "row_count": result.row_count,
        "time_window": {**result.time_window, "description": window_description},
        "chunks": retrieval.as_dicts(chunks),
        "citations": [c.citation() for c in chunks],
        "used_fixture_data": result.used_fixture_data,
    }
    return state


def gather_unclear(state: State) -> State:
    """No tools, no generation. The clarifying question is built from what the
    classifier said was missing, so it asks for something specific."""
    missing = state["classification"].missing_context or ["what you are asking about"]
    asked_for = ", ".join(missing)
    state["evidence"] = {}
    state["payload"] = {
        "question": state["question"],
        "question_class": QuestionClass.UNCLEAR,
        "answer": f"I need one more detail before I can answer: {asked_for}.",
        "clarifying_question": f"Could you tell me {asked_for}?",
        "candidate_interpretations": [],
        "used_fixture_data": False,
        # No model is called on this path even in production, but in stub mode
        # nothing classified the question either - which is the part the reader
        # needs to know.
        "stub_mode": stub.enabled(),
    }
    return state


# ---------------------------------------------------------------------------
# Generation and enforcement
# ---------------------------------------------------------------------------

_SCHEMA_HINTS: dict[QuestionClass, str] = {
    QuestionClass.HISTORICAL: (
        '{"answer": str}  - the figures, the window, and what the rows show'
    ),
    QuestionClass.REFERENCE: '{"answer": str}',
    # `procedure` is now attached ENTIRELY from evidence and config by
    # contracts.procedure_identity() - identity from the confirmed header,
    # controlled_copy_location from settings. The model supplies the prose and
    # the verbatim prerequisites only. Asking it for identity fields it had
    # never been shown is how it came to return "" for all of them.
    QuestionClass.PROCEDURAL: (
        '{"answer": str, "prerequisites_verbatim": [str]}'
    ),
    # documented_limits carries source_file, not a citation: the Citation is
    # attached from the evidence in contracts.documented_limits(). An empty []
    # here told the model nothing about the shape, so it invented
    # {"description", "citation": "WRPS-DEMO-003, Section 4"} and every
    # advisory answer failed its contract.
    QuestionClass.ADVISORY: (
        '{"answer": str, "evidence": [{"description": str, "metric": str, '
        '"value": number, "unit": str, "sample_size": int}], '
        '"documented_limits": [{"description": str, "value": number|str, '
        '"unit": str, "source_file": str}], "deferral": str}'
    ),
}


def generate(state: State) -> State:
    """Generate prose, then enforce the contract. Retry once, then error."""
    klass = state["classification"].question_class
    if klass is QuestionClass.UNCLEAR:
        return state  # gather_unclear already built the payload

    cfg = settings()
    trace = state.get("trace")
    evidence = state["evidence"]
    stubbed = stub.enabled()
    client = None if stubbed else _client()

    def call(attempt: int, previous: ContractViolation | None) -> dict[str, Any]:
        if stubbed:
            # The model's half of the payload, supplied as constants. Note what
            # happens on a retry: the stub is deterministic, so a contract
            # failure here fails again identically and surfaces as a 422. That
            # is correct - it means the contract genuinely rejects what this
            # mode produces, and no amount of retrying changes it.
            generated = stub.generated_fields(klass, state["question"], evidence)
        else:
            generated = _generate_prose(
                klass, state["question"], evidence, previous, client=client, cfg=cfg
            )
        return _assemble(klass, state["question"], evidence, generated, stubbed=stubbed)

    result = enforce_contract(call, klass, trace=trace)
    state["answer"] = result.answer
    state["payload"] = result.answer.model_dump()
    return state


def _generate_prose(
    klass: QuestionClass,
    question: str,
    evidence: dict[str, Any],
    previous: ContractViolation | None,
    *,
    client,
    cfg,
) -> dict[str, Any]:
    """The model's half of the payload: prose, and the fields only it can fill."""
    user = {
        "question": question,
        "evidence": evidence,
        "return_schema": _SCHEMA_HINTS[klass],
    }
    if previous is not None:
        # Tell it what it broke. One retry only - a model that fails a
        # safety contract twice is not going to be argued into compliance.
        user["previous_attempt_rejected"] = {
            "rule": previous.rule,
            "detail": previous.detail,
        }
    response = client.chat.completions.create(
        model=cfg.chat_deployment,
        messages=[
            {"role": "system", "content": PROMPTS[klass]},
            {"role": "user", "content": json.dumps(user, default=str)},
        ],
        temperature=0.1,
        max_tokens=cfg.max_output_tokens,
        response_format={"type": "json_object"},
    )
    return json.loads(response.choices[0].message.content)


def _assemble(
    klass: QuestionClass,
    question: str,
    evidence: dict[str, Any],
    generated: dict[str, Any],
    *,
    stubbed: bool = False,
) -> dict[str, Any]:
    """Attach the factual fields to whatever produced the prose.

    Everything factual - rows, counts, citations, the fixture flag - is attached
    HERE, from the evidence, so the thing that wrote the prose cannot alter it.
    That is true whether the prose came from the model or from stub.py, which is
    what makes the stub a fair test of this path rather than a mock of it.
    """
    payload: dict[str, Any] = {
        "question": question,
        "question_class": klass,
        "answer": generated.get("answer", ""),
        "used_fixture_data": evidence.get("used_fixture_data", False),
        "stub_mode": stubbed,
    }
    if klass is QuestionClass.HISTORICAL:
        payload.update(
            query=evidence["query"],
            rows=evidence["rows"],
            row_count=evidence["row_count"],
            time_window=evidence["time_window"],
            citations=evidence.get("citations", []),
        )
    elif klass is QuestionClass.REFERENCE:
        payload.update(
            citations=evidence.get("citations", []),
            tags_referenced=[t["tag_id"] for t in evidence.get("tags", [])],
        )
    elif klass is QuestionClass.PROCEDURAL:
        payload.update(
            procedure=procedure_identity(evidence),
            prerequisites_verbatim=generated.get("prerequisites_verbatim", []),
            steps_provided=False,
            citations=evidence.get("citations", []),
        )
    elif klass is QuestionClass.ADVISORY:
        payload.update(
            evidence=generated.get("evidence", []),
            documented_limits=documented_limits(evidence, generated),
            recommendation_given=False,
            deferral=generated.get("deferral", ""),
            citations=evidence.get("citations", []),
        )
    return payload


def validate_unclear(state: State) -> State:
    from contracts import validate_answer

    if state["classification"].question_class is QuestionClass.UNCLEAR:
        state["answer"] = validate_answer(state["payload"], QuestionClass.UNCLEAR)
    return state


# ---------------------------------------------------------------------------
# Graph
# ---------------------------------------------------------------------------


def classify_node(state: State) -> State:
    if stub.enabled():
        # No model, so nothing is classified - the caller said which class this
        # is. apply_safety_rules() still runs over the result inside
        # stub.classify(), because those rules are the part that has to hold
        # whoever chose the label.
        state["classification"] = stub.classify(
            state["question"], state.get("forced_class")
        )
        return state
    state["classification"] = classifier.classify(
        state["question"], client=_client(), trace=state.get("trace")
    )
    return state


def route(state: State) -> str:
    return state["classification"].question_class.value


def build_graph():
    graph = StateGraph(State)
    graph.add_node("classify", classify_node)
    graph.add_node("historical", gather_historical)
    graph.add_node("reference", gather_reference)
    graph.add_node("procedural", gather_procedural)
    graph.add_node("advisory", gather_advisory)
    graph.add_node("unclear", gather_unclear)
    graph.add_node("generate", generate)
    graph.add_node("finalise", validate_unclear)

    graph.set_entry_point("classify")
    graph.add_conditional_edges(
        "classify",
        route,
        {
            "historical": "historical",
            "reference": "reference",
            "procedural": "procedural",
            "advisory": "advisory",
            "unclear": "unclear",
        },
    )
    for node in ("historical", "reference", "procedural", "advisory", "unclear"):
        graph.add_edge(node, "generate")
    graph.add_edge("generate", "finalise")
    graph.add_edge("finalise", END)
    return graph.compile()


_GRAPH = None


def answer(question: str, *, trace=None, forced_class: str | None = None):
    """Answer one question. Raises ContractViolation if the contract cannot be
    met - the caller returns an error, never a partial answer.

    forced_class is ignored unless NO_LLM_STUB is on. Letting a caller pick its
    own contract is precisely what the classifier exists to prevent, so the
    check is here rather than being left to the caller to remember.
    """
    global _GRAPH
    if _GRAPH is None:
        _GRAPH = build_graph()
    if forced_class and not stub.enabled():
        log.warning("ignoring forced class %r - NO_LLM_STUB is off", forced_class)
        forced_class = None
    final = _GRAPH.invoke(
        {"question": question, "trace": trace, "forced_class": forced_class}
    )
    return final["answer"], final["classification"]