recon-driven sql composer; pick → compose → execute; llm out of structural sql
This commit is contained in:
95
api/analyses/_narrow.py
Normal file
95
api/analyses/_narrow.py
Normal file
@@ -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)
|
||||
144
api/analyses/_pick.py
Normal file
144
api/analyses/_pick.py
Normal file
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
)
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
26
api/composer/__init__.py
Normal file
26
api/composer/__init__.py
Normal file
@@ -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",
|
||||
]
|
||||
288
api/composer/compose.py
Normal file
288
api/composer/compose.py
Normal file
@@ -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)
|
||||
73
api/composer/dates.py
Normal file
73
api/composer/dates.py
Normal file
@@ -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 `<qualified_col> 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"
|
||||
)
|
||||
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]
|
||||
@@ -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).
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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),
|
||||
)
|
||||
|
||||
|
||||
35
api/prompts/pick.system.txt
Normal file
35
api/prompts/pick.system.txt
Normal file
@@ -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": "<name from the candidate metrics list>",
|
||||
"group_by": ["<column ref>", ...],
|
||||
"where": [<filter>, ...],
|
||||
"order_by": {"by": "metric"|"dimension", "direction": "asc"|"desc", "dimension": "<col ref>"?},
|
||||
"limit": <integer>
|
||||
}
|
||||
|
||||
FILTER SHAPES — exactly one value field per filter:
|
||||
|
||||
{"column": "<col ref>", "equals": "<value>"}
|
||||
{"column": "<col ref>", "in_values": ["<v1>", "<v2>", ...]}
|
||||
{"column": "<col ref>", "between": [<lo>, <hi>]}
|
||||
{"column": "<col ref>", "date_range": "<YYYY|YYYY-Q1..Q4>"}
|
||||
|
||||
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": "<one short sentence on what's missing>"}
|
||||
- No SQL fragments, no raw expressions, no CASE, no subqueries.
|
||||
- Output ONLY the JSON object. No explanation, no markdown.
|
||||
10
api/prompts/pick.user.txt
Normal file
10
api/prompts/pick.user.txt
Normal file
@@ -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.
|
||||
@@ -1,6 +1,7 @@
|
||||
Question: {question}
|
||||
|
||||
Warehouse tables: {table_names}
|
||||
Warehouse tables:
|
||||
{tables_block}
|
||||
|
||||
Metric catalog:
|
||||
{metrics_block}
|
||||
|
||||
@@ -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.
|
||||
@@ -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}
|
||||
@@ -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: <spec>}` or the nested form
|
||||
`{table: {col: <spec>}}`. Each `<spec>` 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),
|
||||
|
||||
@@ -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)"
|
||||
|
||||
@@ -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 []
|
||||
Reference in New Issue
Block a user