| """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 |
|
|
| 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) |
|
|
|
|
| |
| _stream_bus: Optional[StreamBus] = None |
|
|
|
|
| def get_stream_bus() -> StreamBus: |
| global _stream_bus |
| if _stream_bus is None: |
| _stream_bus = StreamBus() |
| return _stream_bus |
|
|
