brain-of-reese/app/rag/retriever.py

"""Hybrid retrieval: pgvector cosine ∪ Postgres FTS, RRF-fused (PLAN §6, A7).

* **Vector list** — top-N chunks by cosine distance (``embedding <=> $1``),
  each carrying its cosine ``1 − distance`` (the honesty-gate input).
* **Lexical list** — top-N chunks matching an OR-``tsquery`` over the
  question's tokens, ordered by ``ts_rank``, UNION the name-hit list:
  documents whose normalized name (title + path stem, alnum-only,
  lowercased) contains a DIGIT-BEARING normalized question token of
  length >= 4 — bare tokens (``1panel``) and the join of adjacent
  question tokens starting with a letter ("Qwen 3.8" → ``qwen38``).
  The digit requirement is the precision guard: plain prose words
  ("server", "arguments" — 4+ chars but no digit) never name-match,
  while versioned product names (the incident's whole point) always
  carry one. The name-hit documents LEAD the lexical list (ranked by
  match count, then total matched length, then catalog order; capped
  at :data:`NAME_HIT_LIMIT`),
  the FTS rows follow. This is what finds name-your-tool questions
  ("gitlab") that vector similarity buries — and, the 2026-09-05
  incident, the versioned-name case the default parser lexes
  incompatibly (question "Qwen 3.8" → ``qwen``/``3``/``8`` can never
  match a document's ``qwen3``/``8``/``27b`` tokens, while every
  unrelated llama.cpp quadlet out-ranks the target on the shared
  ``llama``/``cpp`` tokens). The name-hit rows carry ``fts_hit=True``
  (they ARE the lexical signal — the A8 honesty gate then answers
  instead of deflecting) and ``cosine=0.0``; the RRF fusion is
  unchanged (same lists, same ``1/(k+rank)`` terms).
* **Fusion** — Reciprocal Rank Fusion (``score = Σ 1/(k + rank)`` over the
  lists a chunk appears in; chunks hit by both lists get both terms). The
  fused score ranks; :meth:`select_documents` and :func:`weak_hit_titles`
  keep working off ``score``.

The product requirement (LOCKED A7, revised 2026-08-24): the LLM receives
the **entire relevant document**, not just the chunk — chunk hits map back
to their parents, dedupe, rank by best fused score, and the full text of
the top-N documents is always fed through, never truncated. If a future KB
ever makes the prompt too large for the model, the ``LLMError`` → SSE
``error`` path surfaces it loudly — no silent partial context.

Deterministic tie-break for equal fused scores:
``(−fused, −cosine, document.path, chunk.position)``.
"""
from __future__ import annotations

import math
import re
import uuid
from collections.abc import Sequence
from dataclasses import dataclass, replace
from datetime import UTC, datetime
from pathlib import Path

from sqlalchemy import select, text
from sqlalchemy.orm import Session

from app.config import get_settings
from app.models import Chunk, Document

#: Shared overflow marker — now used by the steering (<tuning>) section
#: only (phase 15; imported by ``app.rag.prompts``). The document context
#: path never truncates (A7 revised, owner permission 2026-08-24).
TRUNCATION_MARKER = "[…truncated…]"

#: Alphanumeric tokens of a question (``to_tsquery`` input, OR-joined),
#: with DOTTED tokens kept whole (``llama.cpp`` → ``llama.cpp``). The
#: default parser lexes a dotted word as ONE lexeme (``to_tsvector`` of
#: "llama.cpp" → ``'llama.cpp'``; of "Qwen 3.8" → ``'3.8'``), so a split
#: token (``llama`` | ``cpp``) can never match the document side — the
#: 2026-09-05 incident's question phrase "llama.cpp" lexed
#: incompatibly on both sides of the query. Dotted tokens are passed
#: through as single ``to_tsquery`` lexemes (verified: Postgres accepts
#: ``'llama.cpp'`` and ``'3.8'`` as lexemes).
_TOKEN_RE = re.compile(r"[a-z0-9]+(?:\.[a-z0-9]+)*")

#: Name-hit policy (the 2026-09-05 incident): a question token shorter
#: than this (normalized) is too weak a name signal ("3", "8", "the") —
#: it would match half the KB by accident.
NAME_TOKEN_MIN_LEN = 4

#: The name-hit list may contribute at most this many documents to the
#: lexical side (before the FTS rows); the rest are dropped.
NAME_HIT_LIMIT = 10

#: Alphanumeric runs of a lowercased string (name normalization).
_ALNUM_RE = re.compile(r"[a-z0-9]+")


def _normalize_name(s: str) -> str:
    """Lowercased, alnum-only form of *s* (``Qwen 3.8`` → ``qwen38``)."""
    return "".join(_ALNUM_RE.findall(s.lower()))


def name_hit_tokens(question: str) -> list[str]:
    """The normalized name-match candidates of one question.

    Only DIGIT-BEARING candidates count (the precision guard — a plain
    prose word like "server" or "arguments" must never name-match a
    document; a versioned product name always carries a digit):

    * the :func:`_normalize_name` form of every whitespace token, kept
      when at least :data:`NAME_TOKEN_MIN_LEN` chars (``1panel``,
      ``qwen38`` from a single written token);
    * the versioned-name JOIN — the normalized concatenation of every
      ADJACENT token pair whose SECOND token is purely numeric (a
      version number: ``"… for Qwen 3.8"`` → ``Qwen`` + ``3.8`` joins
      to ``qwen38``; prose joins like ``correct`` + ``llama`` and the
      word-after-version boundaries like ``3.8`` + ``show`` are
      dropped). A purely-numeric second token means the join always
      starts with the first token's text (a letter in practice), so no
      digit-leading artifact (``38show``) can survive.

    Order of first appearance, de-duplicated.
    """
    tokens = question.split()
    out: list[str] = []
    seen: set[str] = set()

    def add(tok: str) -> None:
        if (
            len(tok) >= NAME_TOKEN_MIN_LEN
            and any(ch.isdigit() for ch in tok)
            and tok not in seen
        ):
            seen.add(tok)
            out.append(tok)

    for i, raw in enumerate(tokens):
        norm = _normalize_name(raw)
        add(norm)
        if i + 1 < len(tokens):
            next_norm = _normalize_name(tokens[i + 1])
            if next_norm and next_norm.isdigit():
                add(norm + next_norm)  # the versioned-name join
    return out


#: One row of the lexical candidate query (all fields needed to build a
#: detached :class:`Document` plus the chunk fields and ``ts_rank``).
_LEXICAL_SQL = text(
    """
    SELECT c.id            AS chunk_id,
           c.position      AS position,
           c.content       AS content,
           d.id            AS doc_id,
           d.source        AS source,
           d.path          AS path,
           d.full_path     AS full_path,
           d.title         AS title,
           d.content       AS doc_content,
           d.content_hash  AS content_hash,
           d.indexed_at    AS indexed_at,
           d.created_at    AS created_at,
           c.is_summary    AS is_summary,
           ts_rank(c.tsv, to_tsquery('english', :tsquery)) AS rank
    FROM chunks c
    JOIN documents d ON d.id = c.document_id
    WHERE c.tsv @@ to_tsquery('english', :tsquery)
    ORDER BY rank DESC, d.path ASC, c.position ASC
    LIMIT :limit
    """
)


@dataclass
class RetrievedChunk:
    """One retrieval candidate: fused rank score + parent document row.

    * ``score`` — RRF fused score (the ranking key for document selection
      and weak-hit titles).
    * ``cosine`` — vector similarity ``1 − distance`` (the honesty-gate
      input; ``0.0`` for lexical-only hits that have no vector rank).
    * ``fts_hit`` — the chunk matched the question's OR-tsquery.
    * ``is_summary`` — True for the lite-model summary chunk (phase 30,
      position −1): its parent *is* the source document, so a summary hit
      resolves to the full source document through the unchanged
      chunk→document mapping (A7 revised). Default ``False`` keeps every
      ordinary content chunk valid.
    """

    chunk_id: uuid.UUID
    position: int
    content: str
    score: float
    document: Document
    cosine: float = 0.0
    fts_hit: bool = False
    is_summary: bool = False


def lexical_tsquery(question: str) -> str | None:
    """OR-joined token string for ``to_tsquery('english', …)``, or ``None``.

    Tokens are lowercased ``[a-z0-9]+`` runs with DOTTED words kept whole
    (see :data:`_TOKEN_RE` — the default parser lexes ``llama.cpp`` as
    one lexeme, so the query must carry it whole too), de-duplicated in
    order of first appearance. Postgres does the lexing/stemming; a
    question whose tokens are all stopwords lexes to an *empty* tsquery
    (which matches nothing), so no special-casing is needed there.
    Pure-symbol questions ("???", "🔧") yield no tokens → ``None`` → no
    lexical query at all.
    """
    seen: set[str] = set()
    tokens: list[str] = []
    for tok in _TOKEN_RE.findall(question.lower()):
        if tok not in seen:
            seen.add(tok)
            tokens.append(tok)
    return " | ".join(tokens) if tokens else None


def fuse(
    vector: Sequence[RetrievedChunk],
    lexical: Sequence[RetrievedChunk],
    k: int,
) -> list[RetrievedChunk]:
    """Reciprocal Rank Fusion over the two ranked candidate lists.

    ``score(chunk) = Σ 1/(k + rank)`` — one term per list the chunk appears
    in (ranks are 1-based; a chunk in both lists gets both terms). Returns
    the union ordered by ``(-score, -cosine, document.path, position)``.

    Lexical-only hits (no vector rank) enter with ``cosine=0.0`` and
    ``fts_hit=True``; vector chunks matched by the lexical list get
    ``fts_hit=True`` in place (the input objects are mutated — callers
    should not reuse them afterwards).
    """
    if k <= 0:
        raise ValueError("rrf k must be > 0")
    by_id: dict[uuid.UUID, RetrievedChunk] = {}
    fused: dict[uuid.UUID, float] = {}
    for rank, rc in enumerate(vector, start=1):
        by_id[rc.chunk_id] = rc
        fused[rc.chunk_id] = fused.get(rc.chunk_id, 0.0) + 1.0 / (k + rank)
    for rank, rc in enumerate(lexical, start=1):
        term = 1.0 / (k + rank)
        if rc.chunk_id in by_id:
            existing = by_id[rc.chunk_id]
            by_id[rc.chunk_id] = replace(existing, fts_hit=True)
            fused[rc.chunk_id] += term
        else:
            rc = replace(rc, fts_hit=True)
            by_id[rc.chunk_id] = rc
            fused[rc.chunk_id] = fused.get(rc.chunk_id, 0.0) + term
    out = [replace(rc, score=fused[rc.chunk_id]) for rc in by_id.values()]
    out.sort(key=lambda rc: (-rc.score, -rc.cosine, rc.document.path, rc.position))
    return out


def apply_recency_boost(
    chunks: Sequence[RetrievedChunk],
    *,
    now: datetime | None = None,
    weight: float | None = None,
    half_life_days: int | None = None,
) -> list[RetrievedChunk]:
    """Additive recency boost on the fused score (phase 106, D6).

    Each chunk's score becomes
    ``score + weight * exp(−age_days / half_life_days)`` where
    ``age_days = max(0, (now − document.created_at))`` in days — a
    zero-age document gets the full *weight* (the MAXIMUM additive
    score), each ``half_life_days`` of age multiplies the remaining
    boost by ``e**-1`` (≈0.37), and a FUTURE ``created_at`` clamps to
    age 0 (the document reads as brand-new — consistent with D3's
    today-folding in :mod:`app.rag.doc_dates`). Defaults: *weight* /
    *half_life_days* from :func:`get_settings` (``recency_boost`` /
    ``recency_half_life_days``), *now* from ``datetime.now(UTC)``.

    Magnitude rationale (the ``0.0007`` default, the k=60 RRF scale):
    rank 1 vs 2 in one list differs by ~0.00026 and rank 1 vs 10 by
    ~0.0021, so the full weight is a bounded 2-3 rank head start —
    enough to break near-ties toward the newer document, far below the
    fused gap between a document that answers and one that merely
    resembles (the phase-106 fine-line battery pins the measured
    margin: 0.00263 ≥ 3× the zero-age boost).

    Pure (the :func:`fuse` convention): the inputs are never mutated —
    every boosted chunk is a ``replace()`` copy — and the result is
    re-sorted with the EXISTING deterministic key
    ``(−score, −cosine, document.path, position)``; with ``weight=0``
    every score is untouched and an already-fused (already-sorted)
    input comes back byte-identical (the kill switch, pinned).

    Untouched by design: the A8 honesty gate and ``query_log.top_score``
    (both read the chunk's ``cosine``, which the boost never modifies),
    :func:`weak_hit_titles` (titles only), and the never-truncated
    top-N contract (:func:`select_documents` still feeds whole
    documents — the boost re-ranks WHICH documents, never truncates).
    SINGLE APPLY SITE: :func:`retrieve()` is the only caller in
    ``app/`` — the chat API and ``scripts/eval_retrieval.py`` inherit
    the boost through it; nothing else may apply it.
    """
    if weight is None or half_life_days is None:
        settings = get_settings()
        if weight is None:
            weight = settings.recency_boost
        if half_life_days is None:
            half_life_days = settings.recency_half_life_days
    if half_life_days <= 0:
        raise ValueError("half_life_days must be > 0")
    if now is None:
        now = datetime.now(UTC)
    boosted: list[RetrievedChunk] = []
    for rc in chunks:
        age_days = max(0.0, (now - rc.document.created_at).total_seconds() / 86400.0)
        boosted.append(
            replace(
                rc,
                score=rc.score + weight * math.exp(-age_days / half_life_days),
            )
        )
    boosted.sort(key=lambda rc: (-rc.score, -rc.cosine, rc.document.path, rc.position))
    return boosted


def _vector_candidates(
    db: Session, question_embedding: list[float], limit: int
) -> list[RetrievedChunk]:
    """Top-*limit* chunks by pgvector cosine distance (``<=>``).

    ``cosine = 1 − distance``. Chunks whose embedding is still NULL
    (two-phase import in progress) are skipped. Each candidate carries
    its ``Chunk.is_summary`` flag (phase 30) so a summary hit stays
    identifiable after fusion.
    """
    distance = Chunk.embedding.cosine_distance(question_embedding)
    rows = db.execute(
        select(Chunk, distance.label("distance"), Document)
        .join(Document, Chunk.document_id == Document.id)
        .where(Chunk.embedding.is_not(None))
        .order_by(distance)
        .limit(limit)
    ).all()
    return [
        RetrievedChunk(
            chunk_id=chunk.id,
            position=chunk.position,
            content=chunk.content,
            score=0.0,  # fused score is filled in by :func:`fuse`
            document=doc,
            cosine=round(1.0 - float(dist), 6),
            is_summary=chunk.is_summary,
        )
        for chunk, dist, doc in rows
    ]


#: One name-hit row: the document fields (detached :class:`Document`)
#: plus the representative chunk — the ``is_summary`` chunk when the
#: document has one (its natural-language summary is the best chunk of
#: a machine file, and the chunk the vector list is likeliest to have
#: ranked too, so the RRF merge dedupes cleanly), else chunk 0.
_NAME_HIT_SQL = text(
    """
    SELECT d.id            AS doc_id,
           d.source        AS source,
           d.path          AS path,
           d.full_path     AS full_path,
           d.title         AS title,
           d.content       AS doc_content,
           d.content_hash  AS content_hash,
           d.indexed_at    AS indexed_at,
           d.created_at    AS created_at,
           c.id            AS chunk_id,
           c.position      AS position,
           c.content       AS content,
           c.is_summary    AS is_summary
    FROM documents d
    LEFT JOIN LATERAL (
        SELECT id, position, content, is_summary
        FROM chunks
        WHERE document_id = d.id
        ORDER BY is_summary DESC, position ASC
        LIMIT 1
    ) c ON true
    WHERE d.id = ANY(:ids) AND c.id IS NOT NULL
    """
)


def _name_hit_chunks(db: Session, question: str) -> list[RetrievedChunk]:
    """The documents whose NAME matches the question (the 2026-09-05
    incident, the versioned-name signal).

    A document is a name hit when its normalized name —
    :func:`_normalize_name` of its ``title`` followed by the normalized
    path stem (``qwen3.8-27b-juggernaut-vulkan.container`` →
    ``qwen3827bjuggernautvulkan``) — contains at least one
    :func:`name_hit_tokens` candidate (``qwen38``). Ranked by
    (distinct matched tokens, total matched length, source, path) —
    catalog order is the final deterministic tie-break — capped at
    :data:`NAME_HIT_LIMIT`. Each hit becomes one lexical
    :class:`RetrievedChunk` (its representative chunk, ``fts_hit=True``,
    ``cosine=0.0``). Two lightweight queries: one projection over
    (id, source, path, title) in catalog order, one LATERAL chunk
    fetch for the ≤ :data:`NAME_HIT_LIMIT` winners (no full-content
    load; the content joins in via the row fetch below).
    """
    tokens = name_hit_tokens(question)
    if not tokens:
        return []
    rows = db.execute(
        select(Document.id, Document.source, Document.path, Document.title).order_by(
            Document.source, Document.path
        )
    ).all()
    scored: list[tuple[int, int, uuid.UUID]] = []
    by_id: dict[uuid.UUID, tuple[str, str]] = {}  # id -> (source, path)
    for doc_id, source, path, title in rows:
        by_id[doc_id] = (source, path)
        name = _normalize_name(title) + _normalize_name(Path(path).stem)
        matched = [t for t in tokens if t in name]
        if matched:
            scored.append((len(matched), sum(len(t) for t in matched), doc_id))
    if not scored:
        return []
    # Ranked by (distinct matched tokens, total matched length), the
    # deterministic catalog tie-break (source, path) last.
    scored.sort(
        key=lambda s: (-s[0], -s[1], by_id[s[2]][0], by_id[s[2]][1])
    )
    ids = [s[2] for s in scored[:NAME_HIT_LIMIT]]
    hit_rows = list(db.execute(_NAME_HIT_SQL, {"ids": ids}).all())
    # The LATERAL query returns winners in id order; re-order by the
    # ranked order computed above so the lexical list is deterministic.
    order = {doc_id: rank for rank, doc_id in enumerate(ids)}
    hit_rows.sort(key=lambda row: order[row.doc_id])
    out: list[RetrievedChunk] = []
    for row in hit_rows:
        doc = Document(
            id=row.doc_id,
            source=row.source,
            path=row.path,
            full_path=row.full_path,
            title=row.title,
            content=row.doc_content,
            content_hash=row.content_hash,
            indexed_at=row.indexed_at,
            created_at=row.created_at,
        )
        out.append(
            RetrievedChunk(
                chunk_id=row.chunk_id,
                position=row.position,
                content=row.content,
                score=0.0,  # filled in by :func:`fuse`
                document=doc,
                cosine=0.0,  # no vector rank — name-only hit
                fts_hit=True,  # the lexical signal — the A8 gate answers
                is_summary=bool(row.is_summary),
            )
        )
    return out


def _lexical_candidates(db: Session, question: str, limit: int) -> list[RetrievedChunk]:
    """The lexical candidate list: name hits, then FTS rows (A7).

    The name-hit documents (:func:`_name_hit_chunks`, the versioned-name
    signal) LEAD the list — they are the strongest lexical evidence for
    a name-your-thing question — followed by the top-*limit* chunks
    matching the question's OR-tsquery, ordered by ``ts_rank`` (with
    ``d.path, c.position`` as the deterministic tie-break). Deduped by
    chunk id (a name-hit representative chunk the FTS list also ranked
    appears once). An empty tsquery (stopword-only question) simply
    contributes no FTS rows; a question with no name tokens contributes
    no name hits — both halves are independent.
    """
    out: list[RetrievedChunk] = []
    seen: set[uuid.UUID] = set()
    for rc in _name_hit_chunks(db, question):
        if rc.chunk_id not in seen:
            seen.add(rc.chunk_id)
            out.append(rc)
    tsquery = lexical_tsquery(question)
    if tsquery is None:
        return out
    rows = db.execute(
        _LEXICAL_SQL, {"tsquery": tsquery, "limit": limit}
    ).all()
    for row in rows:
        if row.chunk_id in seen:
            continue
        seen.add(row.chunk_id)
        doc = Document(
            id=row.doc_id,
            source=row.source,
            path=row.path,
            full_path=row.full_path,
            title=row.title,
            content=row.doc_content,
            content_hash=row.content_hash,
            indexed_at=row.indexed_at,
            created_at=row.created_at,
        )
        out.append(
            RetrievedChunk(
                chunk_id=row.chunk_id,
                position=row.position,
                content=row.content,
                score=0.0,  # filled in by :func:`fuse`
                document=doc,
                cosine=0.0,  # no vector rank — lexical-only hit
                fts_hit=True,
                is_summary=row.is_summary,
            )
        )
    return out


def retrieve(
    db: Session,
    question: str,
    question_embedding: list[float],
    vector_candidates: int | None = None,
    lexical_candidates: int | None = None,
) -> list[RetrievedChunk]:
    """Hybrid retrieval (A7): vector top-N ∪ FTS top-N, RRF-fused.

    Returns the fused candidate list in rank order (best first). Each
    :class:`RetrievedChunk` carries the fused ``score`` (ranking), the
    ``cosine`` similarity (honesty gate) and the ``fts_hit`` flag.
    """
    settings = get_settings()
    v_n = (
        settings.hybrid_vector_candidates if vector_candidates is None else vector_candidates
    )
    l_n = (
        settings.hybrid_lexical_candidates if lexical_candidates is None else lexical_candidates
    )
    if v_n <= 0:
        raise ValueError("vector_candidates must be >= 1")
    if l_n <= 0:
        raise ValueError("lexical_candidates must be >= 1")
    vector = _vector_candidates(db, question_embedding, v_n)
    lexical = _lexical_candidates(db, question, l_n)
    fused = fuse(vector, lexical, settings.rrf_k)
    if settings.recency_boost > 0:
        # Phase 106, D6 — the SINGLE recency-boost apply site: an
        # additive post-fusion re-rank (see :func:`apply_recency_boost`).
        # ``0`` = off: the pre-phase ranking returns byte-identical (the
        # kill switch) and weight-0 callers pay nothing.
        return apply_recency_boost(fused)
    return fused


def weak_hit_titles(chunks: Sequence[RetrievedChunk]) -> list[str]:
    """Distinct parent-document titles of *chunks*, best fused score first.

    Deflection mode (PLAN §6, A8) is built from these *titles only* — the
    LOW prompt and the "Maybe try" chips never see document content.
    """
    titles: list[str] = []
    seen: set[uuid.UUID] = set()
    for rc in sorted(chunks, key=lambda c: c.score, reverse=True):
        if rc.document.id in seen:
            continue
        seen.add(rc.document.id)
        titles.append(rc.document.title)
    return titles


def select_documents(
    chunks: Sequence[RetrievedChunk],
    n: int | None = None,
) -> list[Document]:
    """Map chunk hits to distinct parent documents, ranked by best fused score.

    At most *n* documents are returned (default ``BOR_TOP_N_DOCS``). The
    returned rows carry the full document content, byte-identical — a
    matched parent document is **never truncated** (A7 revised, owner
    permission 2026-08-24). There is deliberately no context budget: an
    oversized prompt must fail loudly through the ``LLMError`` → SSE
    ``error`` path, never arrive as silent partial context.
    """
    top_n = n if n is not None else get_settings().top_n_docs

    docs: list[Document] = []
    seen: set[uuid.UUID] = set()
    for rc in sorted(chunks, key=lambda c: c.score, reverse=True):
        if rc.document.id in seen:
            continue
        seen.add(rc.document.id)
        docs.append(rc.document)
    return docs[:top_n]