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

api/analyses/drill_down/__init__.py

@@ -0,0 +1,3 @@
+from api.analyses.drill_down.analysis import DrillDown
+
+__all__ = ["DrillDown"]

api/analyses/drill_down/analysis.py

@@ -0,0 +1,89 @@
+"""drill_down Analysis — ReAct loop.
+
+Pattern: bounded loop that picks a dimension to slice by, generates SQL,
+executes it, and decides whether to continue. From the runtime's view this
+is just `(args, question) → Finding`; internally it's a small ReAct agent.
+
+This file owns the orchestration. The decision step, the slice execution,
+the interpretation, and the prompt-context formatters all live in
+`helpers.py`. Types and constants live in `types.py`.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+from api import langfuse_client as lf
+from api.analyses.base import Analysis
+from api.analyses.drill_down.helpers import decide_next, execute_slice, interpret
+from api.analyses.drill_down.types import ARGS_SCHEMA, DrillDownArgs, Slice
+from api.analyses.types import Finding
+
+logger = logging.getLogger("nvi.analyses.drill_down")
+
+
+class DrillDown(Analysis):
+    name = "drill_down"
+    description = (
+        "Iteratively slice a metric by candidate dimensions to find which "
+        "ones explain the most variance. ReAct loop: pick a dimension, "
+        "query, decide whether to continue. Best for open-ended 'which "
+        "factors explain X' or 'why are some segments different' questions."
+    )
+    args_schema = ARGS_SCHEMA
+
+    async def run(self, args: dict[str, Any], question: str) -> Finding:
+        a = DrillDownArgs.from_raw(args, default_question=question)
+
+        with lf.span("analysis.drill_down", input={"question": a.question, "metric": a.metric}) as span:
+            try:
+                slices = await self._loop(a)
+                if not slices:
+                    return Finding(
+                        analysis=self.name, summary="",
+                        error="drill_down stopped before producing any slice",
+                    )
+                summary = await interpret(a.question, a.metric, slices)
+            except Exception as e:
+                logger.exception("drill_down failed")
+                span.update(output={"error": str(e)})
+                return Finding(analysis=self.name, summary="", error=str(e))
+
+            span.update(output={"summary": summary, "iterations": len(slices)})
+            return self._finalise(a, slices, summary)
+
+    async def _loop(self, a: DrillDownArgs) -> list[Slice]:
+        """Bounded ReAct loop. Each iteration: decide → slice. Stops when the
+        LLM says so, when the budget runs out, or when a repeat is detected."""
+        slices: list[Slice] = []
+        tried: set[str] = set()
+        for it in range(a.max_iters):
+            remaining = a.max_iters - it
+            decision = await decide_next(a.question, a.metric, a.dimensions, slices, remaining)
+            if decision.get("action") == "stop":
+                break
+            dim = decision.get("dimension")
+            if not dim or dim in tried:
+                break
+            tried.add(dim)
+            slices.append(await execute_slice(
+                a.question, a.metric, dim, decision.get("reason", ""),
+            ))
+        return slices
+
+    def _finalise(self, a: DrillDownArgs, slices: list[Slice], summary: str) -> Finding:
+        rows_combined = [
+            {"dimension": s.dimension, **r} for s in slices for r in s.rows
+        ]
+        return Finding(
+            analysis=self.name,
+            summary=summary,
+            rows=rows_combined,
+            sql=[s.sql for s in slices],
+            metadata={
+                "metric": a.metric,
+                "iterations": len(slices),
+                "dimensions_tried": [s.dimension for s in slices],
+            },
+        )

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

api/analyses/drill_down/types.py

@@ -0,0 +1,61 @@
+"""Public data shapes + constants for the drill_down Analysis."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+DEFAULT_DIMENSIONS: list[str] = ["A2", "A3", "A11", "A12", "A13", "A14"]
+DEFAULT_METRIC: str = "loan_default_rate"
+DEFAULT_MAX_ITERS: int = 3
+
+# JSON-schema-ish surface the planner reads to know which args this
+# Analysis accepts. Kept as a plain dict so `registry.catalog()` can dump
+# it straight to JSON.
+ARGS_SCHEMA: dict[str, dict[str, Any]] = {
+    "question": {
+        "type": "string",
+        "description": "Refined question this Analysis should answer.",
+    },
+    "metric": {
+        "type": "string",
+        "description": "Metric or column to slice (e.g. 'loan_default_rate', 'amount').",
+    },
+    "dimensions": {
+        "type": "array",
+        "items": "string",
+        "description": "Candidate dimensions to drill into. Defaults to district demographics.",
+    },
+    "max_iters": {
+        "type": "integer",
+        "description": f"Cap on slice iterations. Default {DEFAULT_MAX_ITERS}.",
+    },
+}
+
+
+@dataclass
+class DrillDownArgs:
+    """Parsed, defaulted args for one drill_down run."""
+    question: str
+    metric: str = DEFAULT_METRIC
+    dimensions: list[str] = field(default_factory=lambda: list(DEFAULT_DIMENSIONS))
+    max_iters: int = DEFAULT_MAX_ITERS
+
+    @classmethod
+    def from_raw(cls, raw: dict[str, Any], default_question: str) -> "DrillDownArgs":
+        return cls(
+            question=raw.get("question") or default_question,
+            metric=raw.get("metric") or DEFAULT_METRIC,
+            dimensions=raw.get("dimensions") or list(DEFAULT_DIMENSIONS),
+            max_iters=int(raw.get("max_iters") or DEFAULT_MAX_ITERS),
+        )
+
+
+@dataclass
+class Slice:
+    """One iteration's slice: the chosen dimension, the SQL it produced,
+    its rows, and the LLM's stated reason for picking the dimension."""
+    dimension: str
+    sql: str
+    rows: list[dict[str, Any]]
+    reason: str = ""