recon + sqlglot validator + drill_down package; guard ReAct dimension picks against candidate list

api/analyses/drill_down/helpers.py

@@ -0,0 +1,173 @@
+"""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", as_type="generation", 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))
+
+        # 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)
+
+
+# ── 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])