"""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 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 similarity: float 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, "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, 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) return 0.75 * chunk.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, ) return rerank(hits, question, top_n=3) 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, 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, 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.""" hits = lexical_search( question, top_k=12, doc_type="procedure", equipment_id=equipment_id, conn=conn, ) return rerank(hits, question, top_n=3) def as_dicts(chunks: list[Chunk]) -> list[dict[str, Any]]: return [asdict(c) for c in chunks]