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
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 as_dicts(chunks: list[Chunk]) -> list[dict[str, Any]]:
    return [asdict(c) for c in chunks]