"""Document retrieval over pg-ai / pgvector. Two rules that are not optional and are implemented as defaults, not as arguments a caller has to remember: * superseded = FALSE is ALWAYS applied. Citing a withdrawn revision of a procedure is worse than finding nothing. Including superseded revisions is possible only through include_superseded, which exists for the ingest tooling and is never set on the answer path. * Every hit carries full citation metadata - document number, revision, effective date, page, section. A chunk without them cannot be cited, and an answer that cannot cite cannot be given. Procedural retrieval is filtered to doc_type = 'procedure'. A manual describing 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 import psycopg from psycopg.rows import dict_row from config import settings log = logging.getLogger("tools.retrieval") DocType = Literal["procedure", "manual", "rationalisation", "design"] @dataclass class Chunk: id: int source_file: str doc_type: str doc_number: str | None revision: str | None effective_date: date | None page: int | None section_title: str | None equipment_id: str | None chunk_text: str # 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, # 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, "section_title": self.section_title, "source_file": self.source_file, "superseded": False, } def _connect() -> psycopg.Connection: cfg = settings() return psycopg.connect( cfg.dsn(), row_factory=dict_row, application_name="ai-api", connect_timeout=5 ) def search( query_embedding: list[float], *, top_k: int = 8, doc_type: DocType | None = None, equipment_id: str | None = None, include_superseded: bool = False, conn: psycopg.Connection | None = None, ) -> list[Chunk]: """Cosine top-k over doc_chunks, filtered and cited. include_superseded exists for ingest verification only. Setting it on the answer path is a defect - the contract will not stop you, because a superseded citation raises at construction, but the failure will look like a contract bug rather than the caller's mistake. """ owned = conn is None conn = conn or _connect() try: where = [] if include_superseded else ["superseded = FALSE"] params: dict[str, Any] = {"embedding": str(query_embedding), "k": top_k} if doc_type: where.append("doc_type = %(doc_type)s") params["doc_type"] = doc_type if equipment_id: # Equipment-specific chunks first, but do not exclude general ones - # the governing procedure for a pump is often written for the class. where.append("(equipment_id = %(equipment_id)s OR equipment_id IS NULL)") params["equipment_id"] = equipment_id clause = f"WHERE {' AND '.join(where)}" if where else "" with conn.cursor() as cur: cur.execute( f""" SELECT id, source_file, doc_type, doc_number, revision, effective_date, page, section_title, equipment_id, chunk_text, doc_title, authorising_role, 1 - (embedding <=> %(embedding)s::vector) AS similarity FROM doc_chunks {clause} ORDER BY embedding <=> %(embedding)s::vector LIMIT %(k)s """, params, ) return [Chunk(**row) for row in cur.fetchall()] finally: if owned: conn.close() def rerank(chunks: list[Chunk], question: str, *, top_n: int = 4) -> list[Chunk]: """Cheap lexical rerank over the vector hits. Deliberately not a model call: this runs on every question and a reranking model is a second inference per request for a marginal gain on a document set this small. Revisit if retrieval accuracy is the eval failure, and fix it here rather than by adding instructions to the prompt. """ terms = {t.lower().strip(".,?") for t in question.split() if len(t) > 3} def score(chunk: Chunk) -> float: text = chunk.chunk_text.lower() overlap = sum(1 for t in terms if t in text) lexical = overlap / max(len(terms), 1) # 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] def find_procedure( query_embedding: list[float], question: str, *, equipment_id: str | None = None, conn: psycopg.Connection | None = None, ) -> list[Chunk]: """Procedural path: procedures only, live revisions only. The chunks that come back are for IDENTIFYING and QUOTING the procedure. They are not raw material for reconstructing it - ProceduralAnswer's contract rejects any response containing instruction language, whatever these chunks happen to contain. """ hits = search( query_embedding, top_k=12, doc_type="procedure", equipment_id=equipment_id, conn=conn, ) 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( question: str, *, top_k: int = 8, doc_type: DocType | None = None, equipment_id: str | None = None, conn: psycopg.Connection | None = None, ) -> list[Chunk]: """Full-text search over chunk_text. NO-LLM STUB MODE ONLY. Embedding the question needs Azure OpenAI, so until that exists there is no vector to search with. This finds chunks by words instead, which is a genuinely different thing: it matches what the operator typed, not what they meant. "The well is going to overflow" finds nothing here and would find the spill procedure with embeddings. Kept beside search() rather than hidden inside it, and never called on the normal answer path, so that nobody can mistake a lexical hit for a semantic one when reading a trace. superseded = FALSE applies here exactly as it does in search(). There is no include_superseded, because this function has one caller and that caller is a demo. """ owned = conn is None conn = conn or _connect() try: where = ["superseded = FALSE"] params: dict[str, Any] = {"q": question, "k": top_k} if doc_type: where.append("doc_type = %(doc_type)s") params["doc_type"] = doc_type if equipment_id: where.append("(equipment_id = %(equipment_id)s OR equipment_id IS NULL)") params["equipment_id"] = equipment_id clause = " AND ".join(where) with conn.cursor() as cur: cur.execute( f""" SELECT id, source_file, doc_type, doc_number, revision, effective_date, page, section_title, equipment_id, chunk_text, doc_title, authorising_role, ts_rank( to_tsvector('english', chunk_text), plainto_tsquery('english', %(q)s) ) AS similarity FROM doc_chunks WHERE {clause} AND to_tsvector('english', chunk_text) @@ plainto_tsquery('english', %(q)s) ORDER BY similarity DESC LIMIT %(k)s """, params, ) rows = cur.fetchall() if not rows: # plainto_tsquery ANDs every term, so one unmatched word returns # nothing at all. Fall back to any-term matching before giving up, # or a demo question phrased as a sentence never finds anything. terms = [t.strip(".,?;:").lower() for t in question.split() if len(t) > 3] if terms: params["q"] = " | ".join(terms) with conn.cursor() as cur: cur.execute( f""" SELECT id, source_file, doc_type, doc_number, revision, effective_date, page, section_title, equipment_id, chunk_text, doc_title, authorising_role, ts_rank( to_tsvector('english', chunk_text), to_tsquery('english', %(q)s) ) AS similarity FROM doc_chunks WHERE {clause} AND to_tsvector('english', chunk_text) @@ to_tsquery('english', %(q)s) ORDER BY similarity DESC LIMIT %(k)s """, params, ) rows = cur.fetchall() return [Chunk(**row) for row in rows] finally: if owned: conn.close() def find_procedure_lexical( question: str, *, equipment_id: str | None = None, conn: psycopg.Connection | None = None, ) -> list[Chunk]: """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, doc_type="procedure", equipment_id=equipment_id, conn=conn, ) 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]]: return [asdict(c) for c in chunks]