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

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)