digitalkin 0.2.25rc0__py3-none-any.whl → 0.3.1.dev2__py3-none-any.whl

digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.2.25rc0"
+    __version__ = "0.3.1.dev2"

digitalkin/core/__init__.py

@@ -0,0 +1 @@
+"""Core of Digitlakin defining the task management and sub-modules."""

digitalkin/core/common/__init__.py

@@ -0,0 +1,9 @@
+"""Common utilities for the core module."""
+
+from digitalkin.core.common.factories import ConnectionFactory, ModuleFactory, QueueFactory
+
+__all__ = [
+    "ConnectionFactory",
+    "ModuleFactory",
+    "QueueFactory",
+]

digitalkin/core/common/factories.py

@@ -0,0 +1,156 @@
+"""Common factory functions for reducing code duplication in core module."""
+
+import asyncio
+import datetime
+
+from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
+from digitalkin.logger import logger
+from digitalkin.modules._base_module import BaseModule
+
+
+class ConnectionFactory:
+    """Factory for creating SurrealDB connections with consistent configuration."""
+
+    @staticmethod
+    async def create_surreal_connection(
+        database: str = "task_manager",
+        timeout: datetime.timedelta = datetime.timedelta(seconds=5),
+        *,
+        auto_init: bool = True,
+    ) -> SurrealDBConnection:
+        """Create and optionally initialize a SurrealDB connection.
+
+        This factory method centralizes the creation of SurrealDB connections
+        to ensure consistent configuration across the codebase.
+
+        Args:
+            database: Database name to connect to
+            timeout: Connection timeout
+            auto_init: Whether to automatically initialize the connection
+
+        Returns:
+            Initialized or uninitialized SurrealDBConnection instance
+
+        Example:
+            # Create and auto-initialize
+            conn = await ConnectionFactory.create_surreal_connection("taskiq_worker")
+
+            # Create without initialization
+            conn = await ConnectionFactory.create_surreal_connection(auto_init=False)
+            await conn.init_surreal_instance()
+        """
+        logger.debug(
+            "Creating SurrealDB connection for database: %s, timeout: %s",
+            database,
+            timeout,
+            extra={"database": database, "timeout": str(timeout)},
+        )
+
+        connection: SurrealDBConnection = SurrealDBConnection(database, timeout)
+
+        if auto_init:
+            await connection.init_surreal_instance()
+            logger.debug("SurrealDB connection initialized for database: %s", database)
+
+        return connection
+
+
+class ModuleFactory:
+    """Factory for creating module instances with consistent configuration."""
+
+    @staticmethod
+    def create_module_instance(
+        module_class: type[BaseModule],
+        job_id: str,
+        mission_id: str,
+        setup_id: str,
+        setup_version_id: str,
+    ) -> BaseModule:
+        """Create a module instance with standard parameters.
+
+        This factory method centralizes module instantiation to ensure
+        consistent parameter passing across the codebase.
+
+        Args:
+            module_class: The module class to instantiate
+            job_id: Unique job identifier
+            mission_id: Mission identifier
+            setup_id: Setup identifier
+            setup_version_id: Setup version identifier
+
+        Returns:
+            Instantiated module
+
+        Raises:
+            ValueError: If job_id or mission_id is empty
+
+        Example:
+            module = ModuleFactory.create_module_instance(
+                MyModule,
+                job_id="job_123",
+                mission_id="mission:test",
+                setup_id="setup:config",
+                setup_version_id="v1.0",
+            )
+        """
+        # Validate parameters
+        if not job_id:
+            msg = "job_id cannot be empty"
+            raise ValueError(msg)
+        if not mission_id:
+            msg = "mission_id cannot be empty"
+            raise ValueError(msg)
+
+        logger.debug(
+            "Creating module instance: %s for job: %s",
+            module_class.__name__,
+            job_id,
+            extra={
+                "module_class": module_class.__name__,
+                "job_id": job_id,
+                "mission_id": mission_id,
+                "setup_id": setup_id,
+                "setup_version_id": setup_version_id,
+            },
+        )
+
+        return module_class(
+            job_id=job_id,
+            mission_id=mission_id,
+            setup_id=setup_id,
+            setup_version_id=setup_version_id,
+        )
+
+
+class QueueFactory:
+    """Factory for creating asyncio queues with consistent configuration."""
+
+    # Default max queue size to prevent unbounded memory growth
+    DEFAULT_MAX_QUEUE_SIZE = 1000
+
+    @staticmethod
+    def create_bounded_queue(maxsize: int = DEFAULT_MAX_QUEUE_SIZE) -> asyncio.Queue:
+        """Create a bounded asyncio queue with standard configuration.
+
+        Args:
+            maxsize: Maximum queue size (default 1000, 0 means unlimited)
+
+        Returns:
+            Bounded asyncio.Queue instance
+
+        Raises:
+            ValueError: If maxsize is negative
+
+        Example:
+            queue = QueueFactory.create_bounded_queue()
+            # or with custom size
+            queue = QueueFactory.create_bounded_queue(maxsize=500)
+            # unlimited queue
+            queue = QueueFactory.create_bounded_queue(maxsize=0)
+        """
+        if maxsize < 0:
+            msg = "maxsize must be >= 0"
+            raise ValueError(msg)
+
+        logger.debug("Creating bounded queue with maxsize: %d", maxsize, extra={"maxsize": maxsize})
+        return asyncio.Queue(maxsize=maxsize)

digitalkin/core/job_manager/__init__.py

@@ -0,0 +1 @@
+"""Job Manager logic."""

digitalkin/{modules → core}/job_manager/base_job_manager.py

@@ -5,17 +5,126 @@ from collections.abc import AsyncGenerator, AsyncIterator, Callable, Coroutine
 from contextlib import asynccontextmanager
 from typing import Any, Generic
 
-from digitalkin.models import ModuleStatus
+from digitalkin.core.task_manager.base_task_manager import BaseTaskManager
+from digitalkin.core.task_manager.task_session import TaskSession
+from digitalkin.models.core.task_monitor import TaskStatus
 from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
+from digitalkin.models.module.module import ModuleCodeModel
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
 
 
-class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT]):
-    """Abstract base class for managing background module jobs."""
+class BaseJobManager(abc.ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
+    """Abstract base class for managing background module jobs.
 
-    async def _start(self) -> None:
+    Uses composition to delegate task lifecycle management to a TaskManager.
+    """
+
+    module_class: type[BaseModule]
+    services_mode: ServicesMode
+    _task_manager: BaseTaskManager
+
+    def __init__(
+        self,
+        module_class: type[BaseModule],
+        services_mode: ServicesMode,
+        task_manager: BaseTaskManager,
+    ) -> None:
+        """Initialize the job manager.
+
+        Args:
+            module_class: The class of the module to be managed.
+            services_mode: The mode of operation for the services (e.g., ASYNC or SYNC).
+            task_manager: The task manager instance to use for task lifecycle management.
+        """
+        self.module_class = module_class
+        self.services_mode = services_mode
+        self._task_manager = task_manager
+
+        services_config = ServicesConfig(
+            services_config_strategies=self.module_class.services_config_strategies,
+            services_config_params=self.module_class.services_config_params,
+            mode=services_mode,
+        )
+        setattr(self.module_class, "services_config", services_config)
+
+    # Properties to expose task manager attributes
+    @property
+    def tasks_sessions(self) -> dict[str, TaskSession]:
+        """Get task sessions from the task manager."""
+        return self._task_manager.tasks_sessions
+
+    @property
+    def tasks(self) -> dict[str, Any]:
+        """Get tasks from the task manager."""
+        return self._task_manager.tasks
+
+    # Delegate task lifecycle methods to task manager
+    async def create_task(
+        self,
+        task_id: str,
+        mission_id: str,
+        module: BaseModule,
+        coro: Coroutine[Any, Any, None],
+        **kwargs: Any,  # noqa: ANN401
+    ) -> None:
+        """Create a task using the task manager.
+
+        Args:
+            task_id: Unique identifier for the task
+            mission_id: Mission identifier
+            module: Module instance
+            coro: Coroutine to execute
+            **kwargs: Additional arguments for task creation
+        """
+        await self._task_manager.create_task(task_id, mission_id, module, coro, **kwargs)
+
+    async def clean_session(self, task_id: str, mission_id: str) -> bool:
+        """Clean a task's session.
+
+        Args:
+            task_id: Unique identifier for the task.
+            mission_id: Mission identifier.
+
+        Returns:
+            bool: True if the task was successfully cancelled, False otherwise.
+        """
+        return await self._task_manager.clean_session(task_id, mission_id)
+
+    async def cancel_task(self, task_id: str, mission_id: str, timeout: float | None = None) -> bool:
+        """Cancel a task.
+
+        Args:
+            task_id: Unique identifier for the task.
+            mission_id: Mission identifier.
+            timeout: Optional timeout in seconds to wait for the cancellation to complete.
+
+        Returns:
+            bool: True if the task was successfully cancelled, False otherwise.
+        """
+        return await self._task_manager.cancel_task(task_id, mission_id, timeout)
+
+    async def send_signal(self, task_id: str, mission_id: str, signal_type: str, payload: dict) -> bool:
+        """Send signal to a task.
+
+        Args:
+            task_id: Unique identifier for the task.
+            mission_id: Mission identifier.
+            signal_type: Type of signal to send.
+            payload: Payload data for the signal.
+
+        Returns:
+            bool: True if the signal was successfully sent, False otherwise.
+        """
+        return await self._task_manager.send_signal(task_id, mission_id, signal_type, payload)
+
+    async def shutdown(self, mission_id: str, timeout: float = 30.0) -> None:
+        """Shutdown all tasks."""
+        await self._task_manager.shutdown(mission_id, timeout)
+
+    @abc.abstractmethod
+    async def start(self) -> None:
         """Start the job manager.
 
         This method initializes any necessary resources or configurations
@@ -24,8 +133,9 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT]):
 
     @staticmethod
     async def job_specific_callback(
-        callback: Callable[[str, OutputModelT], Coroutine[Any, Any, None]], job_id: str
-    ) -> Callable[[OutputModelT], Coroutine[Any, Any, None]]:
+        callback: Callable[[str, OutputModelT | ModuleCodeModel], Coroutine[Any, Any, None]],
+        job_id: str,
+    ) -> Callable[[OutputModelT | ModuleCodeModel], Coroutine[Any, Any, None]]:
         """Generate a job-specific callback function.
 
         Args:
@@ -36,7 +146,7 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT]):
             Callable: A wrapped callback function that includes the job ID.
         """
 
-        def callback_wrapper(output_data: OutputModelT) -> Coroutine[Any, Any, None]:
+        def callback_wrapper(output_data: OutputModelT | ModuleCodeModel) -> Coroutine[Any, Any, None]:
             """Wrapper for the callback function.
 
             Args:
@@ -49,26 +159,6 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT]):
 
         return callback_wrapper
 
-    def __init__(
-        self,
-        module_class: type[BaseModule],
-        services_mode: ServicesMode,
-    ) -> None:
-        """Initialize the job manager.
-
-        Args:
-            module_class: The class of the module to be managed.
-            services_mode: The mode of operation for the services (e.g., ASYNC or SYNC).
-        """
-        self.module_class = module_class
-
-        services_config = ServicesConfig(
-            services_config_strategies=self.module_class.services_config_strategies,
-            services_config_params=self.module_class.services_config_params,
-            mode=services_mode,
-        )
-        setattr(self.module_class, "services_config", services_config)
-
     @abc.abstractmethod  # type: ignore
     @asynccontextmanager  # type: ignore
     async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:
@@ -104,7 +194,7 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT]):
         """
 
     @abc.abstractmethod
-    async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT:
+    async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT | ModuleCodeModel:
         """Generate a stream consumer for a module's output data.
 
         This method creates an asynchronous generator that streams output data
@@ -115,7 +205,7 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT]):
             job_id: The unique identifier of the job.
 
         Returns:
-            SetupModelT: the SetupModelT object fully processed.
+            SetupModelT | ModuleCodeModel: the SetupModelT object fully processed, or an error code.
         """
 
     @abc.abstractmethod
@@ -156,14 +246,30 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT]):
         """
 
     @abc.abstractmethod
-    async def get_module_status(self, job_id: str) -> ModuleStatus | None:
+    async def get_module_status(self, job_id: str) -> TaskStatus:
         """Retrieve the status of a module job.
 
         Args:
             job_id: The unique identifier of the job.
 
         Returns:
-            ModuleStatus | None: The status of the job, or None if the job does not exist.
+            ModuleStatu: The status of the job.
+        """
+
+    @abc.abstractmethod
+    async def wait_for_completion(self, job_id: str) -> None:
+        """Wait for a task to complete.
+
+        This method blocks until the specified job has reached a terminal state.
+        The implementation varies by job manager type:
+        - SingleJobManager: Awaits the asyncio.Task directly
+        - TaskiqJobManager: Polls task status from SurrealDB
+
+        Args:
+            job_id: The unique identifier of the job to wait for.
+
+        Raises:
+            KeyError: If the job_id is not found.
         """
 
     @abc.abstractmethod