File size: 2,175 Bytes
6b62834 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | """StreamBus — internal pub/sub event bus for streaming events to clients."""
import asyncio
from typing import AsyncIterator, Optional
from agentic_rag.data.models import AgentEvent
class StreamBus:
"""Async pub/sub event bus for agent streaming.
Agents publish events to topics (keyed by turn_id).
Entry points subscribe to relay events to SSE/WebSocket clients.
"""
def __init__(self):
self._subscribers: dict[str, list[asyncio.Queue]] = {}
def publish(self, turn_id: str, event: AgentEvent) -> None:
"""Publish an event to all subscribers of a turn."""
event.turn_id = turn_id
queues = self._subscribers.get(turn_id, [])
for q in queues:
try:
q.put_nowait(event)
except asyncio.QueueFull:
pass # Drop event if subscriber is too slow
async def subscribe(self, turn_id: str, buffer_size: int = 100) -> AsyncIterator[AgentEvent]:
"""Subscribe to events for a turn. Yields events as they arrive."""
queue: asyncio.Queue = asyncio.Queue(maxsize=buffer_size)
if turn_id not in self._subscribers:
self._subscribers[turn_id] = []
self._subscribers[turn_id].append(queue)
try:
while True:
try:
event = await asyncio.wait_for(queue.get(), timeout=60.0)
yield event
if event.event_type.value == "done":
break
except asyncio.TimeoutError:
yield AgentEvent(event_type="error", data={"message": "Stream timeout"})
break
finally:
self._subscribers[turn_id].remove(queue)
if not self._subscribers[turn_id]:
del self._subscribers[turn_id]
def cleanup(self, turn_id: str) -> None:
"""Remove all subscribers for a turn."""
self._subscribers.pop(turn_id, None)
# Global instance
_stream_bus: Optional[StreamBus] = None
def get_stream_bus() -> StreamBus:
global _stream_bus
if _stream_bus is None:
_stream_bus = StreamBus()
return _stream_bus
|