134 lines
4.0 KiB
Python
134 lines
4.0 KiB
Python
"""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}
|