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
from ...middlewares.logging import get_logger

logger = get_logger("planner_inputs")

# Ceilings on what `CatalogSummary.render` emits into the planner prompt (F-12,
# 2026-07-23). See that method's docstring for why these are safety nets set high
# rather than tight caps. Whichever binds first wins:
#   _MAX_TABLES        guards many-small-tables catalogs
#   _MAX_CATALOG_CHARS guards few-very-wide-tables catalogs (a 200-column fact table
#                      blows the budget long before the table count does)
# ~250k chars is roughly 62k tokens. Sizing rationale, measured this session:
#   100 tables x 30 cols  ~121k chars  -> renders IN FULL (a plausible real warehouse)
#   150 tables x 30 cols  ~182k chars  -> renders IN FULL
#   400 tables x 30 cols  ~482k chars  -> truncated to ~250k (was ~241k TOKENS/call)
# The point is to make the catastrophic case survivable without touching anyone real.
# These are provisional: nobody knows the largest actual customer catalog yet (review
# open question #4), so the truncation log is what tells us when to revisit them.
_MAX_TABLES = 150
_MAX_CATALOG_CHARS = 250_000


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.

        **Bounded since 2026-07-23 (F-12).** This used to emit one line per column
        across every table of every source with no truncation anywhere, and
        `PlannerService.plan` rebuilds the whole prompt on each of its 3 retries.
        Measured: 200 tables x 30 cols renders ~481k chars (~120k tokens) per call,
        ~361k tokens across the retries — past the context window, so the Azure call
        fails and the never-throw path degrades it to "Analysis failed". A 400-table
        warehouse is simply unusable, and the symptom looks like a data problem.

        `_MAX_TABLES` is a SAFETY NET, not a tight cap. A low cap would be actively
        harmful: there is no relevance ordering here, so dropping tables can drop the
        very table the question is about. It is set far above any catalog seen so far,
        every truncation is logged (that log is the signal a real customer is
        approaching it), and the planner is TOLD it saw a subset so it can ask the
        user to name a table rather than silently assume it saw everything.

        The durable fix is a relevance-ranked, question-keyed subset — which must not
        be attempted without an eval proving the planner still picks the right table.
        """
        if not self.structured_sources and not self.unstructured_sources:
            return "(catalog is empty — the user has not registered any data yet)"

        rendered_tables = 0
        omitted_tables = 0
        budget_used = 0
        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:
                # Two ceilings, whichever binds first. Table count alone is not
                # enough: 150 tables of 200 columns is just as unbounded as 400
                # tables of 30, and wide fact tables are common in warehouses.
                if rendered_tables >= _MAX_TABLES or budget_used >= _MAX_CATALOG_CHARS:
                    omitted_tables += 1
                    continue
                rendered_tables += 1
                table_start = len(lines)
                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})"
                    )
                budget_used += sum(len(ln) + 1 for ln in lines[table_start:])
            lines.append("")

        if omitted_tables:
            # The planner must know it saw a subset — otherwise "that table isn't in
            # the catalog" becomes a confident, wrong answer.
            lines.append(
                f"... and {omitted_tables} more table(s) not shown (catalog too large "
                f"to render in full). If the question names a table you cannot see "
                f"here, say so and ask the user to name it explicitly."
            )
            logger.warning(
                "catalog render truncated",
                rendered_tables=rendered_tables,
                omitted_tables=omitted_tables,
                chars_used=budget_used,
                max_tables=_MAX_TABLES,
                max_chars=_MAX_CATALOG_CHARS,
            )

        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