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