src/traceability/schemas.py

"""Traceability payload schemas (KM-691).

User-facing provenance for one assistant turn: what the AI planned, which tools it
called (with real inputs + outputs), and which data sources it read (with the
executed query). One `TraceabilityPayload` is built per assistant `message_id`,
stored as a JSONB row, and served by `GET /api/v1/traceability`.

Distinct from Langfuse *observability* (`src/observability/langfuse/`): that is
engineering-only, PII-masked (arg keys + row counts). Traceability shows the user
their own data's provenance — real args, output previews, executed SQL — so it is
NOT masked. Truncation caps (in `scratchpad.py`) bound the payload size instead.

`thinking` is always `null` in v1 (our agents are plain chat completions with no
native reasoning output; synthesizing it post-hoc would be unfaithful). The field
stays in the payload so adding it later is contract-compatible.
"""

from __future__ import annotations

from datetime import datetime
from typing import Any, Literal

from pydantic import BaseModel, ConfigDict, Field


class PlanStep(BaseModel):
    """One CRISP-DM step the Planner ran (derived from an AnalysisRecord task)."""

    step: int
    stage: str
    objective: str
    status: str
    tools_used: list[str] = Field(default_factory=list)


class PlanningInfo(BaseModel):
    """Planner output for a slow-path turn; `null` on every other turn type."""

    goal_restated: str
    assumptions: list[str] = Field(default_factory=list)
    steps: list[PlanStep] = Field(default_factory=list)


class ToolCallInfo(BaseModel):
    """One tool invocation with its real input + output (both truncation-capped).

    The field is spelled `input` in the wire contract (the FE reads it), which
    shadows the `input` builtin — hence the alias + `populate_by_name` so callers
    can pass `input=` on construction and `model_dump(by_alias=True)` emits `input`.

    `summary` is a plain-English one-liner built from a fixed per-tool template
    (never an LLM call) — the FE headline for this step; `input`/`output` stay raw
    for the collapsible "technical details" layer.
    """

    model_config = ConfigDict(populate_by_name=True)

    order: int
    task_id: str | None = None
    name: str
    summary: str | None = None
    input_: dict[str, Any] = Field(default_factory=dict, alias="input")
    output: dict[str, Any] = Field(default_factory=dict)
    status: Literal["success", "error"] = "success"
    error: str | None = None


# ---------------------------------------------------------------------------
# `data_used` — the user-facing "what data did this analysis touch" layer.
# All names are resolved from the catalog at build time (deterministic lookup,
# no LLM). Every `id` field is MACHINE-ONLY — the FE must never render it; it
# exists for click-through/linking, reconciliation, and audit. One DataUsed
# entry per `retrieve_data` call.
# ---------------------------------------------------------------------------


class SourceRef(BaseModel):
    """The data source a pull read from. `id` is machine-only (FE must not render)."""

    id: str
    name: str
    type: str | None = None


class TableRef(BaseModel):
    """A table the query touched. `id` is machine-only. `role`: base | joined."""

    id: str
    name: str
    role: str


class JoinRef(BaseModel):
    """A join, rendered in real names, e.g. condition='order_items.order_id = orders.id'."""

    type: str
    condition: str


class ColumnRef(BaseModel):
    """A real catalog column the query read. `id` is machine-only (FE must not render).

    `roles` records why it was used: selected | aggregated | filtered | grouped |
    joined | ordered. `table` is the real table name for qualification.
    """

    id: str
    name: str
    table: str
    data_type: str | None = None
    pii: bool = False
    roles: list[str] = Field(default_factory=list)


class OutputColumn(BaseModel):
    """A column in the result set.

    `kind='column'` — read straight from the data (has a real `from`).
    `kind='computed'` — calculated; carries a `formula`, and has NO catalog id
    because it is not a stored column (e.g. total_revenue = SUM(line_total)).
    `from` is spelled with an alias (Python keyword) — populate_by_name lets
    callers pass `from_=`.
    """

    model_config = ConfigDict(populate_by_name=True)

    name: str
    kind: Literal["column", "computed"]
    from_: str | None = Field(default=None, alias="from")
    formula: str | None = None


class FilterRef(BaseModel):
    """A filter, resolved to a real column plus a plain-language `description`."""

    column: str
    op: str
    value: Any = None
    description: str


class OrderByRef(BaseModel):
    """A sort. `target` is a real column name (kind='column') OR a computed
    output alias (kind='computed' — the IR stores the alias in `column_id`)."""

    target: str
    kind: Literal["column", "computed"]
    dir: str = "asc"


class DataUsed(BaseModel):
    """One `retrieve_data` pull, fully resolved for the user. Names for display,
    ids for machine linkage only (FE must not render any `id`)."""

    source: SourceRef
    tables: list[TableRef] = Field(default_factory=list)
    joins: list[JoinRef] = Field(default_factory=list)
    columns_read: list[ColumnRef] = Field(default_factory=list)
    output_columns: list[OutputColumn] = Field(default_factory=list)
    filters: list[FilterRef] = Field(default_factory=list)
    group_by: list[str] = Field(default_factory=list)
    order_by: list[OrderByRef] = Field(default_factory=list)
    limit: int | None = None
    rows_returned: int | None = None
    query: str | None = None


class TraceabilityPayload(BaseModel):
    """The full provenance record for one assistant `message_id`."""

    model_config = ConfigDict(populate_by_name=True)

    analysis_id: str
    message_id: str
    user_id: str  # ownership column on the row; the FE may ignore it
    intent: str
    generated_at: datetime
    planning: PlanningInfo | None = None
    thinking: str | None = None
    tool_calls: list[ToolCallInfo] = Field(default_factory=list)
    data_used: list[DataUsed] = Field(default_factory=list)
    sources: list[dict[str, Any]] = Field(default_factory=list)