diff --git a/api/analyses/_narrow.py b/api/analyses/_narrow.py new file mode 100644 index 0000000..2be4801 --- /dev/null +++ b/api/analyses/_narrow.py @@ -0,0 +1,95 @@ +"""Candidate narrowing for the Pick LLM call. + +The composer covers L1's schema-pollution problem — the LLM never sees raw +SQL. But the LLM still needs to choose from *some* candidate set when +picking a metric / dimension / filter column. Sending the full schema for +every question is wasteful (cost, latency) and noisier (more candidates +the model can pick wrong from). + +v1: **structural neighborhood**. Given the candidate metrics, compute the +union of: + - columns of each candidate metric's `from_table` + - columns of every table reachable within `max_hops` via declared + relationships. + +This covers every column that could legitimately group/filter that metric +without hand-curating a list. Cheap, deterministic, easy to reason about. + +v2 (future): embedding-based retrieval against question + metric/column +descriptions, with usage-derived signal (pgvector). Out of scope here. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +from api.recon.types import Metric, Recon + + +@dataclass +class Candidates: + metrics: list[Metric] + # column_refs: each entry is either a bare name (when unambiguous) or + # `table.column` (when the name lives on more than one reachable table). + column_refs: list[str] + + +def neighborhood_tables(recon: Recon, seeds: list[str], max_hops: int = 2) -> set[str]: + """BFS from `seeds` over recon.relationships, returning every table + reachable within `max_hops`. Edges are undirected (FKs go both ways).""" + adj: dict[str, set[str]] = {t: set() for t in recon.tables} + for r in recon.relationships: + adj.setdefault(r.from_table, set()).add(r.to_table) + adj.setdefault(r.to_table, set()).add(r.from_table) + visited: set[str] = set() + frontier: set[str] = {s for s in seeds if s in recon.tables} + for _ in range(max_hops + 1): # +1 to include the seed itself at hop 0 + visited |= frontier + next_frontier: set[str] = set() + for t in frontier: + next_frontier |= adj.get(t, set()) + frontier = next_frontier - visited + if not frontier: + break + return visited + + +def narrow_candidates( + recon: Recon, + *, + allowed_metrics: list[str] | None = None, + max_hops: int = 2, +) -> Candidates: + """Return the Pick-candidate set for the given metrics. + + `allowed_metrics`: if provided, restricts the candidate metrics to this + list (intersected with recon.metrics). If None, every metric is in scope. + `max_hops`: how far to walk from each candidate metric's from_table. + """ + metric_names = ( + [m for m in allowed_metrics if m in recon.metrics] + if allowed_metrics is not None + else list(recon.metrics) + ) + metrics = [recon.metrics[m] for m in metric_names] + seeds = [m.from_table for m in metrics] + tables = neighborhood_tables(recon, seeds, max_hops=max_hops) + + # Build the column ref list. A column name appearing in more than one + # reachable table is exposed as `table.column` for each owner to keep + # the LLM's pick unambiguous. + appearances: dict[str, list[str]] = {} + for tname in tables: + for c in recon.tables[tname].columns: + appearances.setdefault(c.name, []).append(tname) + + refs: list[str] = [] + for col, owners in appearances.items(): + if len(owners) == 1: + refs.append(col) + else: + for t in owners: + refs.append(f"{t}.{col}") + refs.sort() + + return Candidates(metrics=metrics, column_refs=refs) diff --git a/api/analyses/_pick.py b/api/analyses/_pick.py new file mode 100644 index 0000000..b7abfee --- /dev/null +++ b/api/analyses/_pick.py @@ -0,0 +1,144 @@ +"""pick_for_question — one constrained LLM call that returns a typed Pick. + +Used by L2 Analyses that need to translate a free-form question into the +composer's input. The model is shown: + - the narrowed metric catalog (name + description + unit) + - the narrowed candidate column refs (bare or `table.column`) + - the filter grammar (via the system prompt) + +It returns JSON. We validate the JSON into a `Pick` and fail fast on any +shape error — no local retry. Per the project's retry feedback, recovery +belongs at a future global layer that emits a visible event when it fires. + +If the model returns `{"error": "..."}` (the prompt's escape hatch for +"this question can't be expressed"), we raise PickValidationError with +the model's reason — the caller surfaces it cleanly. +""" + +from __future__ import annotations + +import json +import re +from dataclasses import dataclass + +from api import langfuse_client as lf +from api.analyses._narrow import Candidates, narrow_candidates +from api.composer.types import Pick, PickValidationError +from api.llm import chat +from api.prompts import load, render +from api.recon import load_recon +from api.recon.types import Recon + + +@dataclass +class PickOutcome: + pick: Pick + candidates: Candidates # the narrowed set the model saw (kept for tracing) + + +def pick_for_question( + question: str, + *, + recon: Recon | None = None, + allowed_metrics: list[str] | None = None, + max_tokens: int = 512, + span_name: str = "pick.gen", +) -> PickOutcome: + """Run one LLM call to translate `question` into a Pick. + + `allowed_metrics`: optional list to restrict the candidate metrics + (used by Analyses that already know which metric the planner chose, + so we don't waste tokens listing every metric in the catalog). + """ + recon = recon or load_recon() + candidates = narrow_candidates(recon, allowed_metrics=allowed_metrics) + + metrics_block = _render_metrics_block(candidates) + columns_block = _render_columns_block(candidates) + + system = load("pick.system") + user = render( + "pick.user", + question=question, + metrics_block=metrics_block, + columns_block=columns_block, + ) + + with lf.span( + "pick_for_question", + input={ + "question": question, + "metric_candidates": [m.name for m in candidates.metrics], + "column_candidate_count": len(candidates.column_refs), + }, + ) as span: + raw = chat(system=system, user=user, max_tokens=max_tokens, span_name=span_name) + payload = _parse_json(raw) + if "error" in payload: + raise PickValidationError( + f"LLM declined to pick: {payload['error']!s}" + ) + pick = Pick.from_dict(payload) + # Bind-check the pick against recon up-front — the composer would + # raise the same errors later, but raising here makes the failure + # event happen at the pick step instead of compose. + _validate_against_recon(pick, recon) + span.update(output=pick.to_dict()) + return PickOutcome(pick=pick, candidates=candidates) + + +def _render_metrics_block(candidates: Candidates) -> str: + if not candidates.metrics: + return "(no candidate metrics)" + lines: list[str] = [] + for m in candidates.metrics: + unit = f" [{m.unit}]" if m.unit else "" + lines.append(f"- {m.name}{unit}: {m.description}") + return "\n".join(lines) + + +def _render_columns_block(candidates: Candidates) -> str: + if not candidates.column_refs: + return "(no candidate columns)" + return "\n".join(f"- {c}" for c in candidates.column_refs) + + +def _parse_json(text: str) -> dict: + """Best-effort: prefer a ```json``` fence, otherwise extract the first + {...} block. Matches the existing pattern in drill_down._parse_json.""" + m = re.search(r"```(?:json)?\s*(\{.*?\})\s*```", text, re.DOTALL) + raw = m.group(1) if m else text + start, end = raw.find("{"), raw.rfind("}") + if start < 0 or end <= start: + raise PickValidationError( + f"pick_for_question returned no JSON object: {text[:200]!r}" + ) + try: + return json.loads(raw[start:end + 1]) + except json.JSONDecodeError as e: + raise PickValidationError( + f"pick_for_question returned invalid JSON: {e!s}; raw={raw[:200]!r}" + ) from e + + +def _validate_against_recon(pick: Pick, recon: Recon) -> None: + """Surface metric/column resolution errors as PickValidationError up + front so the trace marks the pick step (not compose) as the failure.""" + if pick.metric not in recon.metrics: + raise PickValidationError( + f"picked metric {pick.metric!r} is not in recon" + ) + for ref in pick.group_by: + try: + recon.resolve_column(ref) + except ValueError as e: + raise PickValidationError( + f"group_by column {ref!r}: {e}" + ) from e + for f in pick.where: + try: + recon.resolve_column(f.column) + except ValueError as e: + raise PickValidationError( + f"where filter on {f.column!r}: {e}" + ) from e diff --git a/api/analyses/compare_periods.py b/api/analyses/compare_periods.py index fbe6b04..7a7d29e 100644 --- a/api/analyses/compare_periods.py +++ b/api/analyses/compare_periods.py @@ -1,24 +1,27 @@ """compare_periods Analysis — CoT with two queries. -Both SQL queries are generated in one LLM call (paired prompt), then both -executed, then a single interpretation pass diffs them. Single round-trip -per LLM step → cheap and bounded latency. +One LLM call picks the shape (metric + group_by + non-time filters); the +Analysis injects each period as a typed date_range filter and the composer +emits two SQLs deterministically. No LLM authors SQL. """ from __future__ import annotations -import json import logging -import re +from dataclasses import replace from typing import Any from api import langfuse_client as lf +from api.analyses._pick import pick_for_question from api.analyses.base import Analysis from api.analyses.types import Finding +from api.composer import compose +from api.composer.types import Filter, Pick from api.llm import chat from api.prompts import load, render -from api.runtime import events from api.recon import load_recon +from api.recon.types import Recon +from api.runtime import events from api.tools.execute_sql import execute_sql logger = logging.getLogger("nvi.analyses.compare_periods") @@ -28,44 +31,74 @@ class ComparePeriods(Analysis): name = "compare_periods" description = ( "Compare a metric across two time windows (e.g. Q2 vs Q3, 1995 vs 1996). " - "Generates two parallel SQL queries and diffs the results." + "Picks one shape, composes two SQLs with different date filters, diffs the results." ) args_schema = { "question": {"type": "string", "description": "The metric / scope to compare."}, - "period_a": {"type": "string", "description": "First period label, e.g. '1995' or 'Q2 1996'."}, - "period_b": {"type": "string", "description": "Second period label, e.g. '1996' or 'Q3 1996'."}, + "period_a": {"type": "string", "description": "First period label, e.g. '1995' or '1996-Q2'."}, + "period_b": {"type": "string", "description": "Second period label, e.g. '1996' or '1996-Q3'."}, + "allowed_metrics": {"type": "array", "items": "string", "description": "Metrics the planner judged relevant (optional)."}, } async def run(self, args: dict[str, Any], question: str) -> Finding: sub_q = args.get("question") or question period_a = args.get("period_a", "") period_b = args.get("period_b", "") + allowed_metrics = args.get("allowed_metrics") with lf.span("analysis.compare_periods", input=args) as span: try: + recon = load_recon() + await events.publish_current(events.tool_call_start( - "generate_pair", input={"question": sub_q, "period_a": period_a, "period_b": period_b}, + "pick_for_question", input={"question": sub_q}, )) - pair = _generate_pair(sub_q, period_a, period_b) + outcome = pick_for_question( + sub_q, recon=recon, allowed_metrics=allowed_metrics, + span_name="compare_periods.pick", + ) + base_pick = outcome.pick await events.publish_current(events.tool_call_end( - "generate_pair", - output={"a": {"label": pair["a"]["label"], "sql": pair["a"]["sql"]}, - "b": {"label": pair["b"]["label"], "sql": pair["b"]["sql"]}}, + "pick_for_question", output={"pick": base_pick.to_dict()}, )) - await events.publish_current(events.tool_call_start("execute_sql", input={"label": pair["a"]["label"], "sql": pair["a"]["sql"]})) - res_a = execute_sql(pair["a"]["sql"]) + date_col_ref = _date_column_for(base_pick.metric, recon) + pick_a = _with_period_filter(base_pick, date_col_ref, period_a) + pick_b = _with_period_filter(base_pick, date_col_ref, period_b) + + await events.publish_current(events.tool_call_start( + "compose_sql", input={"label": period_a, "pick": pick_a.to_dict()}, + )) + composed_a = compose(pick_a, recon) + await events.publish_current(events.tool_call_end( + "compose_sql", output={"label": period_a, "sql": composed_a.sql}, + )) + + await events.publish_current(events.tool_call_start( + "compose_sql", input={"label": period_b, "pick": pick_b.to_dict()}, + )) + composed_b = compose(pick_b, recon) + await events.publish_current(events.tool_call_end( + "compose_sql", output={"label": period_b, "sql": composed_b.sql}, + )) + + await events.publish_current(events.tool_call_start( + "execute_sql", input={"label": period_a, "sql": composed_a.sql}, + )) + res_a = execute_sql(composed_a.sql) await events.publish_current(events.tool_call_end( "execute_sql", - output={"label": pair["a"]["label"], "row_count": res_a.row_count, + output={"label": period_a, "row_count": res_a.row_count, "preview": res_a.as_dicts()[:5]}, )) - await events.publish_current(events.tool_call_start("execute_sql", input={"label": pair["b"]["label"], "sql": pair["b"]["sql"]})) - res_b = execute_sql(pair["b"]["sql"]) + await events.publish_current(events.tool_call_start( + "execute_sql", input={"label": period_b, "sql": composed_b.sql}, + )) + res_b = execute_sql(composed_b.sql) await events.publish_current(events.tool_call_end( "execute_sql", - output={"label": pair["b"]["label"], "row_count": res_b.row_count, + output={"label": period_b, "row_count": res_b.row_count, "preview": res_b.as_dicts()[:5]}, )) @@ -73,8 +106,8 @@ class ComparePeriods(Analysis): interpret_user = render( "compare_periods.interpret.user", question=sub_q, - label_a=pair["a"]["label"], - label_b=pair["b"]["label"], + label_a=period_a, + label_b=period_b, rows_a=repr(res_a.as_dicts()[:20]), rows_b=repr(res_b.as_dicts()[:20]), ) @@ -98,48 +131,46 @@ class ComparePeriods(Analysis): analysis=self.name, summary=summary, rows=[ - {"period": pair["a"]["label"], **r} for r in res_a.as_dicts()[:20] + {"period": period_a, **r} for r in res_a.as_dicts()[:20] ] + [ - {"period": pair["b"]["label"], **r} for r in res_b.as_dicts()[:20] + {"period": period_b, **r} for r in res_b.as_dicts()[:20] ], - sql=[pair["a"]["sql"], pair["b"]["sql"]], + sql=[composed_a.sql, composed_b.sql], metadata={ - "label_a": pair["a"]["label"], - "label_b": pair["b"]["label"], + "metric": base_pick.metric, + "label_a": period_a, + "label_b": period_b, "row_count_a": res_a.row_count, "row_count_b": res_b.row_count, }, ) -def _generate_pair(question: str, period_a: str, period_b: str) -> dict[str, dict[str, str]]: - schema = load_recon() - text = chat( - system=load("compare_periods.pair"), - user=render( - "compare_periods.pair.user", - question=question, - period_a=period_a or "(infer)", - period_b=period_b or "(infer)", - schema_block=schema.render_tables(), - metrics_block=schema.render_metrics(), - ), - max_tokens=1024, - span_name="compare_periods.pair", +def _date_column_for(metric_name: str, recon: Recon) -> str: + """Return the date column ref (qualified) for the metric's from_table. + + The composer needs to know which column carries time for the period + filter. We look at the metric's source table and pick the first + column whose sql_type is DATE / TIMESTAMP. Returned qualified + (`table.column`) so resolve_column never trips on ambiguity.""" + metric = recon.metrics[metric_name] + table = recon.tables[metric.from_table] + for c in table.columns: + sql_type = c.sql_type.upper() + if sql_type.startswith("DATE") or sql_type.startswith("TIMESTAMP"): + return f"{table.name}.{c.name}" + raise ValueError( + f"no date/timestamp column on table {table.name!r} for metric {metric_name!r}" ) - obj = json.loads(_extract_json(text)) - for k in ("a", "b"): - if k not in obj or "sql" not in obj[k] or "label" not in obj[k]: - raise ValueError(f"pair generator returned malformed payload: missing {k}") - obj[k]["sql"] = obj[k]["sql"].strip().rstrip(";").strip() - return obj -def _extract_json(text: str) -> str: - m = re.search(r"```(?:json)?\s*(\{.*\})\s*```", text, re.DOTALL) - if m: - return m.group(1) - start, end = text.find("{"), text.rfind("}") - if start >= 0 and end > start: - return text[start:end + 1] - return text +def _with_period_filter(pick: Pick, col_ref: str, period: str) -> Pick: + """Return a new Pick with a date_range filter on `col_ref` set to + `period`. Strips any existing date_range filter on the same column + so the Analysis-injected one takes precedence.""" + where = [ + f for f in pick.where + if not (f.column == col_ref and f.date_range is not None) + ] + where.append(Filter(column=col_ref, date_range=period)) + return replace(pick, where=where) diff --git a/api/analyses/direct_answer.py b/api/analyses/direct_answer.py index ab7d9d0..fe3c95d 100644 --- a/api/analyses/direct_answer.py +++ b/api/analyses/direct_answer.py @@ -10,13 +10,15 @@ import logging from typing import Any from api import langfuse_client as lf +from api.analyses._pick import pick_for_question from api.analyses.base import Analysis from api.analyses.types import Finding +from api.composer import compose from api.llm import chat from api.prompts import load, render +from api.recon import load_recon from api.runtime import events from api.tools.execute_sql import execute_sql -from api.tools.text_to_sql import text_to_sql logger = logging.getLogger("nvi.analyses.direct_answer") @@ -29,23 +31,39 @@ class DirectAnswer(Analysis): ) args_schema = { "question": {"type": "string", "description": "Refined question this Analysis should answer."}, - "hint_tables": {"type": "array", "items": "string", "description": "Tables the planner thinks are relevant (optional)."}, + "allowed_metrics": {"type": "array", "items": "string", "description": "Metrics the planner judged relevant (optional; defaults to all)."}, } async def run(self, args: dict[str, Any], question: str) -> Finding: sub_q = args.get("question") or question - hint_tables = args.get("hint_tables") + allowed_metrics = args.get("allowed_metrics") with lf.span("analysis.direct_answer", input={"question": sub_q}) as span: try: - await events.publish_current(events.tool_call_start("text_to_sql", input={"question": sub_q})) - t2s = text_to_sql(sub_q, hint_tables=hint_tables) + recon = load_recon() + + await events.publish_current(events.tool_call_start( + "pick_for_question", input={"question": sub_q}, + )) + outcome = pick_for_question( + sub_q, recon=recon, allowed_metrics=allowed_metrics, + span_name="direct_answer.pick", + ) await events.publish_current(events.tool_call_end( - "text_to_sql", output={"sql": t2s.sql, "tables": t2s.used_tables}, + "pick_for_question", output={"pick": outcome.pick.to_dict()}, )) - await events.publish_current(events.tool_call_start("execute_sql", input={"sql": t2s.sql})) - result = execute_sql(t2s.sql) + await events.publish_current(events.tool_call_start( + "compose_sql", input={"pick": outcome.pick.to_dict()}, + )) + composed = compose(outcome.pick, recon) + await events.publish_current(events.tool_call_end( + "compose_sql", + output={"sql": composed.sql, "tables": composed.used_tables}, + )) + + await events.publish_current(events.tool_call_start("execute_sql", input={"sql": composed.sql})) + result = execute_sql(composed.sql) await events.publish_current(events.tool_call_end( "execute_sql", output={"row_count": result.row_count, "truncated": result.truncated, @@ -56,7 +74,7 @@ class DirectAnswer(Analysis): interpret_user = render( "direct_answer.interpret.user", question=sub_q, - sql=t2s.sql, + sql=composed.sql, rows=repr(result.as_dicts()[:20]), ) await events.publish_current(events.llm_call( @@ -64,7 +82,10 @@ class DirectAnswer(Analysis): system_len=len(interpret_system), user_len=len(interpret_user), )) - summary = chat(system=interpret_system, user=interpret_user, max_tokens=512) + summary = chat( + system=interpret_system, user=interpret_user, max_tokens=512, + span_name="direct_answer.interpret", + ) except Exception as e: logger.exception("direct_answer failed") await events.publish_current(events.tool_call_end("direct_answer", error=str(e))) @@ -76,6 +97,10 @@ class DirectAnswer(Analysis): analysis=self.name, summary=summary, rows=result.as_dicts()[:20], - sql=[t2s.sql], - metadata={"row_count": result.row_count, "truncated": result.truncated}, + sql=[composed.sql], + metadata={ + "row_count": result.row_count, + "truncated": result.truncated, + "metric": outcome.pick.metric, + }, ) diff --git a/api/analyses/drill_down/helpers.py b/api/analyses/drill_down/helpers.py index 47e3e8f..0584805 100644 --- a/api/analyses/drill_down/helpers.py +++ b/api/analyses/drill_down/helpers.py @@ -13,12 +13,12 @@ from typing import Any from api import langfuse_client as lf from api.analyses.drill_down.types import Slice +from api.composer import Pick, compose from api.llm import chat from api.prompts import load, render from api.recon import load_recon from api.runtime import events from api.tools.execute_sql import execute_sql -from api.tools.text_to_sql import text_to_sql logger = logging.getLogger("nvi.analyses.drill_down.helpers") @@ -69,54 +69,25 @@ async def decide_next(question: str, metric: str, dimensions: list[str], # ── Slice execution ── -def build_slice_question(question: str, metric: str, dim: str) -> str: - """Construct a slice question with explicit table/join hints from the - recon graph. Stops the LLM from inventing FROM clauses that omit the - table that owns the dimension column. - """ - recon = load_recon() - dim_owners = recon.owning_tables(dim) - metric_def = recon.metrics.get(metric) - metric_table = metric_def.from_table if metric_def else None - - hints: list[str] = [] - if dim_owners: - hints.append(f"- The dimension column `{dim}` lives in table: {', '.join(dim_owners)}.") - if metric_table: - hints.append(f"- The metric `{metric}` is defined over table `{metric_table}`.") - if metric_def and metric_def.filter: - hints.append(f" Apply this filter for the metric: {metric_def.filter}.") - if metric_def: - hints.append(f" Compute the metric as: {metric_def.sql} (use this expression verbatim).") - if dim_owners and metric_table and dim_owners[0] != metric_table: - path = recon.join_path(metric_table, dim_owners[0]) - if path: - hints.append(f"- Required JOIN path: {' → '.join(path)}.") - - hint_block = ("\n".join(hints) + "\n\n") if hints else "" - return ( - f"{question}\n\n" - f"Slice the metric `{metric}` by `{dim}` and return the top rows by metric value.\n\n" - f"{hint_block}" - f"GROUP BY the dimension. ORDER BY the metric DESC. Limit to top 10." - ) - - async def execute_slice(question: str, metric: str, dim: str, reason: str) -> Slice: - """Generate SQL for one slice, execute it, emit tool-call events - around both, and return the Slice.""" - slice_q = build_slice_question(question, metric, dim) + """Compose SQL for one slice deterministically (metric + dim are already + chosen by `decide_next`), execute it, emit tool-call events around both, + and return the Slice. Zero LLM calls — recon authors the SQL.""" + recon = load_recon() + pick = Pick(kind="aggregate", metric=metric, group_by=[dim], limit=10) - await events.publish_current(events.tool_call_start("text_to_sql", input={"question": slice_q})) - t2s = text_to_sql(slice_q) + await events.publish_current(events.tool_call_start( + "compose_sql", input={"pick": pick.to_dict()}, + )) + composed = compose(pick, recon) await events.publish_current(events.tool_call_end( - "text_to_sql", output={"sql": t2s.sql, "tables": t2s.used_tables}, + "compose_sql", output={"sql": composed.sql, "tables": composed.used_tables}, )) await events.publish_current(events.tool_call_start( - "execute_sql", input={"sql": t2s.sql, "dimension": dim}, + "execute_sql", input={"sql": composed.sql, "dimension": dim}, )) - result = execute_sql(t2s.sql) + result = execute_sql(composed.sql) await events.publish_current(events.tool_call_end( "execute_sql", output={"dimension": dim, "row_count": result.row_count, @@ -125,7 +96,7 @@ async def execute_slice(question: str, metric: str, dim: str, reason: str) -> Sl return Slice( dimension=dim, - sql=t2s.sql, + sql=composed.sql, rows=result.as_dicts()[:20], reason=reason, ) diff --git a/api/composer/__init__.py b/api/composer/__init__.py new file mode 100644 index 0000000..054067d --- /dev/null +++ b/api/composer/__init__.py @@ -0,0 +1,26 @@ +"""Recon-driven SQL composer. + +The LLM picks a *shape* (which metric, which dimensions, which filters); +the composer walks recon and emits the SQL deterministically. Column–table +binding, joins, and quoting are graph traversals, not LLM guesses. + +See `def/schema-as-constraint.md` for the architectural argument. +""" + +from api.composer.types import ( + ComposeResult, + Filter, + OrderBy, + Pick, + PickValidationError, +) +from api.composer.compose import compose + +__all__ = [ + "compose", + "ComposeResult", + "Filter", + "OrderBy", + "Pick", + "PickValidationError", +] diff --git a/api/composer/compose.py b/api/composer/compose.py new file mode 100644 index 0000000..d21b7b5 --- /dev/null +++ b/api/composer/compose.py @@ -0,0 +1,288 @@ +"""compose(pick, recon) → ComposeResult. + +Deterministic SQL emission from a typed Pick. The composer: + 1. Resolves the metric → from_table + sql expression + default filter. + 2. Resolves each group_by column → owning table (via recon.resolve_column). + 3. Computes the join graph: join_path(metric.from_table, owner) per dim, + deduped, walked in declared order to pick the right relationship edge. + 4. Renders SELECT / FROM + JOINs / WHERE / GROUP BY / ORDER BY / LIMIT + with every identifier quoted via sqlglot's `identify=True` round-trip + as a paranoid post-check. + +No LLM call anywhere in this module. If the Pick can't be expressed against +recon, the composer raises `PickValidationError` — the caller surfaces it. +""" + +from __future__ import annotations + +import sqlglot + +from api.composer.dates import resolve_date_range +from api.composer.types import ( + ComposeResult, + Filter, + OrderBy, + Pick, + PickValidationError, +) +from api.recon.types import Recon, Relationship + + +# ── Alias allocation ──────────────────────────────────────── + +def _alias_for(table: str, used: dict[str, str]) -> str: + """Pick a short alias for `table` that's unique within the query. + + First letter, then first-two, then numbered. `used` maps alias→table so + we can grow on collision. + """ + for base in (table[:1], table[:2], table[:3], table): + if base and base not in used: + return base + i = 1 + while True: + cand = f"{table[:1]}{i}" + if cand not in used: + return cand + i += 1 + + +# ── Join graph ────────────────────────────────────────────── + +def _find_edge(recon: Recon, a: str, b: str) -> Relationship: + """Return the relationship that joins tables a and b (either direction). + Raises if no declared FK connects them.""" + for r in recon.relationships: + if (r.from_table == a and r.to_table == b) or (r.from_table == b and r.to_table == a): + return r + raise PickValidationError( + f"no declared relationship between {a!r} and {b!r}" + ) + + +def _build_joins(recon: Recon, base: str, targets: list[str], + aliases: dict[str, str]) -> list[str]: + """Compute the union of join paths from `base` to each target table, + emit JOIN clauses in walk order. Allocates aliases for any intermediate + table that wasn't already in `aliases`.""" + visited: set[str] = {base} + clauses: list[str] = [] + for target in targets: + path = recon.join_path(base, target) + if path is None: + raise PickValidationError( + f"no join path from {base!r} to {target!r} in recon" + ) + for i in range(1, len(path)): + prev, cur = path[i - 1], path[i] + if cur in visited: + continue + edge = _find_edge(recon, prev, cur) + # Edge direction tells us which side has which column. + if edge.from_table == prev: + left_t, left_c, right_t, right_c = edge.from_table, edge.from_column, edge.to_table, edge.to_column + else: + left_t, left_c, right_t, right_c = edge.to_table, edge.to_column, edge.from_table, edge.from_column + # Ensure both sides have aliases. + for t in (left_t, right_t): + if t not in aliases: + aliases[t] = _alias_for(t, {a: t for t, a in aliases.items()}) + la, ra = aliases[left_t], aliases[right_t] + clauses.append( + f'JOIN "{cur}" AS "{aliases[cur]}" ' + f'ON "{la}"."{left_c}" = "{ra}"."{right_c}"' + ) + visited.add(cur) + return clauses + + +# ── Filter rendering ──────────────────────────────────────── + +def _render_literal(v) -> str: + """Render a Python value as a SQL literal. Strings are single-quoted with + embedded quotes escaped. Numbers pass through.""" + if isinstance(v, str): + return "'" + v.replace("'", "''") + "'" + if isinstance(v, bool): + return "TRUE" if v else "FALSE" + if v is None: + return "NULL" + return str(v) + + +def _render_filter(f: Filter, recon: Recon, aliases: dict[str, str]) -> str: + """Render one Filter to a SQL predicate. Resolves the column's owning + table and aliases it (allocating an alias if needed — and a join later + if that introduces a new table).""" + table, col = recon.resolve_column(f.column) + if table not in aliases: + aliases[table] = _alias_for(table, {a: t for t, a in aliases.items()}) + qualified = f'"{aliases[table]}"."{col.name}"' + + kind = f.kind() + if kind == "equals": + return f"{qualified} = {_render_literal(f.equals)}" + if kind == "in_values": + if not f.in_values: + raise PickValidationError( + f"filter on {f.column!r} has empty in_values" + ) + rendered = ", ".join(_render_literal(v) for v in f.in_values) + return f"{qualified} IN ({rendered})" + if kind == "between": + lo, hi = f.between + return f"{qualified} BETWEEN {_render_literal(lo)} AND {_render_literal(hi)}" + if kind == "date_range": + return resolve_date_range(f.date_range, col, qualified) + raise PickValidationError(f"unknown filter kind {kind!r}") + + +# ── ORDER BY rendering ────────────────────────────────────── + +def _render_order_by(ob: OrderBy, metric_name: str, dim_select: dict[str, str]) -> str: + """`metric_name` is the alias of the metric SELECT expression; + `dim_select` maps dimension ref (column ref) → output alias.""" + if ob.by == "metric": + return f'"{metric_name}" {ob.direction.upper()}' + # by == "dimension" + if ob.dimension not in dim_select: + raise PickValidationError( + f"order_by.dimension={ob.dimension!r} is not in group_by ({list(dim_select)})" + ) + return f'"{dim_select[ob.dimension]}" {ob.direction.upper()}' + + +# ── Metric expression rewriting ───────────────────────────── + +def _qualify_metric_sql(sql_fragment: str, table_alias: str) -> str: + """Rewrite bare column refs inside the metric's SQL expression so they + point at the metric table's alias. + + Metric.sql in metrics.yaml is written as an unaliased expression (e.g. + `AVG(CASE WHEN status='B' ... END)`). We need it qualified to the + metric table's alias so the composer can introduce other tables via + joins without ambiguity. Uses sqlglot to parse and rewrite identifier + refs that don't already have a table prefix. + """ + tree = sqlglot.parse_one(sql_fragment, dialect="postgres") + for col in tree.find_all(sqlglot.exp.Column): + if col.table: + continue + # Inject the alias as the table qualifier. + col.set("table", sqlglot.exp.Identifier(this=table_alias, quoted=True)) + return tree.sql(dialect="postgres", identify=True) + + +def _qualify_metric_filter(filter_fragment: str, table_alias: str) -> str: + """Same as _qualify_metric_sql but for the metric's optional WHERE filter. + metric.filter is written as a bare boolean expression.""" + return _qualify_metric_sql(filter_fragment, table_alias) + + +# ── Compose ───────────────────────────────────────────────── + +def compose(pick: Pick, recon: Recon) -> ComposeResult: + """Render `pick` to SQL against `recon`. Raises PickValidationError on + any reference the recon can't resolve.""" + + # 1. Metric. + if pick.metric not in recon.metrics: + raise PickValidationError( + f"metric {pick.metric!r} not in recon (known: {sorted(recon.metrics)})" + ) + metric = recon.metrics[pick.metric] + if metric.from_table not in recon.tables: + raise PickValidationError( + f"metric {pick.metric!r} declares from_table={metric.from_table!r} " + f"but no such table in recon" + ) + + aliases: dict[str, str] = {} + aliases[metric.from_table] = _alias_for(metric.from_table, {}) + base_alias = aliases[metric.from_table] + + # 2. Resolve every group_by column to (table, Column) and collect target + # tables that aren't the metric's table — those need joins. + group_bindings: list[tuple[str, str, str, str]] = [] # (ref, table, col_name, output_alias) + extra_tables: list[str] = [] + dim_select_alias: dict[str, str] = {} + for ref in pick.group_by: + table, col = recon.resolve_column(ref) + if table not in aliases: + if table != metric.from_table: + extra_tables.append(table) + aliases[table] = _alias_for(table, {a: t for t, a in aliases.items()}) + out_alias = col.name if ref == col.name else ref.replace(".", "_") + group_bindings.append((ref, table, col.name, out_alias)) + dim_select_alias[ref] = out_alias + + # 3. Resolve filter columns first (they may also introduce new tables we + # need to join). Rendered separately so we can keep their predicates + # in WHERE. + filter_clauses: list[str] = [] + for f in pick.where: + table, _ = recon.resolve_column(f.column) + if table not in aliases: + if table != metric.from_table: + extra_tables.append(table) + aliases[table] = _alias_for(table, {a: t for t, a in aliases.items()}) + filter_clauses.append(_render_filter(f, recon, aliases)) + + # 4. Build JOINs for the union of extra tables. + join_clauses = _build_joins(recon, metric.from_table, extra_tables, aliases) + + # 5. SELECT list: dimension columns first (in group_by order), then the + # metric expression. + select_parts: list[str] = [] + for ref, table, col_name, out_alias in group_bindings: + select_parts.append(f'"{aliases[table]}"."{col_name}" AS "{out_alias}"') + metric_expr = _qualify_metric_sql(metric.sql, base_alias) + select_parts.append(f'{metric_expr} AS "{metric.name}"') + + # 6. WHERE: metric's default filter (if any) ANDed with the Pick's filters. + where_parts: list[str] = [] + if metric.filter: + where_parts.append(_qualify_metric_filter(metric.filter, base_alias)) + where_parts.extend(filter_clauses) + + # 7. GROUP BY (just the dim output aliases) + ORDER BY + LIMIT. + group_by_clause = "" + if group_bindings: + group_by_clause = "GROUP BY " + ", ".join( + f'"{out_alias}"' for _, _, _, out_alias in group_bindings + ) + + order_by_clause = "" + if pick.order_by is not None: + order_by_clause = "ORDER BY " + _render_order_by( + pick.order_by, metric.name, dim_select_alias + ) + elif group_bindings: + # Sensible default: largest metric first. + order_by_clause = f'ORDER BY "{metric.name}" DESC' + + limit_clause = f"LIMIT {pick.limit}" if pick.limit is not None else "" + + # 8. Assemble. + sql_parts = [ + "SELECT " + ", ".join(select_parts), + f'FROM "{metric.from_table}" AS "{base_alias}"', + ] + sql_parts.extend(join_clauses) + if where_parts: + sql_parts.append("WHERE " + " AND ".join(where_parts)) + if group_by_clause: + sql_parts.append(group_by_clause) + if order_by_clause: + sql_parts.append(order_by_clause) + if limit_clause: + sql_parts.append(limit_clause) + raw = "\n".join(sql_parts) + + # 9. Paranoid re-parse: catches composer bugs by round-tripping through + # sqlglot. identify=True keeps every identifier quoted. + parsed = sqlglot.parse_one(raw, dialect="postgres") + sql = parsed.sql(dialect="postgres", identify=True) + + used = sorted(aliases) + return ComposeResult(sql=sql, used_tables=used) diff --git a/api/composer/dates.py b/api/composer/dates.py new file mode 100644 index 0000000..d870160 --- /dev/null +++ b/api/composer/dates.py @@ -0,0 +1,73 @@ +"""Date-range resolver for typed Filters. + +Given a `Filter` with `date_range="YYYY"` or `date_range="YYYY-Q[1-4]"` and +the target column, render a SQL fragment that evaluates correctly under the +column's actual storage type. Dispatched in two layers: + +1. `Column.semantic_type` if set (extension point — e.g. `"date_yymmdd"` for + datasets that ship YYMMDD-encoded integers). +2. Otherwise `Column.sql_type` (`DATE`, `TIMESTAMP`, …) — the default path + for warehouses that store dates as real date columns. + +Inclusive on both ends. Bounds quoted with single quotes (Postgres parses +date literals from strings). +""" + +from __future__ import annotations + +import re + +from api.composer.types import PickValidationError +from api.recon.types import Column + +_YEAR_ONLY = re.compile(r"^\d{4}$") +_YEAR_QUARTER = re.compile(r"^(\d{4})-Q([1-4])$") + + +_QUARTER_BOUNDS = { + 1: ("01-01", "03-31"), + 2: ("04-01", "06-30"), + 3: ("07-01", "09-30"), + 4: ("10-01", "12-31"), +} + + +def _bounds(spec: str) -> tuple[str, str]: + """Parse a date_range spec to (lo, hi) calendar dates as ISO strings.""" + if _YEAR_ONLY.match(spec): + year = spec + return f"{year}-01-01", f"{year}-12-31" + m = _YEAR_QUARTER.match(spec) + if m: + year, q = m.group(1), int(m.group(2)) + lo_mmdd, hi_mmdd = _QUARTER_BOUNDS[q] + return f"{year}-{lo_mmdd}", f"{year}-{hi_mmdd}" + raise PickValidationError( + f"unsupported date_range {spec!r}; expected 'YYYY' or 'YYYY-Q1..Q4'" + ) + + +def resolve_date_range(spec: str, column: Column, qualified_col: str) -> str: + """Render the WHERE fragment for ` BETWEEN ...`. + + `qualified_col` is the already-quoted reference the composer wants in + the SQL (e.g. `"l"."date"`). Returns just the predicate, no `WHERE`. + """ + lo, hi = _bounds(spec) + semantic = (column.semantic_type or "").lower() + + if semantic == "date_yymmdd": + # Column stores YYMMDD as integer (some BIRD datasets ship this way). + return ( + f"TO_DATE(LPAD({qualified_col}::text, 6, '0'), 'YYMMDD') " + f"BETWEEN '{lo}' AND '{hi}'" + ) + + sql_type = column.sql_type.upper() + if sql_type.startswith("DATE") or sql_type.startswith("TIMESTAMP"): + return f"{qualified_col} BETWEEN '{lo}' AND '{hi}'" + + raise PickValidationError( + f"date_range on column with sql_type={column.sql_type!r} and " + f"semantic_type={column.semantic_type!r} is not supported" + ) diff --git a/api/composer/types.py b/api/composer/types.py new file mode 100644 index 0000000..279178d --- /dev/null +++ b/api/composer/types.py @@ -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: ` = ` + - 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] diff --git a/api/datasets/financial/schema_docs.yaml b/api/datasets/financial/schema_docs.yaml index c38940f..84f2c87 100644 --- a/api/datasets/financial/schema_docs.yaml +++ b/api/datasets/financial/schema_docs.yaml @@ -27,7 +27,7 @@ tables: columns: account.district_id: Geographic district the account belongs to (joins to district). account.frequency: Statement issuance frequency. "POPLATEK MESICNE" = monthly, "POPLATEK TYDNE" = weekly, "POPLATEK PO OBRATU" = on transaction. - account.date: Date the account was opened (YYMMDD as integer). + account.date: Date the account was opened. client.gender: '"M" or "F". Derived from birth_number.' client.birth_date: Date of birth, normalised. Original birth_number encoded gender by adding 50 to the month field for women. @@ -54,7 +54,7 @@ columns: district.A16: Number of crimes committed in 1996. loan.account_id: The account the loan was granted to. - loan.date: Loan grant date (YYMMDD as integer). + loan.date: Date the loan was granted. loan.amount: Total amount of the loan (CZK). loan.duration: Loan duration in months. loan.payments: Monthly payment (CZK). @@ -65,7 +65,7 @@ columns: card.issued: Date the card was issued. trans.account_id: The account the transaction belongs to. - trans.date: Transaction date (YYMMDD as integer). + trans.date: Transaction date. trans.type: '"PRIJEM" (credit) or "VYDAJ" (debit).' trans.operation: Mode of operation, e.g. "VKLAD" (cash credit), "VYBER" (cash withdrawal), "PREVOD Z UCTU" (transfer from another bank). trans.amount: Amount (CZK). diff --git a/api/llm.py b/api/llm.py index 76ae7d4..1b3bb2b 100644 --- a/api/llm.py +++ b/api/llm.py @@ -9,7 +9,8 @@ Provider selection is via `settings.llm_provider`. Supported: Every call opens a Langfuse `generation` span automatically, tagged with the model name and populated with token usage from the provider's response. Spans nest under whatever observation the caller has open, so the trace -hierarchy shows `text_to_sql > llm.chat` with usage on the inner node. +hierarchy shows the wrapping step (e.g. `pick_for_question`) with the +`llm.chat` generation as a child carrying the usage. """ from __future__ import annotations @@ -119,7 +120,7 @@ def _active_model() -> str: def chat(*, system: str, user: str, max_tokens: int = 1024, span_name: str = "llm.chat") -> str: """Run a single chat completion. Opens a Langfuse generation span around the call with the model name + token usage. `span_name` is the label - shown in Langfuse — pass something descriptive (e.g. "text_to_sql.gen") + shown in Langfuse — pass something descriptive (e.g. "direct_answer.pick") so traces are scannable.""" provider = get_settings().llm_provider.lower() impl = _PROVIDERS.get(provider) diff --git a/api/plan/planner.py b/api/plan/planner.py index 03dbb33..0c4ad26 100644 --- a/api/plan/planner.py +++ b/api/plan/planner.py @@ -25,8 +25,8 @@ def plan(question: str) -> Plan: user = render( "planner.user", question=question, - table_names=", ".join(schema.table_names()), - metrics_block=schema.render_metrics(), + tables_block=schema.render_tables_brief(), + metrics_block=schema.render_metrics_brief(), catalog=json.dumps(catalog(), indent=2), ) diff --git a/api/prompts/pick.system.txt b/api/prompts/pick.system.txt new file mode 100644 index 0000000..31d5da3 --- /dev/null +++ b/api/prompts/pick.system.txt @@ -0,0 +1,35 @@ +You translate a business question into a typed Pick that an SQL composer +will execute deterministically. You do NOT write SQL — you choose from +finite, known sets and the composer builds the query. + +OUTPUT FORMAT — JSON, no prose, no markdown fences: + +{ + "kind": "aggregate", + "metric": "", + "group_by": ["", ...], + "where": [, ...], + "order_by": {"by": "metric"|"dimension", "direction": "asc"|"desc", "dimension": ""?}, + "limit": +} + +FILTER SHAPES — exactly one value field per filter: + +{"column": "", "equals": ""} +{"column": "", "in_values": ["", "", ...]} +{"column": "", "between": [, ]} +{"column": "", "date_range": ""} + +RULES: + +- `metric` MUST be one of the candidate metric names. Don't invent. +- Every column ref in `group_by` and in `where[*].column` MUST appear in + the candidate columns list. If the same name lives on multiple tables, + it will be listed as `table.column` — pick the qualified form. +- Don't include `group_by` or `where` if not needed (empty list / omit). +- Only ONE filter per column. Don't repeat. +- If you cannot express the question using the candidate metrics + + columns + filter shapes, return: + {"error": ""} +- No SQL fragments, no raw expressions, no CASE, no subqueries. +- Output ONLY the JSON object. No explanation, no markdown. diff --git a/api/prompts/pick.user.txt b/api/prompts/pick.user.txt new file mode 100644 index 0000000..8d1b370 --- /dev/null +++ b/api/prompts/pick.user.txt @@ -0,0 +1,10 @@ +QUESTION: +{question} + +CANDIDATE METRICS (pick one for `metric`): +{metrics_block} + +CANDIDATE COLUMNS (pick from these for `group_by` and `where[*].column`): +{columns_block} + +Return the Pick as JSON (or the error object). No other text. diff --git a/api/prompts/planner.user.txt b/api/prompts/planner.user.txt index e2982ee..2719116 100644 --- a/api/prompts/planner.user.txt +++ b/api/prompts/planner.user.txt @@ -1,6 +1,7 @@ Question: {question} -Warehouse tables: {table_names} +Warehouse tables: +{tables_block} Metric catalog: {metrics_block} diff --git a/api/prompts/text_to_sql.system.txt b/api/prompts/text_to_sql.system.txt deleted file mode 100644 index c019c59..0000000 --- a/api/prompts/text_to_sql.system.txt +++ /dev/null @@ -1,16 +0,0 @@ -Task: translate the analyst's question into a single Postgres SELECT (or WITH … SELECT) against the BIRD `financial` schema (PKDD'99 Czech bank). - -Output: exactly one fenced ```sql block, nothing else — no prose, no second block, no trailing semicolon required. - -Constraints: -- search_path is set to `financial`; don't qualify table names with the schema. -- **Postgres folds unquoted identifiers to lowercase.** Always double-quote any column or table name that contains an uppercase letter — including all the `district` columns: `"A2"`, `"A3"`, `"A4"`, `"A5"`, `"A6"`, `"A7"`, `"A8"`, `"A9"`, `"A10"`, `"A11"`, `"A12"`, `"A13"`, `"A14"`, `"A15"`, `"A16"`. Mixing quoted and unquoted forms of the same column (`"A3"` and `A13`) will fail at runtime — be consistent. -- Quote identifiers that collide with SQL keywords (notably `"order"`). -- Dates in the source are YYMMDD integers (e.g. 950315 → 1995-03-15). Convert with `TO_DATE(LPAD(date::text, 6, '0'), 'YYMMDD')` whenever you need a real date for filtering, comparison, or display. -- Loan `status` values: 'A' finished OK, 'B' finished defaulted, 'C' running OK, 'D' running in debt. -- Transaction `type` values: 'PRIJEM' credit, 'VYDAJ' debit. -- Read-only: never emit DDL or DML. - -When the question maps to a metric in the catalog you're given, copy the metric's SQL fragment verbatim — don't paraphrase it. - -If the question is ambiguous, pick the most defensible interpretation and proceed; don't ask for clarification. diff --git a/api/prompts/text_to_sql.user.txt b/api/prompts/text_to_sql.user.txt deleted file mode 100644 index 2aa5244..0000000 --- a/api/prompts/text_to_sql.user.txt +++ /dev/null @@ -1,10 +0,0 @@ -Question: -{question} - -Schema (only the tables that look relevant; you may use any of these): -{schema_block} - -Metric catalog (copy the SQL verbatim when the question maps to one): -{metrics_block} - -Return only the SQL in a ```sql fenced block.{retry_hint} diff --git a/api/recon/build.py b/api/recon/build.py index 907a5bb..5aa287b 100644 --- a/api/recon/build.py +++ b/api/recon/build.py @@ -80,16 +80,25 @@ def _read_yaml(path: Path) -> dict[str, Any]: return yaml.safe_load(path.read_text()) or {} -def _parse_column_descs(raw: dict[str, Any]) -> dict[str, str]: - """Accept either the flat form `{table.col: desc}` or the nested form - `{table: {col: desc}}`. Returns flat form.""" - out: dict[str, str] = {} - for k, v in raw.items(): +def _parse_column_specs(raw: dict[str, Any]) -> dict[str, dict[str, Any]]: + """Accept either the flat form `{table.col: }` or the nested form + `{table: {col: }}`. Each `` is either a bare description + string OR a dict `{desc?: str, type?: str}` where `type` is a domain- + level semantic type (e.g. "date_yymmdd"). Returns flat form: each value + is normalised to `{desc, type}`.""" + def norm(v: Any) -> dict[str, Any]: if isinstance(v, dict): - for col, desc in v.items(): - out[f"{k}.{col}"] = desc + return {"desc": v.get("desc") or v.get("description"), "type": v.get("type")} + return {"desc": v, "type": None} + + out: dict[str, dict[str, Any]] = {} + for k, v in raw.items(): + if isinstance(v, dict) and not ("desc" in v or "description" in v or "type" in v): + # Nested: {table: {col: spec}}. + for col, sub in v.items(): + out[f"{k}.{col}"] = norm(sub) else: - out[k] = v + out[k] = norm(v) return out @@ -101,19 +110,20 @@ def merge_into_recon(dataset: str, extracted: dict[str, Any]) -> Recon: metrics_yaml = _read_yaml(DATASETS_DIR / dataset / "metrics.yaml") table_descs: dict[str, str] = aug.get("tables", {}) or {} - col_descs = _parse_column_descs(aug.get("columns", {}) or {}) + col_specs = _parse_column_specs(aug.get("columns", {}) or {}) tables: dict[str, Table] = {} for tname, t_data in extracted["tables"].items(): - cols = [ - Column( + cols = [] + for c in t_data["columns"]: + spec = col_specs.get(f"{tname}.{c['name']}", {}) + cols.append(Column( name=c["name"], sql_type=c["sql_type"], nullable=c["nullable"], - description=col_descs.get(f"{tname}.{c['name']}"), - ) - for c in t_data["columns"] - ] + description=spec.get("desc"), + semantic_type=spec.get("type"), + )) tables[tname] = Table( name=tname, description=table_descs.get(tname), diff --git a/api/recon/types.py b/api/recon/types.py index 9377759..e61556c 100644 --- a/api/recon/types.py +++ b/api/recon/types.py @@ -23,6 +23,10 @@ class Column: sql_type: str nullable: bool description: str | None = None + # Domain-level type that the composer dispatches on for filter rendering. + # Examples: "date_yymmdd" (BIRD's int-encoded YYMMDD dates), "date" (real + # postgres date). None = treat literally with the column's sql_type. + semantic_type: str | None = None @classmethod def from_inspector(cls, info: dict[str, Any], description: str | None = None) -> "Column": @@ -40,12 +44,16 @@ class Column: "sql_type": self.sql_type, "nullable": self.nullable, "description": self.description, + "semantic_type": self.semantic_type, } @classmethod def from_dict(cls, d: dict[str, Any]) -> "Column": - return cls(name=d["name"], sql_type=d["sql_type"], - nullable=d["nullable"], description=d.get("description")) + return cls( + name=d["name"], sql_type=d["sql_type"], + nullable=d["nullable"], description=d.get("description"), + semantic_type=d.get("semantic_type"), + ) @dataclass @@ -144,7 +152,7 @@ class Recon: column_to_tables: dict[str, list[str]] = field(default_factory=dict) relationships: list[Relationship] = field(default_factory=list) - # ── Read-side helpers used by Analyses + text_to_sql ── + # ── Read-side helpers used by Analyses + the composer ── def table_names(self) -> list[str]: return sorted(self.tables) @@ -156,6 +164,37 @@ class Recon: """Tables that have a column with this name (case-sensitive).""" return self.column_to_tables.get(column, []) + def resolve_column(self, ref: str) -> tuple[str, "Column"]: + """Resolve a `column` or `table.column` reference to its (table_name, + Column) pair. Raises ValueError if the column doesn't exist or is + ambiguous (lives on multiple tables) without an explicit `table.` prefix. + + Used by the composer to bind every Pick.group_by entry to a single + owning table before SQL composition. Disambiguation is on the caller: + if you mean `account.district_id`, say so.""" + if "." in ref: + table, col = ref.split(".", 1) + if table not in self.tables: + raise ValueError(f"unknown table {table!r} in column ref {ref!r}") + for c in self.tables[table].columns: + if c.name == col: + return table, c + raise ValueError(f"column {col!r} not found on table {table!r}") + owners = self.owning_tables(ref) + if not owners: + raise ValueError(f"no table has a column named {ref!r}") + if len(owners) > 1: + raise ValueError( + f"column {ref!r} is ambiguous (lives in: {', '.join(owners)}); " + f"qualify it as 'table.{ref}'" + ) + table = owners[0] + for c in self.tables[table].columns: + if c.name == ref: + return table, c + # Index says it's there but the table doesn't — shouldn't happen. + raise ValueError(f"column {ref!r} indexed under {table!r} but not found") + def join_path(self, src: str, dst: str) -> list[str] | None: """Shortest sequence of tables from src to dst via declared relationships. Returns None if no path exists. The path includes both endpoints.""" @@ -220,6 +259,29 @@ class Recon: ) return "\n".join(lines) + def render_metrics_brief(self) -> str: + """One line per metric — name, unit, description. No SQL, no table. + + For the planner: it picks WHICH analyses to run; it doesn't need + implementation detail. Saves prompt tokens and cuts the noise the + model has to filter through.""" + if not self.metrics: + return "(no metrics defined)" + return "\n".join( + f"- {m.name} [{m.unit or 'unitless'}]: {m.description}" + for m in self.metrics.values() + ) + + def render_tables_brief(self) -> str: + """One line per table — name + description. No columns, no types.""" + if not self.tables: + return "(no tables)" + lines: list[str] = [] + for t in self.tables.values(): + desc = t.description or "(no description)" + lines.append(f"- {t.name}: {desc}") + return "\n".join(lines) + def render_relationships(self) -> str: if not self.relationships: return "(no relationships declared)" diff --git a/api/tools/text_to_sql.py b/api/tools/text_to_sql.py deleted file mode 100644 index a5b6591..0000000 --- a/api/tools/text_to_sql.py +++ /dev/null @@ -1,76 +0,0 @@ -"""LLM-driven NL → SQL with sqlglot validation and one retry.""" - -from __future__ import annotations - -import logging -import re - -import sqlglot - -from api import langfuse_client as lf -from api.llm import chat -from api.prompts import load, render -from api.recon import load_recon -from api.recon.validate import ReconValidationError, validate_sql -from api.tools.types import T2SResult - -logger = logging.getLogger("nvi.tools.text_to_sql") - - -def text_to_sql(question: str, *, hint_tables: list[str] | None = None) -> T2SResult: - recon = load_recon() - tables = hint_tables or recon.table_names() - system = load("text_to_sql.system") - - def _user(retry_hint: str = "") -> str: - return render( - "text_to_sql.user", - question=question, - schema_block=recon.render_tables(tables), - metrics_block=recon.render_metrics(), - retry_hint=retry_hint, - ) - - with lf.span( - "text_to_sql", - as_type="generation", - input={"question": question, "hint_tables": hint_tables}, - ) as span: - raw = _extract_sql(chat(system=system, user=_user())) - sql = _normalize(raw) # raises on parse error — fail fast - validate_sql(sql, recon) # raises ReconValidationError — fail fast - - result = T2SResult(sql=sql, used_tables=_extract_tables(sql)) - span.update(output={"sql": sql, "used_tables": result.used_tables}) - return result - - -def _extract_sql(text: str) -> str: - m = re.search(r"```sql\s*(.*?)```", text, re.DOTALL | re.IGNORECASE) - if m: - return m.group(1).strip().rstrip(";").strip() - return text.strip().rstrip(";").strip() - - -def _normalize(sql: str) -> str: - """Parse the LLM's SQL, then re-render with every identifier quoted. - - Postgres folds unquoted identifiers to lowercase, so `d.A13` becomes - `d.a13` and breaks against case-preserving columns like district."A13". - Re-rendering with `identify=True` puts double quotes around every - identifier — case is preserved, and Postgres treats a quoted lowercase - identifier the same as the unquoted form, so this is safe in both - directions. - """ - parsed = sqlglot.parse_one(sql, dialect="postgres") - if parsed is None: - raise ValueError("sqlglot returned no parse tree") - return parsed.sql(dialect="postgres", identify=True) - - -def _extract_tables(sql: str) -> list[str]: - try: - parsed = sqlglot.parse_one(sql, dialect="postgres") - return sorted({t.name for t in parsed.find_all(sqlglot.exp.Table)}) - except Exception: - return [] diff --git a/tests/test_composer_compose.py b/tests/test_composer_compose.py new file mode 100644 index 0000000..8c58b49 --- /dev/null +++ b/tests/test_composer_compose.py @@ -0,0 +1,229 @@ +"""Composer SQL emission — golden-string tests against a representative recon. + +The fixture mirrors enough of the financial dataset to exercise: metric +with default filter, group_by with single hop and multi-hop joins, typed +date_range filter on a real DATE column, and ORDER BY by both metric and +dimension. Adjust goldens cautiously — if a string changes, confirm the +new SQL still semantically matches before updating the test. +""" + +from api.composer import compose +from api.composer.types import Filter, OrderBy, Pick, PickValidationError +from api.recon.types import Column, Metric, Recon, Relationship, Table + +import pytest + + +def _fixture_recon() -> Recon: + cols_account = [ + Column("account_id", "BIGINT", False), + Column("district_id", "BIGINT", True), + ] + cols_district = [ + Column("district_id", "BIGINT", False), + Column("A2", "TEXT", True), + Column("A3", "TEXT", True), + ] + cols_loan = [ + Column("loan_id", "BIGINT", False), + Column("account_id", "BIGINT", True), + Column("amount", "BIGINT", True), + Column("status", "TEXT", True), + Column("date", "DATE", True), + ] + rels = [ + Relationship("account", "district_id", "district", "district_id"), + Relationship("loan", "account_id", "account", "account_id"), + ] + c2t: dict[str, list[str]] = {} + for t in [Table("account", "accounts", cols_account), + Table("district", "districts", cols_district), + Table("loan", "loans", cols_loan)]: + for c in t.columns: + c2t.setdefault(c.name, []).append(t.name) + for k in c2t: + c2t[k] = sorted(c2t[k]) + + metrics = { + "loan_default_rate": Metric( + name="loan_default_rate", + description="Default rate over finished loans.", + sql="AVG(CASE WHEN status = 'B' THEN 1.0 WHEN status = 'A' THEN 0.0 END)", + from_table="loan", + filter="status IN ('A','B')", + unit="ratio", + ), + "loan_volume": Metric( + name="loan_volume", + description="Total CZK lent.", + sql="SUM(amount)", + from_table="loan", + filter=None, + unit="CZK", + ), + } + return Recon( + schema="financial", + tables={ + "account": Table("account", "accounts", cols_account), + "district": Table("district", "districts", cols_district), + "loan": Table("loan", "loans", cols_loan), + }, + metrics=metrics, + column_to_tables=c2t, + relationships=rels, + ) + + +# ── Aggregate, no group_by ────────────────────────────────── + +def test_metric_only_no_group_by(): + r = _fixture_recon() + pick = Pick(kind="aggregate", metric="loan_volume") + out = compose(pick, r) + assert out.sql == 'SELECT SUM("l"."amount") AS "loan_volume" FROM "loan" AS "l"' + assert out.used_tables == ["loan"] + + +def test_metric_with_default_filter_no_group_by(): + r = _fixture_recon() + pick = Pick(kind="aggregate", metric="loan_default_rate") + out = compose(pick, r) + expected = ( + "SELECT AVG(CASE WHEN \"l\".\"status\" = 'B' THEN 1.0 " + "WHEN \"l\".\"status\" = 'A' THEN 0.0 END) AS \"loan_default_rate\" " + "FROM \"loan\" AS \"l\" " + "WHERE \"l\".\"status\" IN ('A', 'B')" + ) + assert out.sql == expected + + +# ── Group by — single hop ─────────────────────────────────── + +def test_group_by_multi_hop_join(): + r = _fixture_recon() + pick = Pick( + kind="aggregate", + metric="loan_default_rate", + group_by=["A2"], + limit=10, + ) + out = compose(pick, r) + # loan → account → district, A2 owned by district. + assert 'FROM "loan" AS "l"' in out.sql + assert 'JOIN "account" AS "a" ON "l"."account_id" = "a"."account_id"' in out.sql + assert ( + 'JOIN "district" AS "d" ON "a"."district_id" = "d"."district_id"' in out.sql + ) + assert 'GROUP BY "A2"' in out.sql + # default ORDER BY metric DESC + assert 'ORDER BY "loan_default_rate" DESC' in out.sql + assert "LIMIT 10" in out.sql + assert set(out.used_tables) == {"loan", "account", "district"} + + +# ── Filters ───────────────────────────────────────────────── + +def test_where_equals_filter(): + r = _fixture_recon() + pick = Pick( + kind="aggregate", + metric="loan_volume", + where=[Filter(column="status", equals="D")], + ) + out = compose(pick, r) + assert "WHERE \"l\".\"status\" = 'D'" in out.sql + + +def test_where_in_values_filter(): + r = _fixture_recon() + pick = Pick( + kind="aggregate", + metric="loan_volume", + where=[Filter(column="status", in_values=["C", "D"])], + ) + out = compose(pick, r) + assert "\"l\".\"status\" IN ('C', 'D')" in out.sql + + +def test_where_date_range_on_real_date_column(): + r = _fixture_recon() + pick = Pick( + kind="aggregate", + metric="loan_default_rate", + where=[Filter(column="date", date_range="1996")], + ) + out = compose(pick, r) + assert "\"l\".\"date\" BETWEEN '1996-01-01' AND '1996-12-31'" in out.sql + # metric's default filter still ANDed in + assert "\"l\".\"status\" IN ('A', 'B')" in out.sql + + +def test_metric_filter_ands_with_pick_filter(): + r = _fixture_recon() + pick = Pick( + kind="aggregate", + metric="loan_default_rate", + where=[Filter(column="date", date_range="1996")], + ) + out = compose(pick, r) + # Both filters in the same WHERE clause, joined by AND. + assert " AND " in out.sql + + +# ── ORDER BY by dimension ─────────────────────────────────── + +def test_order_by_dimension(): + r = _fixture_recon() + pick = Pick( + kind="aggregate", + metric="loan_volume", + group_by=["A3"], + order_by=OrderBy(by="dimension", direction="asc", dimension="A3"), + ) + out = compose(pick, r) + assert 'ORDER BY "A3" ASC' in out.sql + + +def test_order_by_dimension_not_in_group_by_raises(): + r = _fixture_recon() + pick = Pick( + kind="aggregate", + metric="loan_volume", + group_by=["A2"], + order_by=OrderBy(by="dimension", direction="asc", dimension="A3"), + ) + with pytest.raises(PickValidationError, match="not in group_by"): + compose(pick, r) + + +# ── Validation rejections ─────────────────────────────────── + +def test_unknown_metric_raises(): + r = _fixture_recon() + pick = Pick(kind="aggregate", metric="moon_phase") + with pytest.raises(PickValidationError, match="not in recon"): + compose(pick, r) + + +def test_unknown_group_by_column_raises(): + r = _fixture_recon() + pick = Pick(kind="aggregate", metric="loan_volume", group_by=["nonsense"]) + with pytest.raises(ValueError, match="no table has a column named 'nonsense'"): + compose(pick, r) + + +def test_ambiguous_group_by_column_raises(): + r = _fixture_recon() + # account_id lives on both account and loan; must be qualified. + pick = Pick(kind="aggregate", metric="loan_volume", group_by=["account_id"]) + with pytest.raises(ValueError, match="ambiguous"): + compose(pick, r) + + +def test_qualified_disambiguates(): + r = _fixture_recon() + pick = Pick(kind="aggregate", metric="loan_volume", group_by=["loan.account_id"]) + out = compose(pick, r) + # Output alias becomes "loan_account_id" since the ref was qualified. + assert '"loan_account_id"' in out.sql diff --git a/tests/test_composer_dates.py b/tests/test_composer_dates.py new file mode 100644 index 0000000..310a498 --- /dev/null +++ b/tests/test_composer_dates.py @@ -0,0 +1,50 @@ +"""Date-range resolver tests — bounds parsing + dispatch on column type.""" + +import pytest + +from api.composer.dates import _bounds, resolve_date_range +from api.composer.types import PickValidationError +from api.recon.types import Column + + +def test_bounds_year(): + assert _bounds("1996") == ("1996-01-01", "1996-12-31") + + +def test_bounds_quarter_q1(): + assert _bounds("1996-Q1") == ("1996-01-01", "1996-03-31") + + +def test_bounds_quarter_q3(): + assert _bounds("1996-Q3") == ("1996-07-01", "1996-09-30") + + +def test_bounds_unsupported_raises(): + with pytest.raises(PickValidationError, match="unsupported"): + _bounds("Q3-1996") + + +def test_resolve_real_date_column(): + col = Column(name="date", sql_type="DATE", nullable=False) + out = resolve_date_range("1996", col, '"l"."date"') + assert out == "\"l\".\"date\" BETWEEN '1996-01-01' AND '1996-12-31'" + + +def test_resolve_timestamp_column(): + col = Column(name="created_at", sql_type="TIMESTAMP WITHOUT TIME ZONE", nullable=True) + out = resolve_date_range("1996-Q2", col, '"x"."created_at"') + assert out == "\"x\".\"created_at\" BETWEEN '1996-04-01' AND '1996-06-30'" + + +def test_resolve_yymmdd_semantic_type(): + col = Column(name="date", sql_type="INTEGER", nullable=False, + semantic_type="date_yymmdd") + out = resolve_date_range("1996", col, '"l"."date"') + assert "TO_DATE(LPAD(" in out + assert "BETWEEN '1996-01-01' AND '1996-12-31'" in out + + +def test_resolve_unsupported_column_type_raises(): + col = Column(name="x", sql_type="TEXT", nullable=True) + with pytest.raises(PickValidationError, match="not supported"): + resolve_date_range("1996", col, '"x"."x"') diff --git a/tests/test_events.py b/tests/test_events.py index 9120ff0..9281c60 100644 --- a/tests/test_events.py +++ b/tests/test_events.py @@ -62,13 +62,13 @@ def test_tool_call_shapes_carry_analysis_field(): """Tool/LLM events must include the active Analysis name from the contextvar — UI groups sub-events by it. None when outside an Analysis.""" # Outside any Analysis: analysis is None. - s = events.tool_call_start("text_to_sql", input={"q": "?"}) + s = events.tool_call_start("compose_sql", input={"q": "?"}) assert s["type"] == "tool_call_start" - assert s["tool"] == "text_to_sql" + assert s["tool"] == "compose_sql" assert s["analysis"] is None assert s["input"] == {"q": "?"} - e = events.tool_call_end("text_to_sql", output={"sql": "..."}) + e = events.tool_call_end("compose_sql", output={"sql": "..."}) assert e["analysis"] is None assert e["output"] == {"sql": "..."} assert e["error"] is None diff --git a/tests/test_imports.py b/tests/test_imports.py index 976779a..5d2ef2f 100644 --- a/tests/test_imports.py +++ b/tests/test_imports.py @@ -21,13 +21,19 @@ MODULES = [ "api.recon.types", "api.recon.build", "api.tools.db", - "api.tools.text_to_sql", "api.tools.execute_sql", "api.tools.types", + "api.composer", + "api.composer.types", + "api.composer.compose", + "api.composer.dates", "api.analyses.base", "api.analyses.types", + "api.analyses._pick", + "api.analyses._narrow", "api.analyses.direct_answer", "api.analyses.compare_periods", + "api.analyses.drill_down", "api.analyses.registry", "api.plan.planner", "api.plan.types", diff --git a/tests/test_pick_shape.py b/tests/test_pick_shape.py new file mode 100644 index 0000000..61c8de3 --- /dev/null +++ b/tests/test_pick_shape.py @@ -0,0 +1,122 @@ +"""Pick / Filter / OrderBy validation — round-trips and rejection cases.""" + +import pytest + +from api.composer.types import Filter, OrderBy, Pick, PickValidationError + + +# ── Filter ────────────────────────────────────────────────── + +def test_filter_equals_roundtrip(): + f = Filter.from_dict({"column": "status", "equals": "B"}) + assert f.kind() == "equals" + assert f.to_dict() == {"column": "status", "equals": "B"} + + +def test_filter_in_values_roundtrip(): + f = Filter.from_dict({"column": "status", "in_values": ["A", "B"]}) + assert f.kind() == "in_values" + assert f.to_dict() == {"column": "status", "in_values": ["A", "B"]} + + +def test_filter_between_normalised_to_tuple(): + f = Filter.from_dict({"column": "amount", "between": [1000, 5000]}) + assert f.kind() == "between" + assert f.between == (1000, 5000) + # round-trip emits a list, not a tuple + assert f.to_dict() == {"column": "amount", "between": [1000, 5000]} + + +def test_filter_date_range_roundtrip(): + f = Filter.from_dict({"column": "date", "date_range": "1996"}) + assert f.kind() == "date_range" + assert f.to_dict() == {"column": "date", "date_range": "1996"} + + +def test_filter_missing_value_raises(): + with pytest.raises(PickValidationError, match="no value field"): + Filter(column="x").kind() + + +def test_filter_multiple_value_fields_raises(): + f = Filter(column="x", equals="A", in_values=["A", "B"]) + with pytest.raises(PickValidationError, match="multiple value fields"): + f.kind() + + +def test_filter_missing_column_raises(): + with pytest.raises(PickValidationError, match="missing 'column'"): + Filter.from_dict({"equals": "B"}) + + +def test_filter_between_wrong_arity_raises(): + with pytest.raises(PickValidationError, match="2-element"): + Filter.from_dict({"column": "x", "between": [1, 2, 3]}) + + +# ── OrderBy ───────────────────────────────────────────────── + +def test_order_by_metric(): + ob = OrderBy.from_dict({"by": "metric", "direction": "desc"}) + assert ob.by == "metric" + assert ob.direction == "desc" + assert ob.to_dict() == {"by": "metric", "direction": "desc"} + + +def test_order_by_dimension_requires_dimension(): + with pytest.raises(PickValidationError, match="requires order_by.dimension"): + OrderBy.from_dict({"by": "dimension"}) + + +def test_order_by_invalid_direction_raises(): + with pytest.raises(PickValidationError, match="direction"): + OrderBy.from_dict({"by": "metric", "direction": "sideways"}) + + +# ── Pick ──────────────────────────────────────────────────── + +def test_pick_minimal(): + p = Pick.from_dict({"kind": "aggregate", "metric": "loan_volume"}) + assert p.metric == "loan_volume" + assert p.group_by == [] + assert p.where == [] + assert p.order_by is None + assert p.limit is None + + +def test_pick_full_roundtrip(): + raw = { + "kind": "aggregate", + "metric": "loan_default_rate", + "group_by": ["A2"], + "where": [{"column": "date", "date_range": "1996"}], + "order_by": {"by": "metric", "direction": "desc"}, + "limit": 10, + } + p = Pick.from_dict(raw) + assert p.to_dict() == raw + + +def test_pick_rejects_non_aggregate_kind(): + with pytest.raises(PickValidationError, match="aggregate"): + Pick.from_dict({"kind": "ratio", "metric": "x"}) + + +def test_pick_missing_metric_raises(): + with pytest.raises(PickValidationError, match="metric"): + Pick.from_dict({"kind": "aggregate"}) + + +def test_pick_group_by_must_be_list_of_strings(): + with pytest.raises(PickValidationError, match="group_by"): + Pick.from_dict({"kind": "aggregate", "metric": "x", "group_by": [1, 2]}) + + +def test_pick_where_must_be_list(): + with pytest.raises(PickValidationError, match="where must be a list"): + Pick.from_dict({"kind": "aggregate", "metric": "x", "where": {"foo": "bar"}}) + + +def test_pick_limit_must_be_int(): + with pytest.raises(PickValidationError, match="limit"): + Pick.from_dict({"kind": "aggregate", "metric": "x", "limit": "ten"}) diff --git a/ui/app/src/composables/useRunStream.ts b/ui/app/src/composables/useRunStream.ts index 81fb18a..ca30ee1 100644 --- a/ui/app/src/composables/useRunStream.ts +++ b/ui/app/src/composables/useRunStream.ts @@ -252,7 +252,8 @@ export function useRunStream() { t_started: elapsed(), }) const detail = d.tool === 'execute_sql' ? (d.input?.sql || '') : - d.tool === 'text_to_sql' ? (d.input?.question || '') : + d.tool === 'pick_for_question' ? (d.input?.question || '') : + d.tool === 'compose_sql' && d.input?.pick ? `metric=${d.input.pick.metric}` : JSON.stringify(d.input) push('debug', d.tool, `→ ${detail}`) }) @@ -271,14 +272,14 @@ export function useRunStream() { push('error', d.tool, `✗ ${d.error}`) return } - if (d.tool === 'text_to_sql' && d.output?.sql) { + if (d.tool === 'compose_sql' && d.output?.sql) { push('info', d.tool, `✓ ${d.output.sql}`) + } else if (d.tool === 'pick_for_question' && d.output?.pick) { + push('info', d.tool, `✓ ${d.output.pick.metric} (${(d.output.pick.group_by || []).join(', ') || 'no group_by'})`) } else if (d.tool === 'execute_sql') { const label = d.output?.dimension ? `${d.output.dimension}: ` : d.output?.label ? `${d.output.label}: ` : '' push('info', d.tool, `✓ ${label}${d.output?.row_count} rows`) - } else if (d.tool === 'generate_pair') { - push('info', d.tool, `✓ a=${d.output?.a?.label} · b=${d.output?.b?.label}`) } else { push('info', d.tool, '✓') } diff --git a/ui/app/src/pages/Ask.vue b/ui/app/src/pages/Ask.vue index b7de85a..10123ae 100644 --- a/ui/app/src/pages/Ask.vue +++ b/ui/app/src/pages/Ask.vue @@ -84,10 +84,10 @@ function subEvents(nodeId: string): ToolCall[] { } function shortPreview(call: ToolCall): string { - if (call.tool === 'text_to_sql' && call.output?.sql) return call.output.sql + if (call.tool === 'compose_sql' && call.output?.sql) return call.output.sql if (call.tool === 'execute_sql' && call.input?.sql) return call.input.sql - if (call.tool === 'generate_pair' && call.output) { - return `A: ${call.output.a?.label}\n${call.output.a?.sql}\n\nB: ${call.output.b?.label}\n${call.output.b?.sql}` + if (call.tool === 'pick_for_question' && call.output?.pick) { + return JSON.stringify(call.output.pick, null, 2) } return '' }