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

"""Formal retrieval HTTP API with server-derived synthetic access grants."""

from __future__ import annotations

import uuid
from collections.abc import AsyncIterator
from typing import Annotated, Any, Literal

from fastapi import APIRouter, Depends, Request
from pydantic import BaseModel, ConfigDict, Field, field_validator

from app.adapters.fake import FakeEmbeddingProvider, FakeReranker
from app.adapters.model_gateway import ModelGatewayAdapter
from app.core.config import Settings, get_settings
from app.core.demo_identity import (
    ACCESS_SCOPE_ID,
    BAILIAN_ACCESS_SCOPE_ID,
    BAILIAN_KNOWLEDGE_BASE_ID,
    KNOWLEDGE_BASE_ID,
)
from app.persistence.retrieval import PostgresRetrievalRepository, RetrievalRepository
from app.services.retrieval import (
    QUERY_MAX_LENGTH,
    RERANK_TOP_N_DEFAULT,
    VECTOR_TOP_K_DEFAULT,
    EffectiveRetrievalParameters,
    RetrievalActor,
    RetrievalGrant,
    RetrievalHit,
    RetrievalResult,
    RetrievalService,
    RetrievalTimings,
)


class RetrievalSearchRequest(BaseModel):
    """Bounded client input. Access-scope fields are intentionally forbidden."""

    model_config = ConfigDict(extra="forbid")

    knowledge_base_id: uuid.UUID
    query: str = Field(min_length=1, max_length=QUERY_MAX_LENGTH)
    vector_top_k: int = Field(default=VECTOR_TOP_K_DEFAULT, ge=1, le=10_000)
    rerank_top_n: int = Field(default=RERANK_TOP_N_DEFAULT, ge=1, le=10_000)

    @field_validator("query")
    @classmethod
    def normalize_query(cls, value: str) -> str:
        normalized = " ".join(value.split())
        if not normalized:
            raise ValueError("query must contain non-whitespace text")
        return normalized


class RetrievalProfileResponse(BaseModel):
    profile_hash: str = Field(pattern=r"^[0-9a-f]{64}$")
    model: str
    dimension: Literal[1024]
    synthetic: bool


class RetrievalParametersResponse(BaseModel):
    vector_top_k: int = Field(ge=1, le=50)
    rerank_top_n: int = Field(ge=1, le=10)


class RetrievalTimingsResponse(BaseModel):
    embedding_ms: float = Field(ge=0, allow_inf_nan=False)
    database_ms: float = Field(ge=0, allow_inf_nan=False)
    rerank_ms: float = Field(ge=0, allow_inf_nan=False)
    total_ms: float = Field(ge=0, allow_inf_nan=False)


class RetrievalHitResponse(BaseModel):
    rank: int = Field(ge=1)
    vector_rank: int = Field(ge=1)
    citation_id: uuid.UUID
    document_id: uuid.UUID
    source_name: str = Field(min_length=1, max_length=240)
    snippet: str = Field(min_length=1, max_length=1_200)
    section_path: list[str]
    page_start: int | None = Field(default=None, ge=1)
    page_end: int | None = Field(default=None, ge=1)
    page_label: str
    vector_score: float = Field(ge=-1, le=1, allow_inf_nan=False)
    rerank_score: float | None = Field(default=None, ge=0, le=1, allow_inf_nan=False)


class RetrievalSearchResponse(BaseModel):
    status: Literal["ok", "empty"]
    trace_id: str
    knowledge_base_id: uuid.UUID
    access_scope_count: int = Field(ge=1)
    profile: RetrievalProfileResponse
    parameters: RetrievalParametersResponse
    rerank_status: Literal["applied", "degraded", "skipped_empty"]
    degradation_reason: Literal["rerank_unavailable"] | None
    embedding_request_id: str | None
    rerank_request_id: str | None
    embedding_model: str
    rerank_model: str | None
    timings: RetrievalTimingsResponse
    results: list[RetrievalHitResponse]


_SYNTHETIC_ACTOR = RetrievalActor(
    subject="synthetic-demo-reader",
    grants=(
        RetrievalGrant(
            knowledge_base_id=KNOWLEDGE_BASE_ID,
            access_scope_ids=(ACCESS_SCOPE_ID,),
        ),
        RetrievalGrant(
            knowledge_base_id=BAILIAN_KNOWLEDGE_BASE_ID,
            access_scope_ids=(BAILIAN_ACCESS_SCOPE_ID,),
        ),
    ),
)


def get_retrieval_actor() -> RetrievalActor:
    """Return the temporary server-owned actor until real authentication replaces it."""

    return _SYNTHETIC_ACTOR


def get_retrieval_repository(
    settings: Annotated[Settings, Depends(get_settings)],
) -> RetrievalRepository:
    return PostgresRetrievalRepository(settings)


async def get_retrieval_model_gateway(
    settings: Annotated[Settings, Depends(get_settings)],
) -> AsyncIterator[ModelGatewayAdapter]:
    adapter = ModelGatewayAdapter.from_settings(settings)
    try:
        yield adapter
    finally:
        await adapter.aclose()


def get_retrieval_service(
    repository: Annotated[RetrievalRepository, Depends(get_retrieval_repository)],
    model_gateway: Annotated[ModelGatewayAdapter, Depends(get_retrieval_model_gateway)],
) -> RetrievalService:
    return RetrievalService(
        repository=repository,
        embedding_provider=model_gateway,
        reranker=model_gateway,
        synthetic_embedding_provider=FakeEmbeddingProvider(1024),
        synthetic_reranker=FakeReranker(),
    )


def _profile(result: RetrievalResult) -> RetrievalProfileResponse:
    return RetrievalProfileResponse(
        profile_hash=result.profile.profile_hash,
        model=result.profile.model,
        dimension=1024,
        synthetic=result.profile.synthetic,
    )


def _parameters(value: EffectiveRetrievalParameters) -> RetrievalParametersResponse:
    return RetrievalParametersResponse(
        vector_top_k=value.vector_top_k,
        rerank_top_n=value.rerank_top_n,
    )


def _timings(value: RetrievalTimings) -> RetrievalTimingsResponse:
    return RetrievalTimingsResponse(
        embedding_ms=value.embedding_ms,
        database_ms=value.database_ms,
        rerank_ms=value.rerank_ms,
        total_ms=value.total_ms,
    )


def _hit(value: RetrievalHit) -> RetrievalHitResponse:
    return RetrievalHitResponse(
        rank=value.rank,
        vector_rank=value.vector_rank,
        citation_id=value.citation_id,
        document_id=value.document_id,
        source_name=value.source_name,
        snippet=value.snippet,
        section_path=list(value.section_path),
        page_start=value.page_start,
        page_end=value.page_end,
        page_label=value.page_label,
        vector_score=value.vector_score,
        rerank_score=value.rerank_score,
    )


router = APIRouter(prefix="/api/v1/retrieval", tags=["retrieval"])


@router.post(
    "/search",
    response_model=RetrievalSearchResponse,
    operation_id="searchRetrievalEvidence",
)
async def retrieval_search(
    payload: RetrievalSearchRequest,
    request: Request,
    service: Annotated[RetrievalService, Depends(get_retrieval_service)],
    actor: Annotated[RetrievalActor, Depends(get_retrieval_actor)],
) -> Any:
    result = await service.search(
        actor=actor,
        knowledge_base_id=payload.knowledge_base_id,
        query=payload.query,
        vector_top_k=payload.vector_top_k,
        rerank_top_n=payload.rerank_top_n,
    )
    return RetrievalSearchResponse(
        status=result.status,
        trace_id=str(getattr(request.state, "trace_id", "unavailable")),
        knowledge_base_id=result.knowledge_base_id,
        access_scope_count=result.access_scope_count,
        profile=_profile(result),
        parameters=_parameters(result.parameters),
        rerank_status=result.rerank_status,
        degradation_reason=result.degradation_reason,
        embedding_request_id=result.embedding_request_id,
        rerank_request_id=result.rerank_request_id,
        embedding_model=result.embedding_model,
        rerank_model=result.rerank_model,
        timings=_timings(result.timings),
        results=[_hit(hit) for hit in result.results],
    )