| """Agent orchestrator: routes question → multi-stage pipeline → UIPlan.
|
|
|
| Pipeline stages activate different slices of the 30-key pool, so a single
|
| question distributes load across 8-12 distinct keys in parallel:
|
|
|
| 1. intent → LIGHT (17-19)
|
| 2. retrieve → RETRIEVER (11-16) ← GraphRAG-lite schema rerank
|
| 3. plan → SQL (7-10)
|
| 4. critic → CRITIC (1-6) ← runs concurrently with primary exec
|
| 4b.verify → VERIFIER (20-23) ← second-opinion on result quality
|
| 5. insights → SUMMARY (11-16) ← fan-out, 1 LLM call per result set
|
| 6. narrative → NARRATIVE (7-10)
|
| 7. followup → FOLLOWUP (24-26) ← smart follow-up generation
|
|
|
| Each stage is resilient: any failure degrades to a heuristic and the pipeline
|
| continues.
|
| """
|
| from __future__ import annotations
|
|
|
| import asyncio
|
| import time
|
| from dataclasses import dataclass, field
|
| from typing import Any, Awaitable, Callable
|
|
|
|
|
| EmitFn = Callable[[dict[str, Any]], Awaitable[None]]
|
|
|
| from ..keys.pool import get_pool
|
| from ..retrieval.duck import run_sql
|
| from ..retrieval.sql_safety_duck import validate_duckdb_sql
|
| from ..ui_plan.schema import UIPlan
|
| from .dynamic import (
|
| ResultSet,
|
| _pick_table,
|
| assemble_plan,
|
| dynamic_offline_plan,
|
| heuristic_secondaries,
|
| )
|
| from .pipeline import (
|
| classify_intent,
|
| critique_sql,
|
| exec_narrative,
|
| generate_followups,
|
| per_result_insights,
|
| plan_sqls,
|
| retrieve_tables,
|
| verify_sql,
|
| )
|
|
|
|
|
| @dataclass
|
| class AgentState:
|
| question: str
|
| intent: str = "exploratory"
|
| sql: str | None = None
|
| rows: list[dict[str, Any]] = field(default_factory=list)
|
| plan: UIPlan | None = None
|
| trace: list[dict[str, Any]] = field(default_factory=list)
|
| mode: str = "offline"
|
| keys_used: list[str] = field(default_factory=list)
|
| confidence: float = 0.7
|
| followups: list[str] = field(default_factory=list)
|
|
|
|
|
| def _safe_run(sql: str) -> tuple[str, list[dict[str, Any]], str | None]:
|
| chk = validate_duckdb_sql(sql)
|
| if not chk.ok:
|
| return sql, [], chk.reason
|
| res = run_sql(chk.rewritten)
|
| return chk.rewritten, res.rows, res.error
|
|
|
|
|
| def _track(state: AgentState, kid: str | None) -> None:
|
| if kid and kid not in state.keys_used:
|
| state.keys_used.append(kid)
|
|
|
|
|
| async def _noop_emit(_evt: dict[str, Any]) -> None:
|
| return None
|
|
|
|
|
|
|
|
|
|
|
|
|
| def _is_all_null_or_zero(rows: list[dict[str, Any]]) -> bool:
|
| """True when every numeric value in every row is None or 0,
|
| meaning the result set carries no useful information.
|
| Non-numeric columns (labels, ids) are ignored."""
|
| if not rows:
|
| return True
|
| found_any_numeric = False
|
| for r in rows:
|
| for v in r.values():
|
| if v is None:
|
| continue
|
| if isinstance(v, (int, float)):
|
| found_any_numeric = True
|
| if v != 0:
|
| return False
|
|
|
| elif isinstance(v, str):
|
| try:
|
| fv = float(v)
|
| found_any_numeric = True
|
| if fv != 0:
|
| return False
|
| except (ValueError, OverflowError):
|
| pass
|
|
|
|
|
| return found_any_numeric
|
|
|
|
|
| def _filter_meaningful(results: list[ResultSet]) -> list[ResultSet]:
|
| """Drop result sets whose numeric data is entirely null/zero."""
|
| return [r for r in results if not _is_all_null_or_zero(r.rows)]
|
|
|
|
|
| def _preview_rows(rows: list[dict[str, Any]], n: int = 3) -> list[dict[str, Any]]:
|
| out = []
|
| for r in rows[:n]:
|
| item = {}
|
| for k, v in r.items():
|
| if isinstance(v, (int, float)):
|
| item[k] = v
|
| elif v is None:
|
| item[k] = None
|
| else:
|
| s = str(v)
|
| item[k] = s if len(s) <= 60 else s[:57] + "…"
|
| out.append(item)
|
| return out
|
|
|
|
|
| async def _exec_secondaries_parallel(
|
| question: str, table: str, specs: list[dict[str, Any]]
|
| ) -> list[ResultSet]:
|
| """Run secondary SQLs concurrently; each DuckDB call is short & cheap."""
|
| def _job(spec: dict[str, Any]) -> ResultSet | None:
|
| sq = (spec.get("sql") or "").strip()
|
| st = spec.get("title") or "Companion view"
|
| if not sq:
|
| return None
|
| rewritten, rows, err = _safe_run(sq)
|
| if err or not rows or _is_all_null_or_zero(rows):
|
| return None
|
| return ResultSet(title=st, sql=rewritten, rows=rows, table=table)
|
|
|
| loop = asyncio.get_running_loop()
|
| tasks = [loop.run_in_executor(None, _job, s) for s in specs[:3]]
|
| out = await asyncio.gather(*tasks, return_exceptions=True)
|
| return [r for r in out if isinstance(r, ResultSet)]
|
|
|
|
|
| async def _run_llm_path(state: AgentState, emit: EmitFn = _noop_emit) -> AgentState:
|
| q = state.question
|
| t0 = time.perf_counter()
|
|
|
| def _ms() -> int:
|
| return int((time.perf_counter() - t0) * 1000)
|
|
|
|
|
| await emit({"stage": "intent", "status": "running", "label": "Classifying intent",
|
| "role": "LIGHT", "ms": _ms()})
|
| await emit({"stage": "retrieve", "status": "running",
|
| "label": "Retrieving relevant tables (GraphRAG-lite)",
|
| "role": "RETRIEVER", "ms": _ms()})
|
|
|
| (intent, k_intent), (tables, k_retr) = await asyncio.gather(
|
| classify_intent(q), retrieve_tables(q, k=3),
|
| )
|
| state.intent = intent
|
| _track(state, k_intent)
|
| _track(state, k_retr)
|
| state.trace.append({"stage": "intent", "value": intent, "key": k_intent})
|
| state.trace.append({"stage": "retrieve", "tables": tables, "key": k_retr})
|
| await emit({"stage": "intent", "status": "done", "key": k_intent,
|
| "value": intent, "ms": _ms()})
|
| await emit({"stage": "retrieve", "status": "done", "key": k_retr,
|
| "tables": tables, "ms": _ms()})
|
|
|
| if not tables:
|
| tables = [_pick_table(q)]
|
|
|
|
|
| await emit({"stage": "plan", "status": "running",
|
| "label": "Generating SQL plan (primary + secondaries)",
|
| "role": "SQL", "ms": _ms()})
|
| plan_data, k_plan = await plan_sqls(q, tables, intent)
|
| _track(state, k_plan)
|
| state.trace.append({"stage": "plan", "key": k_plan,
|
| "has_primary": bool(plan_data.get("primary"))})
|
|
|
| table = plan_data.get("table") or tables[0]
|
| primary = plan_data.get("primary") or {}
|
| secondaries = plan_data.get("secondaries") or []
|
|
|
| p_sql = (primary.get("sql") or "").strip()
|
| p_title = primary.get("title") or "Result"
|
|
|
| await emit({"stage": "plan", "status": "done", "key": k_plan,
|
| "table": table, "primary_title": p_title,
|
| "primary_sql": p_sql,
|
| "secondary_count": len(secondaries),
|
| "secondary_titles": [s.get("title") for s in secondaries[:3]],
|
| "ms": _ms()})
|
|
|
|
|
| primary_result: ResultSet | None = None
|
| if p_sql:
|
| await emit({"stage": "critic", "status": "running",
|
| "label": "Auditing primary SQL", "role": "CRITIC", "ms": _ms()})
|
| await emit({"stage": "exec_primary", "status": "running",
|
| "label": "Executing primary query on DuckDB",
|
| "role": "DUCKDB", "ms": _ms()})
|
|
|
| async def _run_primary() -> tuple[str, list[dict], str | None]:
|
| loop = asyncio.get_running_loop()
|
| return await loop.run_in_executor(None, _safe_run, p_sql)
|
|
|
| critic_fut = asyncio.create_task(critique_sql(q, table, p_sql))
|
| exec_fut = asyncio.create_task(_run_primary())
|
| (rewritten, rows, err), (fix, k_crit) = await asyncio.gather(exec_fut, critic_fut)
|
| _track(state, k_crit)
|
| state.trace.append({"stage": "critic", "key": k_crit,
|
| "applied": bool(fix) and bool(err)})
|
|
|
|
|
| all_null = (not err and rows and _is_all_null_or_zero(rows))
|
|
|
| await emit({"stage": "exec_primary", "status": "done",
|
| "row_count": len(rows), "error": err,
|
| "all_null": bool(all_null), "ms": _ms()})
|
| await emit({"stage": "critic", "status": "done", "key": k_crit,
|
| "fix_provided": bool(fix),
|
| "will_apply": bool(fix) and (bool(err) or bool(all_null)),
|
| "ms": _ms()})
|
|
|
| if (err or all_null) and fix:
|
| await emit({"stage": "critic_retry", "status": "running",
|
| "label": "Re-running corrected SQL", "ms": _ms()})
|
| rewritten, rows, err = _safe_run(fix)
|
| state.trace.append({"stage": "critic_retry", "error": err,
|
| "fix_len": len(fix)})
|
| await emit({"stage": "critic_retry", "status": "done",
|
| "row_count": len(rows), "error": err, "ms": _ms()})
|
|
|
| state.sql = rewritten
|
| state.rows = rows
|
| if not err and rows:
|
| primary_result = ResultSet(title=p_title, sql=rewritten,
|
| rows=rows, table=table)
|
|
|
|
|
| verify_task = None
|
| if primary_result and primary_result.rows:
|
| await emit({"stage": "verify", "status": "running",
|
| "label": "Verifying result quality (second opinion)",
|
| "role": "VERIFIER", "ms": _ms()})
|
| verify_task = asyncio.create_task(
|
| verify_sql(q, table, primary_result.sql, primary_result.rows)
|
| )
|
|
|
|
|
| results: list[ResultSet] = []
|
| if primary_result:
|
| results.append(primary_result)
|
|
|
| if secondaries:
|
| await emit({"stage": "secondaries", "status": "running",
|
| "label": f"Running {min(len(secondaries),3)} companion queries in parallel",
|
| "role": "DUCKDB", "ms": _ms()})
|
| sec_results = await _exec_secondaries_parallel(q, table, secondaries)
|
| results.extend(sec_results)
|
| if secondaries:
|
| await emit({"stage": "secondaries", "status": "done",
|
| "succeeded": len(sec_results),
|
| "attempted": min(len(secondaries), 3),
|
| "ms": _ms()})
|
|
|
|
|
| if verify_task:
|
| try:
|
| v_fix, v_conf, k_ver = await verify_task
|
| _track(state, k_ver)
|
| state.confidence = v_conf
|
| state.trace.append({"stage": "verify", "key": k_ver,
|
| "confidence": v_conf,
|
| "fix_provided": bool(v_fix)})
|
| await emit({"stage": "verify", "status": "done", "key": k_ver,
|
| "confidence": v_conf,
|
| "fix_provided": bool(v_fix), "ms": _ms()})
|
|
|
| if v_fix and v_conf < 0.5:
|
| await emit({"stage": "verify_retry", "status": "running",
|
| "label": "Re-running verifier-corrected SQL",
|
| "ms": _ms()})
|
| rewritten, rows, err = _safe_run(v_fix)
|
| if not err and rows and not _is_all_null_or_zero(rows):
|
| state.sql = rewritten
|
| state.rows = rows
|
| results[0] = ResultSet(title=p_title, sql=rewritten,
|
| rows=rows, table=table)
|
| state.confidence = max(v_conf + 0.2, 0.6)
|
| state.trace.append({"stage": "verify_retry", "error": err,
|
| "row_count": len(rows)})
|
| await emit({"stage": "verify_retry", "status": "done",
|
| "row_count": len(rows), "error": err, "ms": _ms()})
|
| except Exception:
|
| state.confidence = 0.5
|
|
|
|
|
| if len(results) < 3:
|
| await emit({"stage": "heuristic_topup", "status": "running",
|
| "label": "Adding heuristic companion views", "ms": _ms()})
|
| added = 0
|
| for ti, sq in heuristic_secondaries(q, table):
|
| if any(r.sql.strip() == sq.strip() for r in results):
|
| continue
|
| rewritten, rows, err = _safe_run(sq)
|
| if err or not rows or _is_all_null_or_zero(rows):
|
| continue
|
| results.append(ResultSet(title=ti, sql=rewritten, rows=rows,
|
| table=table))
|
| added += 1
|
| if len(results) >= 4:
|
| break
|
| await emit({"stage": "heuristic_topup", "status": "done",
|
| "added": added, "ms": _ms()})
|
|
|
|
|
| results = _filter_meaningful(results)
|
|
|
| if not results:
|
| state.trace.append({"stage": "llm_empty_fallback"})
|
| state.plan = dynamic_offline_plan(q)
|
| await emit({"stage": "fallback", "status": "done",
|
| "label": "No results — using offline dynamic plan", "ms": _ms()})
|
| return state
|
|
|
|
|
| await emit({"stage": "result_sets", "status": "done",
|
| "count": len(results),
|
| "previews": [
|
| {"title": r.title, "rows": _preview_rows(r.rows),
|
| "row_count": len(r.rows),
|
| "columns": list(r.rows[0].keys()) if r.rows else []}
|
| for r in results
|
| ],
|
| "ms": _ms()})
|
|
|
|
|
| await emit({"stage": "insights", "status": "running",
|
| "label": f"Generating per-result insights ({len(results)} parallel)",
|
| "role": "SUMMARY", "ms": _ms()})
|
| await emit({"stage": "narrative", "status": "running",
|
| "label": "Drafting executive narrative",
|
| "role": "NARRATIVE", "ms": _ms()})
|
| await emit({"stage": "followups", "status": "running",
|
| "label": "Generating smart follow-up questions",
|
| "role": "FOLLOWUP", "ms": _ms()})
|
|
|
| (insights_pair, narrative_pair, followup_pair) = await asyncio.gather(
|
| per_result_insights(q, [(r.title, r.rows) for r in results]),
|
| exec_narrative(q, [(r.title, r.rows) for r in results]),
|
| generate_followups(q, [(r.title, r.rows) for r in results], intent),
|
| )
|
| insights, k_insights = insights_pair
|
| narrative, k_nar = narrative_pair
|
| followups, k_fol = followup_pair
|
| for k in k_insights:
|
| _track(state, k)
|
| _track(state, k_nar)
|
| _track(state, k_fol)
|
| state.followups = followups
|
| state.trace.append({"stage": "insights", "count": len(insights),
|
| "keys": k_insights})
|
| state.trace.append({"stage": "narrative", "key": k_nar})
|
| state.trace.append({"stage": "followups", "key": k_fol,
|
| "count": len(followups)})
|
|
|
| await emit({"stage": "insights", "status": "done", "keys": k_insights,
|
| "count": len([i for i in insights if i]), "ms": _ms()})
|
| await emit({"stage": "narrative", "status": "done", "key": k_nar,
|
| "preview": (narrative or "")[:200], "ms": _ms()})
|
| await emit({"stage": "followups", "status": "done", "key": k_fol,
|
| "count": len(followups), "ms": _ms()})
|
|
|
| await emit({"stage": "synth", "status": "running",
|
| "label": "Assembling adaptive UI plan", "ms": _ms()})
|
| plan = assemble_plan(q, results, intent=intent)
|
| if narrative:
|
| plan.narrative = narrative
|
|
|
| if insights:
|
| tbl_comps = [c for c in plan.components if c.type == "data-table"]
|
| for comp, ins in zip(tbl_comps, insights):
|
| if ins:
|
| comp.why = ins
|
|
|
| state.plan = plan
|
| state.trace.append({
|
| "stage": "synth",
|
| "components": len(plan.components),
|
| "result_sets": len(results),
|
| "keys_used": state.keys_used,
|
| })
|
| await emit({"stage": "synth", "status": "done",
|
| "components": len(plan.components),
|
| "result_sets": len(results),
|
| "keys_used": state.keys_used, "ms": _ms()})
|
| return state
|
|
|
|
|
| async def run_agent(question: str, emit: EmitFn = _noop_emit) -> AgentState:
|
| state = AgentState(question=question)
|
| if not get_pool().available():
|
| state.mode = "offline"
|
| plan = dynamic_offline_plan(question)
|
| state.plan = plan
|
| state.intent = "descriptive"
|
| for c in plan.citations:
|
| if c.snippet:
|
| state.sql = c.snippet
|
| break
|
| state.trace.append({"stage": "offline_dynamic",
|
| "components": len(plan.components),
|
| "result_sets": len(plan.citations)})
|
| await emit({"stage": "offline", "status": "done",
|
| "label": "No keys available — offline plan",
|
| "components": len(plan.components)})
|
| return state
|
|
|
| state.mode = "groq"
|
| try:
|
| return await _run_llm_path(state, emit=emit)
|
| except Exception as e:
|
| state.trace.append({"stage": "llm_exception", "error": str(e)[:200]})
|
| state.plan = dynamic_offline_plan(question)
|
| state.intent = "descriptive"
|
| await emit({"stage": "exception", "status": "error",
|
| "error": str(e)[:200]})
|
| return state
|
|
|
