nvi/api/analyses/compare_periods.py

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

from __future__ import annotations

import json
import logging
import re
from typing import Any

from api import langfuse_client as lf
from api.analyses.base import Analysis
from api.analyses.types import Finding
from api.llm import chat
from api.prompts import load, render
from api.runtime import events
from api.tools.execute_sql import execute_sql
from api.tools.schema import load_schema

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). "
        "Generates two parallel SQL queries and 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'."},
    }

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

        with lf.span("analysis.compare_periods", input=args) as span:
            try:
                await events.publish_current(events.tool_call_start(
                    "generate_pair", input={"question": sub_q, "period_a": period_a, "period_b": period_b},
                ))
                pair = _generate_pair(sub_q, period_a, period_b)
                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"]}},
                ))

                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"])
                await events.publish_current(events.tool_call_end(
                    "execute_sql",
                    output={"label": pair["a"]["label"], "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_end(
                    "execute_sql",
                    output={"label": pair["b"]["label"], "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=pair["a"]["label"],
                    label_b=pair["b"]["label"],
                    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)
            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": pair["a"]["label"], **r} for r in res_a.as_dicts()[:20]
                ] + [
                    {"period": pair["b"]["label"], **r} for r in res_b.as_dicts()[:20]
                ],
                sql=[pair["a"]["sql"], pair["b"]["sql"]],
                metadata={
                    "label_a": pair["a"]["label"],
                    "label_b": pair["b"]["label"],
                    "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_schema()
    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,
    )
    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