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

digitalkin/{grpc → grpc_servers}/_base_server.py

@@ -11,15 +11,15 @@ from typing import Any, cast
 import grpc
 from grpc import aio as grpc_aio
 
-from digitalkin.grpc.utils.exceptions import (
+from digitalkin.grpc_servers.utils.exceptions import (
     ConfigurationError,
     ReflectionError,
     SecurityError,
     ServerStateError,
     ServicerError,
 )
-from digitalkin.grpc.utils.models import SecurityMode, ServerConfig, ServerMode
-from digitalkin.grpc.utils.types import GrpcServer, ServiceDescriptor, T
+from digitalkin.grpc_servers.utils.models import SecurityMode, ServerConfig, ServerMode
+from digitalkin.grpc_servers.utils.types import GrpcServer, ServiceDescriptor, T
 
 logger = logging.getLogger(__name__)
 

digitalkin/{grpc → grpc_servers}/module_server.py

@@ -5,13 +5,16 @@ import uuid
 
 import grpc
 
-from digitalkin.grpc._base_server import BaseServer
-from digitalkin.grpc.module_servicer import ModuleServicer
-from digitalkin.grpc.utils.exceptions import ServerError
-from digitalkin.grpc.utils.models import ModuleServerConfig, SecurityMode
+from digitalkin.grpc_servers._base_server import BaseServer
+from digitalkin.grpc_servers.module_servicer import ModuleServicer
+from digitalkin.grpc_servers.utils.exceptions import ServerError
+from digitalkin.grpc_servers.utils.models import ModuleServerConfig, SecurityMode
 from digitalkin.modules._base_module import BaseModule
 
-from digitalkin_proto.digitalkin.module.v2 import module_service_pb2, module_service_pb2_grpc
+from digitalkin_proto.digitalkin.module.v2 import (
+    module_service_pb2,
+    module_service_pb2_grpc,
+)
 from digitalkin_proto.digitalkin.module_registry.v2 import (
     metadata_pb2,
     module_registry_service_pb2_grpc,
@@ -111,7 +114,9 @@ class ModuleServer(BaseServer):
         Raises:
             ServerError: If communication with the registry server fails.
         """
-        logger.info("Registering module with registry at %s", self.config.registry_address)
+        logger.info(
+            "Registering module with registry at %s", self.config.registry_address
+        )
 
         # Create appropriate channel based on security mode
         channel = self._create_registry_channel()
@@ -125,11 +130,16 @@ class ModuleServer(BaseServer):
 
             metadata = metadata_pb2.Metadata(
                 name=self.module_class.metadata["name"],
-                tags=[metadata_pb2.Tag(tag=tag) for tag in self.module_class.metadata["tags"]],
+                tags=[
+                    metadata_pb2.Tag(tag=tag)
+                    for tag in self.module_class.metadata["tags"]
+                ],
                 description=self.module_class.metadata["description"],
             )
 
-            self.module_class.metadata["module_id"] = f"{self.module_class.metadata['name']}:{uuid.uuid4()}"
+            self.module_class.metadata["module_id"] = (
+                f"{self.module_class.metadata['name']}:{uuid.uuid4()}"
+            )
             # Create registration request
             request = registration_pb2.RegisterRequest(
                 module_id=self.module_class.metadata["module_id"],
@@ -162,7 +172,9 @@ class ModuleServer(BaseServer):
         Raises:
             ServerError: If communication with the registry server fails.
         """
-        logger.info("Deregistering module from registry at %s", self.config.registry_address)
+        logger.info(
+            "Deregistering module from registry at %s", self.config.registry_address
+        )
 
         # Create appropriate channel based on security mode
         channel = self._create_registry_channel()
@@ -204,13 +216,19 @@ class ModuleServer(BaseServer):
 
             root_certificates = None
             if self.config.credentials.root_cert_path:
-                with open(self.config.credentials.root_cert_path, "rb") as root_cert_file:  # noqa: FURB101
+                with open(
+                    self.config.credentials.root_cert_path, "rb"
+                ) as root_cert_file:  # noqa: FURB101
                     root_certificates = root_cert_file.read()
 
             # Create channel credentials
-            channel_credentials = grpc.ssl_channel_credentials(root_certificates=root_certificates or certificate_chain)
+            channel_credentials = grpc.ssl_channel_credentials(
+                root_certificates=root_certificates or certificate_chain
+            )
 
-            return grpc.secure_channel(self.config.registry_address, channel_credentials)
+            return grpc.secure_channel(
+                self.config.registry_address, channel_credentials
+            )
         # Insecure channel
         return grpc.insecure_channel(self.config.registry_address)
 

digitalkin/{grpc → grpc_servers}/module_servicer.py

@@ -14,9 +14,15 @@ from digitalkin_proto.digitalkin.module.v2 import (
 )
 from google.protobuf import json_format, struct_pb2
 
+from digitalkin.grpc_servers.utils.exceptions import ServicerError
+from digitalkin.models.module import OutputModelT
 from digitalkin.models.module.module import ModuleStatus
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.modules.job_manager import JobManager
+from digitalkin.services.services_models import ServicesMode
+from digitalkin.services.setup.default_setup import DefaultSetup
+from digitalkin.services.setup.grpc_setup import GrpcSetup
+from digitalkin.services.setup.setup_strategy import SetupStrategy
 
 logger = logging.getLogger(__name__)
 
@@ -31,6 +37,8 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         active_jobs: Dictionary tracking active module jobs.
     """
 
+    setup: SetupStrategy
+
     def __init__(self, module_class: type[BaseModule]) -> None:
         """Initialize the module servicer.
 
@@ -41,8 +49,9 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         self.queue: asyncio.Queue = asyncio.Queue()
         self.module_class = module_class
         self.job_manager = JobManager(module_class)
+        self.setup = GrpcSetup() if self.job_manager.args.services_mode == ServicesMode.REMOTE else DefaultSetup()
 
-    async def add_to_queue(self, job_id: str, output_data: dict[str, Any]) -> None:
+    async def add_to_queue(self, job_id: str, output_data: OutputModelT) -> None:
         """Callback used to add the output data to the queue of messages."""
         logger.info("JOB: %s added an output_data: %s", job_id, output_data)
         await self.queue.put({job_id: output_data})
@@ -60,19 +69,26 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
 
         Yields:
             Responses during module execution.
+
+        Raises:
+            ServicerError: the necessary query didn't work.
         """
         logger.info("StartModule called for module: '%s'", self.module_class.__name__)
         # Process the module input
-        input_data = dict(request.input.items())
-        setup_data = self.module_class.storage.get(
-            table="setups",
-            data={
-                "keys": [
-                    # remove prefix 'setups:'
-                    request.setup_id[7:],
-                ],
-            },
-        )[0]
+        # TODO: Check failure of input data format
+        input_data = self.module_class.input_format(**dict(request.input.items()))
+        setup_data_class = self.setup.get_setup(
+            setup_dict={
+                "setup_id": request.setup_id,
+                "mission_id": "missions:test_demo",
+            }
+        )
+
+        if not setup_data_class:
+            msg = "No setup data returned."
+            raise ServicerError(msg)
+        # TODO: Check failure of setup data format
+        setup_data = self.module_class.setup_format(**setup_data_class.current_setup_version.content)
 
         # setup_id should be use to request a precise setup from the module
         # Create a job for this execution
@@ -84,9 +100,9 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
 
         while module.status == ModuleStatus.RUNNING or not self.queue.empty():
             output_data = await self.queue.get()
-            if output_data.get("error", None) is not None:
-                context.set_code(output_data["error"]["code"])
-                context.set_details(output_data["error"]["error_message"])
+            if output_data[job_id].get("error", None) is not None:
+                context.set_code(output_data[job_id]["error"]["code"])
+                context.set_details(output_data[job_id]["error"]["error_message"])
                 yield lifecycle_pb2.StartModuleResponse(success=False)
                 return
             else:

digitalkin/{grpc → grpc_servers}/registry_server.py

@@ -2,11 +2,13 @@
 
 import logging
 
-from digitalkin_proto.digitalkin.module_registry.v2 import module_registry_service_pb2_grpc
+from digitalkin_proto.digitalkin.module_registry.v2 import (
+    module_registry_service_pb2_grpc,
+)
 
-from digitalkin.grpc._base_server import BaseServer
-from digitalkin.grpc.registry_servicer import RegistryModule, RegistryServicer
-from digitalkin.grpc.utils.models import RegistryServerConfig
+from digitalkin.grpc_servers._base_server import BaseServer
+from digitalkin.grpc_servers.registry_servicer import RegistryModule, RegistryServicer
+from digitalkin.grpc_servers.utils.models import RegistryServerConfig
 
 logger = logging.getLogger(__name__)
 

digitalkin/{grpc → grpc_servers}/registry_servicer.py

@@ -88,7 +88,9 @@ class Metadata(BaseModel):
             metadata_pb2.Metadata: The protobuf representation of this metadata.
         """
         return metadata_pb2.Metadata(
-            name=self.name, tags=(t.to_proto() for t in self.tags), description=self.description
+            name=self.name,
+            tags=(t.to_proto() for t in self.tags),
+            description=self.description,
         )
 
     @classmethod
@@ -361,7 +363,11 @@ class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceSer
         Returns:
             status_pb2.ListModulesStatusResponse: A response containing a list of module statuses.
         """
-        logger.info("Getting registered modules with offset %d and limit %d", request.offset, request.list_size)
+        logger.info(
+            "Getting registered modules with offset %d and limit %d",
+            request.offset,
+            request.list_size,
+        )
         if request.offset > len(self.registered_modules):
             message = f"Out of range {request.offset}"
             logger.warning(message)

digitalkin/{grpc → grpc_servers}/utils/factory.py

@@ -3,9 +3,9 @@
 from pathlib import Path
 from typing import Any
 
-from digitalkin.grpc.module_server import ModuleServer
-from digitalkin.grpc.registry_server import RegistryServer
-from digitalkin.grpc.utils.models import (
+from digitalkin.grpc_servers.module_server import ModuleServer
+from digitalkin.grpc_servers.registry_server import RegistryServer
+from digitalkin.grpc_servers.utils.models import (
     ModuleServerConfig,
     RegistryServerConfig,
     SecurityMode,
@@ -167,7 +167,9 @@ def _create_server_config(
     # Add credentials if secure mode
     if security == "secure":
         if not server_key_path or not server_cert_path:
-            raise ValueError("Server key and certificate paths are required for secure mode")
+            raise ValueError(
+                "Server key and certificate paths are required for secure mode"
+            )
 
         config_params["credentials"] = ServerCredentials(
             server_key_path=Path(server_key_path),

digitalkin/grpc_servers/utils/grpc_client_wrapper.py

@@ -0,0 +1,68 @@
+"""Client wrapper to ease channel creation with specific ServerConfig."""
+
+import logging
+from typing import Any
+
+import grpc
+
+from digitalkin.grpc_servers.utils.exceptions import ServerError
+from digitalkin.grpc_servers.utils.models import SecurityMode, ServerConfig
+
+logger = logging.getLogger(__name__)
+
+
+class GrpcClientWrapper:
+    """gRPC client shared by the different services."""
+
+    stub: Any
+
+    @staticmethod
+    def _init_channel(config: ServerConfig) -> grpc.Channel:
+        """Create an appropriate channel to the registry server.
+
+        Returns:
+            A gRPC channel for communication with the registry.
+
+        Raises:
+            ValueError: If credentials are required but not provided.
+        """
+        if config.security == SecurityMode.SECURE and config.credentials:
+            # Secure channel
+            with open(config.credentials.server_cert_path, "rb") as cert_file:  # noqa: FURB101
+                certificate_chain = cert_file.read()
+
+            root_certificates = None
+            if config.credentials.root_cert_path:
+                with open(config.credentials.root_cert_path, "rb") as root_cert_file:  # noqa: FURB101
+                    root_certificates = root_cert_file.read()
+
+            # Create channel credentials
+            channel_credentials = grpc.ssl_channel_credentials(root_certificates=root_certificates or certificate_chain)
+
+            return grpc.secure_channel(f"{config.host}:{config.port}", channel_credentials)
+        # Insecure channel
+        return grpc.insecure_channel(f"{config.host}:{config.port}")
+
+    def exec_grpc_query(self, query_endpoint: str, request: Any) -> Any:  # noqa: ANN401
+        """Execute a gRPC query with from the query's rpc endpoint name.
+
+        Arguments:
+            query_endpoint: rpc query name
+            request: gRPC object to match the rpc query
+
+        Returns:
+            corresponding gRPC reponse.
+
+        Raises:
+            ServerError: gRPC error catching
+        """
+        try:
+            # Call the register method
+            logger.warning("send request to %s", query_endpoint)
+            response = getattr(self.stub, query_endpoint)(request)
+            logger.warning("recive response from request to registry: %s", response)
+        except grpc.RpcError:
+            logger.exception("RPC error during registration:")
+            raise ServerError
+        else:
+            return response

digitalkin/{grpc → grpc_servers}/utils/models.py

@@ -6,7 +6,7 @@ from typing import Any
 
 from pydantic import BaseModel, Field, ValidationInfo, field_validator
 
-from digitalkin.grpc.utils.exceptions import ConfigurationError, SecurityError
+from digitalkin.grpc_servers.utils.exceptions import ConfigurationError, SecurityError
 
 
 class ServerMode(str, Enum):

digitalkin/models/__init__.py

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

digitalkin/models/module/__init__.py

@@ -1,5 +1,11 @@
 """This module contains the models for the modules."""
 
-from .module import Module, ModuleStatus
+from digitalkin.models.module.module import Module, ModuleStatus
+from digitalkin.models.module.module_types import (
+    InputModelT,
+    OutputModelT,
+    SecretModelT,
+    SetupModelT,
+)
 
-__all__ = ["Module", "ModuleStatus"]
+__all__ = ["InputModelT", "Module", "ModuleStatus", "OutputModelT", "SecretModelT", "SetupModelT"]

digitalkin/models/module/module_types.py

@@ -0,0 +1,10 @@
+"""Types for module models."""
+
+from typing import TypeVar
+
+from pydantic import BaseModel
+
+InputModelT = TypeVar("InputModelT", bound=BaseModel)
+OutputModelT = TypeVar("OutputModelT", bound=BaseModel)
+SetupModelT = TypeVar("SetupModelT", bound=BaseModel)
+SecretModelT = TypeVar("SecretModelT", bound=BaseModel)

digitalkin/models/services/__init__.py

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

digitalkin/modules/__init__.py

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

digitalkin/modules/_base_module.py

@@ -4,44 +4,64 @@ 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 collections.abc import Callable, Coroutine
+from typing import Any, ClassVar, Generic
 
 from digitalkin.logger import logger
-from digitalkin.models.module import ModuleStatus
-from digitalkin.services.service_provider import ServiceProvider
+from digitalkin.models.module import InputModelT, ModuleStatus, OutputModelT, SecretModelT, SetupModelT
+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.services_config import ServicesConfig, ServicesStrategy
+from digitalkin.services.snapshot.snapshot_strategy import SnapshotStrategy
 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]):
+class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretModelT]):
     """BaseModule is the abstract base for all modules in the DigitalKin SDK."""
 
+    name: str
+    description: str
     input_format: type[InputModelT]
     output_format: type[OutputModelT]
     setup_format: type[SetupModelT]
+    secret_format: type[SecretModelT]
     metadata: ClassVar[dict[str, Any]]
 
-    local_services: type[ServiceProvider]
-    dev_services: type[ServiceProvider]
-
+    # service config params
+    services_config_strategies: ClassVar[dict[str, ServicesStrategy | None]]
+    services_config_params: ClassVar[dict[str, dict[str, Any | None] | None]]
+    services_config: ServicesConfig
+
+    # services list
+    agent: AgentStrategy
+    cost: CostStrategy
+    filesystem: FilesystemStrategy
+    identity: IdentityStrategy
+    registry: RegistryStrategy
+    snapshot: SnapshotStrategy
     storage: StorageStrategy
 
+    def _init_strategies(self) -> None:
+        """Initialize the services configuration."""
+        for service_name in self.services_config.valid_strategy_names():
+            service = self.services_config.init_strategy(service_name, self.mission_id)
+            setattr(self, service_name, service)
+
     def __init__(
         self,
         job_id: str,
-        name: str | None = None,
+        mission_id: str,
     ) -> None:
         """Initialize the module."""
         self.job_id: str = job_id
-        self.name = name or self.__class__.__name__
+        self.mission_id: str = mission_id
         self._status = ModuleStatus.CREATED
         self._task: asyncio.Task | None = None
+        # Initialize services configuration
+        self._init_strategies()
 
     @property
     def status(self) -> ModuleStatus:
@@ -52,6 +72,23 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
         """
         return self._status
 
+    @classmethod
+    def get_secret_format(cls, llm_format: bool) -> str:  # noqa: FBT001
+        """Get the JSON schema of the secret format model.
+
+        Raises:
+            NotImplementedError: If the `secret_format` is not defined.
+
+        Returns:
+            The JSON schema of the secret format as a string.
+        """
+        if cls.secret_format is not None:
+            if llm_format:
+                return json.dumps(cls.secret_format, indent=2)
+            return json.dumps(cls.secret_format.model_json_schema(), indent=2)
+        msg = f"{cls.__name__}' class does not define a 'secret_format'."
+        raise NotImplementedError(msg)
+
     @classmethod
     def get_input_format(cls, llm_format: bool) -> str:  # noqa: FBT001
         """Get the JSON schema of the input format model.
@@ -62,7 +99,7 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
         Returns:
             The JSON schema of the input format as a string.
         """
-        if cls.output_format is not None:
+        if cls.input_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)
@@ -104,16 +141,16 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
         raise NotImplementedError(msg)
 
     @abstractmethod
-    async def initialize(self, setup_data: dict[str, Any]) -> None:
+    async def initialize(self, setup_data: SetupModelT) -> None:
         """Initialize the module."""
         raise NotImplementedError
 
     @abstractmethod
     async def run(
         self,
-        input_data: dict[str, Any],
-        setup_data: dict[str, Any],
-        callback: Callable,
+        input_data: InputModelT,
+        setup_data: SetupModelT,
+        callback: Callable[[OutputModelT], Coroutine[Any, Any, None]],
     ) -> None:
         """Run the module."""
         raise NotImplementedError
@@ -125,9 +162,9 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
 
     async def _run_lifecycle(
         self,
-        input_data: dict[str, Any],
-        setup_data: dict[str, Any],
-        callback: Callable,
+        input_data: InputModelT,
+        setup_data: SetupModelT,
+        callback: Callable[[OutputModelT], Coroutine[Any, Any, None]],
     ) -> None:
         """Run the module lifecycle.
 
@@ -147,9 +184,9 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT]):
 
     async def start(
         self,
-        input_data: dict[str, Any],
-        setup_data: dict[str, Any],
-        callback: Callable,
+        input_data: InputModelT,
+        setup_data: SetupModelT,
+        callback: Callable[[OutputModelT], Coroutine[Any, Any, None]],
     ) -> None:
         """Start the module."""
         try:

digitalkin/modules/archetype_module.py

@@ -2,13 +2,9 @@
 
 from abc import ABC
 
+from digitalkin.models.module import InputModelT, OutputModelT, SecretModelT, SetupModelT
 from digitalkin.modules._base_module import BaseModule
 
 
-class ArchetypeModule(BaseModule, ABC):
+class ArchetypeModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT], 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"]