digitalkin 0.3.0rc0__py3-none-any.whl → 0.3.0rc1__py3-none-any.whl

digitalkin/{modules/job_manager → core/task_manager}/task_session.py

@@ -1,13 +1,13 @@
-"""."""
+"""Task session easing task lifecycle management."""
 
 import asyncio
 import datetime
 from collections.abc import AsyncGenerator
 
+from digitalkin.core.task_manager.surrealdb_repository import SurrealDBConnection
 from digitalkin.logger import logger
-from digitalkin.models.module.task_monitor import HeartbeatMessage, SignalMessage, SignalType, TaskStatus
+from digitalkin.models.core.task_monitor import HeartbeatMessage, SignalMessage, SignalType, TaskStatus
 from digitalkin.modules._base_module import BaseModule
-from digitalkin.modules.job_manager.surrealdb_repository import SurrealDBConnection
 
 
 class TaskSession:
@@ -23,6 +23,7 @@ class TaskSession:
     signal_queue: AsyncGenerator | None
 
     task_id: str
+    mission_id: str
     signal_record_id: str | None
     heartbeat_record_id: str | None
 
@@ -37,11 +38,12 @@ class TaskSession:
     def __init__(
         self,
         task_id: str,
+        mission_id: str,
         db: SurrealDBConnection,
         module: BaseModule,
         heartbeat_interval: datetime.timedelta = datetime.timedelta(seconds=2),
     ) -> None:
-        """."""
+        """Initialize Task Session."""
         self.db = db
         self.module = module
 
@@ -49,6 +51,8 @@ class TaskSession:
         self.queue: asyncio.Queue = asyncio.Queue()
 
         self.task_id = task_id
+        self.mission_id = mission_id
+
         self.heartbeat = None
         self.started_at = None
         self.completed_at = None
@@ -63,17 +67,17 @@ class TaskSession:
         logger.info(
             "TaskContext initialized for task: '%s'",
             task_id,
-            extra={"task_id": task_id, "heartbeat_interval": heartbeat_interval},
+            extra={"task_id": task_id, "mission_id": mission_id, "heartbeat_interval": heartbeat_interval},
         )
 
     @property
     def cancelled(self) -> bool:
-        """."""
+        """Task cancellation status."""
         return self.is_cancelled.is_set()
 
     @property
     def paused(self) -> bool:
-        """."""
+        """Task paused status."""
         return self._paused.is_set()
 
     async def send_heartbeat(self) -> bool:
@@ -84,6 +88,7 @@ class TaskSession:
         """
         heartbeat = HeartbeatMessage(
             task_id=self.task_id,
+            mission_id=self.mission_id,
             timestamp=datetime.datetime.now(datetime.timezone.utc),
         )
 
@@ -142,7 +147,11 @@ class TaskSession:
             logger.debug(f"Heartbeat tick for task: '{self.task_id}' | {self.cancelled=}")
             success = await self.send_heartbeat()
             if not success:
-                logger.error("Heartbeat failed, cancelling task: '%s'", self.task_id, extra={"task_id": self.task_id})
+                logger.error(
+                    "Heartbeat failed, cancelling task: '%s'",
+                    self.task_id,
+                    extra={"task_id": self.task_id},
+                )
                 await self._handle_cancel()
                 break
             await asyncio.sleep(self._heartbeat_interval.total_seconds())
@@ -150,7 +159,11 @@ class TaskSession:
     async def wait_if_paused(self) -> None:
         """Block execution if task is paused."""
         if self._paused.is_set():
-            logger.info("Task paused, waiting for resume: '%s'", self.task_id, extra={"task_id": self.task_id})
+            logger.info(
+                "Task paused, waiting for resume: '%s'",
+                self.task_id,
+                extra={"task_id": self.task_id},
+            )
             await self._paused.wait()
 
     async def listen_signals(self) -> None:  # noqa: C901
@@ -159,7 +172,11 @@ class TaskSession:
         Raises:
             CancelledError: Asyncio when task cancelling
         """
-        logger.info("Signal listener started for task: '%s'", self.task_id, extra={"task_id": self.task_id})
+        logger.info(
+            "Signal listener started for task: '%s'",
+            self.task_id,
+            extra={"task_id": self.task_id},
+        )
         if self.signal_record_id is None:
             self.signal_record_id = (await self.db.select_by_task_id("tasks", self.task_id)).get("id")
 
@@ -183,7 +200,11 @@ class TaskSession:
                     await self._handle_status_request()
 
         except asyncio.CancelledError:
-            logger.debug("Signal listener cancelled for task: '%s'", self.task_id, extra={"task_id": self.task_id})
+            logger.debug(
+                "Signal listener cancelled for task: '%s'",
+                self.task_id,
+                extra={"task_id": self.task_id},
+            )
             raise
         except Exception as e:
             logger.error(
@@ -194,18 +215,28 @@ class TaskSession:
             )
         finally:
             await self.db.stop_live(live_id)
-            logger.info("Signal listener stopped for task: '%s'", self.task_id, extra={"task_id": self.task_id})
+            logger.info(
+                "Signal listener stopped for task: '%s'",
+                self.task_id,
+                extra={"task_id": self.task_id},
+            )
 
     async def _handle_cancel(self) -> None:
         """Idempotent cancellation with acknowledgment."""
         logger.debug("Handle cancel called")
         if self.is_cancelled.is_set():
             logger.debug(
-                "Cancel signal ignored - task already cancelled: '%s'", self.task_id, extra={"task_id": self.task_id}
+                "Cancel signal ignored - task already cancelled: '%s'",
+                self.task_id,
+                extra={"task_id": self.task_id},
             )
             return
 
-        logger.info("Cancelling task: '%s'", self.task_id, extra={"task_id": self.task_id})
+        logger.info(
+            "Cancelling task: '%s'",
+            self.task_id,
+            extra={"task_id": self.task_id},
+        )
 
         self.status = TaskStatus.CANCELLED
         self.is_cancelled.set()
@@ -219,6 +250,7 @@ class TaskSession:
             self.signal_record_id,  # type: ignore
             SignalMessage(
                 task_id=self.task_id,
+                mission_id=self.mission_id,
                 action=SignalType.ACK_CANCEL,
                 status=self.status,
             ).model_dump(),
@@ -227,7 +259,11 @@ class TaskSession:
     async def _handle_pause(self) -> None:
         """Pause task execution."""
         if not self._paused.is_set():
-            logger.info("Pausing task: '%s'", self.task_id, extra={"task_id": self.task_id})
+            logger.info(
+                "Pausing task: '%s'",
+                self.task_id,
+                extra={"task_id": self.task_id},
+            )
             self._paused.set()
 
         await self.db.update(
@@ -235,6 +271,7 @@ class TaskSession:
             self.signal_record_id,  # type: ignore
             SignalMessage(
                 task_id=self.task_id,
+                mission_id=self.mission_id,
                 action=SignalType.ACK_PAUSE,
                 status=self.status,
             ).model_dump(),
@@ -243,7 +280,11 @@ class TaskSession:
     async def _handle_resume(self) -> None:
         """Resume paused task."""
         if self._paused.is_set():
-            logger.info("Resuming task: '%s'", self.task_id, extra={"task_id": self.task_id})
+            logger.info(
+                "Resuming task: '%s'",
+                self.task_id,
+                extra={"task_id": self.task_id},
+            )
             self._paused.clear()
 
         await self.db.update(
@@ -251,6 +292,7 @@ class TaskSession:
             self.signal_record_id,  # type: ignore
             SignalMessage(
                 task_id=self.task_id,
+                mission_id=self.mission_id,
                 action=SignalType.ACK_RESUME,
                 status=self.status,
             ).model_dump(),
@@ -262,9 +304,10 @@ class TaskSession:
             "tasks",
             self.signal_record_id,  # type: ignore
             SignalMessage(
-                action=SignalType.ACK_STATUS,
+                mission_id=self.mission_id,
                 task_id=self.task_id,
                 status=self.status,
+                action=SignalType.ACK_STATUS,
             ).model_dump(),
         )
 

digitalkin/grpc_servers/__init__.py

@@ -1,19 +1 @@
-"""This package contains the gRPC server and client implementations.
-
-```shell
-digitalkin/grpc/
-├── __init__.py
-├── base_server.py        # Base server implementation with common functionality
-├── module_server.py      # Module-specific server implementation
-├── registry_server.py    # Registry-specific server implementation
-├── module_servicer.py    # gRPC servicer for Module service
-├── registry_servicer.py  # gRPC servicer for Registry service
-├── client/               # Client libraries for connecting to servers
-│   ├── __init__.py
-│   ├── module_client.py
-│   └── registry_client.py
-└── utils/                # Utility functions
-    ├── __init__.py
-    └── server_utils.py   # Common server utilities
-```
-"""
+"""This package contains the gRPC server and client implementations."""

digitalkin/grpc_servers/_base_server.py

@@ -17,9 +17,9 @@ from digitalkin.grpc_servers.utils.exceptions import (
     ServerStateError,
     ServicerError,
 )
-from digitalkin.grpc_servers.utils.models import SecurityMode, ServerConfig, ServerMode
-from digitalkin.grpc_servers.utils.types import GrpcServer, ServiceDescriptor, T
 from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import SecurityMode, ServerConfig, ServerMode
+from digitalkin.models.grpc_servers.types import GrpcServer, ServiceDescriptor, T
 
 
 class BaseServer(abc.ABC):

digitalkin/grpc_servers/module_server.py

@@ -17,12 +17,12 @@ from digitalkin_proto.digitalkin.module_registry.v2 import (
 from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.module_servicer import ModuleServicer
 from digitalkin.grpc_servers.utils.exceptions import ServerError
-from digitalkin.grpc_servers.utils.models import (
+from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import (
     ClientConfig,
     ModuleServerConfig,
     SecurityMode,
 )
-from digitalkin.logger import logger
 from digitalkin.modules._base_module import BaseModule
 
 

digitalkin/grpc_servers/module_servicer.py

@@ -13,12 +13,12 @@ from digitalkin_proto.digitalkin.module.v2 import (
 )
 from google.protobuf import json_format, struct_pb2
 
+from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 from digitalkin.grpc_servers.utils.exceptions import ServicerError
 from digitalkin.logger import logger
+from digitalkin.models.core.job_manager_models import JobManagerMode
 from digitalkin.models.module.module import ModuleStatus
 from digitalkin.modules._base_module import BaseModule
-from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
-from digitalkin.modules.job_manager.job_manager_models import JobManagerMode
 from digitalkin.services.services_models import ServicesMode
 from digitalkin.services.setup.default_setup import DefaultSetup
 from digitalkin.services.setup.grpc_setup import GrpcSetup
@@ -226,7 +226,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
                     yield lifecycle_pb2.StartModuleResponse(success=True, output=proto, job_id=job_id)
         finally:
             await self.job_manager.tasks[job_id]
-            await self.job_manager.clean_session(job_id)
+            await self.job_manager.clean_session(job_id, mission_id=request.mission_id)
 
         logger.info("Job %s finished", job_id)
 

digitalkin/grpc_servers/registry_server.py

@@ -7,8 +7,8 @@ from digitalkin_proto.digitalkin.module_registry.v2 import (
 
 from digitalkin.grpc_servers._base_server import BaseServer
 from digitalkin.grpc_servers.registry_servicer import RegistryModule, RegistryServicer
-from digitalkin.grpc_servers.utils.models import RegistryServerConfig
 from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import RegistryServerConfig
 
 
 class RegistryServer(BaseServer):

digitalkin/grpc_servers/utils/__init__.py

@@ -0,0 +1 @@
+"""gRPC servers utilities package."""

digitalkin/grpc_servers/utils/exceptions.py

@@ -27,11 +27,3 @@ class ServerStateError(ServerError):
 
 class ReflectionError(ServerError):
     """Error related to gRPC reflection service."""
-
-
-class HealthCheckError(ServerError):
-    """Error related to gRPC health check service."""
-
-
-class OptionalFeatureNotImplementedError(NotImplementedError):
-    """Raised when an optional feature is not implemented, but was requested."""

digitalkin/grpc_servers/utils/grpc_client_wrapper.py

@@ -6,8 +6,8 @@ from typing import Any
 import grpc
 
 from digitalkin.grpc_servers.utils.exceptions import ServerError
-from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode
 from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import ClientConfig, SecurityMode
 
 
 class GrpcClientWrapper:

digitalkin/mixins/chat_history_mixin.py

@@ -33,6 +33,8 @@ class ChatHistoryMixin(UserMessageMixin, StorageMixin, LoggerMixin, Generic[Inpu
         Returns:
             Unique history key for the current session
         """
+        # TODO: define mission-specific chat history key not dependant on mission_id
+        # or need customization by user
         mission_id = getattr(context.session, "mission_id", None) or "default"
         return f"{self.CHAT_HISTORY_RECORD_ID}_{mission_id}"
 

digitalkin/mixins/file_history_mixin.py

@@ -30,6 +30,8 @@ class FileHistoryMixin(StorageMixin, LoggerMixin):
         Returns:
             Unique history key for the current session
         """
+        # TODO: define mission-specific chat history key not dependant on mission_id
+        # or need customization by user
         mission_id = getattr(context.session, "mission_id", None) or "default"
         return f"{self.FILE_HISTORY_RECORD_ID}_{mission_id}"
 
@@ -62,38 +64,30 @@ class FileHistoryMixin(StorageMixin, LoggerMixin):
 
         Args:
             context: Module context containing storage strategy
-            role: Message role (user, assistant, system)
             files: list of files model
 
         Raises:
             StorageServiceError: If history update fails
         """
         history_key = self._get_history_key(context)
+        file_history = self.load_file_history(context)
 
-        try:
-            if not self.file_history_front.files:
-                self.file_history_front = self.load_file_history(context)
-            self.file_history_front.files.extend(files)
-        except Exception as e:
-            self.log_error(context, f"Failed to append message to File history: {e}")
-
-        try:
-            self.log_debug(context, f"Updating File history for session: {history_key}")
-            self.update_storage(
+        file_history.files.extend(files)
+        if len(file_history.files) == len(files):
+            # Create new record
+            self.log_debug(context, f"Creating new file history for session: {history_key}")
+            self.store_storage(
                 context,
                 self.FILE_HISTORY_COLLECTION,
                 history_key,
-                self.file_history_front.model_dump(),
+                file_history.model_dump(),
+                data_type="OUTPUT",
             )
-        except Exception as e:
-            self.log_error(context, f"Updating File history for session: {history_key} with error: {e}")
-
-            # Create new record
-            self.log_debug(context, f"Creating new File history for session: {history_key}")
-            self.store_storage(
+        else:
+            self.log_debug(context, f"Updating file history for session: {history_key}")
+            self.update_storage(
                 context,
                 self.FILE_HISTORY_COLLECTION,
                 history_key,
-                self.file_history_front.model_dump(),
-                data_type="OUTPUT",
+                file_history.model_dump(),
             )

digitalkin/mixins/filesystem_mixin.py

@@ -30,13 +30,12 @@ class FilesystemMixin:
         return context.filesystem.upload_files(files)
 
     @staticmethod
-    def get_file(context: ModuleContext, file_id: str) -> tuple[FilesystemRecord, bytes | None]:
+    def get_file(context: ModuleContext, file_id: str) -> FilesystemRecord:
         """Retrieve a file by ID with the content.
 
         Args:
             context: Module context containing the filesystem strategy
             file_id: Unique identifier for the file
-            include_content: Whether to include file content in response
 
         Returns:
             File object with metadata and optionally content

digitalkin/mixins/logger_mixin.py

@@ -17,10 +17,8 @@ class LoggerMixin:
         Args:
             context: Module context containing the callbacks strategy
             message: Debug message to log
-            *args: Positional arguments for message formatting
-            **kwargs: Keyword arguments for logger
         """
-        return context.callbacks.logger.debug(message)
+        return context.callbacks.logger.debug(message, extra=context.session.current_ids())
 
     @staticmethod
     def log_info(context: ModuleContext, message: str) -> None:
@@ -29,10 +27,8 @@ class LoggerMixin:
         Args:
             context: Module context containing the callbacks strategy
             message: Info message to log
-            *args: Positional arguments for message formatting
-            **kwargs: Keyword arguments for logger
         """
-        return context.callbacks.logger.info(message)
+        return context.callbacks.logger.info(message, extra=context.session.current_ids())
 
     @staticmethod
     def log_warning(context: ModuleContext, message: str) -> None:
@@ -41,10 +37,8 @@ class LoggerMixin:
         Args:
             context: Module context containing the callbacks strategy
             message: Warning message to log
-            *args: Positional arguments for message formatting
-            **kwargs: Keyword arguments for logger
         """
-        return context.callbacks.logger.warning(message)
+        return context.callbacks.logger.warning(message, extra=context.session.current_ids())
 
     @staticmethod
     def log_error(context: ModuleContext, message: str) -> None:
@@ -53,7 +47,5 @@ class LoggerMixin:
         Args:
             context: Module context containing the callbacks strategy
             message: Error message to log
-            *args: Positional arguments for message formatting
-            **kwargs: Keyword arguments for logger
         """
-        return context.callbacks.logger.error(message)
+        return context.callbacks.logger.error(message, extra=context.session.current_ids())

digitalkin/models/core/__init__.py

@@ -0,0 +1 @@
+"""Core models."""

digitalkin/{modules/job_manager → models/core}/job_manager_models.py

@@ -4,7 +4,7 @@ from enum import Enum
 
 from pydantic import BaseModel
 
-from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
+from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 
 
 class StreamCodeModel(BaseModel):
@@ -35,10 +35,10 @@ class JobManagerMode(Enum):
         """
         match self:
             case JobManagerMode.SINGLE:
-                from digitalkin.modules.job_manager.single_job_manager import SingleJobManager  # noqa: PLC0415
+                from digitalkin.core.job_manager.single_job_manager import SingleJobManager  # noqa: PLC0415
 
                 return SingleJobManager
             case JobManagerMode.TASKIQ:
-                from digitalkin.modules.job_manager.taskiq_job_manager import TaskiqJobManager  # noqa: PLC0415
+                from digitalkin.core.job_manager.taskiq_job_manager import TaskiqJobManager  # noqa: PLC0415
 
                 return TaskiqJobManager

digitalkin/models/{module → core}/task_monitor.py

@@ -1,4 +1,4 @@
-"""."""
+"""Task monitoring models for signaling and heartbeat messages."""
 
 from datetime import datetime, timezone
 from enum import Enum
@@ -8,7 +8,7 @@ from pydantic import BaseModel, Field
 
 
 class TaskStatus(Enum):
-    """."""
+    """Task status enumeration."""
 
     PENDING = "pending"
     RUNNING = "running"
@@ -18,7 +18,7 @@ class TaskStatus(Enum):
 
 
 class SignalType(Enum):
-    """."""
+    """Signal type enumeration."""
 
     START = "start"
     STOP = "stop"
@@ -34,9 +34,10 @@ class SignalType(Enum):
 
 
 class SignalMessage(BaseModel):
-    """."""
+    """Signal message model for task monitoring."""
 
     task_id: str = Field(..., description="Unique identifier for the task")
+    mission_id: str = Field(..., description="Identifier for the mission")
     status: TaskStatus = Field(..., description="Current status of the task")
     action: SignalType = Field(..., description="Type of signal action")
     timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
@@ -45,7 +46,8 @@ class SignalMessage(BaseModel):
 
 
 class HeartbeatMessage(BaseModel):
-    """."""
+    """Heartbeat message model for task monitoring."""
 
     task_id: str = Field(..., description="Unique identifier for the task")
+    mission_id: str = Field(..., description="Identifier for the mission")
     timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

digitalkin/models/grpc_servers/__init__.py

@@ -0,0 +1 @@
+"""Base gRPC server and client models."""

digitalkin/models/module/module_context.py

@@ -15,32 +15,64 @@ from digitalkin.services.storage.storage_strategy import StorageStrategy
 class Session(SimpleNamespace):
     """Session data container with mandatory setup_id and mission_id."""
 
+    job_id: str
     mission_id: str
+    setup_id: str
     setup_version_id: str
 
-    def __init__(self, mission_id: str, setup_version_id: str, **kwargs: dict[str, Any]) -> None:
+    def __init__(
+        self,
+        job_id: str,
+        mission_id: str,
+        setup_id: str,
+        setup_version_id: str,
+        **kwargs: dict[str, Any],
+    ) -> None:
         """Init Module Session.
 
         Args:
+            job_id: current job_id.
             mission_id: current mission_id.
+            setup_id: used setup config.
             setup_version_id: used setup config.
             kwargs: user defined session variables.
 
         Raises:
             ValueError: If mandatory args are missing
         """
+        if not setup_id:
+            msg = "setup_id is mandatory and cannot be empty"
+            raise ValueError(msg)
         if not setup_version_id:
             msg = "setup_version_id is mandatory and cannot be empty"
             raise ValueError(msg)
         if not mission_id:
             msg = "mission_id is mandatory and cannot be empty"
             raise ValueError(msg)
+        if not job_id:
+            msg = "job_id is mandatory and cannot be empty"
+            raise ValueError(msg)
 
+        self.job_id = job_id
         self.mission_id = mission_id
+        self.setup_id = setup_id
         self.setup_version_id = setup_version_id
 
         super().__init__(**kwargs)
 
+    def current_ids(self) -> dict[str, str]:
+        """Return current session ids as a dictionary.
+
+        Returns:
+            A dictionary containing the current session ids.
+        """
+        return {
+            "job_id": self.job_id,
+            "mission_id": self.mission_id,
+            "setup_id": self.setup_id,
+            "setup_version_id": self.setup_version_id,
+        }
+
 
 class ModuleContext:
     """ModuleContext provides a container for strategies and resources used by a module.

digitalkin/models/services/cost.py

@@ -38,6 +38,7 @@ class CostConfig(BaseModel):
 class CostEvent(BaseModel):
     """Pydantic model that represents a cost event registered during service execution.
 
+    # DEPRECATED
     :param cost_name: Identifier for the cost configuration.
     :param cost_type: The type of cost.
     :param usage: The amount or units consumed.