190 lines
7.4 KiB
Python
190 lines
7.4 KiB
Python
"""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: `<col> = <value>`
|
|
- column + in_values: `<col> IN (<values>)`
|
|
- column + between: `<col> BETWEEN <a> AND <b>` (inclusive, both ends).
|
|
- column + date_range: `<col> BETWEEN '<yyyy>-01-01' AND '<yyyy>-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]
|