Agentic-Service-Data-Eyond-Catalog

Running

Agentic-Service-Data-Eyond-Catalog / src /agents /planner /inputs.py

Rifqi Hafizuddin

[NOTICKET] planner: surface numeric min/max, so thresholds aren't chosen blind

9155dcb about 2 months ago

14.6 kB

"""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
    #: The column holds ONE value in every row (`stats.distinct_count == 1`).
    #:
    #: The catalog has always known this; the summary used to drop it, so the planner
    #: could not see it. In the 2026-09-17 dry-run it targeted `Total_Hours`
    #: (`distinct_count=1, min=24, max=24` — a fixed 24-hour day) with
    #: `analyze_importance`, and the tool correctly refused: nothing can explain a
    #: number that never moves. (A23 Bug 3.)
    #:
    #: A bool, not the raw count: constancy is the only fact that changes a planning
    #: decision, and rendering a count for every column would cost tokens on every
    #: call to say something true but useless about the other 35 columns.
    #:
    #: Safe for PII columns — a cardinality is not a value, so unlike `top_values`
    #: this discloses nothing and is emitted regardless of `pii_flag`.
    is_constant: bool = False
    #: Observed range, straight from `stats.min`/`stats.max` (A23 Bug 14).
    #:
    #: The catalog has carried these all along — Go populates them on 98 of 98
    #: numeric columns — and the summary dropped them, so the planner has been
    #: choosing thresholds blind. Asked to band `pa` into low/medium/high it used
    #: thirds of 0-100; the column actually runs 68.8-95.9, so all 42 rows landed
    #: in "high" and the segmentation said nothing. It did not guess badly, it
    #: could not see the range.
    #:
    #: Unlike `is_constant`, these ARE values out of the customer's data, so they
    #: follow the same rule as `sample_values` and are suppressed for PII columns.
    #:
    #: Date columns currently arrive as None (Go's introspection does not populate
    #: them — handed to Harry 2026-09-17). This renders them the moment they do,
    #: which is the other half of that fix.
    value_min: Any | None = None
    value_max: 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
                            ),
                            is_constant=(
                                col.stats is not None and col.stats.distinct_count == 1
                            ),
                            value_min=(
                                None
                                if col.pii_flag or col.stats is None
                                else col.stats.min
                            ),
                            value_max=(
                                None
                                if col.pii_flag or col.stats is None
                                else col.stats.max
                            ),
                        )
                        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 ""
                    # Emitted ONLY when true, so the normal column costs nothing and
                    # the marker is conspicuous where it matters. (A23 Bug 3.)
                    const = (
                        " [CONSTANT: one value in every row — cannot be a target or a "
                        "grouping key]"
                        if col.is_constant
                        else ""
                    )
                    # Skipped for a constant column: its marker already says the range
                    # is a single point, and "range=24..24" would only repeat it.
                    rng = (
                        f", range={col.value_min}..{col.value_max}"
                        if not col.is_constant
                        and col.value_min is not None
                        and col.value_max is not None
                        else ""
                    )
                    lines.append(
                        f"    - {col.name} [{col.data_type}]: "
                        f"samples={samples}{top}{rng}{const} — 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
    # True since A15 (2026-09-17) activated forecast/anomaly/importance/cluster.
    # Kept as a knob rather than deleted: a caller that wants a strictly
    # descriptive run can still switch modeling off for that turn.
    modeling_allowed: bool = True
    token_budget: int | None = None
    time_budget_seconds: int | None = None
    row_budget: int = 10_000

	"""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
	#: The column holds ONE value in every row (`stats.distinct_count == 1`).
	#:
	#: The catalog has always known this; the summary used to drop it, so the planner
	#: could not see it. In the 2026-09-17 dry-run it targeted `Total_Hours`
	#: (`distinct_count=1, min=24, max=24` — a fixed 24-hour day) with
	#: `analyze_importance`, and the tool correctly refused: nothing can explain a
	#: number that never moves. (A23 Bug 3.)
	#:
	#: A bool, not the raw count: constancy is the only fact that changes a planning
	#: decision, and rendering a count for every column would cost tokens on every
	#: call to say something true but useless about the other 35 columns.
	#:
	#: Safe for PII columns — a cardinality is not a value, so unlike `top_values`
	#: this discloses nothing and is emitted regardless of `pii_flag`.
	is_constant: bool = False
	#: Observed range, straight from `stats.min`/`stats.max` (A23 Bug 14).
	#:
	#: The catalog has carried these all along — Go populates them on 98 of 98
	#: numeric columns — and the summary dropped them, so the planner has been
	#: choosing thresholds blind. Asked to band `pa` into low/medium/high it used
	#: thirds of 0-100; the column actually runs 68.8-95.9, so all 42 rows landed
	#: in "high" and the segmentation said nothing. It did not guess badly, it
	#: could not see the range.
	#:
	#: Unlike `is_constant`, these ARE values out of the customer's data, so they
	#: follow the same rule as `sample_values` and are suppressed for PII columns.
	#:
	#: Date columns currently arrive as None (Go's introspection does not populate
	#: them — handed to Harry 2026-09-17). This renders them the moment they do,
	#: which is the other half of that fix.
	value_min: Any \| None = None
	value_max: 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
	),
	is_constant=(
	col.stats is not None and col.stats.distinct_count == 1
	),
	value_min=(
	None
	if col.pii_flag or col.stats is None
	else col.stats.min
	),
	value_max=(
	None
	if col.pii_flag or col.stats is None
	else col.stats.max
	),
	)
	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 ""
	# Emitted ONLY when true, so the normal column costs nothing and
	# the marker is conspicuous where it matters. (A23 Bug 3.)
	const = (
	" [CONSTANT: one value in every row — cannot be a target or a "
	"grouping key]"
	if col.is_constant
	else ""
	)
	# Skipped for a constant column: its marker already says the range
	# is a single point, and "range=24..24" would only repeat it.
	rng = (
	f", range={col.value_min}..{col.value_max}"
	if not col.is_constant
	and col.value_min is not None
	and col.value_max is not None
	else ""
	)
	lines.append(
	f" - {col.name} [{col.data_type}]: "
	f"samples={samples}{top}{rng}{const} — 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
	# True since A15 (2026-09-17) activated forecast/anomaly/importance/cluster.
	# Kept as a knob rather than deleted: a caller that wants a strictly
	# descriptive run can still switch modeling off for that turn.
	modeling_allowed: bool = True
	token_budget: int \| None = None
	time_budget_seconds: int \| None = None
	row_budget: int = 10_000