digitalkin 0.3.2.dev2__py3-none-any.whl

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

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

digitalkin/models/core/__init__.py

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

digitalkin/models/core/job_manager_models.py

@@ -0,0 +1,36 @@
+"""Job manager models."""
+
+from enum import Enum
+
+from digitalkin.core.job_manager.base_job_manager import BaseJobManager
+
+
+class JobManagerMode(Enum):
+    """Job manager mode."""
+
+    SINGLE = "single"
+    TASKIQ = "taskiq"
+
+    def __str__(self) -> str:
+        """Get the string representation of the job manager mode.
+
+        Returns:
+            str: job manager mode name.
+        """
+        return self.value
+
+    def get_manager_class(self) -> type[BaseJobManager]:
+        """Get the job manager class based on the mode.
+
+        Returns:
+            type: The job manager class.
+        """
+        match self:
+            case JobManagerMode.SINGLE:
+                from digitalkin.core.job_manager.single_job_manager import SingleJobManager  # noqa: PLC0415
+
+                return SingleJobManager
+            case JobManagerMode.TASKIQ:
+                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/models/grpc_servers/models.py

@@ -0,0 +1,275 @@
+"""Data models for gRPC server configurations."""
+
+from enum import Enum
+from pathlib import Path
+from typing import Any
+
+from pydantic import BaseModel, Field, ValidationInfo, field_validator
+
+from digitalkin.grpc_servers.utils.exceptions import ConfigurationError, SecurityError
+
+
+class ServerMode(str, Enum):
+    """Enum for server operation mode."""
+
+    SYNC = "sync"
+    ASYNC = "async"
+
+
+class SecurityMode(str, Enum):
+    """Enum for server security mode."""
+
+    SECURE = "secure"
+    INSECURE = "insecure"
+
+
+class ServerCredentials(BaseModel):
+    """Model for server credentials in secure mode.
+
+    Attributes:
+        server_key_path: Path to the server private key
+        server_cert_path: Path to the server certificate
+        root_cert_path: Optional path to the root certificate
+    """
+
+    server_key_path: Path = Field(..., description="Path to the server private key")
+    server_cert_path: Path = Field(..., description="Path to the server certificate")
+    root_cert_path: Path | None = Field(None, description="Path to the root certificate")
+
+    # Enable __slots__ for memory efficiency
+    model_config = {
+        "extra": "forbid",
+        "arbitrary_types_allowed": True,
+        "validate_assignment": True,
+        "use_enum_values": True,
+        "frozen": True,  # Make immutable
+    }
+
+    @field_validator("server_key_path", "server_cert_path", "root_cert_path")
+    @classmethod
+    def check_path_exists(cls, v: Path | None) -> Path | None:
+        """Validate that the file path exists.
+
+        Args:
+            v: Path to validate
+
+        Returns:
+            The validated path
+
+        Raises:
+            SecurityError: If the path does not exist
+        """
+        if v is not None and not v.exists():
+            msg = f"File not found: {v}"
+            raise SecurityError(msg)
+        return v
+
+
+class ClientCredentials(BaseModel):
+    """Model for client credentials in secure mode.
+
+    Attributes:
+        root_cert_path: path to the root certificate
+        client_key_path: Path to the client private key
+        client_cert_path: Path to the client certificate
+    """
+
+    root_cert_path: Path = Field(..., description="Path to the root certificate")
+    client_key_path: Path | None = Field(None, description="Path to the client private key | mTLS enable")
+    client_cert_path: Path | None = Field(None, description="Path to the client certificate | mTLS enable")
+
+    # Enable __slots__ for memory efficiency
+    model_config = {
+        "extra": "forbid",
+        "arbitrary_types_allowed": True,
+        "validate_assignment": True,
+        "use_enum_values": True,
+        "frozen": True,  # Make immutable
+    }
+
+    @field_validator("client_key_path", "client_cert_path", "root_cert_path")
+    @classmethod
+    def check_path_exists(cls, v: Path | None) -> Path | None:
+        """Validate that the file path exists.
+
+        Args:
+            v: Path to validate
+
+        Returns:
+            The validated path
+
+        Raises:
+            SecurityError: If the path does not exist
+        """
+        if v is not None and not v.exists():
+            msg = f"File not found: {v}"
+            raise SecurityError(msg)
+        return v
+
+
+class ChannelConfig(BaseModel):
+    """Base configuration for gRPC channels.
+
+    Attributes:
+        host: Host address
+        port: Port to listen on
+        mode: communication operation mode (sync/async)
+        security: Security mode (secure/insecure)
+        credentials: Client credentials for secure mode
+    """
+
+    host: str = Field("0.0.0.0", description="Host address to bind the client to")  # noqa: S104
+    port: int = Field(50051, description="Port to listen on")
+    mode: ServerMode = Field(ServerMode.SYNC, description="Client operation mode (sync/async)")
+    security: SecurityMode = Field(SecurityMode.INSECURE, description="Security mode (secure/insecure)")
+
+    # Enable __slots__ for memory efficiency
+    model_config = {
+        "extra": "forbid",
+        "arbitrary_types_allowed": True,
+        "validate_assignment": True,
+        "use_enum_values": True,
+    }
+
+    @field_validator("port")
+    @classmethod
+    def validate_port(cls, v: int) -> int:
+        """Validate that the port is in a valid range.
+
+        Args:
+            v: Port number to validate
+
+        Returns:
+            The validated port number
+
+        Raises:
+            ConfigurationError: If port is outside valid range
+        """
+        if not 0 < v < 65536:  # noqa: PLR2004
+            msg = f"Port must be between 1 and 65535, got {v}"
+            raise ConfigurationError(msg)
+        return v
+
+    @property
+    def address(self) -> str:
+        """Get the server address.
+
+        Returns:
+            The formatted address string
+        """
+        return f"{self.host}:{self.port}"
+
+
+class ClientConfig(ChannelConfig):
+    """Base configuration for gRPC clients.
+
+    Attributes:
+        host: Host address to bind the client to
+        port: Port to listen on
+        mode: Client operation mode (sync/async)
+        security: Security mode (secure/insecure)
+        credentials: Client credentials for secure mode
+        channel_options: Additional channel options
+    """
+
+    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", 100 * 1024 * 1024),  # 100MB
+            ("grpc.max_send_message_length", 100 * 1024 * 1024),  # 100MB
+        ],
+        description="Additional channel options",
+    )
+
+    @field_validator("credentials")
+    @classmethod
+    def validate_credentials(cls, v: ClientCredentials | None, info: ValidationInfo) -> ClientCredentials | None:
+        """Validate that credentials are provided when in secure mode.
+
+        Args:
+            v: The credentials value
+            info: ValidationInfo containing other field values
+
+        Returns:
+            The validated credentials
+
+        Raises:
+            ConfigurationError: If credentials are missing in secure mode
+        """
+        # Access security mode from the info.data dictionary
+        security = info.data.get("security")
+
+        if security == SecurityMode.SECURE and v is None:
+            msg = "Credentials must be provided when using secure mode"
+            raise ConfigurationError(msg)
+        return v
+
+
+class ServerConfig(ChannelConfig):
+    """Base configuration for gRPC servers.
+
+    Attributes:
+        host: Host address to bind the server to
+        port: Port to listen on
+        max_workers: Maximum number of workers for sync mode
+        mode: Server operation mode (sync/async)
+        security: Security mode (secure/insecure)
+        credentials: Server credentials for secure mode
+        server_options: Additional server options
+        enable_reflection: Enable reflection for the server
+    """
+
+    max_workers: int = Field(10, description="Maximum number of workers for sync mode")
+    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", 100 * 1024 * 1024),  # 100MB
+            ("grpc.max_send_message_length", 100 * 1024 * 1024),  # 100MB
+        ],
+        description="Additional server options",
+    )
+    enable_reflection: bool = Field(default=True, description="Enable reflection for the server")
+    enable_health_check: bool = Field(default=True, description="Enable health check service")
+
+    @field_validator("credentials")
+    @classmethod
+    def validate_credentials(cls, v: ServerCredentials | None, info: ValidationInfo) -> ServerCredentials | None:
+        """Validate that credentials are provided when in secure mode.
+
+        Args:
+            v: The credentials value
+            info: ValidationInfo containing other field values
+
+        Returns:
+            The validated credentials
+
+        Raises:
+            ConfigurationError: If credentials are missing in secure mode
+        """
+        # Access security mode from the info.data dictionary
+        security = info.data.get("security")
+
+        if security == SecurityMode.SECURE and v is None:
+            msg = "Credentials must be provided when using secure mode"
+            raise ConfigurationError(msg)
+        return v
+
+
+class ModuleServerConfig(ServerConfig):
+    """Configuration for Module gRPC server.
+
+    Attributes:
+        registry_address: Address of the registry server
+    """
+
+    registry_address: str = Field(..., description="Address of the registry server")
+
+
+class RegistryServerConfig(ServerConfig):
+    """Configuration for Registry gRPC server.
+
+    Attributes:
+        database_url: Database URL for registry data storage
+    """
+
+    database_url: str | None = Field(None, description="Database URL for registry data storage")

digitalkin/models/grpc_servers/types.py

@@ -0,0 +1,24 @@
+"""Type definitions for gRPC utilities."""
+
+from typing import Protocol, TypeVar
+
+import grpc
+from grpc import aio as grpc_aio
+
+GrpcServer = grpc.Server | grpc_aio.Server
+
+# Create a type variable for servicer implementations
+T = TypeVar("T")
+
+
+class ServiceObject(Protocol):
+    """Protocol for individual services in a gRPC descriptor."""
+
+    full_name: str
+
+
+# Create a protocol for service descriptors
+class ServiceDescriptor(Protocol):
+    """Protocol for gRPC service descriptors."""
+
+    services_by_name: dict[str, ServiceObject]

digitalkin/models/module/__init__.py

@@ -0,0 +1,25 @@
+"""This module contains the models for the modules."""
+
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.models.module.module_types import (
+    DataModel,
+    DataTrigger,
+    SetupModel,
+)
+from digitalkin.models.module.utility import (
+    EndOfStreamOutput,
+    UtilityProtocol,
+    UtilityRegistry,
+)
+
+__all__ = [
+    # Core types (used by all SDK users)
+    "DataModel",
+    "DataTrigger",
+    # Utility (commonly used)
+    "EndOfStreamOutput",
+    "ModuleContext",
+    "SetupModel",
+    "UtilityProtocol",
+    "UtilityRegistry",
+]

digitalkin/models/module/module.py

@@ -0,0 +1,40 @@
+"""Module model."""
+
+from enum import Enum, auto
+
+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):
+    """Possible module's state."""
+
+    CREATED = auto()  # Module created but not started
+    STARTING = auto()  # Module is starting
+    RUNNING = auto()  # Module do run
+    STOPPING = auto()  # Module is stopping
+    STOPPED = auto()  # Module stop successfuly
+    FAILED = auto()  # Module stopped due to internal error
+    CANCELLED = auto()  # Module stopped due to internal error
+    NOT_FOUND = auto()
+
+
+class Module(BaseModel):
+    """Module model."""
+
+    name: str
+    cost_schema: str
+    input_schema: str
+    output_schema: str
+    setup_schema: str
+    secret_schema: str
+    type: str
+    version: str
+    description: str