RAG/backend/app/persistence/retrieval.py

"""PostgreSQL/pgvector persistence boundary for formal retrieval.

Authorization, active-model selection, and document lifecycle checks belong in
the candidate SQL.  They must never be applied as an in-memory post-filter.
"""

from __future__ import annotations

import math
import re
import uuid
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Any, Protocol, cast

import psycopg
from pgvector.psycopg import register_vector
from pgvector.vector import Vector
from psycopg.rows import dict_row

from app.core.config import Settings
from app.core.secrets import SecretFileError

_PROFILE_HASH_PATTERN = re.compile(r"^[0-9a-f]{64}$")


@dataclass(frozen=True, slots=True)
class ActiveEmbeddingProfile:
    """The enabled embedding profile selected by a knowledge base."""

    profile_hash: str
    model: str
    dimension: int
    synthetic: bool = False


@dataclass(frozen=True, slots=True)
class RetrievalCandidate:
    """Approved and authorized projection returned by the vector query."""

    citation_id: uuid.UUID
    document_id: uuid.UUID
    source_name: str
    cloud_text: str
    section_path: tuple[str, ...]
    page_start: int | None
    page_end: int | None
    vector_score: float


class RetrievalPersistenceError(RuntimeError):
    """Sanitized storage failure safe for translation at the service boundary."""

    def __init__(self) -> None:
        super().__init__("retrieval persistence unavailable")


class RetrievalRepository(Protocol):
    """Synchronous repository port; the async service runs calls in a thread."""

    def resolve_active_profile(
        self,
        knowledge_base_id: uuid.UUID,
        *,
        allowed_scope_ids: Sequence[uuid.UUID],
    ) -> ActiveEmbeddingProfile | None: ...

    def search_candidates(
        self,
        knowledge_base_id: uuid.UUID,
        *,
        allowed_scope_ids: Sequence[uuid.UUID],
        profile_hash: str,
        query_vector: tuple[float, ...],
        limit: int,
    ) -> list[RetrievalCandidate]: ...


ACTIVE_PROFILE_SQL = """
SELECT
  profile.profile_hash,
  profile.model,
  profile.dimension,
  profile.synthetic
FROM rag.knowledge_bases AS knowledge_base
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = knowledge_base.active_embedding_profile_hash
 AND profile.kind = knowledge_base.active_embedding_profile_kind
WHERE knowledge_base.id = %s
  AND profile.kind = 'embedding'
  AND profile.enabled IS TRUE
  AND profile.dimension = 1024
  AND EXISTS (
    SELECT 1
    FROM rag.access_scopes AS access_scope
    WHERE access_scope.knowledge_base_id = knowledge_base.id
      AND access_scope.id = ANY(%s::uuid[])
  )
LIMIT 1
"""


CANDIDATE_SEARCH_SQL = """
WITH query_input AS (
  SELECT %s::vector AS embedding
)
SELECT
  chunk.citation_id::text AS citation_id,
  document.id::text AS document_id,
  document.filename AS source_name,
  chunk.cloud_text,
  chunk.section_path,
  chunk.page_start,
  chunk.page_end,
  1 - (chunk.embedding <=> query_input.embedding) AS vector_score
FROM rag.chunks AS chunk
JOIN rag.knowledge_bases AS knowledge_base
  ON knowledge_base.id = chunk.knowledge_base_id
JOIN rag.model_profiles AS profile
  ON profile.profile_hash = knowledge_base.active_embedding_profile_hash
 AND profile.kind = knowledge_base.active_embedding_profile_kind
JOIN rag.access_scopes AS access_scope
  ON access_scope.id = chunk.access_scope_id
 AND access_scope.knowledge_base_id = chunk.knowledge_base_id
JOIN rag.documents AS document
  ON document.id = chunk.document_id
 AND document.knowledge_base_id = chunk.knowledge_base_id
 AND document.access_scope_id = chunk.access_scope_id
JOIN rag.document_versions AS document_version
  ON document_version.id = chunk.document_version_id
 AND document_version.document_id = chunk.document_id
CROSS JOIN query_input
WHERE chunk.knowledge_base_id = %s
  AND chunk.access_scope_id = ANY(%s::uuid[])
  AND knowledge_base.active_embedding_profile_hash = %s
  AND knowledge_base.active_embedding_profile_kind = 'embedding'
  AND profile.kind = 'embedding'
  AND profile.enabled IS TRUE
  AND profile.dimension = 1024
  AND chunk.embedding_profile_hash = knowledge_base.active_embedding_profile_hash
  AND chunk.embedding_model = profile.model
  AND chunk.embedding_dimension = profile.dimension
  AND chunk.embedding IS NOT NULL
  AND chunk.embedded_text_sha256 = chunk.embedding_text_sha256
  AND chunk.searchable IS TRUE
  AND chunk.index_status = 'READY'
  AND chunk.approval_status = 'CLOUD_APPROVED'
  AND chunk.deleted_at IS NULL
  AND document.status = 'READY'
  AND document.deleted_at IS NULL
  AND document.active_version_id = chunk.document_version_id
  AND document_version.status = 'READY'
  AND document_version.review_state = 'CLOUD_APPROVED'
  AND document_version.embedding_profile_hash = knowledge_base.active_embedding_profile_hash
  AND document_version.outbound_manifest_sha256 = chunk.outbound_manifest_sha256
ORDER BY chunk.embedding <=> query_input.embedding, chunk.citation_id
LIMIT %s
"""


class PostgresRetrievalRepository:
    """Read-only PostgreSQL implementation with filtered HNSW candidate search."""

    def __init__(self, settings: Settings) -> None:
        self._settings = settings

    def _dsn(self) -> str:
        return (
            self._settings.database_url()
            .set(drivername="postgresql")
            .render_as_string(hide_password=False)
        )

    def resolve_active_profile(
        self,
        knowledge_base_id: uuid.UUID,
        *,
        allowed_scope_ids: Sequence[uuid.UUID],
    ) -> ActiveEmbeddingProfile | None:
        if not allowed_scope_ids:
            return None
        try:
            with psycopg.connect(
                self._dsn(),
                connect_timeout=2,
                row_factory=dict_row,
            ) as connection:
                connection.execute("SET LOCAL statement_timeout = '3000ms'")
                row = connection.execute(
                    ACTIVE_PROFILE_SQL,
                    (knowledge_base_id, list(allowed_scope_ids)),
                ).fetchone()
        except (OSError, SecretFileError, psycopg.Error) as exc:
            raise RetrievalPersistenceError from exc

        if row is None:
            return None
        try:
            profile_hash = cast(str, row["profile_hash"])
            model = cast(str, row["model"])
            dimension = cast(int, row["dimension"])
            synthetic = cast(bool, row["synthetic"])
            if (
                _PROFILE_HASH_PATTERN.fullmatch(profile_hash) is None
                or not model
                or model != model.strip()
                or isinstance(dimension, bool)
                or dimension != 1024
                or not isinstance(synthetic, bool)
            ):
                raise ValueError
        except (KeyError, TypeError, ValueError) as exc:
            raise RetrievalPersistenceError from exc
        return ActiveEmbeddingProfile(
            profile_hash=profile_hash,
            model=model,
            dimension=dimension,
            synthetic=synthetic,
        )

    def search_candidates(
        self,
        knowledge_base_id: uuid.UUID,
        *,
        allowed_scope_ids: Sequence[uuid.UUID],
        profile_hash: str,
        query_vector: tuple[float, ...],
        limit: int,
    ) -> list[RetrievalCandidate]:
        if (
            not allowed_scope_ids
            or _PROFILE_HASH_PATTERN.fullmatch(profile_hash) is None
            or len(query_vector) != 1024
            or isinstance(limit, bool)
            or not 1 <= limit <= 50
        ):
            raise RetrievalPersistenceError

        try:
            vector = Vector(list(query_vector))
            with psycopg.connect(
                self._dsn(),
                connect_timeout=2,
                row_factory=dict_row,
            ) as connection:
                register_vector(connection)
                connection.execute("SET LOCAL statement_timeout = '3000ms'")
                connection.execute("SET LOCAL hnsw.iterative_scan = strict_order")
                connection.execute("SET LOCAL hnsw.ef_search = 100")
                rows = connection.execute(
                    CANDIDATE_SEARCH_SQL,
                    (
                        vector,
                        knowledge_base_id,
                        list(allowed_scope_ids),
                        profile_hash,
                        limit,
                    ),
                ).fetchall()
        except (OSError, SecretFileError, psycopg.Error) as exc:
            raise RetrievalPersistenceError from exc

        try:
            return [self._candidate(row) for row in rows]
        except (KeyError, TypeError, ValueError) as exc:
            raise RetrievalPersistenceError from exc

    @staticmethod
    def _candidate(row: dict[str, Any]) -> RetrievalCandidate:
        raw_section_path = row["section_path"]
        if not isinstance(raw_section_path, list) or any(
            not isinstance(part, str) or not part.strip() for part in raw_section_path
        ):
            raise ValueError
        source_name = row["source_name"]
        cloud_text = row["cloud_text"]
        raw_score = row["vector_score"]
        if (
            not isinstance(source_name, str)
            or not source_name.strip()
            or not isinstance(cloud_text, str)
            or not cloud_text.strip()
            or isinstance(raw_score, bool)
            or not isinstance(raw_score, (int, float))
            or not math.isfinite(float(raw_score))
        ):
            raise ValueError

        page_start = row["page_start"]
        page_end = row["page_end"]
        if (page_start is None) != (page_end is None) or (
            page_start is not None
            and (
                isinstance(page_start, bool)
                or isinstance(page_end, bool)
                or not isinstance(page_start, int)
                or not isinstance(page_end, int)
                or page_start < 1
                or page_end < page_start
            )
        ):
            raise ValueError

        return RetrievalCandidate(
            citation_id=uuid.UUID(cast(str, row["citation_id"])),
            document_id=uuid.UUID(cast(str, row["document_id"])),
            source_name=source_name.strip(),
            cloud_text=cloud_text.strip(),
            section_path=tuple(part.strip() for part in raw_section_path),
            page_start=cast(int | None, page_start),
            page_end=cast(int | None, page_end),
            vector_score=float(raw_score),
        )