From 54c5053454fecf796223112884b4f7f76a6c18b4 Mon Sep 17 00:00:00 2001
From: Gabriel Luiz Freitas Almeida <gabriel@logspace.ai>
Date: Wed, 31 Jan 2024 23:50:13 -0300
Subject: [PATCH] Add SocketIO service and initialize it

---
 src/backend/langflow/services/deps.py   |  8 ++++++--
 src/backend/langflow/services/schema.py |  1 +
 src/backend/langflow/services/utils.py  | 13 ++++++++++---
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/src/backend/langflow/services/deps.py b/src/backend/langflow/services/deps.py
index dea2ee461..19f3dcbf0 100644
--- a/src/backend/langflow/services/deps.py
+++ b/src/backend/langflow/services/deps.py
@@ -3,8 +3,6 @@ from typing import TYPE_CHECKING, Generator
 from langflow.services import ServiceType, service_manager
 
 if TYPE_CHECKING:
-    from sqlmodel import Session
-
     from langflow.services.cache.service import BaseCacheService
     from langflow.services.chat.service import ChatService
     from langflow.services.credentials.service import CredentialService
@@ -13,9 +11,15 @@ if TYPE_CHECKING:
     from langflow.services.plugins.service import PluginService
     from langflow.services.session.service import SessionService
     from langflow.services.settings.service import SettingsService
+    from langflow.services.socket.service import SocketIOService
     from langflow.services.storage.service import StorageService
     from langflow.services.store.service import StoreService
     from langflow.services.task.service import TaskService
+    from sqlmodel import Session
+
+
+def get_socket_service() -> "SocketIOService":
+    return service_manager.get(ServiceType.SOCKET_IO_SERVICE)  # type: ignore
 
 
 def get_storage_service() -> "StorageService":
diff --git a/src/backend/langflow/services/schema.py b/src/backend/langflow/services/schema.py
index 8a348f081..7e78fc178 100644
--- a/src/backend/langflow/services/schema.py
+++ b/src/backend/langflow/services/schema.py
@@ -19,3 +19,4 @@ class ServiceType(str, Enum):
     CREDENTIAL_SERVICE = "credential_service"
     STORAGE_SERVICE = "storage_service"
     MONITOR_SERVICE = "monitor_service"
+    SOCKET_IO_SERVICE = "socket_io_service"
diff --git a/src/backend/langflow/services/utils.py b/src/backend/langflow/services/utils.py
index b6e0350a0..a3cd6fd3e 100644
--- a/src/backend/langflow/services/utils.py
+++ b/src/backend/langflow/services/utils.py
@@ -1,10 +1,12 @@
+from loguru import logger
+from sqlmodel import Session, select
+
 from langflow.services.auth.utils import create_super_user, verify_password
 from langflow.services.database.utils import initialize_database
 from langflow.services.manager import service_manager
 from langflow.services.schema import ServiceType
 from langflow.services.settings.constants import DEFAULT_SUPERUSER, DEFAULT_SUPERUSER_PASSWORD
-from loguru import logger
-from sqlmodel import Session, select
+from langflow.services.socket.utils import set_socketio_server
 
 from .deps import get_db_service, get_session, get_settings_service
 
@@ -19,6 +21,7 @@ def get_factories_and_deps():
     from langflow.services.plugins import factory as plugins_factory
     from langflow.services.session import factory as session_service_factory  # type: ignore
     from langflow.services.settings import factory as settings_factory
+    from langflow.services.socket import factory as socket_factory
     from langflow.services.storage import factory as storage_factory
     from langflow.services.store import factory as store_factory
     from langflow.services.task import factory as task_factory
@@ -48,6 +51,7 @@ def get_factories_and_deps():
         (credentials_factory.CredentialServiceFactory(), [ServiceType.SETTINGS_SERVICE]),
         (storage_factory.StorageServiceFactory(), [ServiceType.SESSION_SERVICE, ServiceType.SETTINGS_SERVICE]),
         (monitor_factory.MonitorServiceFactory(), [ServiceType.SETTINGS_SERVICE]),
+        (socket_factory.SocketIOFactory(), [ServiceType.CACHE_SERVICE]),
     ]
 
 
@@ -189,7 +193,7 @@ def initialize_session_service():
     )
 
 
-def initialize_services(fix_migration: bool = False):
+def initialize_services(fix_migration: bool = False, socketio_server=None):
     """
     Initialize all the services needed.
     """
@@ -214,3 +218,6 @@ def initialize_services(fix_migration: bool = False):
     except Exception as exc:
         logger.error(f"Error migrating flows: {exc}")
         raise RuntimeError("Error migrating flows") from exc
+
+    # Initialize the SocketIO service
+    set_socketio_server(socketio_server)