147 lines
5.0 KiB
Python
147 lines
5.0 KiB
Python
"""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.composer import Pick, compose
|
|
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
|
|
|
|
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 ──
|
|
|
|
async def execute_slice(question: str, metric: str, dim: str, reason: str) -> Slice:
|
|
"""Compose SQL for one slice deterministically (metric + dim are already
|
|
chosen by `decide_next`), execute it, emit tool-call events around both,
|
|
and return the Slice. Zero LLM calls — recon authors the SQL."""
|
|
recon = load_recon()
|
|
pick = Pick(kind="aggregate", metric=metric, group_by=[dim], limit=10)
|
|
|
|
await events.publish_current(events.tool_call_start(
|
|
"compose_sql", input={"pick": pick.to_dict()},
|
|
))
|
|
composed = compose(pick, recon)
|
|
await events.publish_current(events.tool_call_end(
|
|
"compose_sql", output={"sql": composed.sql, "tables": composed.used_tables},
|
|
))
|
|
|
|
await events.publish_current(events.tool_call_start(
|
|
"execute_sql", input={"sql": composed.sql, "dimension": dim},
|
|
))
|
|
result = execute_sql(composed.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=composed.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])
|