"""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), )