"""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()