digitalkin 0.3.1.dev2__py3-none-any.whl → 0.3.2.dev14__py3-none-any.whl

services/filesystem_module.py

@@ -3,16 +3,16 @@
 import asyncio
 import datetime
 from collections.abc import Callable
-from typing import TYPE_CHECKING, Any
+from typing import Any
 
 from pydantic import BaseModel, Field
 
 from digitalkin.logger import logger
 from digitalkin.models.module import ModuleStatus
 from digitalkin.modules.archetype_module import ArchetypeModule
+from digitalkin.services.filesystem.filesystem_strategy import FileFilter, UploadFileData
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
-from digitalkin.services.filesystem.filesystem_strategy import FilesystemRecord, FileFilter, UploadFileData
 
 
 class ExampleInput(BaseModel):
@@ -62,7 +62,10 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
 
     # Define services_config_params with default values
     services_config_strategies = {}
-    services_config_params = {"cost": {"config": {}}, "storage": {"config": {}}} # Filesystem has no config but it's enabled
+    services_config_params = {
+        "cost": {"config": {}},
+        "storage": {"config": {}},
+    }  # Filesystem has no config but it's enabled
 
     def __init__(self, job_id: str, mission_id: str, setup_version_id: str) -> None:
         """Initialize the example module.
@@ -128,7 +131,6 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
             callback(record.model_dump())
         # Call the callback with the output data
 
-
         # Wait a bit to simulate processing time
         await asyncio.sleep(1)
 
@@ -173,7 +175,7 @@ async def test_module() -> None:
 
     # Check the storage
     if module.status == ModuleStatus.STOPPED:
-        files, nb_results = module.filesystem.get_files(
+        files, _nb_results = module.filesystem.get_files(
             filters=FileFilter(name="example_output.txt", context="test-mission-123"),
         )
         for file in files:

services/storage_module.py

@@ -64,7 +64,7 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
 
     # Define services_config_params with default values
     services_config_strategies = {}
-    services_config_params = {"storage": {"config": {"example": ExampleOutput}},"cost": {"config":{}}}
+    services_config_params = {"storage": {"config": {"example": ExampleOutput}}, "cost": {"config": {}}}
 
     def __init__(self, job_id: str, mission_id: str, setup_version_id: str) -> None:
         """Initialize the example module.
@@ -184,7 +184,9 @@ async def test_module() -> None:
 def test_storage_directly() -> None:
     """Test the storage service directly."""
     # Initialize storage service
-    storage = ServicesConfig().storage(mission_id="test-mission",setup_version_id="test-setup-123", config={"example": ExampleStorage})
+    storage = ServicesConfig().storage(
+        mission_id="test-mission", setup_version_id="test-setup-123", config={"example": ExampleStorage}
+    )
 
     # Create a test record
     storage.store("example", "test_table", {"test_key": "test_value"}, "OUTPUT")

digitalkin/grpc_servers/registry_server.py

@@ -1,65 +0,0 @@
-"""Registry gRPC server implementation for DigitalKin."""
-
-from digitalkin_proto.agentic_mesh_protocol.module_registry.v1 import (
-    module_registry_service_pb2,
-    module_registry_service_pb2_grpc,
-)
-
-from digitalkin.grpc_servers._base_server import BaseServer
-from digitalkin.grpc_servers.registry_servicer import RegistryModule, RegistryServicer
-from digitalkin.logger import logger
-from digitalkin.models.grpc_servers.models import RegistryServerConfig
-
-
-class RegistryServer(BaseServer):
-    """gRPC server for DigitalKin module registry.
-
-    This server implements the ModuleRegistryService which allows modules to register
-    themselves and be discovered by other components in the system.
-
-    Attributes:
-        config: Server configuration.
-        registry_servicer: The gRPC servicer handling registry requests.
-    """
-
-    def __init__(
-        self,
-        config: RegistryServerConfig,
-    ) -> None:
-        """Initialize the registry server.
-
-        Args:
-            config: Server configuration.
-        """
-        super().__init__(config)
-        self.config = config
-        self.registry_servicer: RegistryServicer | None = None
-
-    def _register_servicers(self) -> None:
-        """Register the registry servicer with the gRPC server.
-
-        Raises:
-            RuntimeError: If server is not registered during server creation
-        """
-        if self.server is None:
-            msg = "Server must be created before registering servicers"
-            raise RuntimeError(msg)
-
-        logger.debug("Registering registry servicer")
-        self.registry_servicer = RegistryServicer()
-        self.register_servicer(
-            self.registry_servicer,
-            module_registry_service_pb2_grpc.add_ModuleRegistryServiceServicer_to_server,
-            service_descriptor=module_registry_service_pb2.DESCRIPTOR,
-        )
-        logger.debug("Registered registry servicer")
-
-    def get_registered_modules(self) -> list[RegistryModule]:
-        """Get a list of all registered modules.
-
-        Returns:
-            A list of module information objects.
-        """
-        if self.registry_servicer:
-            return list(self.registry_servicer.registered_modules.values())
-        return []

digitalkin/grpc_servers/registry_servicer.py

@@ -1,456 +0,0 @@
-"""Registry servicer implementation for DigitalKin.
-
-This module provides the gRPC service implementation for the Module Registry,
-which handles registration, deregistration, discovery, and status management
-of DigitalKin modules.
-"""
-
-from collections.abc import Iterator
-from enum import Enum
-
-import grpc
-from digitalkin_proto.agentic_mesh_protocol.module_registry.v1 import (
-    discover_pb2,
-    metadata_pb2,
-    module_registry_service_pb2_grpc,
-    registration_pb2,
-    status_pb2,
-)
-from pydantic import BaseModel
-from typing_extensions import Self
-
-from digitalkin.logger import logger
-
-
-class ExtendedEnum(Enum):
-    """Tool enum class."""
-
-    @classmethod
-    def list(cls) -> list:
-        """Classmethod to generate a list of enum values.
-
-        Returns:
-            list: Enum members' values
-        """
-        return [c.value for c in cls]
-
-
-class ModuleStatus(ExtendedEnum):
-    """Describe a Module current status.
-
-    Represents the possible states a module can be in during its lifecycle.
-    """
-
-    # RUNNING: Module alive.
-    RUNNING = 0
-    # IDLE: Module waiting for an event / update.
-    IDLE = 1
-    # ENDED: Module signals the end of task or have been killed.
-    ENDED = 2
-
-
-class Tag(BaseModel):
-    """Words representing a module capabilities.
-
-    Used for module discovery and categorization.
-    """
-
-    # tag: Describe a Module function.
-    tag: str
-
-    def to_proto(self) -> metadata_pb2.Tag:
-        """Convert Tag object from Pydantic to Proto.
-
-        Returns:
-            metadata_pb2.Tag: The protobuf representation of this tag.
-        """
-        return metadata_pb2.Tag(tag=self.tag)
-
-
-class Metadata(BaseModel):
-    """Different informations to index and describe a module.
-
-    Contains human-readable information about a module and its capabilities.
-    """
-
-    # name: Module's name
-    name: str
-    # tags: List of tag to describe a module functionalities.
-    tags: list[Tag]
-    # description: Module's description for search and indexing
-    description: str | None
-
-    def to_proto(self) -> metadata_pb2.Metadata:
-        """Convert Metadata object from Pydantic to Proto.
-
-        Returns:
-            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,
-        )
-
-    @classmethod
-    def from_proto(cls, request_metadata: metadata_pb2.Metadata) -> Self:
-        """Create Metadata object from Proto message.
-
-        Args:
-            request_metadata: The protobuf metadata to convert.
-
-        Returns:
-            Metadata: The Pydantic model representation of the metadata.
-        """
-        return cls(
-            name=request_metadata.name,
-            tags=[Tag(tag=t.tag) for t in request_metadata.tags],
-            description=request_metadata.description,
-        )
-
-
-class RegistryModule(BaseModel):
-    """Module's technical representation to index, search and monitor.
-
-    Contains all the information needed to identify, locate and communicate
-    with a module in the system.
-    """
-
-    # module_id: Id of the module
-    module_id: str
-    # module_type: Type of the module (trigger, tool, kin, view)
-    module_type: str
-    # address: Address used to communicate with the module
-    address: str
-    # port: Port used to communicate with the module
-    port: int
-    # version: Current module version.
-    version: str
-    # metadata: user defined module name, description and tags
-    metadata: Metadata
-    # status: Representation of the Module current state (running, idle, ended...).
-    status: ModuleStatus
-    # message: (Optional) Details about the status.
-    message: str | None
-
-    def to_proto(self) -> discover_pb2.DiscoverInfoResponse:
-        """Convert RegistryModule object from Pydantic to Proto.
-
-        Returns:
-            metadata_pb2.Metadata: The protobuf representation of this metadata.
-        """
-        return discover_pb2.DiscoverInfoResponse(
-            module_id=self.module_id,
-            module_type=self.module_type,
-            address=self.address,
-            port=self.port,
-            version=self.version,
-            metadata=self.metadata.to_proto(),
-        )
-
-
-class RegistryServicer(module_registry_service_pb2_grpc.ModuleRegistryServiceServicer):
-    """Implementation of the ModuleRegistryService.
-
-    This servicer handles the registration, deregistration, and discovery of modules.
-    It maintains an in-memory registry of all active modules and their metadata.
-
-    Attributes:
-        registered_modules: Dictionary mapping module_id to RegistryModule objects.
-    """
-
-    registered_modules: dict[str, RegistryModule]
-
-    def __init__(self) -> None:
-        """Initialize the registry servicer with an empty module registry."""
-        self.registered_modules = {}  # TODO replace with a database
-
-    def RegisterModule(  # noqa: N802
-        self,
-        request: registration_pb2.RegisterRequest,
-        context: grpc.ServicerContext,
-    ) -> registration_pb2.RegisterResponse:
-        """Register a module with the registry.
-
-        Adds a new module to the registry with its connection information and metadata.
-        Fails if a module with the same ID is already registered.
-
-        Args:
-            request: The register request containing module info and address.
-            context: The gRPC context for setting status codes and details.
-
-        Returns:
-            registration_pb2.RegisterResponse: A response indicating success or failure.
-        """
-        module_id = request.module_id
-        logger.debug("Registering module: %s", module_id)
-
-        # Check if module is already registered
-        if module_id in self.registered_modules:
-            message = f"Module '{module_id}' already registered"
-            logger.warning(message)
-
-            context.set_code(grpc.StatusCode.ALREADY_EXISTS)
-            context.set_details(message)
-            return registration_pb2.RegisterResponse(success=False)
-
-        # Store the module info with address
-        self.registered_modules[module_id] = RegistryModule(
-            module_id=request.module_id,
-            module_type=request.module_type,
-            address=request.address,
-            port=request.port,
-            version=request.version,
-            metadata=Metadata.from_proto(request.metadata),
-            status=ModuleStatus.RUNNING,
-            message=None,
-        )
-
-        logger.debug("Module %s registered at %s:%d", module_id, request.address, request.port)
-        return registration_pb2.RegisterResponse(success=True)
-
-    def DeregisterModule(  # noqa: N802
-        self,
-        request: registration_pb2.DeregisterRequest,
-        context: grpc.ServicerContext,
-    ) -> registration_pb2.DeregisterResponse:
-        """Deregister a module from the registry.
-
-        Removes a module from the registry based on its ID.
-        Fails if the specified module is not found in the registry.
-
-        Args:
-            request: The deregister request containing the module ID.
-            context: The gRPC context for setting status codes and details.
-
-        Returns:
-            registration_pb2.DeregisterResponse: A response indicating success or failure.
-        """
-        module_id = request.module_id
-        logger.debug("Deregistering module: %s", module_id)
-
-        # Check if module exists in registry
-        if module_id not in self.registered_modules:
-            message = f"Module {module_id} not found in registry"
-            logger.warning(message)
-
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            context.set_details(message)
-            return registration_pb2.DeregisterResponse()
-
-        # Remove the module
-        del self.registered_modules[module_id]
-
-        logger.debug("Module %s deregistered", module_id)
-        return registration_pb2.DeregisterResponse(success=True)
-
-    def DiscoverInfoModule(  # noqa: N802
-        self,
-        request: discover_pb2.DiscoverInfoRequest,
-        context: grpc.ServicerContext,
-    ) -> discover_pb2.DiscoverInfoResponse:
-        """Discover detailed information about a specific module.
-
-        Retrieves complete information about a module based on its ID.
-
-        Args:
-            request: The discover request containing the module ID.
-            context: The gRPC context (unused).
-
-        Returns:
-            discover_pb2.DiscoverInfoResponse: A response containing the module's information.
-        """
-        logger.debug("Discovering module: %s", request.module_id)
-
-        # Check if module exists in registry
-        if request.module_id not in self.registered_modules:
-            message = f"Module {request.module_id} not found in registry"
-            logger.warning(message)
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            context.set_details(message)
-            return discover_pb2.DiscoverInfoResponse()
-        return self.registered_modules[request.module_id].to_proto()
-
-    def DiscoverSearchModule(  # noqa: N802
-        self,
-        request: discover_pb2.DiscoverSearchRequest,
-        context: grpc.ServicerContext,  # noqa: ARG002
-    ) -> discover_pb2.DiscoverSearchResponse:
-        """Discover modules based on the specified criteria.
-
-        Searches for modules that match the provided filters such as name,
-        type, and tags.
-
-        Args:
-            request: The discover request containing search criteria.
-            context: The gRPC context (unused).
-
-        Returns:
-            discover_pb2.DiscoverSearchResponse: A response containing matching modules.
-        """
-        logger.debug("Discovering modules with criteria:")
-
-        # Start with all modules
-        results = list(self.registered_modules.values())
-        logger.debug("%s", list(results))
-        # Filter by name if specified
-        if request.name:
-            logger.debug("\tname %s", request.name)
-            results = [m for m in results if request.name in m.metadata.name]
-
-        # Filter by type if specified
-        if request.module_type:
-            logger.debug("\tmodule_type %s", request.module_type)
-            results = [m for m in results if m.module_type == request.module_type]
-
-        # Filter by tags if specified
-        if request.tags:
-            logger.debug("\ttags %s", request.tags)
-            results = [m for m in results if any(tag in m.metadata.tags for tag in request.tags)]
-
-        # Filter by description if specified
-        """
-        if request.description:
-            results = [m for m in results if request.description in m.metadata.description]
-        """
-
-        logger.debug("Found %d matching modules", len(results))
-        return discover_pb2.DiscoverSearchResponse(modules=[r.to_proto() for r in results])
-
-    def GetModuleStatus(  # noqa: N802
-        self,
-        request: status_pb2.ModuleStatusRequest,
-        context: grpc.ServicerContext,
-    ) -> status_pb2.ModuleStatusResponse:
-        """Query a specific module's status.
-
-        Retrieves the current status of a module based on its ID.
-
-        Args:
-            request: The status request containing the module ID.
-            context: The gRPC context (unused).
-
-        Returns:
-            status_pb2.ModuleStatusResponse: A response containing the module's status.
-        """
-        logger.debug("Getting status for module: %s", request.module_id)
-
-        # Check if module exists in registry
-        if request.module_id not in self.registered_modules:
-            message = f"Module {request.module_id} not found in registry"
-            logger.warning(message)
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            context.set_details(message)
-            return status_pb2.ModuleStatusResponse()
-
-        module = self.registered_modules[request.module_id]
-        return status_pb2.ModuleStatusResponse(module_id=module.module_id, status=module.status.value)
-
-    def ListModuleStatus(  # noqa: N802
-        self,
-        request: status_pb2.ListModulesStatusRequest,
-        context: grpc.ServicerContext,
-    ) -> status_pb2.ListModulesStatusResponse:
-        """Get a paginated list of registered modules and their statuses.
-
-        Returns a subset of registered modules based on pagination parameters.
-
-        Args:
-            request: The request containing offset and list_size for pagination.
-            context: The gRPC context (unused).
-
-        Returns:
-            status_pb2.ListModulesStatusResponse: A response containing a list of module statuses.
-        """
-        logger.debug(
-            "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)
-            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
-            context.set_details(message)
-            return status_pb2.ListModulesStatusResponse()
-
-        list_size = request.list_size
-        if list_size == 0:
-            list_size = len(self.registered_modules)
-
-        modules_statuses = [
-            status_pb2.ModuleStatusResponse(module_id=module.module_id, status=module.status.value)
-            for module in list(self.registered_modules.values())[request.offset : request.offset + list_size]
-        ]
-
-        logger.debug("Found %d registered modules", len(modules_statuses))
-        return status_pb2.ListModulesStatusResponse(
-            list_size=len(modules_statuses),
-            modules_statuses=modules_statuses,
-        )
-
-    def GetAllModuleStatus(  # noqa: N802
-        self,
-        request: status_pb2.GetAllModulesStatusRequest,  # noqa: ARG002
-        context: grpc.ServicerContext,  # noqa: ARG002
-    ) -> Iterator[status_pb2.ModuleStatusResponse]:
-        """Get all registered modules via a stream.
-
-        Streams the status of all registered modules one by one.
-
-        Args:
-            request: The get all modules request (unused).
-            context: The gRPC context (unused).
-
-        Yields:
-            status_pb2.ModuleStatusResponse: Responses containing individual module statuses.
-        """
-        logger.debug("Streaming all %d registered modules", len(self.registered_modules))
-        for module in self.registered_modules.values():
-            yield status_pb2.ModuleStatusResponse(
-                module_id=module.module_id,
-                status=module.status.value,
-            )
-
-    def UpdateModuleStatus(  # noqa: N802
-        self,
-        request: status_pb2.UpdateStatusRequest,
-        context: grpc.ServicerContext,
-    ) -> status_pb2.UpdateStatusResponse:
-        """Update the status of a registered module.
-
-        Changes the current status of a module based on the provided request.
-        Fails if the specified module is not found in the registry.
-
-        Args:
-            request: The update status request with module ID and new status.
-            context: The gRPC context (unused).
-
-        Returns:
-            status_pb2.UpdateStatusResponse: A response indicating success or failure.
-        """
-        module_id = request.module_id
-        logger.debug("Updating status for module: %s to %s", module_id, request.status)
-
-        # Check if module exists in registry
-        if request.module_id not in self.registered_modules:
-            message = f"Module {request.module_id} not found in registry"
-            logger.warning(message)
-            context.set_code(grpc.StatusCode.NOT_FOUND)
-            context.set_details(message)
-            return status_pb2.UpdateStatusResponse()
-
-        # Check if module status is correct
-        if request.status not in ModuleStatus.list() or request.status is None:
-            message = f"ModuleStatus {request.status} is unknonw, please check the requested status"
-            logger.warning(message)
-            context.set_code(grpc.StatusCode.INVALID_ARGUMENT)
-            context.set_details(message)
-            return status_pb2.UpdateStatusResponse(success=False)
-
-        # Update module status
-        module_info = self.registered_modules[module_id]
-        module_info.status = ModuleStatus(request.status)
-
-        logger.debug("Status for module %s updated to %s", module_id, request.status)
-        return status_pb2.UpdateStatusResponse(success=True)