Add SocketIO service and initialize it

This commit is contained in:
Gabriel Luiz Freitas Almeida 2024-01-31 23:50:13 -03:00
commit 54c5053454
3 changed files with 17 additions and 5 deletions

View file

@ -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":

View file

@ -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"

View file

@ -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)