YAU.POC/yau-plant-assistant
2
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Identify the procedure first, then withhold its steps

Browse source 
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>
This commit is contained in:
Claude 2026-08-28 10:54:13 +10:00
parent dcfb411c3b
commit b3e47506b0
2 changed files with 180 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

81
api/tests/test_rerank_null_similarity.py Normal file
View file

@ -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("")

108
api/tools/retrieval.py
View file

@ -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]]: