verbose live UI + tool-level SSE events + Groq default + regression tests

api/runtime/context.py

@@ -0,0 +1,12 @@
+"""Per-run context — exposes the active run_id to code deep in the call stack
+without threading it through every signature.
+
+Set by the runtime at run entry; read by Analyses and helper functions so they
+can publish trace events without needing the run_id passed in explicitly.
+"""
+
+from __future__ import annotations
+
+from contextvars import ContextVar
+
+current_run_id: ContextVar[str | None] = ContextVar("nvi.current_run_id", default=None)

api/runtime/events.py

@@ -0,0 +1,133 @@
+"""Per-run event bus + factories + SSE formatting.
+
+Tiny in-memory queue keyed by run_id. The runtime publishes on each
+transition; the FastAPI SSE endpoint subscribes and forwards.
+
+Event payloads are constructed via the module-level factory functions below
+so the dict shape lives in one place and is grep-able.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import json
+from typing import Any
+
+from api.analyses.types import Finding
+from api.runtime.context import current_run_id
+
+# run_id -> queue of events. Events are plain dicts.
+_queues: dict[str, asyncio.Queue[dict[str, Any] | None]] = {}
+
+
+# ── Queue management ──
+
+def open_run(run_id: str) -> asyncio.Queue:
+    q: asyncio.Queue = asyncio.Queue()
+    _queues[run_id] = q
+    return q
+
+
+async def publish(run_id: str, event: dict[str, Any]) -> None:
+    q = _queues.get(run_id)
+    if q is not None:
+        await q.put(event)
+
+
+async def publish_current(event: dict[str, Any]) -> None:
+    """Publish to the run identified by the `current_run_id` contextvar.
+    No-op when called outside a run."""
+    rid = current_run_id.get()
+    if rid is not None:
+        await publish(rid, event)
+
+
+async def close(run_id: str) -> None:
+    q = _queues.get(run_id)
+    if q is not None:
+        await q.put(None)
+
+
+def drop(run_id: str) -> None:
+    _queues.pop(run_id, None)
+
+
+# ── SSE formatting + streaming ──
+
+def sse_format(event: dict[str, Any]) -> str:
+    kind = event.get("type", "message")
+    payload = json.dumps({k: v for k, v in event.items() if k != "type"}, default=str)
+    return f"event: {kind}\ndata: {payload}\n\n"
+
+
+async def stream(run_id: str):
+    """Async generator producing SSE-formatted strings until the run ends."""
+    q = _queues.get(run_id)
+    if q is None:
+        yield sse_format({"type": "error", "message": f"unknown run_id {run_id}"})
+        return
+    try:
+        while True:
+            event = await q.get()
+            if event is None:
+                yield sse_format({"type": "end"})
+                return
+            yield sse_format(event)
+    finally:
+        drop(run_id)
+
+
+# ── Event factories ──
+# Each returns the dict payload to pass into `publish()` /
+# `publish_current()`. Names mirror the `type` field for grep-ability.
+
+def run_start(question: str) -> dict[str, Any]:
+    return {"type": "run_start", "question": question}
+
+
+def run_end(*, answer: str | None = None, error: str | None = None) -> dict[str, Any]:
+    return {"type": "run_end", "answer": answer or "", "error": error}
+
+
+def plan_ready(rationale: str, steps: list[dict[str, Any]]) -> dict[str, Any]:
+    return {"type": "plan_ready", "rationale": rationale, "steps": steps}
+
+
+def node_update(node: str, update: Any) -> dict[str, Any]:
+    return {
+        "type": "node_update",
+        "node": node,
+        "keys": sorted(update.keys()) if isinstance(update, dict) else [],
+    }
+
+
+def analysis_start(name: str, args: dict[str, Any], why: str) -> dict[str, Any]:
+    return {"type": "analysis_start", "analysis": name, "args": args, "why": why}
+
+
+def analysis_end(name: str, finding: Finding) -> dict[str, Any]:
+    return {
+        "type": "analysis_end",
+        "analysis": name,
+        "summary": finding.summary,
+        "error": finding.error,
+        "row_count": len(finding.rows),
+    }
+
+
+def synth_start() -> dict[str, Any]:
+    return {"type": "synth_start"}
+
+
+def tool_call_start(tool: str, *, input: Any = None) -> dict[str, Any]:
+    return {"type": "tool_call_start", "tool": tool, "input": input}
+
+
+def tool_call_end(tool: str, *, output: Any = None, error: str | None = None) -> dict[str, Any]:
+    return {"type": "tool_call_end", "tool": tool, "output": output, "error": error}
+
+
+def llm_call(label: str, *, system_len: int, user_len: int) -> dict[str, Any]:
+    """Lightweight breadcrumb for an LLM call. Doesn't carry the prompt
+    contents (those go through Langfuse) — just the fact and rough size."""
+    return {"type": "llm_call", "label": label, "system_chars": system_len, "user_chars": user_len}

api/runtime/runner.py

@@ -0,0 +1,139 @@
+"""langgraph wiring — the one place the framework appears.
+
+Graph: START → plan → execute → synthesize → END
+
+`execute` iterates Plan steps and dispatches each to its Analysis. Sequential
+for now; swap to a fan-out with parallel nodes once we want to parallelise.
+"""
+
+from __future__ import annotations
+
+import dataclasses
+import logging
+from typing import Any
+
+from langgraph.graph import END, START, StateGraph
+
+from api import langfuse_client as lf
+from api.analyses.registry import REGISTRY as ANALYSES
+from api.analyses.types import Finding
+from api.llm import chat
+from api.plan.planner import plan as run_planner
+from api.prompts import load, render
+from api.runtime import events
+from api.runtime.context import current_run_id
+from api.runtime.state import RunState
+
+logger = logging.getLogger("nvi.runtime")
+
+
+def _plan_node(state: RunState) -> dict[str, Any]:
+    plan_obj = run_planner(state["question"])
+    return {
+        "plan_rationale": plan_obj.rationale,
+        "plan_steps": [s.to_dict() for s in plan_obj.steps],
+    }
+
+
+async def _execute_node(state: RunState) -> dict[str, Any]:
+    findings: list[dict[str, Any]] = []
+    for step in state.get("plan_steps", []):
+        name = step["analysis"]
+        analysis = ANALYSES.get(name)
+        await events.publish(
+            state["run_id"],
+            events.analysis_start(name, step.get("args", {}), step.get("why", "")),
+        )
+        if analysis is None:
+            f = Finding(analysis=name, summary="", error=f"unknown analysis {name!r}")
+        else:
+            try:
+                f = await analysis.run(step.get("args", {}), state["question"])
+            except Exception as e:
+                logger.exception("analysis %s raised", name)
+                f = Finding(analysis=name, summary="", error=str(e))
+        findings.append(dataclasses.asdict(f))
+        await events.publish(state["run_id"], events.analysis_end(name, f))
+    return {"findings": findings}
+
+
+async def _synthesize_node(state: RunState) -> dict[str, Any]:
+    await events.publish(state["run_id"], events.synth_start())
+
+    findings = state.get("findings", [])
+    if not findings or all(f.get("error") for f in findings):
+        return {
+            "answer": "I couldn't answer this question — see the per-analysis errors above.",
+            "error": "all_analyses_failed",
+        }
+
+    findings_block = "\n\n".join(
+        f"[{f['analysis']}]\n"
+        f"summary: {f.get('summary') or '(none)'}\n"
+        f"error: {f.get('error') or '(none)'}\n"
+        f"sql: {f.get('sql')}"
+        for f in findings
+    )
+
+    with lf.span("synthesize", as_type="generation", input={"findings": len(findings)}) as span:
+        answer = chat(
+            system=load("synthesize.system"),
+            user=render(
+                "synthesize.user",
+                question=state["question"],
+                findings_block=findings_block,
+            ),
+            max_tokens=512,
+        )
+        span.update(output={"answer": answer})
+    return {"answer": answer}
+
+
+def build_graph():
+    g: StateGraph = StateGraph(RunState)
+    g.add_node("plan", _plan_node)
+    g.add_node("execute", _execute_node)
+    g.add_node("synthesize", _synthesize_node)
+    g.add_edge(START, "plan")
+    g.add_edge("plan", "execute")
+    g.add_edge("execute", "synthesize")
+    g.add_edge("synthesize", END)
+    return g.compile()
+
+
+GRAPH = build_graph()
+
+
+async def run(run_id: str, question: str) -> RunState:
+    """Drive the graph end-to-end, publishing SSE events along the way."""
+    initial: RunState = {"run_id": run_id, "question": question}
+    final: RunState = initial
+    token = current_run_id.set(run_id)
+    try:
+        await events.publish(run_id, events.run_start(question))
+        async for event in GRAPH.astream(initial, stream_mode="updates"):
+            for node, update in event.items():
+                await events.publish(run_id, events.node_update(node, update))
+                if isinstance(update, dict):
+                    final = {**final, **update}
+                if node == "plan" and isinstance(update, dict):
+                    await events.publish(
+                        run_id,
+                        events.plan_ready(
+                            update.get("plan_rationale", ""),
+                            update.get("plan_steps", []),
+                        ),
+                    )
+        await events.publish(
+            run_id,
+            events.run_end(answer=final.get("answer", ""), error=final.get("error")),
+        )
+    except Exception as e:
+        logger.exception("run %s failed", run_id)
+        await events.publish(run_id, events.run_end(error=str(e)))
+        final = {**final, "error": str(e)}
+    finally:
+        current_run_id.reset(token)
+        await events.close(run_id)
+        lf.flush()
+    return final

api/runtime/state.py

@@ -0,0 +1,20 @@
+"""Shared state passed between runtime nodes.
+
+This is the one type langgraph sees. Domain code in api/plan, api/analyses,
+api/tools never imports it — they take/return their own dataclasses.
+"""
+
+from __future__ import annotations
+
+from typing import Annotated, Any, TypedDict
+from operator import add
+
+
+class RunState(TypedDict, total=False):
+    run_id: str
+    question: str
+    plan_rationale: str
+    plan_steps: list[dict[str, Any]]      # serialised PlanStep
+    findings: Annotated[list[dict[str, Any]], add]   # serialised Finding; concatenated by langgraph
+    answer: str
+    error: str