src/api/dto/query.py

"""DTOs for ``/query/*`` endpoints (Stage 3 covers the read side only).

Stage 3 exposes ``GET /query/runs/{runId}`` — the persisted envelope for
a committed run. The ``POST /query`` request shape + live SSE land in
Stage 4. This module models the response envelope as a **discriminated
union by ``mode``** so the FE's typed codegen produces three concrete TS
types (``LookupEnvelope`` / ``ComparisonEnvelope`` / ``ConceptualEnvelope``)
the user can switch over.

Why ``mode: RouterModeEnum`` is repeated on each leaf rather than living
on ``_Base``: Pydantic v2's discriminated unions need the discriminator
to be a literal-typed field on every member so the OpenAPI ``oneOf``
emits a ``discriminator.mapping`` block the codegen can consume. We
declare it on the leaves with a default = the right enum value.

The ``run_id`` and ``cost`` fields are required because Stage-3 reads
only return persisted (committed) runs — the cost is known by the time
the envelope is materialised. During live runs the cost streams via SSE
``progress.costSoFarUsd`` (CONTRACT.md §5).
"""

from __future__ import annotations

from enum import Enum
from typing import Annotated, Union

from pydantic import Field, model_validator

from .common import ApiModel

# Server-side bounds on the cost amplifiers. These cap how much a single
# request can fan out to paid providers — the judge flagged top_k,
# comparison_spec fan-out, and query length as unbounded. The daily spend cap
# (assert_within_daily_cap) is the real backstop on cost, so this ceiling is
# generous (up to 1k results for deep research; default stays low).
_MAX_TOP_K = 1000
_MAX_GROUPS = 10
_MAX_BOOKS_PER_GROUP = 50
_MAX_COMPARISON_BOOKS = 60  # union across groups (per_book = 1 paid call each)
_MAX_FILTER_IDS = 500
_MAX_QUERY_CHARS = 4000


# ---------- Enums ---------------------------------------------------------
class RouterModeEnum(str, Enum):
    """Final mode the router classified the query into.

    LOOKUP = single citation-grounded answer.
    COMPARISON = side-by-side across groups (per_book / per_group / head_to_head).
    CONCEPTUAL = open-ended exploration with per-candidate judging.

    Locked vocabulary; matches the strings stored in
    ``query_history.mode`` so the wire round-trips persisted runs without
    translation.
    """

    LOOKUP = "LOOKUP"
    COMPARISON = "COMPARISON"
    CONCEPTUAL = "CONCEPTUAL"


class ComparisonStrategyEnum(str, Enum):
    """How COMPARISON mode partitions retrieval across books/groups."""

    per_book = "per_book"
    per_group = "per_group"
    head_to_head = "head_to_head"


# ---------- Shared shapes -------------------------------------------------
class Citation(ApiModel):
    """One retrieved-and-cited chunk projected onto the wire.

    ``page_label`` is **server-formatted** (e.g. ``"p. 42 (PDF p. 51)"``)
    so the FE never has to know about printed-page mapping. Matches the
    behavior of :func:`src.stage9_ui.shared.format_page_label`.
    """

    book_id: str
    book_title: str
    author: str | None = None
    pdf_page: int
    printed_page: int | None = None
    page_label: str
    excerpt: str
    score: float | None = None


class CostBreakdown(ApiModel):
    """Aggregate cost for one query run.

    ``total_usd`` is the sum of every ``llm_calls.cost_usd`` row that
    falls within the run's ``llm_call_id_range`` (CONTRACT.md §12).
    """

    total_usd: float
    duration_ms: int
    n_calls: int
    by_stage: dict[str, float]


class LlmCallIdRange(ApiModel):
    """Inclusive id range of ``llm_calls`` rows spanned by this run.

    Backs the cost panel's ``GET /costs/llm-calls?sinceId=`` query —
    the FE asks for everything billed strictly after ``firstId`` (or
    inclusive of it; both ends survive the wire as separate fields so
    the UI can pick).
    """

    first_id: int | None = None
    last_id: int | None = None


class ConsultedBook(ApiModel):
    """LOOKUP mode: one book that contributed citations."""

    book_id: str
    title: str
    n_citations: int


class ComparisonRow(ApiModel):
    """One row of a COMPARISON envelope.

    Shape stays free-form per ``BACKEND_BUILD.md §8.2``; the Stage-4
    query router will populate concrete fields. Stage 3 only persists
    what the existing pipeline emits.
    """

    group_id: str | None = None
    group_label: str | None = None
    book_id: str | None = None
    book_title: str | None = None
    answer: str | None = None
    citations: list[Citation] | None = None
    cost_usd: float | None = None


class AllusionMatch(ApiModel):
    """One per-candidate verdict from CONCEPTUAL mode's judging stage."""

    book_id: str
    book_title: str
    pdf_page: int
    printed_page: int | None = None
    page_label: str
    excerpt: str
    verdict: str
    reasoning: str
    score: float | None = None


# ---------- Envelopes (discriminated by mode) -----------------------------
class _BaseEnvelope(ApiModel):
    """Fields shared by every mode's envelope.

    ``router`` carries the classifier's reasoning blob (kept as ``dict``
    because the pipeline emits free-form JSON here and we don't want to
    couple the wire DTO to internal classifier evolution).
    """

    run_id: int
    query: str
    router: dict
    cost: CostBreakdown
    llm_call_id_range: LlmCallIdRange
    created_at: str


class LookupEnvelope(_BaseEnvelope):
    """LOOKUP mode response: one grounded answer + citations."""

    mode: RouterModeEnum = RouterModeEnum.LOOKUP
    answer: str
    citations: list[Citation]
    n_retrieved: int
    books_consulted: list[ConsultedBook]
    confidence: float | None = None
    ungrounded_claims: list[str] | None = None


class ComparisonEnvelope(_BaseEnvelope):
    """COMPARISON mode: per-row partition with optional cross-cutting synthesis."""

    mode: RouterModeEnum = RouterModeEnum.COMPARISON
    strategy: ComparisonStrategyEnum
    groups: list[dict] | None = None
    rows: list[ComparisonRow]
    cross_cutting: str | None = None


class ConceptualEnvelope(_BaseEnvelope):
    """CONCEPTUAL mode: open-ended exploration with per-candidate verdicts."""

    mode: RouterModeEnum = RouterModeEnum.CONCEPTUAL
    matches: list[AllusionMatch]
    n_candidates_judged: int


# Pydantic v2 discriminated union — the OpenAPI ``oneOf`` emits a
# ``discriminator.mapping`` so codegen produces three concrete TS types.
QueryResponseEnvelope = Annotated[
    Union[LookupEnvelope, ComparisonEnvelope, ConceptualEnvelope],
    Field(discriminator="mode"),
]


# ---------- Stage 4: POST /query request shape ----------------------------
# Locked field names per CONTRACT.md §11. The FE references these by name
# in its OpenAPI codegen; renaming any of them breaks the wire silently.


class PageEnum(str, Enum):
    """Which research page kicked off the query.

    Tagged onto ``query_history.page`` so the history list can scope by
    page. ``search`` and ``compare`` are user-facing; ``allusions`` runs
    CONCEPTUAL mode under the hood.
    """

    search = "search"
    compare = "compare"
    allusions = "allusions"


class Filters(ApiModel):
    """Optional retrieval scope (CONTRACT.md §11).

    All fields nullable; omitted fields mean "no scope constraint on that
    axis". The router treats an all-null filter set as the unrestricted
    corpus (matches Streamlit parity).
    """

    tradition: str | None = None
    era: str | None = None
    language: str | None = None
    author_id: str | None = None
    book_ids: list[str] | None = Field(default=None, max_length=_MAX_FILTER_IDS)
    label_ids: list[str] | None = Field(default=None, max_length=_MAX_FILTER_IDS)


class ComparisonGroup(ApiModel):
    """One named group of books used by COMPARISON mode."""

    id: str
    label: str
    book_ids: list[str] = Field(..., min_length=1, max_length=_MAX_BOOKS_PER_GROUP)


class ComparisonSpec(ApiModel):
    """How COMPARISON mode partitions retrieval across books/groups.

    ``per_book`` ignores ``groups``; ``per_group`` requires it; ``head_to_head``
    expects exactly two groups. Validation happens server-side in the router.
    """

    strategy: ComparisonStrategyEnum
    top_k_per_book: int | None = Field(default=None, ge=1, le=_MAX_TOP_K)
    groups: list[ComparisonGroup] | None = Field(default=None, max_length=_MAX_GROUPS)

    @model_validator(mode="after")
    def _bound_total_books(self) -> "ComparisonSpec":
        """Cap the union of books across groups.

        ``per_book`` strategy issues one paid generation call per unique
        book in the union, so an unbounded group list is a cost amplifier.
        Per-group and per-request bounds alone (10 × 50) would still allow
        500 calls; this caps the total that actually fans out.
        """
        if self.groups:
            total = sum(len(g.book_ids) for g in self.groups)
            if total > _MAX_COMPARISON_BOOKS:
                raise ValueError(
                    f"comparison_spec spans {total} books; the maximum is "
                    f"{_MAX_COMPARISON_BOOKS} (each book is a separate paid "
                    f"generation call)."
                )
        return self


class QueryModelOverrides(ApiModel):
    """Per-stage model overrides from the AdvancedModels expander.

    Each field is the model id (``"gemini-2.5-flash"``, ``"claude-haiku-4-5"``,
    etc.) — the FE picks them from ``GET /tools?category=llm``. None on
    any field falls back to the config default for that stage.
    """

    generation: str | None = None
    classify: str | None = None
    judge: str | None = None


class QueryRequest(ApiModel):
    """``POST /query`` body (BACKEND_BUILD.md §8.1b).

    Field names are LOCKED — the FE typed client references them by name.
    The runner strips snake_case off the dict on its way into
    ``src.stage8_router.router.answer``; that function takes Python
    snake_case kwargs (which is why ``populate_by_name=True`` on the
    :class:`ApiModel` config matters).
    """

    query: str = Field(..., min_length=1, max_length=_MAX_QUERY_CHARS)
    page: PageEnum
    force_mode: RouterModeEnum | None = None
    top_k: int | None = Field(default=None, ge=1, le=_MAX_TOP_K)
    filters: Filters | None = None
    comparison_spec: ComparisonSpec | None = None
    models: QueryModelOverrides | None = None
    use_premium: bool = False


# ---------- History row ---------------------------------------------------
class HistoryRowDTO(ApiModel):
    """Compact projection used by ``GET /history``.

    The full envelope is fetched on-demand via ``GET /query/runs/{runId}``.
    This row is what fits in the recent-queries table without overloading
    the wire — the ``query_preview`` is server-trimmed to 120 chars.
    """

    id: int
    timestamp: str
    page: str
    mode: RouterModeEnum | None = None
    query: str
    query_preview: str
    cost_usd: float
    duration_ms: int
    n_calls: int
    pinned: bool
    note: str | None = None
    username: str | None = None


class HistoryListResponse(ApiModel):
    """Paginated wrapper for ``GET /history``.

    ``cursor`` is the smallest ``id`` returned in this page — the FE
    passes it back as ``?cursor=`` to fetch older rows. ``null`` means
    "no more pages".
    """

    items: list[HistoryRowDTO]
    next_cursor: int | None = None