digitalkin 0.1.1__py3-none-any.whl

digitalkin/grpc/utils/models.py

@@ -0,0 +1,169 @@
+"""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.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 ServerConfig(BaseModel):
+    """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
+    """
+
+    host: str = Field("0.0.0.0", description="Host address to bind the server to")  # noqa: S104
+    port: int = Field(50051, description="Port to listen on")
+    max_workers: int = Field(10, description="Maximum number of workers for sync mode")
+    mode: ServerMode = Field(ServerMode.SYNC, description="Server operation mode (sync/async)")
+    security: SecurityMode = Field(SecurityMode.INSECURE, description="Security mode (secure/insecure)")
+    credentials: ServerCredentials | None = Field(None, description="Server credentials for secure mode")
+    server_options: list[tuple[str, Any]] = Field(default_factory=list, 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")
+
+    # Enable __slots__ for memory efficiency
+    model_config = {
+        "extra": "forbid",
+        "arbitrary_types_allowed": True,
+        "validate_assignment": True,
+        "use_enum_values": True,
+    }
+
+    @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
+
+    @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 ModuleServerConfig(ServerConfig):
+    """Configuration for Module gRPC server.
+
+    Attributes:
+        registry_address: Address of the registry server
+    """
+
+    registry_address: str | None = Field(None, 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/grpc/utils/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/logger.py

@@ -0,0 +1,17 @@
+"""This module sets up a logger."""
+
+import logging
+import sys
+
+logging.basicConfig(
+    level=logging.DEBUG,
+    stream=sys.stdout,
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+
+logging.getLogger("grpc").setLevel(logging.DEBUG)
+logging.getLogger("asyncio").setLevel(logging.DEBUG)
+
+
+logger = logging.getLogger("digitalkin")

digitalkin/models/__init__.py

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

digitalkin/models/module/__init__.py

@@ -0,0 +1,5 @@
+"""This module contains the models for the modules."""
+
+from .module import Module, ModuleStatus
+
+__all__ = ["Module", "ModuleStatus"]

digitalkin/models/module/module.py

@@ -0,0 +1,31 @@
+"""Module model."""
+
+from enum import Enum, auto
+
+from pydantic import BaseModel
+
+
+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
+    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

digitalkin/models/services/__init__.py

@@ -0,0 +1,6 @@
+"""This module contains the models for the services."""
+
+from .cost import CostEvent
+from .storage import StorageModel
+
+__all__ = ["CostEvent", "StorageModel"]

digitalkin/models/services/cost.py

@@ -0,0 +1,53 @@
+"""Pydantic models for cost service."""
+
+from datetime import datetime, timezone
+from enum import Enum
+from typing import Any
+
+from pydantic import BaseModel, Field
+
+
+class CostTypeEnum(Enum):
+    """Enumeration of supported cost types."""
+
+    TOKEN_INPUT = "token_input"
+    TOKEN_OUTPUT = "token_output"
+    API_CALL = "api_call"
+    STORAGE = "storage"
+    TIME = "time"
+    CUSTOM = "custom"
+
+
+class CostConfig(BaseModel):
+    """Pydantic model that defines a cost configuration.
+
+    :param cost_name: Name of the cost (unique identifier in the service).
+    :param cost_type: The type/category of the cost.
+    :param description: A short description of the cost.
+    :param unit: The unit of measurement (e.g. token, call, MB).
+    :param rate: The cost per unit (e.g. dollars per token).
+    """
+
+    name: str
+    type: CostTypeEnum
+    description: str | None = None
+    unit: str
+    rate: float
+
+
+class CostEvent(BaseModel):
+    """Pydantic model that represents a cost event registered during service execution.
+
+    :param cost_name: Identifier for the cost configuration.
+    :param cost_type: The type of cost.
+    :param usage: The amount or units consumed.
+    :param cost_amount: The computed cost amount; if not provided it is computed as usage*rate.
+    :param timestamp: The time when the cost event was recorded.
+    :param metadata: Additional contextual information about the cost event.
+    """
+
+    name: str
+    usage: float
+    amount: float
+    timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
+    metadata: dict[str, Any] | None = None

digitalkin/models/services/storage.py

@@ -0,0 +1,10 @@
+"""Storage model."""
+
+from pydantic import BaseModel
+
+
+class StorageModel(BaseModel):
+    """Storage model."""
+
+    data: object
+    type: str

digitalkin/modules/__init__.py

@@ -0,0 +1,7 @@
+"""Module package for DigitalKin."""
+
+from .archetype_module import ArchetypeModule
+from .tool_module import ToolModule
+from .trigger_module import TriggerModule
+
+__all__ = ["ArchetypeModule", "ToolModule", "TriggerModule"]

digitalkin/modules/_base_module.py

@@ -0,0 +1,177 @@
+"""BaseModule is the abstract base for all modules in the DigitalKin SDK."""
+
+import asyncio
+import contextlib
+import json
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from typing import Any, ClassVar, Generic, TypeVar
+
+from pydantic import BaseModel
+
+from digitalkin.logger import logger
+from digitalkin.models.module import ModuleStatus
+from digitalkin.services.service_provider import ServiceProvider
+from digitalkin.services.storage.storage_strategy import StorageStrategy
+
+InputModelT = TypeVar("InputModelT", bound=BaseModel)
+OutputModelT = TypeVar("OutputModelT", bound=BaseModel)
+SetupModelT = TypeVar("SetupModelT", bound=BaseModel)
+
+
+class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
+    """BaseModule is the abstract base for all modules in the DigitalKin SDK."""
+
+    input_format: type[InputModelT]
+    output_format: type[OutputModelT]
+    setup_format: type[SetupModelT]
+    metadata: ClassVar[dict[str, Any]]
+
+    local_services: type[ServiceProvider]
+    dev_services: type[ServiceProvider]
+
+    storage: StorageStrategy
+
+    def __init__(
+        self,
+        job_id: str,
+        name: str | None = None,
+    ) -> None:
+        """Initialize the module."""
+        self.job_id: str = job_id
+        self.name = name or self.__class__.__name__
+        self._status = ModuleStatus.CREATED
+        self._task: asyncio.Task | None = None
+
+    @property
+    def status(self) -> ModuleStatus:
+        """Get the module status.
+
+        Returns:
+            The module status
+        """
+        return self._status
+
+    @classmethod
+    def get_input_format(cls, llm_format: bool) -> str:  # noqa: FBT001
+        """Get the JSON schema of the input format model.
+
+        Raises:
+            NotImplementedError: If the `input_format` is not defined.
+
+        Returns:
+            The JSON schema of the input format as a string.
+        """
+        if cls.output_format is not None:
+            if llm_format:
+                return json.dumps(cls.input_format, indent=2)
+            return json.dumps(cls.input_format.model_json_schema(), indent=2)
+        msg = f"{cls.__name__}' class does not define an 'input_format'."
+        raise NotImplementedError(msg)
+
+    @classmethod
+    def get_output_format(cls, llm_format: bool) -> str:  # noqa: FBT001
+        """Get the JSON schema of the output format model.
+
+        Raises:
+            NotImplementedError: If the `output_format` is not defined.
+
+        Returns:
+            The JSON schema of the output format as a string.
+        """
+        if cls.output_format is not None:
+            if llm_format:
+                return json.dumps(cls.output_format, indent=2)
+            return json.dumps(cls.output_format.model_json_schema(), indent=2)
+        msg = "'%s' class does not define an 'output_format'."
+        raise NotImplementedError(msg)
+
+    @classmethod
+    def get_setup_format(cls, llm_format: bool) -> str:  # noqa: FBT001
+        """Gets the JSON schema of the setup format model.
+
+        Raises:
+            NotImplementedError: If the `setup_format` is not defined.
+
+        Returns:
+            The JSON schema of the setup format as a string.
+        """
+        if cls.setup_format is not None:
+            if llm_format:
+                return json.dumps(cls.setup_format, indent=2)
+            return json.dumps(cls.setup_format.model_json_schema(), indent=2)
+        msg = "'%s' class does not define an 'setup_format'."
+        raise NotImplementedError(msg)
+
+    @abstractmethod
+    async def initialize(self, setup_data: dict[str, Any]) -> None:
+        """Initialize the module."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def run(
+        self,
+        input_data: dict[str, Any],
+        setup_data: dict[str, Any],
+        callback: Callable,
+    ) -> None:
+        """Run the module."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def cleanup(self) -> None:
+        """Run the module."""
+        raise NotImplementedError
+
+    async def _run_lifecycle(
+        self,
+        input_data: dict[str, Any],
+        setup_data: dict[str, Any],
+        callback: Callable,
+    ) -> None:
+        """Run the module lifecycle.
+
+        Raises:
+            asyncio.CancelledError: If the module is cancelled
+        """
+        try:
+            await self.run(input_data, setup_data, callback)
+            await self.stop()
+        except asyncio.CancelledError:
+            logger.info(f"Module {self.name} cancelled")
+        except Exception:
+            self._status = ModuleStatus.FAILED
+            logger.exception("Error inside module %s", self.name)
+        else:
+            self._status = ModuleStatus.STOPPED
+
+    async def start(
+        self,
+        input_data: dict[str, Any],
+        setup_data: dict[str, Any],
+        callback: Callable,
+    ) -> None:
+        """Start the module."""
+        try:
+            await self.initialize(setup_data=setup_data)
+            self._status = ModuleStatus.RUNNING
+            self._task = asyncio.create_task(self._run_lifecycle(input_data, setup_data, callback))
+        except Exception:
+            self._status = ModuleStatus.FAILED
+            logger.exception("Error starting module")
+
+    async def stop(self) -> None:
+        """Stop the module."""
+        if self._status != ModuleStatus.RUNNING:
+            return
+
+        try:
+            self._status = ModuleStatus.STOPPING
+            if self._task and not self._task.done():
+                self._task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await self._task
+            await self.cleanup()
+        except Exception:
+            self._status = ModuleStatus.FAILED
+            logger.exception("Error stopping module")

digitalkin/modules/archetype_module.py

@@ -0,0 +1,14 @@
+"""ArchetypeModule extends BaseModule to implement specific module types."""
+
+from abc import ABC
+
+from digitalkin.modules._base_module import BaseModule
+
+
+class ArchetypeModule(BaseModule, ABC):
+    """ArchetypeModule extends BaseModule to implement specific module types."""
+
+    def __init__(self, name: str | None = None) -> None:
+        """Initialize the module with the given metadata."""
+        super().__init__(self.job_id, name=name)
+        self.capabilities = ["archetype"]

digitalkin/modules/job_manager.py

@@ -0,0 +1,158 @@
+"""Background module manager."""
+
+import asyncio
+import uuid
+from argparse import ArgumentParser, Namespace
+from collections.abc import Callable
+from typing import Any
+
+from digitalkin.logger import logger
+from digitalkin.models import ModuleStatus
+from digitalkin.modules._base_module import BaseModule
+from digitalkin.utils.arg_parser import ArgParser, DevelopmentModeMappingAction
+
+
+class JobManager(ArgParser):
+    """Background module manager."""
+
+    args: Namespace
+
+    def _add_parser_args(self, parser: ArgumentParser) -> None:
+        class_mapping = {
+            "local": self.module_class.local_services,
+            "development": self.module_class.dev_services,
+        }
+
+        super()._add_parser_args(parser)
+        parser.add_argument(
+            "-d",
+            "--dev-mode",
+            env_var="SERVICE_PROVIDER",
+            class_mapping=class_mapping,
+            choices=class_mapping.keys(),
+            default="local",
+            action=DevelopmentModeMappingAction,
+            dest="service_providers",
+            help="Define Module Service configurations for endpoints",
+        )
+
+    def __init__(self, module_class: type[BaseModule]) -> None:
+        """Initialize the job manager."""
+        self.module_class = module_class
+        self.modules: dict[str, BaseModule] = {}
+        self._lock = asyncio.Lock()
+        super().__init__()
+
+        explicit_fields = {
+            name: self.args.service_providers.__dict__[name]
+            for name in self.args.service_providers.__class_vars__
+            if name in self.args.service_providers.__dict__
+        }
+
+        # services are now available as class vars.
+        # init the services provided allowing cold start during module creation
+        for service_name in explicit_fields:
+            service_type = getattr(self.args.service_providers, service_name)
+            setattr(self.module_class, service_name, service_type)
+
+    async def create_job(  # noqa: D417
+        self,
+        input_data: dict[str, Any],
+        setup_data: dict[str, Any],
+        callback: Callable,
+        *args: tuple,
+        **kwargs: dict,
+    ) -> tuple[str, BaseModule]:
+        """Start new module job in background (asyncio).
+
+        Args:
+            module_class: Classe du module à instancier
+            *args: Arguments à passer au constructeur du module
+            **kwargs: Arguments à passer au constructeur du module
+
+        Returns:
+            str: job_id of the module entity
+        """
+        job_id = str(uuid.uuid4())
+        """TODO: check uniqueness of the job_id"""
+        # Création et démarrage du module
+        module = self.module_class(job_id, *args, **kwargs)  # type: ignore
+        self.modules[job_id] = module
+        try:
+            await module.start(input_data, setup_data, callback)
+            logger.info("Module %s (%s) started successfully", job_id, module.name)
+        except Exception:
+            # En cas d'erreur, supprimer le module du gestionnaire
+            del self.modules[job_id]
+            logger.exception("Échec du démarrage du module %s: %s", job_id)
+            raise
+        else:
+            return job_id, module
+
+    async def stop_module(self, job_id: str) -> bool:
+        """Arrête un module en cours d'exécution.
+
+        Args:
+            job_id: Identifiant du module à arrêter
+
+        Returns:
+            True si le module a été arrêté, False s'il n'existe pas.
+        """
+        async with self._lock:
+            module = self.modules.get(job_id)
+            if not module:
+                logger.warning(f"Module {job_id} introuvable")
+                return False
+            try:
+                await module.stop()
+                logger.info(f"Module {job_id} ({module.name}) arrêté avec succès")
+            except Exception as e:
+                logger.error(f"Erreur lors de l'arrêt du module {job_id}: {e}")
+                raise
+            else:
+                return True
+
+    def get_module_status(self, job_id: str) -> ModuleStatus | None:
+        """Obtient le statut d'un module.
+
+        Args:
+            job_id: Identifiant du module
+
+        Returns:
+            Le statut du module ou None si le module n'existe pas.
+        """
+        module = self.modules.get(job_id)
+        return module.status if module else None
+
+    def get_module(self, job_id: str) -> BaseModule | None:
+        """Récupère une référence au module.
+
+        Args:
+            job_id: Identifiant du module
+
+        Returns:
+            Le module ou None s'il n'existe pas.
+        """
+        return self.modules.get(job_id)
+
+    async def stop_all_modules(self) -> None:
+        """Arrête tous les modules en cours d'exécution."""
+        async with self._lock:
+            stop_tasks = [self.stop_module(job_id) for job_id in list(self.modules.keys())]
+            if stop_tasks:
+                await asyncio.gather(*stop_tasks, return_exceptions=True)
+
+    def list_modules(self) -> dict[str, dict[str, Any]]:
+        """Liste tous les modules avec leur statut.
+
+        Returns:
+            Dictionnaire des modules avec leurs informations.
+        """
+        return {
+            job_id: {
+                "name": module.name,
+                "status": module.status,
+                "class": module.__class__.__name__,
+            }
+            for job_id, module in self.modules.items()
+        }