RAG/backend/app/api/v1/demo.py

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