"""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