yau-plant-assistant/api/tools/retrieval.py

"""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]