225 lines
8 KiB
Python
225 lines
8 KiB
Python
# Copyright (c) 2025 Patrick Motsch
|
|
# All rights reserved.
|
|
"""
|
|
Generic streaming event manager for real-time updates.
|
|
Manages event queues for SSE streaming across all features (chatbot, workflows, documents, etc.).
|
|
Supports event-driven streaming instead of polling.
|
|
"""
|
|
|
|
import logging
|
|
import asyncio
|
|
from typing import Dict, Optional, Any, List, AsyncIterator, Set
|
|
from datetime import datetime
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class StreamingEventManager:
|
|
"""
|
|
Generic event manager for real-time streaming across all features.
|
|
Supports multiple event types and contexts (workflows, documents, tasks, etc.).
|
|
Thread-safe event emission and queue management.
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""Initialize the event manager."""
|
|
self._queues: Dict[str, asyncio.Queue] = {}
|
|
self._locks: Dict[str, asyncio.Lock] = {}
|
|
self._cleanup_tasks: Dict[str, asyncio.Task] = {}
|
|
self._subscribers: Dict[str, Set[str]] = {} # context_id -> set of queue_ids (for future multi-subscriber support)
|
|
|
|
def create_queue(self, context_id: str) -> asyncio.Queue:
|
|
"""
|
|
Create a new event queue for a context.
|
|
|
|
Args:
|
|
context_id: Context ID (workflow_id, document_id, task_id, etc.)
|
|
|
|
Returns:
|
|
Event queue for the context
|
|
"""
|
|
if context_id not in self._queues:
|
|
self._queues[context_id] = asyncio.Queue()
|
|
self._locks[context_id] = asyncio.Lock()
|
|
self._subscribers[context_id] = set()
|
|
logger.debug(f"Created event queue for context {context_id}")
|
|
return self._queues[context_id]
|
|
|
|
def get_queue(self, context_id: str) -> Optional[asyncio.Queue]:
|
|
"""
|
|
Get existing event queue for a context.
|
|
|
|
Args:
|
|
context_id: Context ID
|
|
|
|
Returns:
|
|
Event queue if exists, None otherwise
|
|
"""
|
|
return self._queues.get(context_id)
|
|
|
|
async def emit_event(
|
|
self,
|
|
context_id: str,
|
|
event_type: str,
|
|
data: Dict[str, Any],
|
|
event_category: str = "default",
|
|
message: Optional[str] = None,
|
|
step: Optional[str] = None
|
|
):
|
|
"""
|
|
Emit an event to the context's event queue.
|
|
|
|
Args:
|
|
context_id: Context ID (workflow_id, document_id, etc.)
|
|
event_type: Type of event ("message", "log", "status", "progress", "complete", "error", "chatdata")
|
|
data: Event data dictionary (will be included in event)
|
|
event_category: Category of event for filtering ("chat", "workflow", "document", etc.)
|
|
message: Optional event message (for backward compatibility)
|
|
step: Optional processing step (for backward compatibility)
|
|
"""
|
|
queue = self.get_queue(context_id)
|
|
if not queue:
|
|
logger.debug(f"No event queue found for context {context_id}, skipping event")
|
|
return
|
|
|
|
event = {
|
|
"type": event_type,
|
|
"category": event_category,
|
|
"timestamp": datetime.now().timestamp(),
|
|
"data": data,
|
|
"message": message, # For backward compatibility
|
|
"step": step # For backward compatibility
|
|
}
|
|
|
|
try:
|
|
await queue.put(event)
|
|
logger.debug(f"Emitted {event_type} event (category: {event_category}) for context {context_id}")
|
|
except Exception as e:
|
|
logger.error(f"Error emitting event for context {context_id}: {e}")
|
|
|
|
async def stream_events(
|
|
self,
|
|
context_id: str,
|
|
event_categories: Optional[List[str]] = None,
|
|
timeout: Optional[float] = None
|
|
) -> AsyncIterator[Dict[str, Any]]:
|
|
"""
|
|
Async generator for streaming events from a context.
|
|
|
|
Args:
|
|
context_id: Context ID to stream events from
|
|
event_categories: Optional list of event categories to filter by
|
|
timeout: Optional timeout in seconds (None = no timeout)
|
|
|
|
Yields:
|
|
Event dictionaries
|
|
"""
|
|
queue = self.get_queue(context_id)
|
|
if not queue:
|
|
logger.warning(f"No queue found for context {context_id}")
|
|
return
|
|
|
|
start_time = asyncio.get_event_loop().time() if timeout else None
|
|
|
|
while True:
|
|
# Check timeout
|
|
if timeout and start_time:
|
|
elapsed = asyncio.get_event_loop().time() - start_time
|
|
if elapsed > timeout:
|
|
logger.debug(f"Stream timeout for context {context_id}")
|
|
break
|
|
|
|
try:
|
|
# Wait for event with timeout
|
|
wait_timeout = 1.0 # Check timeout every second
|
|
if timeout and start_time:
|
|
remaining = timeout - (asyncio.get_event_loop().time() - start_time)
|
|
if remaining <= 0:
|
|
break
|
|
wait_timeout = min(wait_timeout, remaining)
|
|
|
|
event = await asyncio.wait_for(queue.get(), timeout=wait_timeout)
|
|
|
|
# Filter by category if specified
|
|
if event_categories and event.get("category") not in event_categories:
|
|
continue
|
|
|
|
yield event
|
|
|
|
except asyncio.TimeoutError:
|
|
# Check if we should continue or timeout
|
|
if timeout and start_time:
|
|
elapsed = asyncio.get_event_loop().time() - start_time
|
|
if elapsed >= timeout:
|
|
break
|
|
continue
|
|
except Exception as e:
|
|
logger.error(f"Error in stream_events for context {context_id}: {e}")
|
|
break
|
|
|
|
async def cleanup(self, context_id: str, delay: float = 60.0):
|
|
"""
|
|
Schedule cleanup of event queue after delay.
|
|
This allows time for any remaining events to be consumed.
|
|
|
|
Args:
|
|
context_id: Context ID
|
|
delay: Delay in seconds before cleanup (default: 60 seconds)
|
|
"""
|
|
if context_id in self._cleanup_tasks:
|
|
# Cancel existing cleanup task
|
|
self._cleanup_tasks[context_id].cancel()
|
|
|
|
async def _cleanup():
|
|
try:
|
|
await asyncio.sleep(delay)
|
|
if context_id in self._queues:
|
|
# Drain remaining events
|
|
queue = self._queues[context_id]
|
|
while not queue.empty():
|
|
try:
|
|
queue.get_nowait()
|
|
except asyncio.QueueEmpty:
|
|
break
|
|
|
|
del self._queues[context_id]
|
|
if context_id in self._locks:
|
|
del self._locks[context_id]
|
|
if context_id in self._subscribers:
|
|
del self._subscribers[context_id]
|
|
logger.info(f"Cleaned up event queue for context {context_id}")
|
|
except asyncio.CancelledError:
|
|
pass
|
|
except Exception as e:
|
|
logger.error(f"Error during cleanup for context {context_id}: {e}")
|
|
finally:
|
|
if context_id in self._cleanup_tasks:
|
|
del self._cleanup_tasks[context_id]
|
|
|
|
self._cleanup_tasks[context_id] = asyncio.create_task(_cleanup())
|
|
|
|
def has_queue(self, context_id: str) -> bool:
|
|
"""
|
|
Check if a queue exists for a context.
|
|
|
|
Args:
|
|
context_id: Context ID
|
|
|
|
Returns:
|
|
True if queue exists, False otherwise
|
|
"""
|
|
return context_id in self._queues
|
|
|
|
|
|
# Backward compatibility: ChatbotEventManager is an alias
|
|
ChatbotEventManager = StreamingEventManager
|
|
|
|
# Global singleton instance
|
|
_event_manager = StreamingEventManager()
|
|
|
|
|
|
def get_event_manager() -> StreamingEventManager:
|
|
"""Get the global event manager instance."""
|
|
return _event_manager
|
|
|
|
|