File size: 7,258 Bytes
81e5fe7 0721bb4 81e5fe7 0721bb4 81e5fe7 3743cfe 81e5fe7 3743cfe 81e5fe7 3743cfe 81e5fe7 3743cfe 81e5fe7 3743cfe 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 | """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
|