Identify the procedure first, then withhold its steps

Two retrieval faults, both only visible once a real model ran.

find_procedure ranked a procedure's chunks by similarity to the question. For
"how do I lift the interlock on Pump 02" the closest chunks ARE the step list -
so the branch whose entire purpose is not reproducing steps was handing the
answer writer nothing but steps, while omitting the header block carrying the
title and the authorising role. The model was being asked for a title it had
never been shown, and returned "".

Once the document is identified, WHICH document it is settles what to send:
the header and the prerequisites, in document order, never the steps.
STEP_SECTION_RE is a second line behind ProceduralAnswer's instruction-language
check, not a replacement for it - the contract still rejects instruction
language whatever arrives here. This removes the temptation rather than relying
on catching it.

Separately, rerank did 0.75 * chunk.similarity where similarity is NULL for a
chunk ingested with --no-embed: `1 - (NULL <=> vec)` is NULL, so it raised
TypeError and 500'd the whole question rather than ranking that chunk last. It
now degrades to the lexical half - an unembedded chunk is still findable, just
not by meaning - and Chunk.similarity is typed honestly as float | None.

find_procedure_lexical takes the same identify-then-expand shape, so the stub
keeps testing the shape it always did.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

api/tests/test_rerank_null_similarity.py

@@ -0,0 +1,81 @@
+"""rerank() must rank an unembedded chunk, not crash on it.
+
+Split from test_model_contract_shapes.py because tools.retrieval imports
+psycopg: this module skips on a bare checkout and runs inside the ai-api image,
+where the whole suite runs with nothing skipped:
+
+    docker run --rm --entrypoint python yau/ai-api:local -m pytest /app/tests -q
+"""
+
+import pytest
+
+pytest.importorskip("psycopg")
+
+from tools.retrieval import Chunk, rerank  # noqa: E402
+
+
+def chunk(cid: int, text: str, similarity: float | None) -> Chunk:
+    return Chunk(
+        id=cid, source_file="procedures/X.md", doc_type="procedure",
+        doc_number="WRPS-X-001", revision="0", effective_date=None, page=1,
+        section_title="S", equipment_id=None, chunk_text=text,
+        similarity=similarity,
+    )
+
+
+def test_rerank_survives_a_null_similarity():
+    """The bug: `0.75 * None` raised TypeError and 500'd the question.
+
+    A chunk ingested with --no-embed has a NULL embedding, so
+    `1 - (NULL <=> vec)` comes back NULL.
+    """
+    chunks = [chunk(1, "interlock bypass pump", None),
+              chunk(2, "unrelated text", 0.4)]
+    assert len(rerank(chunks, "interlock bypass procedure", top_n=2)) == 2
+
+
+def test_unembedded_chunk_still_ranks_on_lexical_overlap():
+    """It should degrade to the lexical half, not vanish and not win."""
+    strong_words = chunk(1, "interlock bypass procedure pump", None)
+    weak_words = chunk(2, "boiler feedwater chemistry", None)
+    assert rerank([weak_words, strong_words],
+                  "interlock bypass procedure", top_n=1) == [strong_words]
+
+
+def test_embedded_chunk_outranks_an_unembedded_one_on_equal_text():
+    text = "interlock bypass procedure"
+    embedded = chunk(1, text, 0.9)
+    unembedded = chunk(2, text, None)
+    assert rerank([unembedded, embedded], text, top_n=1) == [embedded]
+
+
+
+
+# ---------------------------------------------------------------------------
+# Step sections are withheld from the procedural branch
+# ---------------------------------------------------------------------------
+
+from tools.retrieval import STEP_SECTION_RE  # noqa: E402
+
+
+def test_step_sections_are_recognised():
+    """These are the sections that ARE the instructions. find_procedure must
+    not hand them to the model - an interlock exists because somebody assessed
+    a hazard, and a bypass reassembled from fragments is a safety document
+    nobody approved."""
+    for title in ["3. Procedure", "Steps", "4. Method", "Instructions",
+                  "5. Execution", "Work Instruction", "6. Restoration"]:
+        assert STEP_SECTION_RE.search(title), title
+
+
+def test_identifying_sections_are_not_withheld():
+    """The header carries the title and authorising role; prerequisites are
+    what a procedural answer quotes. Withholding either was the bug."""
+    for title in ["DEMO DOCUMENT — NOT A CONTROLLED DOCUMENT",
+                  "1. Purpose", "2. Prerequisites", "Scope",
+                  "Hazards", "References"]:
+        assert not STEP_SECTION_RE.search(title), title
+
+
+def test_a_missing_section_title_is_not_treated_as_a_step():
+    assert not STEP_SECTION_RE.search("")

api/tools/retrieval.py

@@ -18,6 +18,7 @@ how an interlock works is not the procedure that authorises lifting it.
 from __future__ import annotations
 
 import logging
+import re
 from dataclasses import dataclass, asdict
 from datetime import date
 from typing import Any, Literal
@@ -44,14 +45,20 @@ class Chunk:
     section_title: str | None
     equipment_id: str | None
     chunk_text: str
-    similarity: float
+    # None when the chunk has no embedding - see rerank().
+    similarity: float | None
+    # Added by migration 007; NULL on rows ingested before it, hence defaults.
+    doc_title: str | None = None
+    authorising_role: str | None = None
 
     def citation(self) -> dict[str, Any]:
         """The citation dict a contract expects. Title falls back to the file
         name, because a chunk with no title is still traceable to a document."""
         return {
             "doc_number": self.doc_number or self.source_file,
-            "title": self.section_title or self.source_file,
+            # doc_title is the document's own title from the confirmed header
+            # (migration 007). The fallbacks are for rows ingested before it.
+            "title": self.doc_title or self.section_title or self.source_file,
             "revision": self.revision or "unknown",
             "effective_date": self.effective_date,
             "page": self.page,
@@ -104,7 +111,7 @@ def search(
                 f"""
                 SELECT id, source_file, doc_type, doc_number, revision,
                        effective_date, page, section_title, equipment_id,
-                       chunk_text,
+                       chunk_text, doc_title, authorising_role,
                        1 - (embedding <=> %(embedding)s::vector) AS similarity
                   FROM doc_chunks
                   {clause}
@@ -133,7 +140,13 @@ def rerank(chunks: list[Chunk], question: str, *, top_n: int = 4) -> list[Chunk]
         text = chunk.chunk_text.lower()
         overlap = sum(1 for t in terms if t in text)
         lexical = overlap / max(len(terms), 1)
-        return 0.75 * chunk.similarity + 0.25 * lexical
+        # similarity is NULL for a chunk ingested with --no-embed, and for any
+        # chunk whose embedding call failed. `1 - (NULL <=> vec)` is NULL, so
+        # this used to raise TypeError and 500 the whole question rather than
+        # rank that chunk last. Degrade to the lexical half instead: an
+        # unembedded chunk is still findable, just not by meaning.
+        similarity = chunk.similarity if chunk.similarity is not None else 0.0
+        return 0.75 * similarity + 0.25 * lexical
 
     return sorted(chunks, key=score, reverse=True)[:top_n]
 
@@ -159,7 +172,76 @@ def find_procedure(
         equipment_id=equipment_id,
         conn=conn,
     )
-    return rerank(hits, question, top_n=3)
+    best = rerank(hits, question, top_n=1)
+    if not best:
+        return []
+    return document_sections(best[0].source_file, conn=conn)
+
+
+# Sections that ARE the instructions. Excluded from what the procedural branch
+# retrieves - see document_sections().
+STEP_SECTION_RE = re.compile(
+    r"\b(procedure|steps?|method|instructions?|execution|restoration|"
+    r"work\s+instruction)\b",
+    re.IGNORECASE,
+)
+
+
+def document_sections(
+    source_file: str,
+    *,
+    limit: int = 6,
+    conn: psycopg.Connection | None = None,
+) -> list[Chunk]:
+    """The identifying sections of ONE document, in document order.
+
+    Ranking a procedure's chunks by similarity to the question was the wrong
+    shape for this branch. It returned the three chunks that most resembled
+    "how do I lift the interlock" - which is the step list - and left out the
+    header block carrying the title and the authorising role, and often the
+    prerequisites too. So the model was asked for a title it had never been
+    shown (and returned "") while being handed the one section it must never
+    reproduce.
+
+    Once the document is identified, WHICH document it is settles what to send:
+    the header and the prerequisites, in the order they appear, never the
+    steps. STEP_SECTION_RE is a second line behind ProceduralAnswer's
+    instruction-language check, not a replacement for it - the contract still
+    rejects instruction language whatever arrives here.
+
+    Ordering by id is document order: ingest.py inserts sections in file order
+    in a single executemany, and replaces every chunk of a source_file in one
+    transaction, so ids within a document are monotonic in the document.
+    """
+    owned = conn is None
+    conn = conn or _connect()
+    try:
+        with conn.cursor() as cur:
+            cur.execute(
+                """
+                SELECT id, source_file, doc_type, doc_number, revision,
+                       effective_date, page, section_title, equipment_id,
+                       chunk_text, doc_title, authorising_role,
+                       NULL::float AS similarity
+                  FROM doc_chunks
+                 WHERE source_file = %(source_file)s
+                   AND superseded = FALSE
+                 ORDER BY id
+                """,
+                {"source_file": source_file},
+            )
+            chunks = [Chunk(**row) for row in cur.fetchall()]
+    finally:
+        if owned:
+            conn.close()
+
+    kept = [c for c in chunks
+            if not STEP_SECTION_RE.search(c.section_title or "")]
+    dropped = len(chunks) - len(kept)
+    if dropped:
+        log.info("%s: withheld %d step section(s) from the procedural branch",
+                 source_file, dropped)
+    return kept[:limit]
 
 
 def lexical_search(
@@ -204,7 +286,7 @@ def lexical_search(
                 f"""
                 SELECT id, source_file, doc_type, doc_number, revision,
                        effective_date, page, section_title, equipment_id,
-                       chunk_text,
+                       chunk_text, doc_title, authorising_role,
                        ts_rank(
                          to_tsvector('english', chunk_text),
                          plainto_tsquery('english', %(q)s)
@@ -232,7 +314,7 @@ def lexical_search(
                         f"""
                         SELECT id, source_file, doc_type, doc_number, revision,
                                effective_date, page, section_title, equipment_id,
-                               chunk_text,
+                               chunk_text, doc_title, authorising_role,
                                ts_rank(
                                  to_tsvector('english', chunk_text),
                                  to_tsquery('english', %(q)s)
@@ -260,7 +342,12 @@ def find_procedure_lexical(
     equipment_id: str | None = None,
     conn: psycopg.Connection | None = None,
 ) -> list[Chunk]:
-    """The find_procedure() shape, without an embedding. Stub mode only."""
+    """The find_procedure() shape, without an embedding. Stub mode only.
+
+    Identify-then-expand exactly as find_procedure() does, so the stub keeps
+    testing the same shape it always did. Only the identification step differs:
+    words rather than meaning.
+    """
     hits = lexical_search(
         question,
         top_k=12,
@@ -268,7 +355,10 @@ def find_procedure_lexical(
         equipment_id=equipment_id,
         conn=conn,
     )
-    return rerank(hits, question, top_n=3)
+    best = rerank(hits, question, top_n=1)
+    if not best:
+        return []
+    return document_sections(best[0].source_file, conn=conn)
 
 
 def as_dicts(chunks: list[Chunk]) -> list[dict[str, Any]]: