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