digitalkin 0.1.1__py3-none-any.whl → 0.2.1__py3-none-any.whl

digitalkin/modules/job_manager.py

@@ -3,12 +3,15 @@
 import asyncio
 import uuid
 from argparse import ArgumentParser, Namespace
-from collections.abc import Callable
+from collections.abc import Callable, Coroutine
 from typing import Any
 
 from digitalkin.logger import logger
 from digitalkin.models import ModuleStatus
+from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
 from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.services_config import ServicesConfig
+from digitalkin.services.services_models import ServicesMode
 from digitalkin.utils.arg_parser import ArgParser, DevelopmentModeMappingAction
 
 
@@ -17,22 +20,43 @@ class JobManager(ArgParser):
 
     args: Namespace
 
-    def _add_parser_args(self, parser: ArgumentParser) -> None:
-        class_mapping = {
-            "local": self.module_class.local_services,
-            "development": self.module_class.dev_services,
-        }
+    @staticmethod
+    async def _job_specific_callback(
+        callback: Callable[[str, OutputModelT], Coroutine[Any, Any, None]], job_id: str
+    ) -> Callable[[OutputModelT], Coroutine[Any, Any, None]]:
+        """Return a callback function for the job.
+
+        Args:
+            callback: Callback function to be called when the job is done
+            job_id: Identifiant du module
+
+        Returns:
+            Callable: Callback function
+        """
 
+        def callback_wrapper(output_data: OutputModelT) -> Coroutine[Any, Any, None]:
+            """Wrapper for the callback function.
+
+            Args:
+                output_data: Output data of the job
+
+            Returns:
+                Coroutine: Callback function
+            """
+            return callback(job_id, output_data)
+
+        return callback_wrapper
+
+    def _add_parser_args(self, parser: ArgumentParser) -> None:
         super()._add_parser_args(parser)
         parser.add_argument(
             "-d",
             "--dev-mode",
-            env_var="SERVICE_PROVIDER",
-            class_mapping=class_mapping,
-            choices=class_mapping.keys(),
+            env_var="SERVICE_MODE",
+            choices=ServicesMode.__members__,
             default="local",
             action=DevelopmentModeMappingAction,
-            dest="service_providers",
+            dest="services_mode",
             help="Define Module Service configurations for endpoints",
         )
 
@@ -43,25 +67,18 @@ class JobManager(ArgParser):
         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)
+        services_config = ServicesConfig(
+            services_config_strategies=self.module_class.services_config_strategies,
+            services_config_params=self.module_class.services_config_params,
+            mode=self.args.services_mode,
+        )
+        setattr(self.module_class, "services_config", services_config)
 
     async def create_job(  # noqa: D417
         self,
-        input_data: dict[str, Any],
-        setup_data: dict[str, Any],
-        callback: Callable,
-        *args: tuple,
-        **kwargs: dict,
+        input_data: InputModelT,
+        setup_data: SetupModelT,
+        callback: Callable[[str, OutputModelT], Coroutine[Any, Any, None]],
     ) -> tuple[str, BaseModule]:
         """Start new module job in background (asyncio).
 
@@ -74,12 +91,13 @@ class JobManager(ArgParser):
             str: job_id of the module entity
         """
         job_id = str(uuid.uuid4())
+        mission_id = "missions:test_demo"
         """TODO: check uniqueness of the job_id"""
         # Création et démarrage du module
-        module = self.module_class(job_id, *args, **kwargs)  # type: ignore
+        module = self.module_class(job_id, mission_id=mission_id)
         self.modules[job_id] = module
         try:
-            await module.start(input_data, setup_data, callback)
+            await module.start(input_data, setup_data, await JobManager._job_specific_callback(callback, job_id))
             logger.info("Module %s (%s) started successfully", job_id, module.name)
         except Exception:
             # En cas d'erreur, supprimer le module du gestionnaire

digitalkin/modules/tool_module.py

@@ -2,13 +2,9 @@
 
 from abc import ABC
 
-from ._base_module import BaseModule
+from digitalkin.models.module import InputModelT, OutputModelT, SecretModelT, SetupModelT
+from digitalkin.modules._base_module import BaseModule # type: ignore
 
 
-class ToolModule(BaseModule, ABC):
+class ToolModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT], ABC):
     """ToolModule extends BaseModule to implement specific module types."""
-
-    def __init__(self, metadata):
-        """Initialize the module with the given metadata."""
-        super().__init__(metadata)
-        self.capabilities = ["tool"]

digitalkin/modules/trigger_module.py

@@ -2,13 +2,8 @@
 
 from abc import ABC
 
-from ._base_module import BaseModule
+from digitalkin.modules._base_module import BaseModule, InputModelT, OutputModelT, SecretModelT, SetupModelT
 
 
-class TriggerModule(BaseModule, ABC):
+class TriggerModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT], ABC):
     """TriggerModule extends BaseModule to implement specific module types."""
-
-    def __init__(self, metadata):
-        """Initialize the module with the given metadata."""
-        super().__init__(metadata)
-        self.capabilities = ["trigger"]

digitalkin/services/__init__.py

@@ -1,13 +1,12 @@
 """This package contains the abstract base class for all services."""
 
-from .agent import AgentStrategy, DefaultAgent
-from .cost import CostStrategy, DefaultCost
-from .filesystem import DefaultFilesystem, FilesystemStrategy
-from .identity import DefaultIdentity, IdentityStrategy
-from .registry import DefaultRegistry, RegistryStrategy
-from .service_provider import ServiceProvider
-from .snapshot import DefaultSnapshot, SnapshotStrategy
-from .storage import DefaultStorage, StorageStrategy
+from digitalkin.services.agent import AgentStrategy, DefaultAgent
+from digitalkin.services.cost import CostStrategy, DefaultCost
+from digitalkin.services.filesystem import DefaultFilesystem, FilesystemStrategy
+from digitalkin.services.identity import DefaultIdentity, IdentityStrategy
+from digitalkin.services.registry import DefaultRegistry, RegistryStrategy
+from digitalkin.services.snapshot import DefaultSnapshot, SnapshotStrategy
+from digitalkin.services.storage import DefaultStorage, StorageStrategy
 
 __all__ = [
     "AgentStrategy",
@@ -22,7 +21,6 @@ __all__ = [
     "FilesystemStrategy",
     "IdentityStrategy",
     "RegistryStrategy",
-    "ServiceProvider",
     "SnapshotStrategy",
     "StorageStrategy",
 ]

digitalkin/services/agent/__init__.py

@@ -1,6 +1,6 @@
 """This module is responsible for handling the agent services."""
 
-from .agent_strategy import AgentStrategy
-from .default_agent import DefaultAgent
+from digitalkin.services.agent.agent_strategy import AgentStrategy
+from digitalkin.services.agent.default_agent import DefaultAgent
 
 __all__ = ["AgentStrategy", "DefaultAgent"]

digitalkin/services/agent/agent_strategy.py

@@ -2,14 +2,11 @@
 
 from abc import ABC, abstractmethod
 
+from digitalkin.services.base_strategy import BaseStrategy
 
-class AgentStrategy(ABC):
-    """Abstract base class for agent strategies."""
 
-    @abstractmethod
-    def __init__(self) -> None:
-        """Initialize the agent strategy."""
-        raise NotImplementedError
+class AgentStrategy(BaseStrategy, ABC):
+    """Abstract base class for agent strategies."""
 
     @abstractmethod
     def start(self) -> None:

digitalkin/services/agent/default_agent.py

@@ -1,14 +1,11 @@
 """Default agent implementation for the agent service."""
 
-from .agent_strategy import AgentStrategy
+from digitalkin.services.agent.agent_strategy import AgentStrategy
 
 
 class DefaultAgent(AgentStrategy):
     """Default agent implementation for the agent service."""
 
-    def __init__(self) -> None:
-        """Initialize the default agent."""
-
     def start(self) -> None:
         """Start the agent."""
 

digitalkin/services/base_strategy.py

@@ -0,0 +1,18 @@
+"""This module contains the abstract base class for storage strategies."""
+
+from abc import ABC
+
+
+class BaseStrategy(ABC):
+    """Abstract base class for all strategies.
+
+    This class defines the interface for all strategies.
+    """
+
+    def __init__(self, mission_id: str) -> None:
+        """Initialize the strategy.
+
+        Args:
+            mission_id: The ID of the mission this strategy is associated with
+        """
+        self.mission_id: str = mission_id

digitalkin/services/cost/__init__.py

@@ -1,6 +1,7 @@
 """This module is responsible for handling the cost services."""
 
-from .cost_strategy import CostStrategy
-from .default_cost import DefaultCost
+from digitalkin.services.cost.cost_strategy import CostStrategy
+from digitalkin.services.cost.default_cost import DefaultCost
+from digitalkin.services.cost.grpc_cost import GrpcCost
 
-__all__ = ["CostStrategy", "DefaultCost"]
+__all__ = ["CostStrategy", "DefaultCost", "GrpcCost"]

digitalkin/services/cost/cost_strategy.py

@@ -1,15 +1,45 @@
 """This module contains the abstract base class for cost strategies."""
 
 from abc import ABC, abstractmethod
+from enum import Enum, auto
+from typing import Any
 
+from pydantic import BaseModel
 
-class CostStrategy(ABC):
+from digitalkin.services.base_strategy import BaseStrategy
+
+
+class CostType(Enum):
+    """."""
+
+    OTHER = auto()
+    TOKEN_INPUT = auto()
+    TOKEN_OUTPUT = auto()
+    API_CALL = auto()
+    STORAGE = auto()
+    TIME = auto()
+
+
+class CostData(BaseModel):
+    """."""
+
+    cost: float
+    mission_id: str
+    name: str
+    type: CostType
+    unit: str
+
+
+class CostStrategy(BaseStrategy, ABC):
     """Abstract base class for cost strategies."""
 
     @abstractmethod
-    def __init__(self) -> None:
-        """Initialize the cost strategy."""
+    def add_cost(self, cost_dict: dict[str, Any]) -> str:
+        """Register a new cost."""
+
+    def __post_init__(self, *args, **kwargs) -> None:  # noqa: ANN002, ANN003
+        """Allow post init configuration."""
 
     @abstractmethod
-    def register_cost(self, data: dict[str, str]) -> None:
-        """Register a new cost."""
+    def get(self, cost_dict: dict[str, Any]) -> list[CostData]:
+        """Get a cost."""

digitalkin/services/cost/default_cost.py

@@ -1,13 +1,30 @@
 """Default cost."""
 
-from .cost_strategy import CostStrategy
+import logging
+from typing import Any
+
+from digitalkin.services.cost.cost_strategy import CostData, CostStrategy
+
+logger = logging.getLogger(__name__)
 
 
 class DefaultCost(CostStrategy):
     """Default cost strategy."""
 
-    def __init__(self) -> None:
-        """Initialize the cost strategy."""
+    def add_cost(self, cost_dict: dict[str, Any]) -> str:  # noqa: PLR6301
+        """Create a new record in the cost database.
+
+        Returns:
+            str: The ID of the new record
+        """
+        logger.info("Cost added with cost_dict: %s", cost_dict)
+        return ""
+
+    def get(self, cost_dict: dict[str, Any]) -> list[CostData]:  # noqa: PLR6301
+        """Get records from the database.
 
-    def register_cost(self, data: dict[str, str]) -> None:
-        """Register a new cost."""
+        Returns:
+            list[CostData]: The list of records
+        """
+        logger.info("Costs querried with cost_dict: %s", cost_dict)
+        return []

digitalkin/services/cost/grpc_cost.py

@@ -0,0 +1,81 @@
+"""This module implements the default Cost strategy."""
+
+import logging
+from typing import Any
+
+from digitalkin_proto.digitalkin.cost.v1 import cost_pb2, cost_service_pb2_grpc
+from google.protobuf import json_format
+from pydantic import ValidationError
+
+from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
+from digitalkin.grpc_servers.utils.models import ServerConfig
+from digitalkin.services.cost.cost_strategy import CostData, CostStrategy
+
+logger = logging.getLogger(__name__)
+
+
+class GrpcCost(CostStrategy, GrpcClientWrapper):
+    """This class implements the default Cost strategy."""
+
+    def _get_costs_by_name(self, cost_dict: dict[str, Any]) -> list[CostData]:
+        request = cost_pb2.GetCostsByNameRequest(
+            mission_id=cost_dict["mission_id"],
+            name=cost_dict["name"],
+        )
+        response = self.exec_grpc_query("GetCostsByName", request)
+        return [CostData(**json_format.MessageToDict(cost)) for cost in response.costs]
+
+    def _get_costs_by_mission(self, cost_dict: dict[str, Any]) -> list[CostData]:
+        request = cost_pb2.GetCostsByMissionRequest(mission_id=cost_dict["mission_id"])
+        response = self.exec_grpc_query("GetCostsByMission", request)
+        return [CostData(**json_format.MessageToDict(cost)) for cost in response.costs]
+
+    def _get_costs_by_type(self, cost_dict: dict[str, Any]) -> list[CostData]:
+        request = cost_pb2.GetCostsByTypeRequest(
+            mission_id=cost_dict["mission_id"],
+            type=cost_dict["type"],
+        )
+        response = self.exec_grpc_query("GetCostsBytype", request)
+        return [CostData(**json_format.MessageToDict(cost)) for cost in response.costs]
+
+    def __init__(self, mission_id: str, config: ServerConfig) -> None:
+        """Initialize the cost."""
+        super().__init__(mission_id)
+        channel = self._init_channel(config)
+        self.stub = cost_service_pb2_grpc.CostServiceStub(channel)
+        logger.info("Channel client 'Cost' initialized succesfully")
+
+    def add_cost(self, cost_dict: dict[str, Any]) -> str:
+        """Create a new record in the cost database.
+
+        Required arguments:
+            data: Object representation of CostData
+
+        Returns:
+            str: The ID of the new record
+        """
+        try:
+            valid_data = CostData.model_validate(cost_dict["data"])
+        except ValidationError:
+            logger.exception("Validation failed for model StorageData")
+            return ""
+        except KeyError:
+            logger.exception("Missing mandatory 'data' in dict.")
+            return ""
+
+        request = cost_pb2.AddCostRequest(**valid_data.model_dump())
+        return self.exec_grpc_query("AddCost", request)
+
+    def get(self, cost_dict: dict[str, Any]) -> list[CostData]:
+        """Get records from the database.
+
+        Returns:
+            list[CostData]: The list of records
+        """
+        if "mission_id" not in cost_dict:
+            return []
+        if "name" in cost_dict:
+            return self._get_costs_by_name(cost_dict)
+        if "type" in cost_dict:
+            return self._get_costs_by_type(cost_dict)
+        return self._get_costs_by_mission(cost_dict)

digitalkin/services/filesystem/__init__.py

@@ -1,6 +1,7 @@
 """This module is responsible for handling the filesystem services."""
 
-from .default_filesystem import DefaultFilesystem
-from .filesystem_strategy import FilesystemStrategy
+from digitalkin.services.filesystem.default_filesystem import DefaultFilesystem
+from digitalkin.services.filesystem.filesystem_strategy import FilesystemStrategy
+from digitalkin.services.filesystem.grpc_filesystem import GrpcFilesystem
 
-__all__ = ["DefaultFilesystem", "FilesystemStrategy"]
+__all__ = ["DefaultFilesystem", "FilesystemStrategy", "GrpcFilesystem"]

digitalkin/services/filesystem/default_filesystem.py

@@ -1,29 +1,209 @@
 """Default filesystem."""
 
-from typing import Any
+import logging
+import os
+import tempfile
+from pathlib import Path
 
-from .filesystem_strategy import FilesystemStrategy
+from digitalkin.services.filesystem.filesystem_strategy import (
+    FilesystemData,
+    FilesystemServiceError,
+    FilesystemStrategy,
+    FileType,
+)
+
+logger = logging.getLogger(__name__)
 
 
 class DefaultFilesystem(FilesystemStrategy):
     """Default state filesystem strategy."""
 
-    def create(self, data: dict[str, Any]) -> str:
-        """Create a new file in the file system."""
-        raise NotImplementedError
+    def __init__(self, mission_id: str, config: dict[str, str]) -> None:
+        """Initialize the default filesystem strategy.
+
+        Args:
+            mission_id: The ID of the mission this strategy is associated with
+            config: A dictionary mapping names to Pydantic model classes
+        """
+        super().__init__(mission_id, config)
+        self.temp_root: str = self.config.get("temp_root", "") or tempfile.gettempdir()
+        os.makedirs(self.temp_root, exist_ok=True)
+        self.db: dict[str, FilesystemData] = {}
+
+    def _get_kin_context_temp_dir(self, kin_context: str) -> str:
+        """Get the temporary directory path for a specific kin_context.
+
+        Args:
+            kin_context: The mission ID or setup ID.
+
+        Returns:
+            str: Path to the kin_context's temporary directory
+        """
+        # Create a kin_context-specific directory to organize files
+        kin_context_dir = os.path.join(self.temp_root, kin_context.replace(":", "_"))
+        os.makedirs(kin_context_dir, exist_ok=True)
+        return kin_context_dir
+
+    def upload(self, content: bytes, name: str, file_type: FileType) -> FilesystemData:
+        """Create a new file in the file system.
+
+        Args:
+            content: The content of the file to be uploaded
+            name: The name of the file to be created
+            file_type: The type of data being uploaded
+
+        Returns:
+            FilesystemData: Metadata about the uploaded file
+
+        Raises:
+            FileExistsError: If the file already exists
+            FilesystemServiceError: If there is an error during upload
+        """
+        if self.db.get(name):
+            msg = f"File with name {name} already exists."
+            logger.error(msg)
+            raise FileExistsError(msg)
+        try:
+            kin_context_dir = self._get_kin_context_temp_dir(self.mission_id)
+            file_path = os.path.join(kin_context_dir, name)
+            Path(file_path).write_bytes(content)
+            url = str(Path(file_path).resolve())
+            return FilesystemData(
+                kin_context=self.mission_id,
+                name=name,
+                file_type=file_type,
+                url=url,
+            )
+        except Exception:
+            msg = f"Error uploading file {name}"
+            logger.exception(msg)
+            raise FilesystemServiceError(msg)
+
+    def get(self, name: str) -> FilesystemData:
+        """Get file from the filesystem.
+
+        Args:
+            name: The name of the file to be retrieved
+
+        Returns:
+            FilesystemData: Metadata about the retrieved file
+
+        Raises:
+            FileNotFoundError: If the file does not exist
+            FilesystemServiceError: If the file does not exist
+        """
+        try:
+            return self.db[name]
+        except KeyError:
+            # If the file does not exist in the database, raise an error
+            msg = f"File with name {name} does not exist."
+            logger.exception(msg)
+            raise FileNotFoundError(msg)
+        except Exception:
+            msg = f"Error getting file {name}"
+            logger.exception(msg)
+            raise FilesystemServiceError(msg)
+
+    def update(self, name: str, content: bytes, file_type: FileType) -> FilesystemData:
+        """Update files in the filesystem.
+
+        Args:
+            content: The new content of the file
+            name: The name of the file to be updated
+            file_type: The type of data being updated
+
+        Returns:
+            FilesystemData: Metadata about the updated file
+
+        Raises:
+            FileNotFoundError: If the file does not exist
+            FilesystemServiceError: If there is an error during update
+        """
+        if name not in self.db:
+            msg = f"File with name {name} does not exist."
+            logger.error(msg)
+            raise FileNotFoundError(msg)
+        try:
+            kin_context_dir = self._get_kin_context_temp_dir(self.mission_id)
+            file_path = os.path.join(kin_context_dir, name)
+            Path(file_path).write_bytes(content)
+            url = str(Path(file_path).resolve())
+            file = FilesystemData(
+                kin_context=self.mission_id,
+                name=name,
+                file_type=file_type,
+                url=url,
+            )
+            self.db[name] = file
+        except Exception:
+            msg = f"Error updating file {name}"
+            logger.exception(msg)
+            raise FilesystemServiceError(msg)
+        else:
+            return file
+
+    def delete(self, name: str) -> bool:
+        """Delete files from the filesystem.
+
+        Args:
+            name: The name of the file to be deleted
+
+        Returns:
+            int: 1 if the file was deleted successfully
+
+        Raises:
+            FileNotFoundError: If the file does not exist
+            FilesystemServiceError: If there is an error during deletion
+        """
+        # First check if the file exists in the database
+        if name not in self.db:
+            msg = f"File with name {name} does not exist in the database."
+            logger.error(msg)
+            raise FileNotFoundError(msg)
+
+        # Get the file path
+        kin_context_dir = self._get_kin_context_temp_dir(self.mission_id)
+        file_path = os.path.join(kin_context_dir, name)
+
+        # Check if the file exists in the filesystem
+        if not os.path.exists(file_path):
+            msg = f"File {name} exists in database but not in filesystem at {file_path}."
+            logger.error(msg)
+            # We could decide to just remove from DB here, but that might hide a larger issue
+            # So we're raising a custom error to alert about the inconsistency
+            raise FilesystemServiceError(msg)
+
+        try:
+            os.remove(file_path)
+            del self.db[name]
+            logger.info("File %s successfully deleted.", name)
+
+        except OSError:
+            msg = f"Error deleting file {name} from filesystem"
+            logger.exception(msg)
+            raise FilesystemServiceError(msg)
+        except Exception:
+            msg = f"Unexpected error deleting file {name}"
+            logger.exception(msg)
+            raise FilesystemServiceError(msg)
+        else:
+            return True
+
+    def get_all(self) -> list[FilesystemData]:
+        """Get all files from the filesystem.
 
-    def get(self, data: dict[str, Any]) -> list[dict[str, Any]]:
-        """Get files from the file system."""
-        raise NotImplementedError
+        Returns:
+            list[FilesystemData]: A list of all files in the filesystem
+        """
+        return list(self.db.values())
 
-    def update(self, data: dict[str, Any]) -> int:
-        """Update files in the file system."""
-        raise NotImplementedError
+    def get_batch(self, names: list[str]) -> dict[str, FilesystemData | None]:
+        """Get files from the filesystem.
 
-    def delete(self, data: dict[str, Any]) -> int:
-        """Delete files from the file system."""
-        raise NotImplementedError
+        Args:
+            names: The names of the files to be retrieved
 
-    def get_all(self) -> list[dict[str, Any]]:
-        """Get all files from the file system."""
-        raise NotImplementedError
+        Returns:
+            dict[FilesystemData | None]: Metadata about the retrieved files
+        """
+        return {name: self.db.get(name, None) for name in names}