177 lines
7.5 KiB
Python
177 lines
7.5 KiB
Python
"""compare_periods Analysis — CoT with two queries.
|
|
|
|
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 logging
|
|
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.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")
|
|
|
|
|
|
class ComparePeriods(Analysis):
|
|
name = "compare_periods"
|
|
description = (
|
|
"Compare a metric across two time windows (e.g. Q2 vs Q3, 1995 vs 1996). "
|
|
"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 '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(
|
|
"pick_for_question", input={"question": sub_q},
|
|
))
|
|
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(
|
|
"pick_for_question", output={"pick": base_pick.to_dict()},
|
|
))
|
|
|
|
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": 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": 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": period_b, "row_count": res_b.row_count,
|
|
"preview": res_b.as_dicts()[:5]},
|
|
))
|
|
|
|
interpret_system = load("compare_periods.interpret")
|
|
interpret_user = render(
|
|
"compare_periods.interpret.user",
|
|
question=sub_q,
|
|
label_a=period_a,
|
|
label_b=period_b,
|
|
rows_a=repr(res_a.as_dicts()[:20]),
|
|
rows_b=repr(res_b.as_dicts()[:20]),
|
|
)
|
|
await events.publish_current(events.llm_call(
|
|
"compare_periods.interpret",
|
|
system_len=len(interpret_system),
|
|
user_len=len(interpret_user),
|
|
))
|
|
summary = chat(
|
|
system=interpret_system, user=interpret_user,
|
|
max_tokens=512, span_name="compare_periods.interpret",
|
|
)
|
|
except Exception as e:
|
|
logger.exception("compare_periods failed")
|
|
await events.publish_current(events.tool_call_end("compare_periods", error=str(e)))
|
|
span.update(output={"error": str(e)})
|
|
return Finding(analysis=self.name, summary="", error=str(e))
|
|
|
|
span.update(output={"summary": summary})
|
|
return Finding(
|
|
analysis=self.name,
|
|
summary=summary,
|
|
rows=[
|
|
{"period": period_a, **r} for r in res_a.as_dicts()[:20]
|
|
] + [
|
|
{"period": period_b, **r} for r in res_b.as_dicts()[:20]
|
|
],
|
|
sql=[composed_a.sql, composed_b.sql],
|
|
metadata={
|
|
"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 _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}"
|
|
)
|
|
|
|
|
|
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)
|