pablan/backend/app/rag/retrieval.py

"""Hybrid retrieval: permission filter BEFORE anything else, then vector +
German full-text candidates merged with Reciprocal Rank Fusion.

There is no search without a user — the permission CTE is part of the one
SQL statement, so an unauthorized chunk is structurally impossible to
retrieve. Queries are user content and are never logged.
"""

import logging
import time
import uuid
from dataclasses import dataclass
from typing import Any

from sqlalchemy import Text, cast, func, literal, select, union
from sqlalchemy.dialects.postgresql import REGCONFIG
from sqlalchemy.ext.asyncio import AsyncSession

from app.llm.client import embed
from app.metrics import metrics
from app.models import Chunk, Document, User
from app.rag.permissions import has_open_review, searchable_documents_filter

logger = logging.getLogger("pablan.rag")

RRF_K = 60
CANDIDATES_PER_SOURCE = 20

# Calibrated against bge-m3 (2026-07): matched top
# hits land at cosine distance ~0.34-0.45, unrelated queries at 0.50+.
NO_ANSWER_MIN_DISTANCE = 0.45


@dataclass
class SearchResult:
    chunk_id: uuid.UUID
    document_id: uuid.UUID
    title: str
    heading_path: str
    content: str
    score: float
    vector_distance: float | None
    fts_match: bool
    # An unanswered request to check this document. Travels with every hit so
    # an answer can mark the source it leaned on as not-yet-settled.
    review_pending: bool = False


async def search(
    db: AsyncSession,
    query: str,
    *,
    user: User,
    top_k: int = 5,
) -> list[SearchResult]:
    started = time.monotonic()
    query_vector = (await embed([query]))[0]

    allowed = (
        select(
            Document.id,
            Document.title,
            has_open_review().label("review_pending"),
        )
        .where(searchable_documents_filter(user))
        .cte("allowed")
    )

    distance = Chunk.embedding.cosine_distance(query_vector)
    vec = (
        select(
            Chunk.id.label("chunk_id"),
            func.row_number().over(order_by=distance).label("rank"),
            distance.label("distance"),
        )
        .join(allowed, allowed.c.id == Chunk.document_id)
        .order_by(distance)
        .limit(CANDIDATES_PER_SOURCE)
        .subquery("vec")
    )

    tsquery = func.websearch_to_tsquery(cast(literal("german"), REGCONFIG), query)
    fts_order = func.ts_rank_cd(Chunk.tsv, tsquery).desc()
    fts = (
        select(
            Chunk.id.label("chunk_id"),
            func.row_number().over(order_by=fts_order).label("rank"),
        )
        .join(allowed, allowed.c.id == Chunk.document_id)
        .where(Chunk.tsv.op("@@")(tsquery))
        .order_by(fts_order)
        .limit(CANDIDATES_PER_SOURCE)
        .subquery("fts")
    )

    candidate_ids = union(select(vec.c.chunk_id), select(fts.c.chunk_id)).subquery(
        "ids"
    )
    score = (
        func.coalesce(1.0 / (RRF_K + vec.c.rank), 0.0)
        + func.coalesce(1.0 / (RRF_K + fts.c.rank), 0.0)
    ).label("score")

    rows = (
        await db.execute(
            select(
                Chunk.id,
                Chunk.document_id,
                Chunk.content,
                Chunk.meta,
                allowed.c.title,
                allowed.c.review_pending,
                score,
                vec.c.distance,
                fts.c.rank.label("fts_rank"),
            )
            .join(candidate_ids, candidate_ids.c.chunk_id == Chunk.id)
            .join(allowed, allowed.c.id == Chunk.document_id)
            .outerjoin(vec, vec.c.chunk_id == Chunk.id)
            .outerjoin(fts, fts.c.chunk_id == Chunk.id)
            .order_by(score.desc())
            .limit(top_k)
        )
    ).all()

    results = [
        SearchResult(
            chunk_id=row.id,
            document_id=row.document_id,
            title=row.title,
            heading_path=heading_path(row.meta),
            content=row.content,
            score=float(row.score),
            vector_distance=float(row.distance) if row.distance is not None else None,
            fts_match=row.fts_rank is not None,
            review_pending=bool(row.review_pending),
        )
        for row in rows
    ]

    duration = time.monotonic() - started
    metrics.inc("retrieval_searches_total")
    metrics.observe("retrieval_seconds", duration)
    metrics.observe("retrieval_results", float(len(results)))
    for result in results:
        source = (
            "both"
            if result.fts_match and result.vector_distance is not None
            else ("fts" if result.fts_match else "vector")
        )
        metrics.inc("retrieval_result_source_total", {"source": source})
    logger.info(
        "retrieval",
        extra={
            "event": "retrieval",
            "duration_ms": round(duration * 1000),
            "result_count": len(results),
            "top_k": top_k,
        },
    )
    return results


def _any_term_tsquery(query: str) -> Any:
    """The fallback's query: the lexemes `websearch_to_tsquery` produces, but
    ORed instead of ANDed.

    With no vector half to carry the recall, an AND query answers a natural
    question ("Wie läuft die Qualitätsprüfung im Wareneingang?") with nothing
    at all unless one single chunk happens to contain every word of it. ORing
    keeps the question usable and leaves the ordering to `ts_rank_cd`, which
    is what ranks a chunk matching more of the terms higher. Only the AND
    operators between groups are rewritten, so quoted phrases and exclusions
    survive. A query of nothing but stop words rewrites to an empty string,
    and NULLIF turns that into a query that matches nothing rather than a
    syntax error.
    """
    websearch = func.websearch_to_tsquery(cast(literal("german"), REGCONFIG), query)
    lexemes = func.nullif(func.replace(cast(websearch, Text), " & ", " | "), "")
    return func.to_tsquery(cast(literal("german"), REGCONFIG), lexemes)


async def text_search(
    db: AsyncSession,
    query: str,
    *,
    user: User,
    top_k: int = 5,
) -> list[SearchResult]:
    """The full-text half of `search()` alone: German `tsvector` matching over
    the GIN index (an inverted index), with no embedding call.

    This is what keeps the knowledge base searchable when no model answers at
    the configured endpoint. It finds less than the hybrid path (keywords, not
    meaning), so it is a fallback the user is told about, never a silent
    substitute. Same permission CTE as everything else — there is no search
    without a user.

    Terms are ORed here while the hybrid path ANDs them (`_any_term_tsquery`):
    alone, a full-sentence question must not come back empty, and in the
    hybrid path the AND is what makes `fts_match` mean "the words are really
    in there" for the no-answer signal.
    """
    started = time.monotonic()
    allowed = (
        select(
            Document.id,
            Document.title,
            has_open_review().label("review_pending"),
        )
        .where(searchable_documents_filter(user))
        .cte("allowed")
    )
    tsquery = _any_term_tsquery(query)
    rank = func.ts_rank_cd(Chunk.tsv, tsquery)
    rows = (
        await db.execute(
            select(
                Chunk.id,
                Chunk.document_id,
                Chunk.content,
                Chunk.meta,
                allowed.c.title,
                allowed.c.review_pending,
                rank.label("rank"),
            )
            .join(allowed, allowed.c.id == Chunk.document_id)
            .where(Chunk.tsv.op("@@")(tsquery))
            .order_by(rank.desc())
            .limit(top_k)
        )
    ).all()

    results = [
        SearchResult(
            chunk_id=row.id,
            document_id=row.document_id,
            title=row.title,
            heading_path=heading_path(row.meta),
            content=row.content,
            score=float(row.rank),
            vector_distance=None,
            fts_match=True,
            review_pending=bool(row.review_pending),
        )
        for row in rows
    ]

    duration = time.monotonic() - started
    metrics.inc("retrieval_text_searches_total")
    metrics.observe("retrieval_seconds", duration)
    logger.info(
        "retrieval (text only)",
        extra={
            "event": "retrieval_text",
            "duration_ms": round(duration * 1000),
            "result_count": len(results),
            "top_k": top_k,
        },
    )
    return results


def heading_path(meta: dict[str, Any] | None) -> str:
    return (meta or {}).get("heading_path", "")


def results_are_low_confidence(results: list[SearchResult]) -> bool:
    """No-answer signal: no keyword match anywhere and the best vector
    candidate is far away. Callers should not present such results as
    grounding."""
    if not results:
        return True
    top = results[0]
    return not top.fts_match and (
        top.vector_distance is None or top.vector_distance >= NO_ANSWER_MIN_DISTANCE
    )