digitalkin 0.2.23__py3-none-any.whl → 0.3.1.dev2__py3-none-any.whl

digitalkin/mixins/cost_mixin.py

@@ -0,0 +1,76 @@
+"""Cost Mixin to ease trigger deveolpment."""
+
+from typing import Literal
+
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.services.cost.cost_strategy import CostData
+
+
+class CostMixin:
+    """Mixin providing cost tracking operations through the cost strategy.
+
+    This mixin wraps cost strategy calls to provide a cleaner API
+    for cost tracking in trigger handlers.
+    """
+
+    @staticmethod
+    def add_cost(context: ModuleContext, name: str, cost_config_name: str, quantity: float) -> None:
+        """Add a cost entry using the cost strategy.
+
+        Args:
+            context: Module context containing the cost strategy
+            name: Name/identifier for this cost entry
+            cost_config_name: Name of the cost configuration to use
+            quantity: Quantity of units consumed
+
+        Raises:
+            CostServiceError: If cost addition fails
+        """
+        return context.cost.add(name, cost_config_name, quantity)
+
+    @staticmethod
+    def get_cost(context: ModuleContext, name: str) -> list[CostData]:
+        """Get cost entries for a specific name.
+
+        Args:
+            context: Module context containing the cost strategy
+            name: Name/identifier to get costs for
+
+        Returns:
+            List of cost data entries
+
+        Raises:
+            CostServiceError: If cost retrieval fails
+        """
+        return context.cost.get(name)
+
+    @staticmethod
+    def get_costs(
+        context: ModuleContext,
+        names: list[str] | None = None,
+        cost_types: list[
+            Literal[
+                "TOKEN_INPUT",
+                "TOKEN_OUTPUT",
+                "API_CALL",
+                "STORAGE",
+                "TIME",
+                "OTHER",
+            ]
+        ]
+        | None = None,
+    ) -> list[CostData]:
+        """Get filtered cost entries.
+
+        Args:
+            context: Module context containing the cost strategy
+            names: Optional list of names to filter by
+            cost_types: Optional list of cost types to filter by
+
+        Returns:
+            List of filtered cost data entries
+
+        Raises:
+            CostServiceError: If cost retrieval fails
+        """
+        return context.cost.get_filtered(names, cost_types)

digitalkin/mixins/file_history_mixin.py

@@ -0,0 +1,93 @@
+"""Context mixins providing ergonomic access to service strategies.
+
+This module provides mixins that wrap service strategy calls with cleaner APIs,
+following Django/FastAPI patterns where context is passed explicitly to each method.
+"""
+
+from digitalkin.mixins.logger_mixin import LoggerMixin
+from digitalkin.mixins.storage_mixin import StorageMixin
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.models.services.storage import FileHistory, FileModel
+
+
+class FileHistoryMixin(StorageMixin, LoggerMixin):
+    """Mixin providing File history operations through storage strategy.
+
+    This mixin provides a higher-level API for managing File history,
+    using the storage strategy as the underlying persistence mechanism.
+    """
+
+    file_history_front: FileHistory = FileHistory(files=[])
+    FILE_HISTORY_COLLECTION = "file_history"
+    FILE_HISTORY_RECORD_ID = "full_file_history"
+
+    def _get_history_key(self, context: ModuleContext) -> str:
+        """Get session-specific history key.
+
+        Args:
+            context: Module context containing session information
+
+        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}"
+
+    def load_file_history(self, context: ModuleContext) -> FileHistory:
+        """Load File history for the current session.
+
+        Args:
+            context: Module context containing storage strategy
+
+        Returns:
+            File history object, empty if none exists or loading fails
+        """
+        history_key = self._get_history_key(context)
+
+        if self.file_history_front is None:
+            try:
+                record = self.read_storage(
+                    context,
+                    self.FILE_HISTORY_COLLECTION,
+                    history_key,
+                )
+                if record and record.data:
+                    return FileHistory.model_validate(record.data)
+            except Exception as e:
+                self.log_warning(context, f"Failed to load File history: {e}")
+        return self.file_history_front
+
+    def append_files_history(self, context: ModuleContext, files: list[FileModel]) -> None:
+        """Append a message to File history.
+
+        Args:
+            context: Module context containing storage strategy
+            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)
+
+        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,
+                file_history.model_dump(),
+                data_type="OUTPUT",
+            )
+        else:
+            self.log_debug(context, f"Updating file history for session: {history_key}")
+            self.update_storage(
+                context,
+                self.FILE_HISTORY_COLLECTION,
+                history_key,
+                file_history.model_dump(),
+            )

digitalkin/mixins/filesystem_mixin.py

@@ -0,0 +1,46 @@
+"""Filesystem Mixin to ease filesystem use."""
+
+from typing import Any
+
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.services.filesystem.filesystem_strategy import FilesystemRecord
+
+
+class FilesystemMixin:
+    """Mixin providing filesystem operations through the filesystem strategy.
+
+    This mixin wraps filesystem strategy calls to provide a cleaner API
+    for file operations in trigger handlers.
+    """
+
+    @staticmethod
+    def upload_files(context: ModuleContext, files: list[Any]) -> tuple[list[FilesystemRecord], int, int]:
+        """Upload files using the filesystem strategy.
+
+        Args:
+            context: Module context containing the filesystem strategy
+            files: List of files to upload
+
+        Returns:
+            Tuple of (all_files, succeeded_files, failed_files)
+
+        Raises:
+            FilesystemServiceError: If upload operation fails
+        """
+        return context.filesystem.upload_files(files)
+
+    @staticmethod
+    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
+
+        Returns:
+            File object with metadata and optionally content
+
+        Raises:
+            FilesystemServiceError: If file retrieval fails
+        """
+        return context.filesystem.get_file(file_id, include_content=True)

digitalkin/mixins/logger_mixin.py

@@ -0,0 +1,51 @@
+"""Logger Mixin to ease and merge every logs."""
+
+from digitalkin.models.module.module_context import ModuleContext
+
+
+class LoggerMixin:
+    """Mixin providing callback operations through the callbacks strategy.
+
+    This mixin wraps callback strategy calls to provide a cleaner API
+    for logging and messaging in trigger handlers.
+    """
+
+    @staticmethod
+    def log_debug(context: ModuleContext, message: str) -> None:
+        """Log debug message using the callbacks strategy.
+
+        Args:
+            context: Module context containing the callbacks strategy
+            message: Debug message to log
+        """
+        return context.callbacks.logger.debug(message, extra=context.session.current_ids())
+
+    @staticmethod
+    def log_info(context: ModuleContext, message: str) -> None:
+        """Log info message using the callbacks strategy.
+
+        Args:
+            context: Module context containing the callbacks strategy
+            message: Info message to log
+        """
+        return context.callbacks.logger.info(message, extra=context.session.current_ids())
+
+    @staticmethod
+    def log_warning(context: ModuleContext, message: str) -> None:
+        """Log warning message using the callbacks strategy.
+
+        Args:
+            context: Module context containing the callbacks strategy
+            message: Warning message to log
+        """
+        return context.callbacks.logger.warning(message, extra=context.session.current_ids())
+
+    @staticmethod
+    def log_error(context: ModuleContext, message: str) -> None:
+        """Log error message using the callbacks strategy.
+
+        Args:
+            context: Module context containing the callbacks strategy
+            message: Error message to log
+        """
+        return context.callbacks.logger.error(message, extra=context.session.current_ids())

digitalkin/mixins/storage_mixin.py

@@ -0,0 +1,79 @@
+"""Storage Mixin to ease storage access in Triggers."""
+
+from typing import Any, Literal
+
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.services.storage.storage_strategy import StorageRecord
+
+
+class StorageMixin:
+    """Mixin providing storage operations through the storage strategy.
+
+    This mixin wraps storage strategy calls to provide a cleaner API
+    for trigger handlers.
+    """
+
+    @staticmethod
+    def store_storage(
+        context: ModuleContext,
+        collection: str,
+        record_id: str | None,
+        data: dict[str, Any],
+        data_type: Literal["OUTPUT", "VIEW", "LOGS", "OTHER"] = "OUTPUT",
+    ) -> StorageRecord:
+        """Store data using the storage strategy.
+
+        Args:
+            context: Module context containing the storage strategy
+            collection: Collection name for the data
+            record_id: Optional record identifier
+            data: Data to store
+            data_type: Type of data being stored
+
+        Returns:
+            Result from the storage strategy
+
+        Raises:
+            StorageServiceError: If storage operation fails
+        """
+        return context.storage.store(collection, record_id, data, data_type=data_type)
+
+    @staticmethod
+    def read_storage(context: ModuleContext, collection: str, record_id: str) -> StorageRecord | None:
+        """Read data from storage.
+
+        Args:
+            context: Module context containing the storage strategy
+            collection: Collection name
+            record_id: Record identifier
+
+        Returns:
+            Retrieved data
+
+        Raises:
+            StorageServiceError: If read operation fails
+        """
+        return context.storage.read(collection, record_id)
+
+    @staticmethod
+    def update_storage(
+        context: ModuleContext,
+        collection: str,
+        record_id: str,
+        data: dict[str, Any],
+    ) -> StorageRecord | None:
+        """Update existing data in storage.
+
+        Args:
+            context: Module context containing the storage strategy
+            collection: Collection name
+            record_id: Record identifier
+            data: Updated data
+
+        Returns:
+            Result from the storage strategy
+
+        Raises:
+            StorageServiceError: If update operation fails
+        """
+        return context.storage.update(collection, record_id, data)

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/core/task_monitor.py

@@ -0,0 +1,70 @@
+"""Task monitoring models for signaling and heartbeat messages."""
+
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class TaskStatus(Enum):
+    """Task status enumeration."""
+
+    PENDING = "pending"
+    RUNNING = "running"
+    CANCELLED = "cancelled"
+    COMPLETED = "completed"
+    FAILED = "failed"
+
+
+class CancellationReason(Enum):
+    """Reason for task cancellation - helps distinguish cleanup vs real cancellation."""
+
+    # Cleanup cancellations (not errors)
+    SUCCESS_CLEANUP = "success_cleanup"  # Main task completed, cleaning up helper tasks
+    FAILURE_CLEANUP = "failure_cleanup"  # Main task failed, cleaning up helper tasks
+
+    # Real cancellations
+    SIGNAL = "signal"  # External signal requested cancellation
+    HEARTBEAT_FAILURE = "heartbeat_failure"  # Heartbeat stopped working
+    TIMEOUT = "timeout"  # Task timed out
+    SHUTDOWN = "shutdown"  # Manager is shutting down
+
+    # Unknown/unset
+    UNKNOWN = "unknown"  # Reason not determined
+
+
+class SignalType(Enum):
+    """Signal type enumeration."""
+
+    START = "start"
+    STOP = "stop"
+    CANCEL = "cancel"
+    PAUSE = "pause"
+    RESUME = "resume"
+    STATUS = "status"
+
+    ACK_CANCEL = "ack_cancel"
+    ACK_PAUSE = "ack_pause"
+    ACK_RESUME = "ack_resume"
+    ACK_STATUS = "ack_status"
+
+
+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))
+    payload: dict[str, Any] = Field(default={}, description="Optional payload for the signal")
+    model_config = {"use_enum_values": True}
+
+
+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/{grpc_servers/utils → models/grpc_servers}/models.py

@@ -175,8 +175,8 @@ class ClientConfig(ChannelConfig):
     credentials: ClientCredentials | None = Field(None, description="Client credentials for secure mode")
     channel_options: list[tuple[str, Any]] = Field(
         default_factory=lambda: [
-            ("grpc.max_receive_message_length", 50 * 1024 * 1024),  # 50MB
-            ("grpc.max_send_message_length", 50 * 1024 * 1024),  # 50MB
+            ("grpc.max_receive_message_length", 100 * 1024 * 1024),  # 100MB
+            ("grpc.max_send_message_length", 100 * 1024 * 1024),  # 100MB
         ],
         description="Additional channel options",
     )
@@ -223,8 +223,8 @@ class ServerConfig(ChannelConfig):
     credentials: ServerCredentials | None = Field(None, description="Server credentials for secure mode")
     server_options: list[tuple[str, Any]] = Field(
         default_factory=lambda: [
-            ("grpc.max_receive_message_length", 50 * 1024 * 1024),  # 50MB
-            ("grpc.max_send_message_length", 50 * 1024 * 1024),  # 50MB
+            ("grpc.max_receive_message_length", 100 * 1024 * 1024),  # 100MB
+            ("grpc.max_send_message_length", 100 * 1024 * 1024),  # 100MB
         ],
         description="Additional server options",
     )
@@ -262,7 +262,7 @@ class ModuleServerConfig(ServerConfig):
         registry_address: Address of the registry server
     """
 
-    registry_address: str | None = Field(None, description="Address of the registry server")
+    registry_address: str = Field(..., description="Address of the registry server")
 
 
 class RegistryServerConfig(ServerConfig):

digitalkin/models/module/__init__.py

@@ -5,6 +5,7 @@ from digitalkin.models.module.module_context import ModuleContext
 from digitalkin.models.module.module_types import (
     DataModel,
     DataTrigger,
+    DataTriggerT,
     InputModelT,
     OutputModelT,
     SecretModelT,
@@ -15,6 +16,7 @@ from digitalkin.models.module.module_types import (
 __all__ = [
     "DataModel",
     "DataTrigger",
+    "DataTriggerT",
     "InputModelT",
     "Module",
     "ModuleContext",

digitalkin/models/module/module.py

@@ -2,7 +2,15 @@
 
 from enum import Enum, auto
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
+
+
+class ModuleCodeModel(BaseModel):
+    """typed error/code model."""
+
+    code: str = Field(...)
+    message: str | None = Field(default=None)
+    short_description: str | None = Field(default=None)
 
 
 class ModuleStatus(Enum):

digitalkin/models/module/module_context.py

@@ -1,24 +1,140 @@
 """Define the module context used in the triggers."""
 
 from types import SimpleNamespace
+from typing import Any
 
+from digitalkin.services.agent.agent_strategy import AgentStrategy
 from digitalkin.services.cost.cost_strategy import CostStrategy
 from digitalkin.services.filesystem.filesystem_strategy import FilesystemStrategy
+from digitalkin.services.identity.identity_strategy import IdentityStrategy
+from digitalkin.services.registry.registry_strategy import RegistryStrategy
+from digitalkin.services.snapshot.snapshot_strategy import SnapshotStrategy
 from digitalkin.services.storage.storage_strategy import StorageStrategy
 
 
-class ModuleContext(SimpleNamespace):
-    """ModuleContext provides a container for strategies and resources used by a module.
+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,
+        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.
 
-    Attributes:
-        cost (CostStrategy): The strategy used to calculate or manage costs within the module.
-        filesystem (FilesystemStrategy): The strategy for interacting with the filesystem.
-        storage (StorageStrategy): The strategy for handling storage operations.
+        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.
 
     This context object is designed to be passed to module components, providing them with
     access to shared strategies and resources. Additional attributes may be set dynamically.
     """
 
+    # services list
+    agent: AgentStrategy
     cost: CostStrategy
     filesystem: FilesystemStrategy
+    identity: IdentityStrategy
+    registry: RegistryStrategy
+    snapshot: SnapshotStrategy
     storage: StorageStrategy
+
+    session: Session
+    callbacks: SimpleNamespace
+    metadata: SimpleNamespace
+    helpers: SimpleNamespace
+    state: SimpleNamespace = SimpleNamespace()
+
+    def __init__(  # noqa: PLR0913, PLR0917
+        self,
+        agent: AgentStrategy,
+        cost: CostStrategy,
+        filesystem: FilesystemStrategy,
+        identity: IdentityStrategy,
+        registry: RegistryStrategy,
+        snapshot: SnapshotStrategy,
+        storage: StorageStrategy,
+        session: dict[str, Any],
+        metadata: dict[str, Any] = {},
+        helpers: dict[str, Any] = {},
+        callbacks: dict[str, Any] = {},
+    ) -> None:
+        """Register mandatory services, session, metadata and callbacks.
+
+        Args:
+            agent: AgentStrategy.
+            cost: CostStrategy.
+            filesystem: FilesystemStrategy.
+            identity: IdentityStrategy.
+            registry: RegistryStrategy.
+            snapshot: SnapshotStrategy.
+            storage: StorageStrategy.
+            metadata: dict defining differents Module metadata.
+            helpers: dict different user defined helpers.
+            session: dict referring the session IDs or informations.
+            callbacks: Functions allowing user to agent interaction.
+        """
+        # Core services
+        self.agent = agent
+        self.cost = cost
+        self.filesystem = filesystem
+        self.identity = identity
+        self.registry = registry
+        self.snapshot = snapshot
+        self.storage = storage
+
+        self.metadata = SimpleNamespace(**metadata)
+        self.session = Session(**session)
+        self.helpers = SimpleNamespace(**helpers)
+        self.callbacks = SimpleNamespace(**callbacks)