src/agents/planner/inputs.py

"""Planner input models — CatalogSummary and Constraints.

`CatalogSummary` is a condensed, PII-safe view of the user's `Catalog`, built
for the planner prompt. It carries every table + column id/type/PII flag + row
counts + low-cardinality top_values, with `sample_values` nulled on PII columns
(INV: no PII sample values into the prompt, see doc §13). It also lists the
available unstructured sources so the planner can plan `retrieve_knowledge`.

The planner *validator* still checks inline `retrieve_data` IRs against the
full `Catalog` via the existing IRValidator — the summary is a prompt input, not
the validation source of truth.

See AGENT_ARCHITECTURE_CONTEXT_new.md §7.3.
"""

from __future__ import annotations

from typing import Any

from pydantic import BaseModel, Field

from ...catalog.models import Catalog, DataType


class ColumnSummary(BaseModel):
    column_id: str
    name: str
    data_type: DataType
    pii_flag: bool = False
    sample_values: list[Any] | None = None  # nulled when pii_flag is True
    top_values: list[Any] | None = None


class ForeignKeySummary(BaseModel):
    """A declared FK edge — the only joins the IR validator accepts.

    Maps directly onto a `retrieve_data` IR join: `column_id` → `left_column_id`,
    `target_table_id` → `target_table_id`, `target_column_id` → `right_column_id`.
    """

    column_id: str
    target_table_id: str
    target_column_id: str


class TableSummary(BaseModel):
    table_id: str
    name: str
    row_count: int | None = None
    columns: list[ColumnSummary] = Field(default_factory=list)
    foreign_keys: list[ForeignKeySummary] = Field(default_factory=list)


class StructuredSourceSummary(BaseModel):
    source_id: str
    name: str
    source_type: str  # "schema" | "tabular"
    tables: list[TableSummary] = Field(default_factory=list)


class UnstructuredSourceSummary(BaseModel):
    source_id: str
    name: str


class CatalogSummary(BaseModel):
    structured_sources: list[StructuredSourceSummary] = Field(default_factory=list)
    unstructured_sources: list[UnstructuredSourceSummary] = Field(default_factory=list)

    @classmethod
    def from_catalog(cls, catalog: Catalog) -> CatalogSummary:
        structured: list[StructuredSourceSummary] = []
        unstructured: list[UnstructuredSourceSummary] = []

        for source in catalog.sources:
            if source.source_type == "unstructured":
                unstructured.append(
                    UnstructuredSourceSummary(source_id=source.source_id, name=source.name)
                )
                continue

            tables = [
                TableSummary(
                    table_id=table.table_id,
                    name=table.name,
                    row_count=table.row_count,
                    columns=[
                        ColumnSummary(
                            column_id=col.column_id,
                            name=col.name,
                            data_type=col.data_type,
                            pii_flag=col.pii_flag,
                            # PII columns leak nothing into the prompt: both
                            # sample_values and (low-cardinality) top_values are
                            # suppressed — top_values are the same class of data.
                            sample_values=None if col.pii_flag else col.sample_values,
                            top_values=(
                                None
                                if col.pii_flag or col.stats is None
                                else col.stats.top_values
                            ),
                        )
                        for col in table.columns
                    ],
                    # The declared FKs — the only joins the validator accepts. FKs
                    # carry no PII (ids only), so they're always surfaced.
                    foreign_keys=[
                        ForeignKeySummary(
                            column_id=fk.column_id,
                            target_table_id=fk.target_table_id,
                            target_column_id=fk.target_column_id,
                        )
                        for fk in table.foreign_keys
                    ],
                )
                for table in source.tables
            ]
            structured.append(
                StructuredSourceSummary(
                    source_id=source.source_id,
                    name=source.name,
                    source_type=source.source_type,
                    tables=tables,
                )
            )

        return cls(structured_sources=structured, unstructured_sources=unstructured)

    def render(self) -> str:
        """Render the summary as compact text for the planner prompt."""
        if not self.structured_sources and not self.unstructured_sources:
            return "(catalog is empty — the user has not registered any data yet)"

        lines: list[str] = []
        for source in self.structured_sources:
            lines.append(f"Source: {source.name} ({source.source_type}) — id={source.source_id}")
            # Name lookups (within a source) so FK edges render with readable
            # table/column names alongside the ids the IR join must copy verbatim.
            table_name_by_id = {t.table_id: t.name for t in source.tables}
            col_name_by_id = {
                c.column_id: c.name for t in source.tables for c in t.columns
            }
            for table in source.tables:
                rc = f" ({table.row_count:,} rows)" if table.row_count is not None else ""
                lines.append(f"  Table: {table.name}{rc} — id={table.table_id}")
                for col in table.columns:
                    samples = "PII (suppressed)" if col.pii_flag else (col.sample_values or [])
                    top = f", top={col.top_values}" if col.top_values else ""
                    lines.append(
                        f"    - {col.name} [{col.data_type}]: "
                        f"samples={samples}{top} — id={col.column_id}"
                    )
                for fk in table.foreign_keys:
                    tgt_table = table_name_by_id.get(fk.target_table_id, fk.target_table_id)
                    tgt_col = col_name_by_id.get(fk.target_column_id, fk.target_column_id)
                    src_col = col_name_by_id.get(fk.column_id, fk.column_id)
                    lines.append(
                        f"    FK: {src_col} → {tgt_table}.{tgt_col} "
                        f"(join: target_table_id={fk.target_table_id}, "
                        f"left_column_id={fk.column_id}, "
                        f"right_column_id={fk.target_column_id})"
                    )
            lines.append("")

        if self.unstructured_sources:
            lines.append("Unstructured sources (for retrieve_knowledge):")
            for src in self.unstructured_sources:
                lines.append(f"  - {src.name} — id={src.source_id}")

        return "\n".join(lines).rstrip()


class Constraints(BaseModel):
    max_tasks: int = 5
    modeling_allowed: bool = False  # no modeling tools in v1
    token_budget: int | None = None
    time_budget_seconds: int | None = None
    row_budget: int = 10_000