diff --git a/src/backend/base/langflow/main.py b/src/backend/base/langflow/main.py index 7aa67b335..697f74c95 100644 --- a/src/backend/base/langflow/main.py +++ b/src/backend/base/langflow/main.py @@ -99,7 +99,7 @@ def get_lifespan(fix_migration=False, socketio_server=None, version=None): raise # Shutdown message rprint("[bold red]Shutting down Langflow...[/bold red]") - teardown_services() + await teardown_services() return lifespan diff --git a/src/backend/base/langflow/services/base.py b/src/backend/base/langflow/services/base.py index 594d697bc..430ecd11e 100644 --- a/src/backend/base/langflow/services/base.py +++ b/src/backend/base/langflow/services/base.py @@ -22,7 +22,7 @@ class Service(ABC): } return schema - def teardown(self): + async def teardown(self): pass def set_ready(self): diff --git a/src/backend/base/langflow/services/database/service.py b/src/backend/base/langflow/services/database/service.py index 32e0b08e2..2fee292a4 100644 --- a/src/backend/base/langflow/services/database/service.py +++ b/src/backend/base/langflow/services/database/service.py @@ -294,7 +294,7 @@ class DatabaseService(Service): logger.debug("Database and tables created successfully") - def teardown(self): + async def teardown(self): logger.debug("Tearing down database") try: settings_service = get_settings_service() diff --git a/src/backend/base/langflow/services/manager.py b/src/backend/base/langflow/services/manager.py index cd27754d1..02bf95571 100644 --- a/src/backend/base/langflow/services/manager.py +++ b/src/backend/base/langflow/services/manager.py @@ -1,3 +1,4 @@ +import asyncio import importlib import inspect from typing import TYPE_CHECKING, Dict, Optional @@ -6,7 +7,6 @@ from loguru import logger if TYPE_CHECKING: from langflow.services.base import Service - from langflow.services.factory import ServiceFactory from langflow.services.schema import ServiceType @@ -93,7 +93,7 @@ class ServiceManager: self.services.pop(service_name, None) self.get(service_name) - def teardown(self): + async def teardown(self): """ Teardown all the services. """ @@ -102,7 +102,9 @@ class ServiceManager: continue logger.debug(f"Teardown service {service.name}") try: - service.teardown() + result = service.teardown() + if asyncio.iscoroutine(result): + await result except Exception as exc: logger.exception(exc) self.services = {} diff --git a/src/backend/base/langflow/services/plugins/service.py b/src/backend/base/langflow/services/plugins/service.py index 98f210fea..4b08eebd5 100644 --- a/src/backend/base/langflow/services/plugins/service.py +++ b/src/backend/base/langflow/services/plugins/service.py @@ -53,7 +53,7 @@ class PluginService(Service): return plugin.get() return None - def teardown(self): + async def teardown(self): for plugin in self.plugins.values(): plugin.teardown() diff --git a/src/backend/base/langflow/services/storage/local.py b/src/backend/base/langflow/services/storage/local.py index b4eb81e22..1c3a67daf 100644 --- a/src/backend/base/langflow/services/storage/local.py +++ b/src/backend/base/langflow/services/storage/local.py @@ -90,6 +90,6 @@ class LocalStorageService(StorageService): else: logger.warning(f"Attempted to delete non-existent file {file_name} in flow {flow_id}.") - def teardown(self): + async def teardown(self): """Perform any cleanup operations when the service is being torn down.""" pass # No specific teardown actions required for local diff --git a/src/backend/base/langflow/services/storage/service.py b/src/backend/base/langflow/services/storage/service.py index ac226bc31..15f51d65c 100644 --- a/src/backend/base/langflow/services/storage/service.py +++ b/src/backend/base/langflow/services/storage/service.py @@ -38,5 +38,5 @@ class StorageService(Service): async def delete_file(self, flow_id: str, file_name: str) -> bool: raise NotImplementedError - def teardown(self): + async def teardown(self): raise NotImplementedError diff --git a/src/backend/base/langflow/services/telemetry/service.py b/src/backend/base/langflow/services/telemetry/service.py index 53a46838a..7a26c9a1a 100644 --- a/src/backend/base/langflow/services/telemetry/service.py +++ b/src/backend/base/langflow/services/telemetry/service.py @@ -61,7 +61,8 @@ class TelemetryService(Service): if path: url = f"{url}/{path}" try: - response = await self.client.get(url, params=payload.model_dump()) + payload_dict = payload.model_dump(exclude_none=True, exclude_unset=True) + response = await self.client.get(url, params=payload_dict) if response.status_code != 200: logger.error(f"Failed to send telemetry data: {response.status_code} {response.text}") else: @@ -134,3 +135,6 @@ class TelemetryService(Service): await self.client.aclose() except Exception as e: logger.error(f"Error stopping tracing service: {e}") + + async def teardown(self): + await self.stop() diff --git a/src/backend/base/langflow/services/utils.py b/src/backend/base/langflow/services/utils.py index 06c45bd8d..558eaf028 100644 --- a/src/backend/base/langflow/services/utils.py +++ b/src/backend/base/langflow/services/utils.py @@ -109,7 +109,7 @@ def teardown_superuser(settings_service, session): raise RuntimeError("Could not remove default superuser.") from exc -def teardown_services(): +async def teardown_services(): """ Teardown all the services. """ @@ -120,7 +120,7 @@ def teardown_services(): try: from langflow.services.manager import service_manager - service_manager.teardown() + await service_manager.teardown() except Exception as exc: logger.exception(exc)