| """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 |
| 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 |
| 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, |
| |
| |
| |
| 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 |
| ], |
| |
| |
| 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}") |
| |
| |
| 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 |
| token_budget: int | None = None |
| time_budget_seconds: int | None = None |
| row_budget: int = 10_000 |
|
|
