recon-driven sql composer; pick → compose → execute; llm out of structural sql
This commit is contained in:
189
api/composer/types.py
Normal file
189
api/composer/types.py
Normal file
@@ -0,0 +1,189 @@
|
||||
"""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]
|
||||
Reference in New Issue
Block a user