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

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}