"""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 # Drop-in for `openai.AzureOpenAI` - identical constructor and call signatures, # and it records every call to Langfuse as a generation carrying the model and # the token usage the plain client discards. That is the only thing Langfuse # can price, so this import is what makes cost appear in the UI. Reads # LANGFUSE_PUBLIC_KEY / LANGFUSE_SECRET_KEY / LANGFUSE_HOST from the # environment, which the api service already has. Falls back to plain # behaviour when they are unset - observability never breaks the answer path. from langfuse.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"]