local-deep-research 0.4.4__py3-none-any.whl → 0.5.2__py3-none-any.whl

local_deep_research/web/services/socket_service.py

@@ -1,210 +1,317 @@
-import logging
+from datetime import datetime
+from threading import Lock
+from typing import Any, NoReturn
 
-# Initialize logger
-logger = logging.getLogger(__name__)
+from flask import Flask, current_app, request
+from flask_socketio import SocketIO
+from loguru import logger
 
-# Make this a module variable to be set by the Flask app on initialization
-socketio = None
-# Socket subscription tracking
-socket_subscriptions = {}
+from ..models.database import get_db_connection
+from ..routes.globals import get_globals
 
 
-def set_socketio(socket_instance):
-    """Set the Socket.IO instance for the service."""
-    global socketio
-    socketio = socket_instance
-    logger.info("Socket.IO instance attached to socket service")
-
-
-def emit_socket_event(event, data, room=None):
+class SocketIOService:
     """
-    Emit a socket event to clients.
-
-    Args:
-        event: The event name to emit
-        data: The data to send with the event
-        room: Optional room ID to send to specific client
-
-    Returns:
-        bool: True if emission was successful, False otherwise
+    Singleton class for managing SocketIO connections and subscriptions.
     """
-    global socketio
-
-    if not socketio:
-        logger.error("Socket.IO not initialized when attempting to emit event")
-        return False
-
-    try:
-        # If room is specified, only emit to that room
-        if room:
-            socketio.emit(event, data, room=room)
-        else:
-            # Otherwise broadcast to all
-            socketio.emit(event, data)
-        return True
-    except Exception as e:
-        logger.error(f"Error emitting socket event {event}: {str(e)}")
-        return False
 
-
-def emit_to_subscribers(event_base, research_id, data):
-    """
-    Emit an event to all subscribers of a specific research.
-
-    Args:
-        event_base: Base event name (will be formatted with research_id)
-        research_id: ID of the research
-        data: The data to send with the event
-
-    Returns:
-        bool: True if emission was successful, False otherwise
-    """
-    global socketio
-
-    if not socketio:
-        logger.error("Socket.IO not initialized when attempting to emit to subscribers")
-        return False
-
-    try:
-        # Emit to the general channel for the research
-        full_event = f"{event_base}_{research_id}"
-        socketio.emit(full_event, data)
-
-        # Emit to specific subscribers
-        if research_id in socket_subscriptions and socket_subscriptions[research_id]:
-            for sid in socket_subscriptions[research_id]:
-                try:
-                    socketio.emit(full_event, data, room=sid)
-                except Exception as sub_err:
-                    logger.error(f"Error emitting to subscriber {sid}: {str(sub_err)}")
-
-        return True
-    except Exception as e:
-        logger.error(
-            f"Error emitting to subscribers for research {research_id}: {str(e)}"
+    _instance = None
+
+    def __new__(cls, *args: Any, app: Flask | None = None, **kwargs: Any):
+        """
+        Args:
+            app: The Flask app to bind this service to. It must be specified
+                the first time this is called and the singleton instance is
+                created, but will be ignored after that.
+            *args: Arguments to pass to the superclass's __new__ method.
+            **kwargs: Keyword arguments to pass to the superclass's __new__ method.
+        """
+        if not cls._instance:
+            if app is None:
+                raise ValueError(
+                    "Flask app must be specified to create a SocketIOService instance."
+                )
+            cls._instance = super(SocketIOService, cls).__new__(
+                cls, *args, **kwargs
+            )
+            cls._instance.__init_singleton(app)
+        return cls._instance
+
+    def __init_singleton(self, app: Flask) -> None:
+        """
+        Initializes the singleton instance.
+
+        Args:
+            app: The app to bind this service to.
+
+        """
+        self.__socketio = SocketIO(
+            app,
+            cors_allowed_origins="*",
+            async_mode="threading",
+            path="/research/socket.io",
+            logger=False,
+            engineio_logger=False,
+            ping_timeout=20,
+            ping_interval=5,
         )
-        return False
 
-
-# Socket event handlers moved from app.py
-def handle_connect(request):
-    """Handle client connection"""
-    logger.info(f"Client connected: {request.sid}")
-
-
-def handle_disconnect(request):
-    """Handle client disconnection"""
-    try:
-        logger.info(f"Client disconnected: {request.sid}")
-        # Clean up subscriptions for this client
-        global socket_subscriptions
-        for research_id, subscribers in list(socket_subscriptions.items()):
-            if request.sid in subscribers:
-                subscribers.remove(request.sid)
-            if not subscribers:
-                socket_subscriptions.pop(research_id, None)
-                logger.info(f"Removed empty subscription for research {research_id}")
-    except Exception as e:
-        logger.error(f"Error handling disconnect: {e}")
-
-
-def handle_subscribe(data, request, active_research=None):
-    """Handle client subscription to research updates"""
-    from datetime import datetime
-
-    from ..models.database import get_db_connection
-
-    research_id = data.get("research_id")
-    if research_id:
-        # First check if this research is still active
-        conn = get_db_connection()
-        cursor = conn.cursor()
-        cursor.execute(
-            "SELECT status FROM research_history WHERE id = ?", (research_id,)
-        )
-        result = cursor.fetchone()
-        conn.close()
-
-        # Only allow subscription to valid research
-        if result:
-            status = result[0]
-
-            # Initialize subscription set if needed
-            global socket_subscriptions
-            if research_id not in socket_subscriptions:
-                socket_subscriptions[research_id] = set()
-
-            # Add this client to the subscribers
-            socket_subscriptions[research_id].add(request.sid)
-            logger.info(f"Client {request.sid} subscribed to research {research_id}")
-
-            # Send current status immediately if available
-            if active_research and research_id in active_research:
-                progress = active_research[research_id]["progress"]
-                latest_log = (
-                    active_research[research_id]["log"][-1]
-                    if active_research[research_id]["log"]
-                    else None
+        # Socket subscription tracking.
+        self.__socket_subscriptions = {}
+        # Set to false to disable logging in the event handlers. This can
+        # be necessary because it will sometimes run the handlers directly
+        # during a call to `emit` that was made in a logging handler.
+        self.__logging_enabled = True
+        # Protects access to shared state.
+        self.__lock = Lock()
+
+        # Register events.
+        @self.__socketio.on("connect")
+        def on_connect():
+            self.__handle_connect(request)
+
+        @self.__socketio.on("disconnect")
+        def on_disconnect(reason: str):
+            self.__handle_disconnect(request, reason)
+
+        @self.__socketio.on("subscribe_to_research")
+        def on_subscribe(data):
+            globals_dict = get_globals()
+            active_research = globals_dict.get("active_research", {})
+            self.__handle_subscribe(data, request, active_research)
+
+        @self.__socketio.on_error
+        def on_error(e):
+            return self.__handle_socket_error(e)
+
+        @self.__socketio.on_error_default
+        def on_default_error(e):
+            return self.__handle_default_error(e)
+
+    def __log_info(self, message: str, *args: Any, **kwargs: Any) -> None:
+        """Log an info message."""
+        if self.__logging_enabled:
+            logger.info(message, *args, **kwargs)
+
+    def __log_error(self, message: str, *args: Any, **kwargs: Any) -> None:
+        """Log an error message."""
+        if self.__logging_enabled:
+            logger.error(message, *args, **kwargs)
+
+    def __log_exception(self, message: str, *args: Any, **kwargs: Any) -> None:
+        """Log an exception."""
+        if self.__logging_enabled:
+            logger.exception(message, *args, **kwargs)
+
+    def emit_socket_event(self, event, data, room=None):
+        """
+        Emit a socket event to clients.
+
+        Args:
+            event: The event name to emit
+            data: The data to send with the event
+            room: Optional room ID to send to specific client
+
+        Returns:
+            bool: True if emission was successful, False otherwise
+        """
+        try:
+            # If room is specified, only emit to that room
+            if room:
+                self.__socketio.emit(event, data, room=room)
+            else:
+                # Otherwise broadcast to all
+                self.__socketio.emit(event, data)
+            return True
+        except Exception as e:
+            logger.error(f"Error emitting socket event {event}: {str(e)}")
+            return False
+
+    def emit_to_subscribers(
+        self, event_base, research_id, data, enable_logging: bool = True
+    ):
+        """
+        Emit an event to all subscribers of a specific research.
+
+        Args:
+            event_base: Base event name (will be formatted with research_id)
+            research_id: ID of the research
+            data: The data to send with the event
+            enable_logging: If set to false, this will disable all logging,
+                which is useful if we are calling this inside of a logging
+                handler.
+
+        Returns:
+            bool: True if emission was successful, False otherwise
+
+        """
+        if not enable_logging:
+            self.__logging_enabled = False
+
+        try:
+            # Emit to the general channel for the research
+            full_event = f"{event_base}_{research_id}"
+            self.__socketio.emit(full_event, data)
+
+            # Emit to specific subscribers
+            with self.__lock:
+                subscriptions = self.__socket_subscriptions.get(research_id)
+            if subscriptions is not None:
+                for sid in subscriptions:
+                    try:
+                        self.__socketio.emit(full_event, data, room=sid)
+                    except Exception:
+                        self.__log_exception(
+                            f"Error emitting to subscriber {sid}"
+                        )
+
+            return True
+        except Exception:
+            self.__log_exception(
+                f"Error emitting to subscribers for research {research_id}"
+            )
+            return False
+        finally:
+            self.__logging_enabled = True
+
+    def __handle_connect(self, request):
+        """Handle client connection"""
+        self.__log_info(f"Client connected: {request.sid}")
+
+    def __handle_disconnect(self, request, reason: str):
+        """Handle client disconnection"""
+        try:
+            self.__log_info(
+                f"Client {request.sid} disconnected because: {reason}"
+            )
+            # Clean up subscriptions for this client
+            with self.__lock:
+                if request.sid in self.__socket_subscriptions:
+                    del self.__socket_subscriptions[request.sid]
+            self.__log_info(f"Removed subscription for client {request.sid}")
+        except Exception as e:
+            self.__log_error(f"Error handling disconnect: {e}")
+
+    def __handle_subscribe(self, data, request, active_research=None):
+        """Handle client subscription to research updates"""
+        research_id = data.get("research_id")
+        if research_id:
+            # First check if this research is still active
+            conn = get_db_connection()
+            cursor = conn.cursor()
+            cursor.execute(
+                "SELECT status FROM research_history WHERE id = ?",
+                (research_id,),
+            )
+            result = cursor.fetchone()
+            conn.close()
+
+            # Only allow subscription to valid research
+            if result:
+                status = result[0]
+
+                # Initialize subscription set if needed
+                with self.__lock:
+                    if research_id not in self.__socket_subscriptions:
+                        self.__socket_subscriptions[research_id] = set()
+
+                        # Add this client to the subscribers
+                        self.__socket_subscriptions[research_id].add(
+                            request.sid
+                        )
+                self.__log_info(
+                    f"Client {request.sid} subscribed to research {research_id}"
                 )
 
-                if latest_log:
-                    emit_socket_event(
+                # Send current status immediately if available
+                if active_research and research_id in active_research:
+                    progress = active_research[research_id]["progress"]
+                    latest_log = (
+                        active_research[research_id]["log"][-1]
+                        if active_research[research_id]["log"]
+                        else None
+                    )
+
+                    if latest_log:
+                        self.emit_socket_event(
+                            f"research_progress_{research_id}",
+                            {
+                                "progress": progress,
+                                "message": latest_log.get(
+                                    "message", "Processing..."
+                                ),
+                                "status": "in_progress",
+                                "log_entry": latest_log,
+                            },
+                            room=request.sid,
+                        )
+                elif status in ["completed", "failed", "suspended"]:
+                    # Send final status for completed research
+                    self.emit_socket_event(
                         f"research_progress_{research_id}",
                         {
-                            "progress": progress,
-                            "message": latest_log.get("message", "Processing..."),
-                            "status": "in_progress",
-                            "log_entry": latest_log,
-                        },
-                        room=request.sid,
-                    )
-            elif status in ["completed", "failed", "suspended"]:
-                # Send final status for completed research
-                emit_socket_event(
-                    f"research_progress_{research_id}",
-                    {
-                        "progress": 100 if status == "completed" else 0,
-                        "message": (
-                            "Research completed successfully"
-                            if status == "completed"
-                            else (
-                                "Research failed"
-                                if status == "failed"
-                                else "Research was suspended"
-                            )
-                        ),
-                        "status": status,
-                        "log_entry": {
-                            "time": datetime.utcnow().isoformat(),
-                            "message": f"Research is {status}",
                             "progress": 100 if status == "completed" else 0,
-                            "metadata": {
-                                "phase": (
-                                    "complete" if status == "completed" else "error"
+                            "message": (
+                                "Research completed successfully"
+                                if status == "completed"
+                                else (
+                                    "Research failed"
+                                    if status == "failed"
+                                    else "Research was suspended"
                                 )
+                            ),
+                            "status": status,
+                            "log_entry": {
+                                "time": datetime.utcnow().isoformat(),
+                                "message": f"Research is {status}",
+                                "progress": 100 if status == "completed" else 0,
+                                "metadata": {
+                                    "phase": (
+                                        "complete"
+                                        if status == "completed"
+                                        else "error"
+                                    )
+                                },
                             },
                         },
-                    },
+                        room=request.sid,
+                    )
+            else:
+                # Research not found
+                self.emit_socket_event(
+                    "error",
+                    {"message": f"Research ID {research_id} not found"},
                     room=request.sid,
                 )
-        else:
-            # Research not found
-            emit_socket_event(
-                "error",
-                {"message": f"Research ID {research_id} not found"},
-                room=request.sid,
-            )
 
+    def __handle_socket_error(self, e):
+        """Handle Socket.IO errors"""
+        self.__log_error(f"Socket.IO error: {str(e)}")
+        # Don't propagate exceptions to avoid crashing the server
+        return False
 
-def handle_socket_error(e):
-    """Handle Socket.IO errors"""
-    logger.error(f"Socket.IO error: {str(e)}")
-    # Don't propagate exceptions to avoid crashing the server
-    return False
-
+    def __handle_default_error(self, e):
+        """Handle unhandled Socket.IO errors"""
+        self.__log_error(f"Unhandled Socket.IO error: {str(e)}")
+        # Don't propagate exceptions to avoid crashing the server
+        return False
 
-def handle_default_error(e):
-    """Handle unhandled Socket.IO errors"""
-    logger.error(f"Unhandled Socket.IO error: {str(e)}")
-    # Don't propagate exceptions to avoid crashing the server
-    return False
+    def run(self, host: str, port: int, debug: bool = False) -> NoReturn:
+        """
+        Runs the SocketIO server.
+
+        Args:
+            host: The hostname to bind the server to.
+            port: The port number to listen on.
+            debug: Whether to run in debug mode. Defaults to False.
+
+        """
+        logger.info(f"Starting web server on {host}:{port} (debug: {debug})")
+        self.__socketio.run(
+            current_app,
+            debug=debug,
+            host=host,
+            port=port,
+            allow_unsafe_werkzeug=True,
+            use_reloader=False,
+        )