34d2ccc576
Build spec and host brief carried in from C:\Claude and WRPS/02-env; the plant model (equipment, tags, alarm bitmask, enums, unit conversions) is derived from WRPS/04-plc/register-map.csv, WRPS/05-scada/modbus/scada-points.csv and WRPS-CTL-003. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
406 lines
15 KiB
Python
406 lines
15 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
|
|
from config import settings
|
|
from contracts import ContractViolation, QuestionClass
|
|
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
|
|
|
|
|
|
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.
|
|
|
|
Give the procedure number, revision, effective date, title, the authorising
|
|
role, and where the controlled copy is. 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.
|
|
|
|
If no controlled procedure was retrieved, say so and stop.
|
|
"""
|
|
|
|
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 with
|
|
their citations. 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:
|
|
client = _client()
|
|
equipment_id, _ = _resolve_entities(state)
|
|
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."""
|
|
client = _client()
|
|
equipment_id, _ = _resolve_entities(state)
|
|
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")
|
|
client = _client()
|
|
result = metrics.run(metrics.pump_down_evidence(days=30), trace=trace)
|
|
_, _, window_description = metrics.rolling_window(30)
|
|
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,
|
|
}
|
|
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}',
|
|
QuestionClass.PROCEDURAL: (
|
|
'{"answer": str, "procedure": {"doc_number": str, "title": str, '
|
|
'"revision": str, "effective_date": "YYYY-MM-DD", '
|
|
'"authorising_role": str, "controlled_copy_location": str}, '
|
|
'"prerequisites_verbatim": [str]}'
|
|
),
|
|
QuestionClass.ADVISORY: (
|
|
'{"answer": str, "evidence": [{"description": str, "metric": str, '
|
|
'"value": number, "unit": str, "sample_size": int}], '
|
|
'"documented_limits": [], "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()
|
|
client = _client()
|
|
trace = state.get("trace")
|
|
evidence = state["evidence"]
|
|
|
|
def call(attempt: int, previous: ContractViolation | None) -> dict[str, Any]:
|
|
user = {
|
|
"question": state["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"},
|
|
)
|
|
generated = json.loads(response.choices[0].message.content)
|
|
|
|
# The model supplies prose and its own structured fields. Everything
|
|
# factual - rows, counts, citations, the fixture flag - is attached
|
|
# here from the evidence, so the model cannot alter it.
|
|
payload: dict[str, Any] = {
|
|
"question": state["question"],
|
|
"question_class": klass,
|
|
"answer": generated.get("answer", ""),
|
|
"used_fixture_data": evidence.get("used_fixture_data", False),
|
|
}
|
|
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=generated.get("procedure"),
|
|
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=generated.get("documented_limits", []),
|
|
recommendation_given=False,
|
|
deferral=generated.get("deferral", ""),
|
|
citations=evidence.get("citations", []),
|
|
)
|
|
return payload
|
|
|
|
result = enforce_contract(call, klass, trace=trace)
|
|
state["answer"] = result.answer
|
|
state["payload"] = result.answer.model_dump()
|
|
return state
|
|
|
|
|
|
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:
|
|
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):
|
|
"""Answer one question. Raises ContractViolation if the contract cannot be
|
|
met - the caller returns an error, never a partial answer."""
|
|
global _GRAPH
|
|
if _GRAPH is None:
|
|
_GRAPH = build_graph()
|
|
final = _GRAPH.invoke({"question": question, "trace": trace})
|
|
return final["answer"], final["classification"]
|