Class ServerSentEventsHandler

rag_system/api_server_with_progress.py:29–61

Handler for Server-Sent Events (SSE) for real-time progress updates

print("✅ RAG Agent initialized successfully.")

class ServerSentEventsHandler:
    """Handler for Server-Sent Events (SSE) for real-time progress updates"""
    
    active_connections: Dict[str, Any] = {}
    
    @classmethod
    def add_connection(cls, session_id: str, response_handler):
        """Add a new SSE connection"""
        cls.active_connections[session_id] = response_handler
        logger.info(f"SSE connection added for session: {session_id}")
    
    @classmethod
    def remove_connection(cls, session_id: str):
        """Remove an SSE connection"""
        if session_id in cls.active_connections:
            del cls.active_connections[session_id]
            logger.info(f"SSE connection removed for session: {session_id}")
    
    @classmethod
    def send_event(cls, session_id: str, event_type: str, data: Dict[str, Any]):
        """Send an SSE event to a specific session"""
        if session_id not in cls.active_connections:
            return
        
        try:
            handler = cls.active_connections[session_id]
            event_data = json.dumps(data)
            message = f"event: {event_type}\ndata: {event_data}\n\n"
            handler.wfile.write(message.encode('utf-8'))
            handler.wfile.flush()
        except Exception as e:
            logger.error(f"Failed to send SSE event: {e}")
            cls.remove_connection(session_id)

class RealtimeProgressTracker(ProgressTracker):
    """Enhanced ProgressTracker that sends updates via Server-Sent Events"""