163 lines
6.1 KiB
Python
163 lines
6.1 KiB
Python
"""SSE event emitter — wraps sse-starlette to emit typed events."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import json
|
|
import uuid
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
|
|
from app.core.logging import get_logger
|
|
|
|
_logger = get_logger(__name__)
|
|
|
|
|
|
class SseEmitter:
|
|
"""Async queue-based SSE emitter.
|
|
|
|
Usage:
|
|
emitter = SseEmitter()
|
|
async with emitter.stream() as stream:
|
|
async for event in stream:
|
|
yield event
|
|
|
|
In a producer task:
|
|
await emitter.emit("tool_call", {...})
|
|
await emitter.done({"result": "ok"})
|
|
|
|
Production mode: when `debug=False`, raw `tool_call` and `llm_call_*`
|
|
events are filtered or transformed into friendly `status` events.
|
|
"""
|
|
|
|
def __init__(self, debug: bool = True) -> None:
|
|
self._queue: asyncio.Queue[tuple[str, str, str] | None] = asyncio.Queue()
|
|
self._event_counter = 0
|
|
self._closed = False
|
|
self._debug = debug
|
|
|
|
async def emit(self, event_type: str, data: Any) -> None:
|
|
if self._closed:
|
|
return
|
|
# In production mode, transform/fiter debug-only events
|
|
if not self._debug:
|
|
event_type, data = self._transform_for_prod(event_type, data)
|
|
if event_type is None:
|
|
return # event filtered out
|
|
self._event_counter += 1
|
|
event_id = f"evt_{self._event_counter}"
|
|
try:
|
|
data_str = json.dumps(data, ensure_ascii=False, default=str)
|
|
except (TypeError, ValueError):
|
|
data_str = json.dumps({"error": "serialization_failed"})
|
|
await self._queue.put((event_type, data_str, event_id))
|
|
|
|
def _transform_for_prod(self, event_type: str, data: Any) -> tuple[str | None, Any]:
|
|
"""Transform debug events into user-friendly status events for production."""
|
|
if event_type == "tool_call":
|
|
# Transform tool_call into a friendly status message
|
|
tool = data.get("tool", "") if isinstance(data, dict) else ""
|
|
result = data.get("result", {}) if isinstance(data, dict) else {}
|
|
is_success = data.get("is_success", True) if isinstance(data, dict) else True
|
|
# Friendly message based on tool type
|
|
friendly = self._friendly_tool_message(tool, result, is_success)
|
|
if friendly:
|
|
return ("status", {"message": friendly, "type": "tool"})
|
|
return (None, None) # filter out
|
|
elif event_type in ("llm_call_start", "llm_call_end"):
|
|
# Filter out raw LLM call events in production
|
|
return (None, None)
|
|
elif event_type == "phase_start":
|
|
# Keep but with friendly name
|
|
phase = data.get("phase") if isinstance(data, dict) else None
|
|
friendly_names = {
|
|
1: "planning",
|
|
2: "writing",
|
|
3: "sending",
|
|
}
|
|
name = friendly_names.get(phase, data.get("name", ""))
|
|
return ("phase_start", {"phase": phase, "name": name, "status": name})
|
|
elif event_type == "phase_end":
|
|
return (event_type, data)
|
|
elif event_type == "warning":
|
|
return (event_type, data)
|
|
else:
|
|
return (event_type, data)
|
|
|
|
def _friendly_tool_message(self, tool: str, result: dict, is_success: bool) -> str | None:
|
|
"""Generate a user-friendly message for a tool call."""
|
|
if not is_success:
|
|
return None # hide failed tool calls in production
|
|
data = result.get("data", {}) if isinstance(result, dict) else {}
|
|
msg = result.get("message", "") if isinstance(result, dict) else ""
|
|
if tool == "entity_create":
|
|
name = data.get("entity_id", "")
|
|
return f"Added new entity" + (f": {name}" if name else "")
|
|
elif tool == "entity_update":
|
|
return "Updated entity"
|
|
elif tool == "entity_delete":
|
|
return "Removed entity"
|
|
elif tool == "env_update":
|
|
return "Updated game state"
|
|
elif tool == "update_plot_rails":
|
|
return "Updated story progress"
|
|
elif tool == "advance_time":
|
|
new_time = data.get("new_time", "")
|
|
return f"Time advanced" + (f" to {new_time}" if new_time else "")
|
|
elif tool == "schedule_trigger":
|
|
return "Scheduled future event"
|
|
elif tool == "rag_query":
|
|
return None # hide RAG queries in production
|
|
elif tool == "rag_add":
|
|
return "Recorded a new fact"
|
|
elif tool == "calc":
|
|
return None # hide calculations
|
|
elif tool == "random_choice":
|
|
return None
|
|
elif tool == "submit_plan":
|
|
return None
|
|
elif tool == "submit_step":
|
|
return None
|
|
elif tool == "suggest_actions":
|
|
return None
|
|
return msg if msg else None
|
|
|
|
async def ping(self) -> None:
|
|
await self.emit("ping", {"ts": _now_iso()})
|
|
|
|
async def done(self, result: Any = None) -> None:
|
|
await self.emit("done", result if result is not None else {})
|
|
await self._queue.put(None) # sentinel
|
|
self._closed = True
|
|
|
|
async def error(self, code: str, message: str, details: Any = None) -> None:
|
|
payload: dict[str, Any] = {"code": code, "message": message}
|
|
if details is not None:
|
|
payload["details"] = details
|
|
await self.emit("error", payload)
|
|
await self._queue.put(None)
|
|
self._closed = True
|
|
|
|
async def stream(self) -> AsyncIterator[dict[str, str]]:
|
|
"""Yield SSE-formatted dicts until the emitter is closed."""
|
|
try:
|
|
while True:
|
|
item = await self._queue.get()
|
|
if item is None:
|
|
break
|
|
event_type, data_str, event_id = item
|
|
yield {
|
|
"event": event_type,
|
|
"data": data_str,
|
|
"id": event_id,
|
|
}
|
|
except asyncio.CancelledError:
|
|
_logger.info("sse_stream_cancelled")
|
|
raise
|
|
|
|
|
|
def _now_iso() -> str:
|
|
from datetime import datetime, timezone
|
|
|
|
return datetime.now(timezone.utc).isoformat()
|