dbd72dccc7
metrics.run() computed outside_retention but nothing consumed it, so the flag was dead and the distinction it exists to carry never reached an answer. Both gather_historical and gather_advisory now pass it, with retention_days, as EVIDENCE - not as an instruction in a prompt. Also corrects eval case H29 to record what it actually is: a deliberately failing case, pinned before the fix per the house convention. It cannot pass yet for a reason that predates this change - gather_historical always queries a rolling 7 days and never parses the window the question asks about, so outside_retention can never be true on that path. Observed today: "No records were found for June 2026. The data provided is for the window from 2026-08-24 to 2026-08-31." Honest, and it states the window, but it leads with "no records were found" and never says the historian keeps only seven days - so an operator cannot tell a retention limit from a quiet month. Question-window parsing is the missing piece and is a separate change. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
529 lines
20 KiB
Python
529 lines
20 KiB
Python
"""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,
|
|
# Zero rows because the historian does not reach that far is NOT the
|
|
# same fact as zero rows because nothing happened, and only one of them
|
|
# is true. Passed as evidence so the answer can say which - not as an
|
|
# instruction in a prompt.
|
|
"outside_retention": result.outside_retention,
|
|
"retention_days": metrics.HISTORY_RETENTION_DAYS,
|
|
}
|
|
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,
|
|
# An advisory answer is only as good as the sample size behind it, so
|
|
# the window's limits are part of the evidence rather than a footnote.
|
|
"outside_retention": result.outside_retention,
|
|
"retention_days": metrics.HISTORY_RETENTION_DAYS,
|
|
}
|
|
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"]
|