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

digitalkin/services/registry/grpc_registry.py

@@ -0,0 +1,306 @@
+"""gRPC Registry client implementation.
+
+This module provides a gRPC-based registry client that communicates with
+the Service Provider's Registry service.
+"""
+
+from typing import Any
+
+from agentic_mesh_protocol.registry.v1 import (
+    registry_enums_pb2,
+    registry_models_pb2,
+    registry_requests_pb2,
+    registry_service_pb2_grpc,
+)
+
+from digitalkin.grpc_servers.utils.exceptions import ServerError
+from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
+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.models.services.registry import (
+    ModuleInfo,
+    RegistryModuleStatus,
+    RegistryModuleType,
+)
+from digitalkin.services.registry.exceptions import (
+    RegistryModuleNotFoundError,
+    RegistryServiceError,
+)
+from digitalkin.services.registry.registry_models import ModuleStatusInfo
+from digitalkin.services.registry.registry_strategy import RegistryStrategy
+
+
+class GrpcRegistry(RegistryStrategy, GrpcClientWrapper, GrpcErrorHandlerMixin):
+    """gRPC-based registry client.
+
+    This client communicates with the Service Provider's Registry service
+    to perform module discovery, registration, and status management operations.
+    """
+
+    def __init__(
+        self,
+        mission_id: str,
+        setup_id: str,
+        setup_version_id: str,
+        client_config: ClientConfig,
+        config: dict[str, Any] | None = None,
+    ) -> None:
+        """Initialize the gRPC registry client."""
+        RegistryStrategy.__init__(self, mission_id, setup_id, setup_version_id, config)
+        self.service_name = "RegistryService"
+        self.stub = registry_service_pb2_grpc.RegistryServiceStub(self._init_channel(client_config))
+        logger.debug("Channel client 'Registry' initialized successfully")
+
+    @staticmethod
+    def _proto_to_module_info(
+        descriptor: registry_models_pb2.ModuleDescriptor,
+    ) -> ModuleInfo:
+        """Convert proto ModuleDescriptor to ModuleInfo.
+
+        Args:
+            descriptor: Proto ModuleDescriptor message.
+
+        Returns:
+            ModuleInfo with mapped fields.
+        """
+        type_name = registry_enums_pb2.ModuleType.Name(descriptor.module_type).removeprefix("MODULE_TYPE_")
+        return ModuleInfo(
+            module_id=descriptor.id,
+            module_type=RegistryModuleType[type_name],
+            address=descriptor.address,
+            port=descriptor.port,
+            version=descriptor.version,
+            name=descriptor.name,
+            documentation=descriptor.documentation or None,
+        )
+
+    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.
+            RegistryServiceError: If gRPC call fails.
+        """
+        logger.debug("Discovering module by ID", extra={"module_id": module_id})
+
+        with self.handle_grpc_errors("GetModule", RegistryServiceError):
+            try:
+                response = self.exec_grpc_query(
+                    "GetModule",
+                    registry_requests_pb2.GetModuleRequest(module_id=module_id),
+                )
+            except ServerError as e:
+                msg = f"Failed to discover module '{module_id}': {e}"
+                logger.error(msg)
+                raise RegistryServiceError(msg) from e
+
+            if not response.id:
+                logger.warning("Module not found in registry", extra={"module_id": module_id})
+                raise RegistryModuleNotFoundError(module_id)
+
+            logger.debug(
+                "Module discovered",
+                extra={
+                    "module_id": response.id,
+                    "address": response.address,
+                    "port": response.port,
+                },
+            )
+            return self._proto_to_module_info(response)
+
+    def search(
+        self,
+        name: str | None = None,
+        module_type: str | None = None,
+        organization_id: str | None = None,
+    ) -> list[ModuleInfo]:
+        """Search for modules by criteria.
+
+        Args:
+            name: Filter by name (partial match via query).
+            module_type: Filter by type (archetype, tool).
+            organization_id: Filter by organization.
+
+        Returns:
+            List of matching modules.
+
+        Raises:
+            RegistryServiceError: If gRPC call fails.
+        """
+        logger.debug(
+            "Searching modules",
+            extra={
+                "name": name,
+                "module_type": module_type,
+                "organization_id": organization_id,
+            },
+        )
+
+        with self.handle_grpc_errors("DiscoverModules", RegistryServiceError):
+            module_types = []
+            if module_type:
+                enum_val = RegistryModuleType[module_type.upper()]
+                module_types.append(getattr(registry_enums_pb2, f"MODULE_TYPE_{enum_val.name}"))
+
+            try:
+                response = self.exec_grpc_query(
+                    "DiscoverModules",
+                    registry_requests_pb2.DiscoverModulesRequest(
+                        query=name or "",
+                        organization_id=organization_id or "",
+                        module_types=module_types,
+                    ),
+                )
+            except ServerError as e:
+                msg = f"Failed to search modules: {e}"
+                logger.error(msg)
+                raise RegistryServiceError(msg) from e
+
+            logger.debug("Search returned %d modules", len(response.modules))
+            return [self._proto_to_module_info(m) for m in response.modules]
+
+    def get_status(self, module_id: str) -> ModuleStatusInfo:
+        """Get module status by fetching the module.
+
+        Args:
+            module_id: The module identifier.
+
+        Returns:
+            ModuleStatusInfo with current status.
+
+        Raises:
+            RegistryModuleNotFoundError: If module not found.
+            RegistryServiceError: If gRPC call fails.
+        """
+        logger.debug("Getting module status", extra={"module_id": module_id})
+
+        with self.handle_grpc_errors("GetModule", RegistryServiceError):
+            try:
+                response = self.exec_grpc_query(
+                    "GetModule",
+                    registry_requests_pb2.GetModuleRequest(module_id=module_id),
+                )
+            except ServerError as e:
+                msg = f"Failed to get module status for '{module_id}': {e}"
+                logger.error(msg)
+                raise RegistryServiceError(msg) from e
+
+            if not response.id:
+                logger.warning("Module not found in registry", extra={"module_id": module_id})
+                raise RegistryModuleNotFoundError(module_id)
+
+            status_name = registry_enums_pb2.ModuleStatus.Name(response.status).removeprefix("MODULE_STATUS_")
+            logger.debug(
+                "Module status retrieved",
+                extra={"module_id": response.id, "status": status_name},
+            )
+            return ModuleStatusInfo(
+                module_id=response.id,
+                status=RegistryModuleStatus[status_name],
+            )
+
+    def register(
+        self,
+        module_id: str,
+        address: str,
+        port: int,
+        version: str,
+    ) -> ModuleInfo | None:
+        """Register a module with the registry.
+
+        Note: The new proto only updates address/port/version for an existing module.
+        The module must already exist in the registry database.
+
+        Args:
+            module_id: Unique module identifier.
+            address: Network address.
+            port: Network port.
+            version: Module version.
+
+        Returns:
+            ModuleInfo if successful, None if module not found.
+
+        Raises:
+            RegistryServiceError: If gRPC call fails.
+        """
+        logger.info(
+            "Registering module with registry",
+            extra={
+                "module_id": module_id,
+                "address": address,
+                "port": port,
+                "version": version,
+            },
+        )
+
+        with self.handle_grpc_errors("RegisterModule", RegistryServiceError):
+            try:
+                response = self.exec_grpc_query(
+                    "RegisterModule",
+                    registry_requests_pb2.RegisterModuleRequest(
+                        module_id=module_id,
+                        address=address,
+                        port=port,
+                        version=version,
+                    ),
+                )
+            except ServerError as e:
+                msg = f"Failed to register module '{module_id}': {e}"
+                logger.error(msg)
+                raise RegistryServiceError(msg) from e
+
+            if not response.module or not response.module.id:
+                logger.warning(
+                    "Registry returned empty response for module registration",
+                    extra={"module_id": module_id},
+                )
+                return None
+
+            logger.info(
+                "Module registered successfully",
+                extra={
+                    "module_id": response.module.id,
+                    "address": response.module.address,
+                    "port": response.module.port,
+                },
+            )
+            return self._proto_to_module_info(response.module)
+
+    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:
+            RegistryServiceError: If gRPC call fails.
+        """
+        logger.debug("Sending heartbeat", extra={"module_id": module_id})
+
+        with self.handle_grpc_errors("Heartbeat", RegistryServiceError):
+            try:
+                response = self.exec_grpc_query(
+                    "Heartbeat",
+                    registry_requests_pb2.HeartbeatRequest(module_id=module_id),
+                )
+            except ServerError as e:
+                msg = f"Failed to send heartbeat for '{module_id}': {e}"
+                logger.error(msg)
+                raise RegistryServiceError(msg) from e
+
+            status_name = registry_enums_pb2.ModuleStatus.Name(response.status).removeprefix("MODULE_STATUS_")
+            logger.debug(
+                "Heartbeat response",
+                extra={"module_id": module_id, "status": status_name},
+            )
+            return RegistryModuleStatus[status_name]

digitalkin/services/registry/registry_models.py

@@ -0,0 +1,15 @@
+"""Registry data models.
+
+This module contains Pydantic models for registry service data structures.
+"""
+
+from pydantic import BaseModel
+
+from digitalkin.models.services.registry import RegistryModuleStatus
+
+
+class ModuleStatusInfo(BaseModel):
+    """Module status response."""
+
+    module_id: str
+    status: RegistryModuleStatus

digitalkin/services/registry/registry_strategy.py

@@ -1,14 +1,98 @@
-"""This module contains the abstract base class for registry strategies."""
+"""Abstract base class for registry strategies."""
 
 from abc import ABC, abstractmethod
+from typing import Any
 
+from digitalkin.models.services.registry import (
+    ModuleInfo,
+    RegistryModuleStatus,
+)
 from digitalkin.services.base_strategy import BaseStrategy
+from digitalkin.services.registry.registry_models import ModuleStatusInfo
 
 
 class RegistryStrategy(BaseStrategy, ABC):
-    """Abstract base class for registry strategies."""
+    """Abstract base class for registry strategies.
+
+    Defines the interface for registry operations including module discovery,
+    registration, and status management.
+    """
+
+    def __init__(
+        self,
+        mission_id: str,
+        setup_id: str,
+        setup_version_id: str,
+        config: dict[str, Any] | None = None,
+    ) -> None:
+        """Initialize the strategy."""
+        super().__init__(mission_id, setup_id, setup_version_id)
+        self.config = config
+
+    @abstractmethod
+    def discover_by_id(self, module_id: str) -> ModuleInfo:
+        """Get module info by ID."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def search(
+        self,
+        name: str | None = None,
+        module_type: str | None = None,
+        organization_id: str | None = None,
+    ) -> list[ModuleInfo]:
+        """Search for modules by criteria.
+
+        Args:
+            name: Filter by name (partial match via query).
+            module_type: Filter by type (archetype, tool).
+            organization_id: Filter by organization.
+
+        Returns:
+            List of matching modules.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_status(self, module_id: str) -> ModuleStatusInfo:
+        """Get module status."""
+        raise NotImplementedError
 
     @abstractmethod
-    def get_by_id(self, module_id: str) -> None:
-        """Get services from the registry."""
+    def register(
+        self,
+        module_id: str,
+        address: str,
+        port: int,
+        version: str,
+    ) -> ModuleInfo | None:
+        """Register a module with the registry.
+
+        Note: The new proto only updates address/port/version for an existing module.
+        The module must already exist in the registry database.
+
+        Args:
+            module_id: Unique module identifier.
+            address: Network address.
+            port: Network port.
+            version: Module version.
+
+        Returns:
+            ModuleInfo if successful, None otherwise.
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    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.
+        """
         raise NotImplementedError

digitalkin/services/services_config.py

@@ -5,13 +5,15 @@ from typing import Any, ClassVar
 from pydantic import BaseModel, Field, PrivateAttr
 
 from digitalkin.services.agent import AgentStrategy, DefaultAgent
+from digitalkin.services.communication import CommunicationStrategy, DefaultCommunication, GrpcCommunication
 from digitalkin.services.cost import CostStrategy, DefaultCost, GrpcCost
 from digitalkin.services.filesystem import DefaultFilesystem, FilesystemStrategy, GrpcFilesystem
 from digitalkin.services.identity import DefaultIdentity, IdentityStrategy
-from digitalkin.services.registry import DefaultRegistry, RegistryStrategy
+from digitalkin.services.registry import DefaultRegistry, GrpcRegistry, RegistryStrategy
 from digitalkin.services.services_models import ServicesMode, ServicesStrategy
 from digitalkin.services.snapshot import DefaultSnapshot, SnapshotStrategy
 from digitalkin.services.storage import DefaultStorage, GrpcStorage, StorageStrategy
+from digitalkin.services.user_profile import DefaultUserProfile, GrpcUserProfile, UserProfileStrategy
 
 
 class ServicesConfig(BaseModel):
@@ -38,9 +40,9 @@ class ServicesConfig(BaseModel):
     )
     _config_snapshot: dict[str, Any | None] = PrivateAttr(default_factory=dict)
     _registry: ServicesStrategy[RegistryStrategy] = PrivateAttr(
-        default_factory=lambda: ServicesStrategy(local=DefaultRegistry, remote=DefaultRegistry)
+        default_factory=lambda: ServicesStrategy(local=DefaultRegistry, remote=GrpcRegistry)
     )
-    _config__registry: dict[str, Any | None] = PrivateAttr(default_factory=dict)
+    _config_registry: dict[str, Any | None] = PrivateAttr(default_factory=dict)
     _filesystem: ServicesStrategy[FilesystemStrategy] = PrivateAttr(
         default_factory=lambda: ServicesStrategy(local=DefaultFilesystem, remote=GrpcFilesystem)
     )
@@ -53,6 +55,14 @@ class ServicesConfig(BaseModel):
         default_factory=lambda: ServicesStrategy(local=DefaultIdentity, remote=DefaultIdentity)
     )
     _config_identity: dict[str, Any | None] = PrivateAttr(default_factory=dict)
+    _communication: ServicesStrategy[CommunicationStrategy] = PrivateAttr(
+        default_factory=lambda: ServicesStrategy(local=DefaultCommunication, remote=GrpcCommunication)
+    )
+    _config_communication: dict[str, Any | None] = PrivateAttr(default_factory=dict)
+    _user_profile: ServicesStrategy[UserProfileStrategy] = PrivateAttr(
+        default_factory=lambda: ServicesStrategy(local=DefaultUserProfile, remote=GrpcUserProfile)
+    )
+    _config_user_profile: dict[str, Any | None] = PrivateAttr(default_factory=dict)
 
     # List of valid strategy names for validation
     _valid_strategy_names: ClassVar[set[str]] = {
@@ -63,6 +73,8 @@ class ServicesConfig(BaseModel):
         "filesystem",
         "agent",
         "identity",
+        "communication",
+        "user_profile",
     }
 
     def __init__(
@@ -169,6 +181,16 @@ class ServicesConfig(BaseModel):
         """Get the identity service strategy class based on the current mode."""
         return self._identity[self.mode.value]
 
+    @property
+    def communication(self) -> type[CommunicationStrategy]:
+        """Get the communication service strategy class based on the current mode."""
+        return self._communication[self.mode.value]
+
+    @property
+    def user_profile(self) -> type[UserProfileStrategy]:
+        """Get the user_profile service strategy class based on the current mode."""
+        return self._user_profile[self.mode.value]
+
     def update_mode(self, mode: ServicesMode) -> None:
         """Update the strategy mode.
 

digitalkin/services/services_models.py

@@ -7,23 +7,27 @@ from pydantic import BaseModel
 
 from digitalkin.logger import logger
 from digitalkin.services.agent import AgentStrategy
+from digitalkin.services.communication import CommunicationStrategy
 from digitalkin.services.cost import CostStrategy
 from digitalkin.services.filesystem import FilesystemStrategy
 from digitalkin.services.identity import IdentityStrategy
 from digitalkin.services.registry import RegistryStrategy
 from digitalkin.services.snapshot import SnapshotStrategy
 from digitalkin.services.storage import StorageStrategy
+from digitalkin.services.user_profile import UserProfileStrategy
 
 # Define type variables
 T = TypeVar(
     "T",
     bound=AgentStrategy
+    | CommunicationStrategy
     | CostStrategy
     | FilesystemStrategy
     | IdentityStrategy
     | RegistryStrategy
     | SnapshotStrategy
-    | StorageStrategy,
+    | StorageStrategy
+    | UserProfileStrategy,
 )
 
 

digitalkin/services/setup/default_setup.py

@@ -7,8 +7,7 @@ from typing import Any
 from pydantic import ValidationError
 
 from digitalkin.logger import logger
-from digitalkin.services.setup.grpc_setup import SetupData, SetupVersionData
-from digitalkin.services.setup.setup_strategy import SetupServiceError, SetupStrategy
+from digitalkin.services.setup.setup_strategy import SetupData, SetupServiceError, SetupStrategy, SetupVersionData
 
 
 class DefaultSetup(SetupStrategy):
@@ -47,7 +46,7 @@ class DefaultSetup(SetupStrategy):
         )
         valid_data.id = setup_id
         self.setups[setup_id] = valid_data
-        logger.debug("CREATE SETUP DATA %s:%s succesfull", setup_id, valid_data)
+        logger.debug("CREATE SETUP DATA %s:%s successful", setup_id, valid_data)
         return setup_id
 
     def get_setup(self, setup_dict: dict[str, Any]) -> SetupData:
@@ -131,7 +130,7 @@ class DefaultSetup(SetupStrategy):
         if setup_version_dict["setup_id"] not in self.setup_versions:
             self.setup_versions[setup_version_dict["setup_id"]] = {}
         self.setup_versions[setup_version_dict["setup_id"]][valid_data.version] = valid_data
-        logger.debug("CREATE SETUP VERSION DATA %s:%s succesfull", setup_version_dict["setup_id"], valid_data)
+        logger.debug("CREATE SETUP VERSION DATA %s:%s successful", setup_version_dict["setup_id"], valid_data)
         return valid_data.version
 
     def get_setup_version(self, setup_version_dict: dict[str, Any]) -> SetupVersionData:
@@ -173,8 +172,8 @@ class DefaultSetup(SetupStrategy):
 
         return [
             value
-            for value in setup_version_dict["setup_id"].values()
-            if setup_version_dict["query_versions"] in value.version or setup_version_dict["name"] in value.name
+            for value in self.setup_versions[setup_version_dict["setup_id"]].values()
+            if setup_version_dict["query_versions"] in value.version
         ]
 
     def update_setup_version(self, setup_version_dict: dict[str, Any]) -> bool:
@@ -190,7 +189,7 @@ class DefaultSetup(SetupStrategy):
             logger.debug("UPDATE setup_id = %s: setup_id DOESN'T EXIST", setup_version_dict["setup_id"])
             return False
 
-        if setup_version_dict["version"] not in self.setup_versions["setup_id"]:
+        if setup_version_dict["version"] not in self.setup_versions[setup_version_dict["setup_id"]]:
             logger.debug("UPDATE setup_id = %s: setup_id DOESN'T EXIST", setup_version_dict["setup_id"])
             return False