apps/api/mbf/graph/agent.py

"""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

# Type for stage event emitter: async function(event_dict) -> None
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


# --------------------------------------------------------------------------- #
# Result-quality helpers                                                      #
# --------------------------------------------------------------------------- #

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
            # strings that look like "0" or "0.00" — common from VARCHARs
            elif isinstance(v, str):
                try:
                    fv = float(v)
                    found_any_numeric = True
                    if fv != 0:
                        return False
                except (ValueError, OverflowError):
                    pass  # genuine text label — skip
    # If there were no numeric columns at all, treat the result as meaningful
    # (it's probably a label/lookup query).
    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)

    # Stages 1 + 2 run in parallel — intent and retrieval are independent.
    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)]

    # Stage 3 — SQL plan
    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()})

    # Stage 4 — SQL critic runs CONCURRENTLY with primary execution.
    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)})

        # Apply critic fix if execution failed OR if results are all-null/zero
        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)

    # Stage 4b — SQL verifier (second opinion, runs concurrently w/ secondaries)
    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)
        )

    # Stage 3b — secondaries fan-out in parallel DuckDB executions
    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()})

    # Collect verifier result (ran concurrently with secondaries)
    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 verifier says results are wrong and provides a fix, re-run
            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:  # noqa: BLE001
            state.confidence = 0.5

    # top-up with heuristic companion views if still sparse
    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()})

    # Drop result sets whose numeric data is entirely null/zero
    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

    # Emit a preview of result sets so the UI can show data right away
    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()})

    # Stages 5 + 6 + 7 in parallel — insights, narrative, AND follow-ups
    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
    # Attach per-slice insights as the `why` on each data-table component
    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:  # noqa: BLE001
        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