mariano/nvi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

recon-driven sql composer; pick → compose → execute; llm out of structural sql

Browse Source
This commit is contained in:
buenosairesam
2026-06-03 11:01:02 -03:00
parent 61494362a3
commit 29c620b2c2
27 changed files with 1516 additions and 249 deletions
Download Patch File Download Diff File Expand all files Collapse all files

95
api/analyses/_narrow.py Normal file
View 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
View 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

141
api/analyses/compare_periods.py
View File

@@ -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)

49
api/analyses/direct_answer.py
View File

@@ -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,
},
)

57
api/analyses/drill_down/helpers.py
View File

@@ -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,
)