nvi/api/analyses/drill_down/helpers.py

"""drill_down helpers — one function per concern.

Kept separate from `analysis.py` so the DrillDown class stays readable as
high-level orchestration: decide → slice → loop → interpret.
"""

from __future__ import annotations

import json
import logging
import re
from typing import Any

from api import langfuse_client as lf
from api.analyses.drill_down.types import Slice
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")


# ── Decision step ──

async def decide_next(question: str, metric: str, dimensions: list[str],
                      slices: list[Slice], budget: int) -> dict[str, Any]:
    """One LLM call to pick the next dimension or stop.

    If the LLM picks something not in the candidate list, drop the choice
    and stop — the planner's fallback (if any) takes over. We don't try to
    coerce the LLM into a valid dimension because it might be hallucinating
    a name (e.g. a metric name) that has no obvious mapping.
    """
    system = load("drill_down.next.system")
    user = render(
        "drill_down.next.user",
        question=question,
        metric=metric,
        dimensions=", ".join(dimensions),
        history=format_history(slices),
        budget=budget,
    )
    with lf.span("drill_down.next", input={"budget": budget}) as span:
        await events.publish_current(events.llm_call(
            "drill_down.next", system_len=len(system), user_len=len(user),
        ))
        decision = _parse_json(chat(
            system=system, user=user, max_tokens=256, span_name="drill_down.next.gen",
        ))

        # Hard guard: the chosen dimension MUST be in the candidate list.
        if decision.get("action") == "drill":
            dim = decision.get("dimension")
            if dim not in dimensions:
                logger.warning(
                    "drill_down.next picked %r (not in candidates %s); stopping",
                    dim, dimensions,
                )
                decision = {
                    "action": "stop",
                    "reason": f"LLM picked {dim!r}, which isn't in the candidate dimensions",
                }
        span.update(output=decision)
        return decision


# ── 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)

    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_end(
        "text_to_sql", output={"sql": t2s.sql, "tables": t2s.used_tables},
    ))

    await events.publish_current(events.tool_call_start(
        "execute_sql", input={"sql": t2s.sql, "dimension": dim},
    ))
    result = execute_sql(t2s.sql)
    await events.publish_current(events.tool_call_end(
        "execute_sql",
        output={"dimension": dim, "row_count": result.row_count,
                "preview": result.as_dicts()[:5]},
    ))

    return Slice(
        dimension=dim,
        sql=t2s.sql,
        rows=result.as_dicts()[:20],
        reason=reason,
    )


# ── Final interpretation ──

async def interpret(question: str, metric: str, slices: list[Slice]) -> str:
    system = load("drill_down.interpret.system")
    user = render(
        "drill_down.interpret.user",
        question=question,
        metric=metric,
        slices_block=format_slices_block(slices),
    )
    await events.publish_current(events.llm_call(
        "drill_down.interpret", system_len=len(system), user_len=len(user),
    ))
    return chat(system=system, user=user, max_tokens=512, span_name="drill_down.interpret")


# ── Prompt-context formatters ──

def format_history(slices: list[Slice]) -> str:
    """Render the 'already tried' block fed to drill_down.next."""
    if not slices:
        return "(none yet)"
    return "\n".join(f"- {s.dimension}: {repr(s.rows[:5])}" for s in slices)


def format_slices_block(slices: list[Slice]) -> str:
    """Render every slice's rows for drill_down.interpret."""
    return "\n\n".join(
        f"dimension: {s.dimension}\nrows: {repr(s.rows)}"
        for s in slices
    )


# ── JSON extraction ──

def _parse_json(text: str) -> dict[str, Any]:
    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 ValueError(f"drill_down.next returned no JSON: {text[:200]!r}")
    return json.loads(raw[start:end + 1])