digitalkin 0.2.25rc0__py3-none-any.whl → 0.3.2.dev14__py3-none-any.whl

digitalkin/services/cost/__init__.py

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

digitalkin/services/cost/grpc_cost.py

@@ -1,16 +1,14 @@
 """This module implements the gRPC Cost strategy."""
 
-from collections.abc import Generator
-from contextlib import contextmanager
-from typing import Any, Literal
+from typing import Literal
 
-from digitalkin_proto.digitalkin.cost.v1 import cost_pb2, cost_service_pb2_grpc
+from agentic_mesh_protocol.cost.v1 import cost_pb2, cost_service_pb2_grpc
 from google.protobuf import json_format
 
-from digitalkin.grpc_servers.utils.exceptions import ServerError
 from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
-from digitalkin.grpc_servers.utils.models import ClientConfig
+from digitalkin.grpc_servers.utils.grpc_error_handler import GrpcErrorHandlerMixin
 from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import ClientConfig
 from digitalkin.services.cost.cost_strategy import (
     CostConfig,
     CostData,
@@ -20,40 +18,9 @@ from digitalkin.services.cost.cost_strategy import (
 )
 
 
-class GrpcCost(CostStrategy, GrpcClientWrapper):
+class GrpcCost(CostStrategy, GrpcClientWrapper, GrpcErrorHandlerMixin):
     """This class implements the default Cost strategy."""
 
-    @staticmethod
-    @contextmanager
-    def _handle_grpc_errors(operation: str) -> Generator[Any, Any, Any]:
-        """Context manager for consistent gRPC error handling.
-
-        Yields:
-            Allow error handling in context.
-
-        Args:
-            operation: Description of the operation being performed.
-
-        Raises:
-            ValueError: Error with the model validation.
-            ServerError: from gRPC Client.
-            CostServiceError: Unexpected error.
-        """
-        try:
-            yield
-        except CostServiceError as e:
-            msg = f"CostServiceError in {operation}: {e}"
-            logger.exception(msg)
-            raise CostServiceError(msg) from e
-        except ServerError as e:
-            msg = f"gRPC {operation} failed: {e}"
-            logger.exception(msg)
-            raise ServerError(msg) from e
-        except Exception as e:
-            msg = f"Unexpected error in {operation}"
-            logger.exception(msg)
-            raise CostServiceError(msg) from e
-
     def __init__(
         self,
         mission_id: str,
@@ -66,7 +33,7 @@ class GrpcCost(CostStrategy, GrpcClientWrapper):
         super().__init__(mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id, config=config)
         channel = self._init_channel(client_config)
         self.stub = cost_service_pb2_grpc.CostServiceStub(channel)
-        logger.debug("Channel client 'Cost' initialized succesfully")
+        logger.debug("Channel client 'Cost' initialized successfully")
 
     def add(
         self,
@@ -84,7 +51,7 @@ class GrpcCost(CostStrategy, GrpcClientWrapper):
         Raises:
             CostServiceError: If the cost config is invalid
         """
-        with self._handle_grpc_errors("AddCost"):
+        with self.handle_grpc_errors("AddCost", CostServiceError):
             cost_config = self.config.get(cost_config_name)
             if cost_config is None:
                 msg = f"Cost config {cost_config_name} not found in the configuration."
@@ -122,7 +89,7 @@ class GrpcCost(CostStrategy, GrpcClientWrapper):
         Returns:
             CostData: The cost data
         """
-        with self._handle_grpc_errors("GetCost"):
+        with self.handle_grpc_errors("GetCost", CostServiceError):
             request = cost_pb2.GetCostRequest(name=name, mission_id=self.mission_id)
             response: cost_pb2.GetCostResponse = self.exec_grpc_query("GetCost", request)
             cost_data_list = [
@@ -150,7 +117,7 @@ class GrpcCost(CostStrategy, GrpcClientWrapper):
         Returns:
             list[CostData]: The cost data
         """
-        with self._handle_grpc_errors("GetCosts"):
+        with self.handle_grpc_errors("GetCosts", CostServiceError):
             request = cost_pb2.GetCostsRequest(
                 mission_id=self.mission_id,
                 filter=cost_pb2.CostFilter(

digitalkin/services/filesystem/default_filesystem.py

@@ -198,8 +198,6 @@ class DefaultFilesystem(FilesystemStrategy):
             end_idx = start_idx + list_size
             paginated_files = filtered_files[start_idx:end_idx]
 
-            logger.critical(f"{filters=} | {paginated_files=}")
-
             if include_content:
                 for file in paginated_files:
                     file.content = Path(file.storage_uri).read_bytes()

digitalkin/services/filesystem/grpc_filesystem.py

@@ -1,17 +1,15 @@
 """gRPC filesystem implementation."""
 
-from collections.abc import Generator
-from contextlib import contextmanager
 from typing import Any, Literal
 
-from digitalkin_proto.digitalkin.filesystem.v1 import filesystem_pb2, filesystem_service_pb2_grpc
+from agentic_mesh_protocol.filesystem.v1 import filesystem_pb2, filesystem_service_pb2_grpc
 from google.protobuf import struct_pb2
 from google.protobuf.json_format import MessageToDict
 
-from digitalkin.grpc_servers.utils.exceptions import ServerError
 from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
-from digitalkin.grpc_servers.utils.models import ClientConfig
+from digitalkin.grpc_servers.utils.grpc_error_handler import GrpcErrorHandlerMixin
 from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import ClientConfig
 from digitalkin.services.filesystem.filesystem_strategy import (
     FileFilter,
     FilesystemRecord,
@@ -21,36 +19,9 @@ from digitalkin.services.filesystem.filesystem_strategy import (
 )
 
 
-class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
+class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper, GrpcErrorHandlerMixin):
     """Default state filesystem strategy."""
 
-    @staticmethod
-    @contextmanager
-    def _handle_grpc_errors(operation: str) -> Generator[Any, Any, Any]:
-        """Context manager for consistent gRPC error handling.
-
-        Yields:
-            Allow error handling in context.
-
-        Args:
-            operation: Description of the operation being performed.
-
-        Raises:
-            ValueError: Error with the model validation.
-            ServerError: from gRPC Client.
-            FilesystemServiceError: Filesystem service internal.
-        """
-        try:
-            yield
-        except ServerError as e:
-            msg = f"gRPC {operation} failed: {e}"
-            logger.exception(msg)
-            raise ServerError(msg) from e
-        except Exception as e:
-            msg = f"Unexpected error in {operation}"
-            logger.exception(msg)
-            raise FilesystemServiceError(msg) from e
-
     @staticmethod
     def _file_type_to_enum(file_type: str) -> filesystem_pb2.FileType:
         """Convert a file type string to a FileType enum.
@@ -148,7 +119,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
         self.service_name = "FilesystemService"
         channel = self._init_channel(client_config)
         self.stub = filesystem_service_pb2_grpc.FilesystemServiceStub(channel)
-        logger.debug("Channel client 'Filesystem' initialized succesfully")
+        logger.debug("Channel client 'Filesystem' initialized successfully")
 
     def upload_files(
         self,
@@ -163,7 +134,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
             tuple[list[FilesystemRecord], int, int]: List of uploaded files, total uploaded count, total failed count
         """
         logger.debug("Uploading %d files", len(files))
-        with GrpcFilesystem._handle_grpc_errors("UploadFiles"):
+        with self.handle_grpc_errors("UploadFiles", FilesystemServiceError):
             upload_files: list[filesystem_pb2.UploadFileData] = []
             for file in files:
                 metadata_struct: struct_pb2.Struct | None = None
@@ -213,7 +184,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
                 context_id = self.setup_id
             case "mission":
                 context_id = self.mission_id
-        with GrpcFilesystem._handle_grpc_errors("GetFile"):
+        with self.handle_grpc_errors("GetFile", FilesystemServiceError):
             request = filesystem_pb2.GetFileRequest(
                 context=context_id,
                 file_id=file_id,
@@ -261,7 +232,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
         Raises:
             FilesystemServiceError: If there is an error during update
         """
-        with GrpcFilesystem._handle_grpc_errors("UpdateFile"):
+        with self.handle_grpc_errors("UpdateFile", FilesystemServiceError):
             request = filesystem_pb2.UpdateFileRequest(
                 context=self.mission_id,
                 file_id=file_id,
@@ -295,7 +266,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
         Returns:
             tuple[dict[str, bool], int, int]: Results per file, total deleted count, total failed count
         """
-        with GrpcFilesystem._handle_grpc_errors("DeleteFiles"):
+        with self.handle_grpc_errors("DeleteFiles", FilesystemServiceError):
             request = filesystem_pb2.DeleteFilesRequest(
                 context=self.mission_id,
                 filters=self._filter_to_proto(filters),
@@ -332,7 +303,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
                 context_id = self.setup_id
             case "mission":
                 context_id = self.mission_id
-        with GrpcFilesystem._handle_grpc_errors("GetFiles"):
+        with self.handle_grpc_errors("GetFiles", FilesystemServiceError):
             request = filesystem_pb2.GetFilesRequest(
                 context=context_id,
                 filters=self._filter_to_proto(filters),

digitalkin/services/registry/__init__.py

@@ -1,6 +1,27 @@
 """This module is responsible for handling the registry service."""
 
+from digitalkin.models.services.registry import (
+    ModuleInfo,
+    RegistryModuleStatus,
+    RegistryModuleType,
+)
 from digitalkin.services.registry.default_registry import DefaultRegistry
+from digitalkin.services.registry.exceptions import (
+    RegistryModuleNotFoundError,
+    RegistryServiceError,
+)
+from digitalkin.services.registry.grpc_registry import GrpcRegistry
+from digitalkin.services.registry.registry_models import ModuleStatusInfo
 from digitalkin.services.registry.registry_strategy import RegistryStrategy
 
-__all__ = ["DefaultRegistry", "RegistryStrategy"]
+__all__ = [
+    "DefaultRegistry",
+    "GrpcRegistry",
+    "ModuleInfo",
+    "ModuleStatusInfo",
+    "RegistryModuleNotFoundError",
+    "RegistryModuleStatus",
+    "RegistryModuleType",
+    "RegistryServiceError",
+    "RegistryStrategy",
+]

digitalkin/services/registry/default_registry.py

@@ -1,10 +1,141 @@
-"""Default registry."""
+"""Default registry implementation."""
 
+from typing import ClassVar
+
+from digitalkin.models.services.registry import (
+    ModuleInfo,
+    RegistryModuleStatus,
+    RegistryModuleType,
+)
+from digitalkin.services.registry.exceptions import RegistryModuleNotFoundError
+from digitalkin.services.registry.registry_models import ModuleStatusInfo
 from digitalkin.services.registry.registry_strategy import RegistryStrategy
 
 
 class DefaultRegistry(RegistryStrategy):
-    """Default registry strategy."""
+    """Default registry strategy using in-memory storage."""
+
+    _modules: ClassVar[dict[str, ModuleInfo]] = {}
+
+    def discover_by_id(self, module_id: str) -> ModuleInfo:
+        """Get module info by ID.
+
+        Args:
+            module_id: The module identifier.
+
+        Returns:
+            ModuleInfo with module details.
+
+        Raises:
+            RegistryModuleNotFoundError: If module not found.
+        """
+        if module_id not in self._modules:
+            raise RegistryModuleNotFoundError(module_id)
+        return self._modules[module_id]
+
+    def search(
+        self,
+        name: str | None = None,
+        module_type: str | None = None,
+        organization_id: str | None = None,  # noqa: ARG002
+    ) -> list[ModuleInfo]:
+        """Search for modules by criteria.
+
+        Args:
+            name: Filter by name (partial match).
+            module_type: Filter by type (archetype, tool).
+            organization_id: Filter by organization (not used in local storage).
+
+        Returns:
+            List of matching modules.
+        """
+        results = list(self._modules.values())
+
+        if name:
+            results = [m for m in results if name in m.name]
+
+        if module_type:
+            results = [m for m in results if m.module_type == module_type]
+
+        return results
+
+    def get_status(self, module_id: str) -> ModuleStatusInfo:
+        """Get module status.
+
+        Args:
+            module_id: The module identifier.
+
+        Returns:
+            ModuleStatusInfo with current status.
+
+        Raises:
+            RegistryModuleNotFoundError: If module not found.
+        """
+        if module_id not in self._modules:
+            raise RegistryModuleNotFoundError(module_id)
+
+        module = self._modules[module_id]
+        return ModuleStatusInfo(
+            module_id=module_id,
+            status=module.status or RegistryModuleStatus.UNSPECIFIED,
+        )
+
+    def register(
+        self,
+        module_id: str,
+        address: str,
+        port: int,
+        version: str,
+    ) -> ModuleInfo | None:
+        """Register a module with the registry.
+
+        Note: Updates existing module or creates new one in local storage.
+
+        Args:
+            module_id: Unique module identifier.
+            address: Network address.
+            port: Network port.
+            version: Module version.
+
+        Returns:
+            ModuleInfo if successful, None otherwise.
+        """
+        existing = self._modules.get(module_id)
+        self._modules[module_id] = ModuleInfo(
+            module_id=module_id,
+            module_type=existing.module_type if existing else RegistryModuleType.UNSPECIFIED,
+            address=address,
+            port=port,
+            version=version,
+            name=existing.name if existing else module_id,
+            status=RegistryModuleStatus.ACTIVE,
+        )
+        return self._modules[module_id]
+
+    def heartbeat(self, module_id: str) -> RegistryModuleStatus:
+        """Send heartbeat to keep module active.
+
+        Args:
+            module_id: The module identifier.
+
+        Returns:
+            Current module status after heartbeat.
+
+        Raises:
+            RegistryModuleNotFoundError: If module not found.
+        """
+        if module_id not in self._modules:
+            raise RegistryModuleNotFoundError(module_id)
 
-    def get_by_id(self, module_id: str) -> None:
-        """Get services from the registry."""
+        module = self._modules[module_id]
+        # Update status to ACTIVE on heartbeat
+        self._modules[module_id] = ModuleInfo(
+            module_id=module.module_id,
+            module_type=module.module_type,
+            address=module.address,
+            port=module.port,
+            version=module.version,
+            name=module.name,
+            status=RegistryModuleStatus.ACTIVE,
+        )
+        return RegistryModuleStatus.ACTIVE

digitalkin/services/registry/exceptions.py

@@ -0,0 +1,47 @@
+"""Registry-specific exceptions.
+
+This module contains custom exceptions for registry service operations.
+"""
+
+
+class RegistryServiceError(Exception):
+    """Base exception for registry service errors."""
+
+
+class RegistryModuleNotFoundError(RegistryServiceError):
+    """Raised when a module is not found in the registry."""
+
+    def __init__(self, module_id: str) -> None:
+        """Initialize the exception.
+
+        Args:
+            module_id: The ID of the module that was not found.
+        """
+        self.module_id = module_id
+        super().__init__(f"Module '{module_id}' not found in registry")
+
+
+class ModuleAlreadyExistsError(RegistryServiceError):
+    """Raised when attempting to register an already-registered module."""
+
+    def __init__(self, module_id: str) -> None:
+        """Initialize the exception.
+
+        Args:
+            module_id: The ID of the module that already exists.
+        """
+        self.module_id = module_id
+        super().__init__(f"Module '{module_id}' already registered")
+
+
+class InvalidStatusError(RegistryServiceError):
+    """Raised when an invalid status is provided."""
+
+    def __init__(self, status: int) -> None:
+        """Initialize the exception.
+
+        Args:
+            status: The invalid status value.
+        """
+        self.status = status
+        super().__init__(f"Invalid module status: {status}")