214 lines
6.4 KiB
Python
214 lines
6.4 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.
|
|
|
|
Many event types carry an `analysis` field that names the active Analysis
|
|
(read from `current_analysis` ContextVar). The UI uses it to group
|
|
sub-events under the right Analysis node without inferring from time order.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, IO
|
|
|
|
from api.analyses.types import Finding
|
|
from api.runtime.context import current_analysis, current_run_id
|
|
|
|
logger = logging.getLogger("nvi.runtime.events")
|
|
|
|
# run_id -> queue of events. Events are plain dicts.
|
|
_queues: dict[str, asyncio.Queue[dict[str, Any] | None]] = {}
|
|
|
|
# run_id -> append-mode JSONL log file. Mirrors every published event so the
|
|
# whole run is replayable from disk. Path: `<LOG_DIR>/<run_id>.jsonl`.
|
|
_log_files: dict[str, IO[str]] = {}
|
|
LOG_DIR = Path(".data/logs")
|
|
|
|
|
|
def log_path(run_id: str) -> Path:
|
|
"""Where this run's JSONL log lives. Resolved relative to cwd of the api
|
|
process — `/app/.data/logs/` in the pod."""
|
|
return LOG_DIR / f"{run_id}.jsonl"
|
|
|
|
|
|
# ── Queue management ──
|
|
|
|
def open_run(run_id: str) -> asyncio.Queue:
|
|
q: asyncio.Queue = asyncio.Queue()
|
|
_queues[run_id] = q
|
|
try:
|
|
LOG_DIR.mkdir(parents=True, exist_ok=True)
|
|
_log_files[run_id] = log_path(run_id).open("a", encoding="utf-8")
|
|
except OSError as e:
|
|
# Filesystem issues shouldn't kill the run; just lose the on-disk log.
|
|
logger.warning("could not open log file for run %s: %s", run_id, e)
|
|
return q
|
|
|
|
|
|
def _write_log(run_id: str, event: dict[str, Any]) -> None:
|
|
f = _log_files.get(run_id)
|
|
if f is None:
|
|
return
|
|
try:
|
|
record = {"t": time.time(), **event}
|
|
f.write(json.dumps(record, default=str) + "\n")
|
|
f.flush() # so `kubectl exec ... -- tail -f` shows live progress
|
|
except (OSError, ValueError) as e:
|
|
logger.warning("log write failed for run %s: %s", run_id, e)
|
|
|
|
|
|
async def publish(run_id: str, event: dict[str, Any]) -> None:
|
|
_write_log(run_id, event)
|
|
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)
|
|
f = _log_files.pop(run_id, None)
|
|
if f is not None:
|
|
try:
|
|
f.close()
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
# ── 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.
|
|
# Factories that fire from inside an Analysis pick up the analysis name from
|
|
# the `current_analysis` ContextVar so callers don't have to pass it.
|
|
|
|
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, *,
|
|
step_id: str | None = None) -> dict[str, Any]:
|
|
return {"type": "analysis_start", "analysis": name, "step_id": step_id,
|
|
"args": args, "why": why}
|
|
|
|
|
|
def analysis_end(name: str, finding: Finding, *,
|
|
step_id: str | None = None) -> dict[str, Any]:
|
|
return {
|
|
"type": "analysis_end",
|
|
"analysis": name,
|
|
"step_id": step_id,
|
|
"summary": finding.summary,
|
|
"error": finding.error,
|
|
"row_count": len(finding.rows),
|
|
}
|
|
|
|
|
|
def analysis_fallback(*, primary: str, fallback: str, reason: str,
|
|
step_id: str | None = None) -> dict[str, Any]:
|
|
return {
|
|
"type": "analysis_fallback",
|
|
"primary": primary,
|
|
"fallback": fallback,
|
|
"reason": reason,
|
|
"step_id": step_id,
|
|
}
|
|
|
|
|
|
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,
|
|
"analysis": current_analysis.get(),
|
|
"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,
|
|
"analysis": current_analysis.get(),
|
|
"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,
|
|
"analysis": current_analysis.get(),
|
|
"system_chars": system_len,
|
|
"user_chars": user_len,
|
|
}
|