digitalkin 0.3.0rc2__tar.gz → 0.3.1.dev0__tar.gz

{digitalkin-0.3.0rc2 → digitalkin-0.3.1.dev0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.3.0rc2
+Version: 0.3.1.dev0
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International
@@ -452,17 +452,17 @@ Classifier: License :: Other/Proprietary License
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: digitalkin-proto>=0.1.16
-Requires-Dist: grpcio-health-checking>=1.71.0
-Requires-Dist: grpcio-reflection>=1.71.0
-Requires-Dist: grpcio-status>=1.71.0
-Requires-Dist: pydantic>=2.12.3
+Requires-Dist: digitalkin-proto>=0.2.0.dev4
+Requires-Dist: grpcio-health-checking>=1.76.0
+Requires-Dist: grpcio-reflection>=1.76.0
+Requires-Dist: grpcio-status>=1.76.0
+Requires-Dist: pydantic>=2.12.4
 Requires-Dist: surrealdb>=1.0.6
 Provides-Extra: taskiq
 Requires-Dist: rstream>=0.40.0; extra == "taskiq"
 Requires-Dist: taskiq-aio-pika>=0.4.4; extra == "taskiq"
 Requires-Dist: taskiq-redis>=1.1.2; extra == "taskiq"
-Requires-Dist: taskiq[reload]>=0.11.19; extra == "taskiq"
+Requires-Dist: taskiq[reload]>=0.11.20; extra == "taskiq"
 Dynamic: license-file
 
 # DigitalKin Python SDK

{digitalkin-0.3.0rc2 → digitalkin-0.3.1.dev0}/pyproject.toml

@@ -12,7 +12,7 @@
 
     keywords = [ "digitalkin", "kin", "agent", "gprc", "sdk" ]
 
-    version = "0.3.0-rc2"
+    version = "0.3.1.dev0"
     classifiers = [
         "Development Status :: 3 - Alpha",
         "Intended Audience :: Developers",
@@ -28,11 +28,11 @@
     ]
 
     dependencies = [
-        "digitalkin-proto>=0.1.16",
-        "grpcio-health-checking>=1.71.0",
-        "grpcio-reflection>=1.71.0",
-        "grpcio-status>=1.71.0",
-        "pydantic>=2.12.3",
+        "digitalkin-proto>=0.2.0.dev4",
+        "grpcio-health-checking>=1.76.0",
+        "grpcio-reflection>=1.76.0",
+        "grpcio-status>=1.76.0",
+        "pydantic>=2.12.4",
         "surrealdb>=1.0.6",
     ]
 
@@ -41,7 +41,7 @@
             "rstream>=0.40.0",
             "taskiq-aio-pika>=0.4.4",
             "taskiq-redis>=1.1.2",
-            "taskiq[reload]>=0.11.19",
+            "taskiq[reload]>=0.11.20",
         ]
 
     [project.urls]
@@ -55,11 +55,11 @@
 
 [dependency-groups]
     dev = [
-        "typos>=1.38.0",
-        "ruff>=0.14.2",
+        "typos>=1.39.2",
+        "ruff>=0.14.5",
         "mypy>=1.18.2",
         "pyright>=1.1.407",
-        "pre-commit>=4.3.0",
+        "pre-commit>=4.4.0",
         "bump-my-version>=1.2.4",
         "build>=1.3.0",
         "twine>=6.2.0",
@@ -68,18 +68,19 @@
     tests = [
         "freezegun>=1.5.3",
         "hdrhistogram>=0.10.3",
-        "grpcio-testing>=1.71.0",
-        "psutil>=7.1.2",
-        "pytest>=8.4.2",
-        "pytest-asyncio>=1.2.0",
+        "grpcio-testing>=1.76.0",
+        "psutil>=7.1.3",
+        "pytest>=9.0.1",
+        "pytest-asyncio>=1.3.0",
         "pytest-cov>=7.0.0",
         "pytest-html==4.1.1",
         "pytest-json-report==1.5.0",
+        "pytest-timeout>=2.3.1",
     ]
     docs = [
         "mike>=2.1.3",
         "markdown-callouts>=0.4",
-        "markdown-exec>=1.8",
+        "markdown-exec>=1.12.1",
         "mkdocs>=1.6",
         "mkdocs-coverage>=2.0.0",
         "mkdocs-llmstxt>=0.2",
@@ -89,18 +90,18 @@
         "mkdocs-autorefs>=1.4.3",
         "mkdocs-awesome-pages-plugin>=2.10.1",
         "mkdocs-git-committers-plugin-2>=2.5.0",
-        "mkdocs-git-revision-date-localized-plugin>=1.4.7",
+        "mkdocs-git-revision-date-localized-plugin>=1.5.0",
         "mkdocs-glightbox>=0.5.2",
         "mkdocs-include-markdown-plugin>=7.2.0",
         "mkdocs-literate-nav>=0.6.2",
-        "mkdocs-material[imaging]>=9.6.22",
+        "mkdocs-material[imaging]>=9.7.0",
         "mkdocs-minify-plugin>=0.8.0",
         "mkdocs-section-index>=0.3.10",
         "mkdocstrings>=0.30.1",
-        "mkdocstrings-python>=1.18.2",
+        "mkdocstrings-python>=1.19.0",
         "mkdocs-open-in-new-tab>=1.0.8",
         "tomli>=2.3.0",
-]
+    ]
 
 [tool.setuptools]
     package-dir = { "" = "src" }
@@ -261,3 +262,20 @@
     testpaths    = [ "tests" ]
     python_files = "test_*.py"
     asyncio_mode = "auto"
+
+    log_cli   = false
+    log_level = "ERROR"
+
+    # Default timeout for tests (in seconds)
+    timeout        = 10
+    timeout_method = "signal" # More reliable for asyncio tests on Linux
+
+    # Custom markers
+    markers = [
+        "integration: marks tests that require real SurrealDB connection (deselect with '-m \"not integration\"')",
+        "grpc: marks tests for gRPC service functionality",
+        "smoke: marks critical path tests that should always pass",
+        "edge_case: marks tests for boundary conditions and edge cases",
+        "validation: marks tests for input validation and schema checking",
+        "regression: marks tests for previously fixed bugs",
+    ]

{digitalkin-0.3.0rc2 → digitalkin-0.3.1.dev0}/src/digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.3.0-rc2"
+    __version__ = "0.3.1.dev0"

digitalkin-0.3.1.dev0/src/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-0.3.1.dev0/src/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-0.3.0rc2 → digitalkin-0.3.1.dev0}/src/digitalkin/core/job_manager/base_job_manager.py

@@ -5,7 +5,8 @@ from collections.abc import AsyncGenerator, AsyncIterator, Callable, Coroutine
 from contextlib import asynccontextmanager
 from typing import Any, Generic
 
-from digitalkin.core.task_manager.task_manager import TaskManager
+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
@@ -14,9 +15,115 @@ from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
 
 
-class BaseJobManager(abc.ABC, TaskManager, Generic[InputModelT, SetupModelT, OutputModelT]):
-    """Abstract base class for managing background module jobs."""
+class BaseJobManager(abc.ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
+    """Abstract base class for managing background module jobs.
 
+    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.
 
@@ -52,29 +159,6 @@ class BaseJobManager(abc.ABC, TaskManager, Generic[InputModelT, SetupModelT, Out
 
         return callback_wrapper
 
-    def __init__(
-        self,
-        module_class: type[BaseModule],
-        services_mode: ServicesMode,
-        **kwargs,  # noqa: ANN003
-    ) -> 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).
-            **kwargs: Additional keyword arguments for the job manager.
-        """
-        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)
-        super().__init__(**kwargs)
-
     @abc.abstractmethod  # type: ignore
     @asynccontextmanager  # type: ignore
     async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:
@@ -110,7 +194,7 @@ class BaseJobManager(abc.ABC, TaskManager, Generic[InputModelT, SetupModelT, Out
         """
 
     @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
@@ -121,7 +205,7 @@ class BaseJobManager(abc.ABC, TaskManager, Generic[InputModelT, SetupModelT, Out
             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
@@ -172,6 +256,22 @@ class BaseJobManager(abc.ABC, TaskManager, Generic[InputModelT, SetupModelT, Out
             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
     async def stop_all_modules(self) -> None:
         """Stop all currently running module jobs.