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

"""Public SSE API for evidence-grounded, single-turn chat."""

from __future__ import annotations

import json
import uuid
from collections.abc import AsyncIterator, Mapping
from typing import Annotated, Literal

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

from app.adapters.model_gateway import ModelGatewayAdapter
from app.api.v1.retrieval import (
    get_retrieval_actor,
    get_retrieval_model_gateway,
    get_retrieval_service,
)
from app.services.chat import (
    CHAT_MAX_TOKENS_DEFAULT,
    CHAT_MAX_TOKENS_LIMIT,
    ChatEvent,
    GroundedChatService,
    PreparedChat,
)
from app.services.retrieval import (
    QUERY_MAX_LENGTH,
    RERANK_TOP_N_DEFAULT,
    VECTOR_TOP_K_DEFAULT,
    RetrievalActor,
    RetrievalService,
)


class StrictDto(BaseModel):
    """Reject unknown fields on every public chat DTO."""

    model_config = ConfigDict(extra="forbid")


class ChatCompletionRequest(StrictDto):
    knowledge_base_id: uuid.UUID
    question: 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)
    max_tokens: int = Field(default=CHAT_MAX_TOKENS_DEFAULT, ge=1, le=CHAT_MAX_TOKENS_LIMIT)

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


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


class ChatMetaEventDto(StrictDto):
    seq: int = Field(ge=1)
    trace_id: str = Field(min_length=1)
    knowledge_base_id: uuid.UUID
    profile: ChatProfileDto
    generation_mode: Literal["synthetic_extractive", "cloud_grounded"]


class ChatEvidenceDto(StrictDto):
    label: str = Field(pattern=r"^S[1-9]\d*$")
    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 ChatTimingsDto(StrictDto):
    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 ChatRetrievalEventDto(StrictDto):
    seq: int = Field(ge=1)
    status: Literal["ok", "empty"]
    rerank_status: Literal["applied", "degraded", "skipped_empty"]
    degradation_reason: Literal["rerank_unavailable"] | None
    evidence: list[ChatEvidenceDto]
    timings: ChatTimingsDto


class ChatDeltaEventDto(StrictDto):
    seq: int = Field(ge=1)
    text: str


class ChatCitationsEventDto(StrictDto):
    seq: int = Field(ge=1)
    citations: list[ChatEvidenceDto]


class ChatUsageEventDto(StrictDto):
    seq: int = Field(ge=1)
    model: str = Field(min_length=1)
    request_id: str | None
    input_tokens: int | None = Field(default=None, ge=0)
    output_tokens: int | None = Field(default=None, ge=0)
    total_tokens: int | None = Field(default=None, ge=0)


class ChatDoneEventDto(StrictDto):
    seq: int = Field(ge=1)
    status: Literal["complete"]
    answer_mode: Literal["grounded", "refused", "retrieval_only"]
    finish_reason: str | None


class ChatErrorEventDto(StrictDto):
    seq: int = Field(ge=1)
    status: Literal["error"]
    code: Literal["CHAT_PROVIDER_UNAVAILABLE", "CHAT_GENERATION_FAILED"]
    title: str
    retryable: bool
    answer_mode: Literal["retrieval_only"]


_EVENT_MODELS: Mapping[str, type[BaseModel]] = {
    "meta": ChatMetaEventDto,
    "retrieval": ChatRetrievalEventDto,
    "delta": ChatDeltaEventDto,
    "citations": ChatCitationsEventDto,
    "usage": ChatUsageEventDto,
    "done": ChatDoneEventDto,
    "error": ChatErrorEventDto,
}


def get_chat_service(
    retrieval_service: Annotated[RetrievalService, Depends(get_retrieval_service)],
    model_gateway: Annotated[ModelGatewayAdapter, Depends(get_retrieval_model_gateway)],
) -> GroundedChatService:
    return GroundedChatService(
        retrieval_service=retrieval_service,
        chat_provider=model_gateway,
    )


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


@router.post(
    "/completions",
    operation_id="streamGroundedChatCompletion",
    response_class=StreamingResponse,
    responses={
        200: {
            "description": "Monotonic grounded-chat event stream",
            "content": {"text/event-stream": {"schema": {"type": "string"}}},
        }
    },
)
async def chat_completion(
    payload: ChatCompletionRequest,
    request: Request,
    service: Annotated[GroundedChatService, Depends(get_chat_service)],
    actor: Annotated[RetrievalActor, Depends(get_retrieval_actor)],
) -> StreamingResponse:
    # Preparation is intentionally awaited before StreamingResponse.  Formal
    # retrieval problems therefore remain normal RFC-style problem JSON.
    prepared = await service.prepare(
        actor=actor,
        knowledge_base_id=payload.knowledge_base_id,
        question=payload.question,
        vector_top_k=payload.vector_top_k,
        rerank_top_n=payload.rerank_top_n,
        max_tokens=payload.max_tokens,
    )
    trace_id = str(getattr(request.state, "trace_id", "unavailable"))
    return StreamingResponse(
        _event_stream(service, prepared, trace_id=trace_id),
        media_type="text/event-stream",
        headers={
            "Cache-Control": "no-store",
            "X-Accel-Buffering": "no",
        },
    )


async def _event_stream(
    service: GroundedChatService,
    prepared: PreparedChat,
    *,
    trace_id: str,
) -> AsyncIterator[bytes]:
    async for event in service.stream(prepared, trace_id=trace_id):
        yield _serialize_event(event)


def _serialize_event(event: ChatEvent) -> bytes:
    model_type = _EVENT_MODELS[event.name]
    payload = model_type.model_validate({"seq": event.seq, **event.data}).model_dump(mode="json")
    serialized = json.dumps(payload, ensure_ascii=False, separators=(",", ":"))
    # JSON strings are plain text, but HTML-sensitive code points are escaped so
    # even an unsafe intermediary cannot turn raw evidence into active markup.
    serialized = serialized.replace("&", "\\u0026").replace("<", "\\u003c").replace(">", "\\u003e")
    return f"event: {event.name}\ndata: {serialized}\n\n".encode()