"""Read-only offline RAG demo endpoints backed by the synthetic dataset.""" from __future__ import annotations import asyncio import hashlib import re from dataclasses import dataclass from typing import Annotated, Any, Literal, Protocol, cast import psycopg from fastapi import APIRouter, Depends, HTTPException, status from pgvector.psycopg import register_vector from pgvector.vector import Vector from psycopg.rows import dict_row from pydantic import BaseModel, Field, field_validator from app.adapters.fake import FakeEmbeddingProvider, FakeReranker from app.core.config import Settings, get_settings from app.core.demo_identity import ( ACCESS_SCOPE_ID, DEMO_EXPECTED_CHUNKS, DEMO_FAKE_EMBEDDING_MODEL, DEMO_SCOPE_NAME, KNOWLEDGE_BASE_ID, offline_embedding_profile_hash, ) from app.core.secrets import SecretFileError QUERY_MAX_LENGTH = 500 SNIPPET_MAX_LENGTH = 600 TITLE_MAX_LENGTH = 120 _SPACE_PATTERN = re.compile(r"\s+") class DemoCounts(BaseModel): """Public aggregate counts for the synthetic dataset.""" chunks: int = Field(ge=0) vectors: int = Field(ge=0) searchable: int = Field(ge=0) class DemoStatusResponse(BaseModel): """Safe readiness summary without database or approval internals.""" status: Literal["ready", "empty_dataset", "incomplete_dataset"] dataset: Literal["synthetic-demo"] = "synthetic-demo" counts: DemoCounts class DemoSearchRequest(BaseModel): """Bounded request accepted by the offline demo search.""" query: str = Field(min_length=1, max_length=QUERY_MAX_LENGTH) top_k: int = Field(default=5, ge=1, le=10) @field_validator("query") @classmethod def normalize_query(cls, value: str) -> str: normalized = _SPACE_PATTERN.sub(" ", value).strip() if not normalized: raise ValueError("query must contain non-whitespace text") return normalized class DemoSearchItem(BaseModel): """Public synthetic result; internal identifiers and hashes are excluded.""" title: str snippet: str page_label: str score: float = Field(ge=0.0, le=1.0) citation_id: str class DemoSearchResponse(BaseModel): """Offline retrieval response with an explicit empty-dataset state.""" status: Literal["ok", "empty_dataset"] dataset: Literal["synthetic-demo"] = "synthetic-demo" results: list[DemoSearchItem] @dataclass(frozen=True, slots=True) class DemoCandidate: """Private retrieval projection used only while constructing safe results.""" source_key: str title: str text: str page_start: int | None page_end: int | None @property def rerank_text(self) -> str: return f"{self.title}\n{self.text}" class DemoRepository(Protocol): """Read-only persistence boundary used by the demo router.""" def counts(self) -> DemoCounts: ... def search(self, query_vector: tuple[float, ...], *, limit: int) -> list[DemoCandidate]: ... class PostgresDemoRepository: """Read-only PostgreSQL/pgvector projection for approved synthetic chunks.""" 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 counts(self) -> DemoCounts: profile_hash = offline_embedding_profile_hash(self._settings.embedding_dimension) with psycopg.connect( self._dsn(), connect_timeout=2, row_factory=dict_row, ) as connection: row = connection.execute( """ SELECT count(*)::integer AS chunks, count(*) FILTER (WHERE chunk.embedding IS NOT NULL)::integer AS vectors, count(*) FILTER (WHERE chunk.searchable)::integer AS searchable FROM rag.chunks AS chunk JOIN rag.access_scopes AS scope ON scope.id = chunk.access_scope_id AND 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 version ON version.id = chunk.document_version_id AND version.document_id = chunk.document_id WHERE chunk.knowledge_base_id = %s AND chunk.access_scope_id = %s AND scope.name = %s AND chunk.metadata ->> 'source_type' = 'synthetic' AND chunk.index_status = 'READY' AND chunk.approval_status = 'CLOUD_APPROVED' AND chunk.deleted_at IS NULL AND chunk.embedding_model = %s AND chunk.embedding_profile_hash = %s AND document.status = 'READY' AND document.deleted_at IS NULL AND document.active_version_id = chunk.document_version_id AND version.status = 'READY' AND version.review_state = 'CLOUD_APPROVED' AND version.outbound_manifest_sha256 = chunk.outbound_manifest_sha256 AND version.embedding_profile_hash = chunk.embedding_profile_hash """, ( KNOWLEDGE_BASE_ID, ACCESS_SCOPE_ID, DEMO_SCOPE_NAME, DEMO_FAKE_EMBEDDING_MODEL, profile_hash, ), ).fetchone() if row is None: return DemoCounts(chunks=0, vectors=0, searchable=0) return DemoCounts( chunks=int(row["chunks"]), vectors=int(row["vectors"]), searchable=int(row["searchable"]), ) def search(self, query_vector: tuple[float, ...], *, limit: int) -> list[DemoCandidate]: if len(query_vector) != self._settings.embedding_dimension: raise ValueError("query vector dimension does not match the demo index") vector = Vector(list(query_vector)) profile_hash = offline_embedding_profile_hash(self._settings.embedding_dimension) 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( """ SELECT chunk.id::text AS source_key, COALESCE(NULLIF(chunk.section_path ->> 0, ''), '合成地质资料') AS title, chunk.cloud_text AS text, chunk.page_start, chunk.page_end FROM rag.chunks AS chunk JOIN rag.access_scopes AS scope ON scope.id = chunk.access_scope_id AND 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 version ON version.id = chunk.document_version_id AND version.document_id = chunk.document_id WHERE chunk.knowledge_base_id = %s AND chunk.access_scope_id = %s AND scope.name = %s AND chunk.metadata ->> 'source_type' = 'synthetic' AND chunk.searchable IS TRUE AND chunk.embedding IS NOT NULL AND chunk.index_status = 'READY' AND chunk.approval_status = 'CLOUD_APPROVED' AND chunk.deleted_at IS NULL AND chunk.embedding_model = %s AND chunk.embedding_profile_hash = %s AND document.status = 'READY' AND document.deleted_at IS NULL AND document.active_version_id = chunk.document_version_id AND version.status = 'READY' AND version.review_state = 'CLOUD_APPROVED' AND version.outbound_manifest_sha256 = chunk.outbound_manifest_sha256 AND version.embedding_profile_hash = chunk.embedding_profile_hash ORDER BY chunk.embedding <=> %s LIMIT %s """, ( KNOWLEDGE_BASE_ID, ACCESS_SCOPE_ID, DEMO_SCOPE_NAME, DEMO_FAKE_EMBEDDING_MODEL, profile_hash, vector, limit, ), ).fetchall() return [ DemoCandidate( source_key=cast(str, row["source_key"]), title=cast(str, row["title"]), text=cast(str, row["text"]), page_start=cast(int | None, row["page_start"]), page_end=cast(int | None, row["page_end"]), ) for row in rows ] def get_demo_repository( settings: Annotated[Settings, Depends(get_settings)], ) -> DemoRepository: """Build the default repository without loading any model credential.""" return PostgresDemoRepository(settings) def _bounded_text(value: str, max_length: int) -> str: normalized = _SPACE_PATTERN.sub(" ", value).strip() if len(normalized) <= max_length: return normalized return f"{normalized[: max_length - 1]}…" def make_citation_id(source_key: str) -> str: """Derive a stable opaque citation without returning an internal UUID.""" digest = hashlib.sha256(f"{DEMO_SCOPE_NAME}:{source_key}".encode()).hexdigest() return f"demo-{digest[:16]}" def make_page_label(page_start: int | None, page_end: int | None) -> str: if page_start is None or page_end is None: return "页码未知" if page_start == page_end: return f"第 {page_start} 页" return f"第 {page_start}-{page_end} 页" def _safe_result(candidate: DemoCandidate, score: float) -> DemoSearchItem: return DemoSearchItem( title=f"合成资料|{_bounded_text(candidate.title, TITLE_MAX_LENGTH)}", snippet=_bounded_text(candidate.text, SNIPPET_MAX_LENGTH), page_label=make_page_label(candidate.page_start, candidate.page_end), score=round(max(0.0, min(1.0, score)), 6), citation_id=make_citation_id(candidate.source_key), ) def _database_unavailable(exc: BaseException) -> HTTPException: return HTTPException( status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="database unavailable", ) router = APIRouter(prefix="/api/v1/demo", tags=["offline-demo"]) @router.get("/status", response_model=DemoStatusResponse) def demo_status(repository: Annotated[DemoRepository, Depends(get_demo_repository)]) -> Any: try: counts = repository.counts() except (OSError, SecretFileError, psycopg.Error) as exc: raise _database_unavailable(exc) from exc state: Literal["ready", "empty_dataset", "incomplete_dataset"] if counts.chunks == 0: state = "empty_dataset" elif ( counts.chunks == DEMO_EXPECTED_CHUNKS and counts.vectors == DEMO_EXPECTED_CHUNKS and counts.searchable == DEMO_EXPECTED_CHUNKS ): state = "ready" else: state = "incomplete_dataset" return DemoStatusResponse(status=state, counts=counts) @router.post("/search", response_model=DemoSearchResponse) async def demo_search( request: DemoSearchRequest, repository: Annotated[DemoRepository, Depends(get_demo_repository)], settings: Annotated[Settings, Depends(get_settings)], ) -> Any: embedder = FakeEmbeddingProvider(settings.embedding_dimension) reranker = FakeReranker() query_result = await embedder.embed_query(request.query) candidate_limit = min(max(request.top_k * 3, 10), 20) try: candidates = await asyncio.to_thread( repository.search, query_result.vectors[0], limit=candidate_limit, ) except (OSError, SecretFileError, psycopg.Error) as exc: raise _database_unavailable(exc) from exc if not candidates: return DemoSearchResponse(status="empty_dataset", results=[]) reranked = await reranker.rerank( request.query, [candidate.rerank_text for candidate in candidates], top_n=min(request.top_k, len(candidates)), ) results = [ _safe_result(candidates[item.index], item.relevance_score) for item in reranked.items ] return DemoSearchResponse(status="ok", results=results)