"""Pick / Filter / OrderBy — the typed shape the LLM emits and the composer
consumes.
Validation is strict: a Pick that references a metric or column not in recon
is rejected before composition. That's the whole point — the composer's
input must be expressible in terms of recon entities. There is no raw-SQL
escape hatch on a Filter; questions outside the supported shapes get a clean
"beyond what I'm built to answer" rather than a fabricated query.
"""
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Any, Literal
class PickValidationError(ValueError):
"""The Pick references a metric/column/filter shape that doesn't fit
recon. Surfaced cleanly to the caller (no retry — see feedback memory)."""
# ── Filter ──────────────────────────────────────────────────
@dataclass
class Filter:
"""One WHERE-clause fragment. Exactly one of the value fields must be set.
- column + equals: ` = `
- column + in_values: ` IN ()`
- column + between: ` BETWEEN AND ` (inclusive, both ends).
- column + date_range: ` BETWEEN '-01-01' AND '-12-31'`
for `"YYYY"`, or quarter bounds for `"YYYY-Q[1-4]"`.
"""
column: str
equals: str | int | float | None = None
in_values: list[Any] | None = None
between: tuple[Any, Any] | None = None
date_range: str | None = None
def kind(self) -> str:
"""Which filter shape is set; raises if none or more than one."""
set_fields = [
name for name in ("equals", "in_values", "between", "date_range")
if getattr(self, name) is not None
]
if len(set_fields) == 0:
raise PickValidationError(
f"filter on column {self.column!r} has no value field set"
)
if len(set_fields) > 1:
raise PickValidationError(
f"filter on column {self.column!r} sets multiple value fields: {set_fields}"
)
return set_fields[0]
def to_dict(self) -> dict[str, Any]:
out: dict[str, Any] = {"column": self.column}
for k in ("equals", "in_values", "between", "date_range"):
v = getattr(self, k)
if v is not None:
out[k] = list(v) if isinstance(v, tuple) else v
return out
@classmethod
def from_dict(cls, d: dict[str, Any]) -> "Filter":
if "column" not in d:
raise PickValidationError(f"filter missing 'column': {d!r}")
between = d.get("between")
if between is not None:
if not isinstance(between, (list, tuple)) or len(between) != 2:
raise PickValidationError(
f"filter 'between' must be a 2-element list, got {between!r}"
)
between = (between[0], between[1])
return cls(
column=d["column"],
equals=d.get("equals"),
in_values=d.get("in_values"),
between=between,
date_range=d.get("date_range"),
)
# ── OrderBy ─────────────────────────────────────────────────
@dataclass
class OrderBy:
"""ORDER BY clause. Either by metric or by a named dimension."""
by: Literal["metric", "dimension"]
direction: Literal["asc", "desc"] = "desc"
dimension: str | None = None
def to_dict(self) -> dict[str, Any]:
out: dict[str, Any] = {"by": self.by, "direction": self.direction}
if self.dimension is not None:
out["dimension"] = self.dimension
return out
@classmethod
def from_dict(cls, d: dict[str, Any]) -> "OrderBy":
by = d.get("by")
if by not in ("metric", "dimension"):
raise PickValidationError(f"order_by.by must be 'metric'|'dimension', got {by!r}")
direction = d.get("direction", "desc")
if direction not in ("asc", "desc"):
raise PickValidationError(f"order_by.direction must be 'asc'|'desc', got {direction!r}")
dim = d.get("dimension")
if by == "dimension" and not dim:
raise PickValidationError("order_by.by='dimension' requires order_by.dimension")
return cls(by=by, direction=direction, dimension=dim)
# ── Pick ────────────────────────────────────────────────────
@dataclass
class Pick:
"""What the LLM emits — fully resolvable by the composer against recon.
Fields:
kind: literal "aggregate" (only supported shape in v1).
metric: must be a name in recon.metrics.
group_by: column refs (`col` or `table.col`); each resolved by
recon.resolve_column.
where: typed filters; the composer renders each against the column's
sql_type / semantic_type.
order_by: optional sort.
limit: optional row limit.
"""
kind: Literal["aggregate"]
metric: str
group_by: list[str] = field(default_factory=list)
where: list[Filter] = field(default_factory=list)
order_by: OrderBy | None = None
limit: int | None = None
def to_dict(self) -> dict[str, Any]:
out: dict[str, Any] = {
"kind": self.kind,
"metric": self.metric,
"group_by": list(self.group_by),
"where": [f.to_dict() for f in self.where],
}
if self.order_by is not None:
out["order_by"] = self.order_by.to_dict()
if self.limit is not None:
out["limit"] = self.limit
return out
@classmethod
def from_dict(cls, d: dict[str, Any]) -> "Pick":
kind = d.get("kind", "aggregate")
if kind != "aggregate":
raise PickValidationError(
f"only kind='aggregate' is supported; got {kind!r}"
)
metric = d.get("metric")
if not isinstance(metric, str) or not metric:
raise PickValidationError(f"pick missing 'metric': {d!r}")
group_by = d.get("group_by", []) or []
if not isinstance(group_by, list) or not all(isinstance(c, str) for c in group_by):
raise PickValidationError(f"pick.group_by must be list[str], got {group_by!r}")
where_raw = d.get("where", []) or []
if not isinstance(where_raw, list):
raise PickValidationError(f"pick.where must be a list, got {where_raw!r}")
where = [Filter.from_dict(f) for f in where_raw]
for f in where:
f.kind() # raises if shape is missing/ambiguous
order_by = OrderBy.from_dict(d["order_by"]) if d.get("order_by") else None
limit = d.get("limit")
if limit is not None and not isinstance(limit, int):
raise PickValidationError(f"pick.limit must be int or None, got {limit!r}")
return cls(
kind=kind,
metric=metric,
group_by=group_by,
where=where,
order_by=order_by,
limit=limit,
)
# ── ComposeResult ───────────────────────────────────────────
@dataclass
class ComposeResult:
"""What `compose()` returns. Same shape as the legacy T2SResult so call
sites can swap with minimal churn."""
sql: str
used_tables: list[str]