File size: 11,044 Bytes
81e5fe7 0721bb4 81e5fe7 0721bb4 81e5fe7 f282b15 81e5fe7 3743cfe 81e5fe7 3743cfe 81e5fe7 3743cfe 81e5fe7 f282b15 81e5fe7 f282b15 81e5fe7 3743cfe 81e5fe7 f282b15 81e5fe7 3743cfe f282b15 81e5fe7 f282b15 81e5fe7 0721bb4 81e5fe7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 | """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
|