digitalkin 0.2.25rc0__py3-none-any.whl → 0.3.2.dev14__py3-none-any.whl

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/__init__.py

@@ -1,6 +1,6 @@
 """This package contains the models for DigitalKin."""
 
-from digitalkin.models.module import Module, ModuleStatus
+from digitalkin.models.module.module import Module, ModuleStatus
 
 __all__ = [
     "Module",

digitalkin/models/core/__init__.py

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

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

@@ -2,15 +2,7 @@
 
 from enum import Enum
 
-from pydantic import BaseModel
-
-from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
-
-
-class StreamCodeModel(BaseModel):
-    """Typed error/code model."""
-
-    code: str
+from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 
 
 class JobManagerMode(Enum):
@@ -35,10 +27,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,74 @@
+"""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")
+    setup_id: str = Field(default="", description="Identifier for the setup")
+    setup_version_id: str = Field(default="", description="Identifier for the setup version")
+    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")
+    setup_id: str = Field(default="", description="Identifier for the setup")
+    setup_version_id: str = Field(default="", description="Identifier for the setup version")
+    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

@@ -65,6 +65,42 @@ class ServerCredentials(BaseModel):
         return v
 
 
+class RetryPolicy(BaseModel):
+    """gRPC retry policy configuration for resilient connections.
+
+    Attributes:
+        max_attempts: Maximum retry attempts including the original call
+        initial_backoff: Initial backoff duration (e.g., "0.1s")
+        max_backoff: Maximum backoff duration (e.g., "10s")
+        backoff_multiplier: Multiplier for exponential backoff
+        retryable_status_codes: gRPC status codes that trigger retry
+    """
+
+    max_attempts: int = Field(default=5, ge=1, le=10, description="Maximum retry attempts including the original call")
+    initial_backoff: str = Field(default="0.1s", description="Initial backoff duration (e.g., '0.1s')")
+    max_backoff: str = Field(default="10s", description="Maximum backoff duration (e.g., '10s')")
+    backoff_multiplier: float = Field(default=2.0, ge=1.0, description="Multiplier for exponential backoff")
+    retryable_status_codes: list[str] = Field(
+        default_factory=lambda: ["UNAVAILABLE", "RESOURCE_EXHAUSTED"],
+        description="gRPC status codes that trigger retry",
+    )
+
+    model_config = {"extra": "forbid", "frozen": True}
+
+    def to_service_config_json(self) -> str:
+        """Serialize to gRPC service config JSON string.
+
+        Returns:
+            JSON string for grpc.service_config channel option.
+        """
+        codes = "[" + ",".join(f'"{c}"' for c in self.retryable_status_codes) + "]"
+        return (
+            f'{{"methodConfig":[{{"name":[{{}}],"retryPolicy":{{"maxAttempts":{self.max_attempts},'
+            f'"initialBackoff":"{self.initial_backoff}","maxBackoff":"{self.max_backoff}",'
+            f'"backoffMultiplier":{self.backoff_multiplier},"retryableStatusCodes":{codes}}}}}]}}'
+        )
+
+
 class ClientCredentials(BaseModel):
     """Model for client credentials in secure mode.
 
@@ -170,15 +206,47 @@ class ClientConfig(ChannelConfig):
         security: Security mode (secure/insecure)
         credentials: Client credentials for secure mode
         channel_options: Additional channel options
+        retry_policy: Retry policy for failed RPCs
     """
 
     credentials: ClientCredentials | None = Field(None, description="Client credentials for secure mode")
+    retry_policy: RetryPolicy = Field(default_factory=lambda: RetryPolicy(), description="Retry policy for failed RPCs")  # noqa: PLW0108
     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),
+            ("grpc.max_send_message_length", 100 * 1024 * 1024),
+            # === DNS Re-resolution (Critical for Container Environments) ===
+            # Minimum milliseconds between DNS re-resolution attempts (500 ms)
+            # When connection fails, gRPC will re-query DNS after this interval
+            # Solves: Container restarts with new IPs causing "No route to host"
+            ("grpc.dns_min_time_between_resolutions_ms", 500),
+            # Initial delay before first reconnection attempt (1 second)
+            ("grpc.initial_reconnect_backoff_ms", 1000),
+            # Maximum delay between reconnection attempts (10 seconds)
+            # Prevents overwhelming the network during extended outages
+            ("grpc.max_reconnect_backoff_ms", 10000),
+            # Minimum delay between reconnection attempts (500ms)
+            # Ensures rapid recovery for brief network glitches
+            ("grpc.min_reconnect_backoff_ms", 500),
+            # === Keepalive Settings (Detect Dead Connections) ===
+            # Send keepalive ping every 60 seconds when connection is idle
+            # Proactively detects dead connections before RPC calls fail
+            ("grpc.keepalive_time_ms", 60000),
+            # Wait 20 seconds for keepalive response before declaring connection dead
+            # Triggers reconnection (with DNS re-resolution) if pong not received
+            ("grpc.keepalive_timeout_ms", 20000),
+            # Send keepalive pings even when no RPCs are in flight
+            # Essential for long-lived connections that may sit idle
+            ("grpc.keepalive_permit_without_calls", True),
+            # Minimum interval between HTTP/2 pings (30 seconds)
+            # Must be >= server's grpc.http2.min_ping_interval_without_data_ms (10s)
+            ("grpc.http2.min_time_between_pings_ms", 30000),
+            # === Retry Configuration ===
+            # Enable automatic retry for failed RPCs (1 = enabled)
+            # Works with retryable status codes: UNAVAILABLE, RESOURCE_EXHAUSTED
+            ("grpc.enable_retries", 1),
         ],
-        description="Additional channel options",
+        description="Resilient gRPC channel options with DNS re-resolution, keepalive, and retries",
     )
 
     @field_validator("credentials")
@@ -204,6 +272,15 @@ class ClientConfig(ChannelConfig):
             raise ConfigurationError(msg)
         return v
 
+    @property
+    def grpc_options(self) -> list[tuple[str, Any]]:
+        """Get channel options with retry policy service config.
+
+        Returns:
+            Full list of gRPC channel options.
+        """
+        return [*self.channel_options, ("grpc.service_config", self.retry_policy.to_service_config_json())]
+
 
 class ServerConfig(ChannelConfig):
     """Base configuration for gRPC servers.
@@ -223,10 +300,18 @@ 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),
+            ("grpc.max_send_message_length", 100 * 1024 * 1024),
+            # === Keepalive Permission (Required for Client Keepalive) ===
+            # Allow clients to send keepalive pings without active RPCs
+            # Without this, server rejects client keepalives with GOAWAY
+            ("grpc.keepalive_permit_without_calls", True),
+            # Minimum interval server allows between client pings (10 seconds)
+            # Prevents "too_many_pings" GOAWAY errors
+            # Must match or be less than client's http2.min_time_between_pings_ms
+            ("grpc.http2.min_ping_interval_without_data_ms", 10000),
         ],
-        description="Additional server options",
+        description="gRPC server options with keepalive support",
     )
     enable_reflection: bool = Field(default=True, description="Enable reflection for the server")
     enable_health_check: bool = Field(default=True, description="Enable health check service")
@@ -262,7 +347,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

@@ -1,26 +1,33 @@
 """This module contains the models for the modules."""
 
-from digitalkin.models.module.module import Module, ModuleStatus
 from digitalkin.models.module.module_context import ModuleContext
 from digitalkin.models.module.module_types import (
     DataModel,
     DataTrigger,
-    InputModelT,
-    OutputModelT,
-    SecretModelT,
     SetupModel,
-    SetupModelT,
+)
+from digitalkin.models.module.tool_reference import (
+    ToolReference,
+    ToolReferenceConfig,
+    ToolSelectionMode,
+)
+from digitalkin.models.module.utility import (
+    EndOfStreamOutput,
+    ModuleStartInfoOutput,
+    UtilityProtocol,
+    UtilityRegistry,
 )
 
 __all__ = [
     "DataModel",
     "DataTrigger",
-    "InputModelT",
-    "Module",
+    "EndOfStreamOutput",
     "ModuleContext",
-    "ModuleStatus",
-    "OutputModelT",
-    "SecretModelT",
+    "ModuleStartInfoOutput",
     "SetupModel",
-    "SetupModelT",
+    "ToolReference",
+    "ToolReferenceConfig",
+    "ToolSelectionMode",
+    "UtilityProtocol",
+    "UtilityRegistry",
 ]

digitalkin/models/module/base_types.py

@@ -0,0 +1,61 @@
+"""Base types for module models."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar
+
+from pydantic import BaseModel, Field
+
+if TYPE_CHECKING:
+    from digitalkin.models.module.setup_types import SetupModel
+
+
+class DataTrigger(BaseModel):
+    """Defines the root input/output model exposing the protocol.
+
+    The mandatory protocol is important to define the module beahvior following the user or agent input/output.
+
+    Example:
+        class MyInput(DataModel):
+            root: DataTrigger
+            user_define_data: Any
+
+        # Usage
+        my_input = MyInput(root=DataTrigger(protocol="message"))
+        print(my_input.root.protocol)  # Output: message
+    """
+
+    protocol: ClassVar[str]
+    created_at: str = Field(
+        default_factory=lambda: datetime.now(tz=timezone.utc).isoformat(),
+        title="Created At",
+        description="Timestamp when the payload was created.",
+    )
+
+
+DataTriggerT = TypeVar("DataTriggerT", bound=DataTrigger)
+
+
+class DataModel(BaseModel, Generic[DataTriggerT]):
+    """Base definition of input/output model showing mandatory root fields.
+
+    The Model define the Module Input/output, usually referring to multiple input/output type defined by an union.
+
+    Example:
+        class ModuleInput(DataModel):
+            root: FileInput | MessageInput
+    """
+
+    root: DataTriggerT
+    annotations: dict[str, str] = Field(
+        default={},
+        title="Annotations",
+        description="Additional metadata or annotations related to the output. ex {'role': 'user'}",
+    )
+
+
+InputModelT = TypeVar("InputModelT", bound=DataModel)
+OutputModelT = TypeVar("OutputModelT", bound=DataModel)
+SecretModelT = TypeVar("SecretModelT", bound=BaseModel)
+SetupModelT = TypeVar("SetupModelT", bound="SetupModel")

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):