diff --git a/api/tests/test_rerank_null_similarity.py b/api/tests/test_rerank_null_similarity.py new file mode 100644 index 0000000..14de3f2 --- /dev/null +++ b/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("") diff --git a/api/tools/retrieval.py b/api/tools/retrieval.py index 513bfa7..3276970 100644 --- a/api/tools/retrieval.py +++ b/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]]: