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

digitalkin/services/registry/grpc_registry.py

@@ -0,0 +1,382 @@
+"""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,
+    RegistrySetupStatus,
+    RegistryVisibility,
+    SetupInfo,
+)
+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,
+            module_name=descriptor.name,
+            documentation=descriptor.documentation or None,
+        )
+
+    @staticmethod
+    def _proto_to_setup_info(descriptor: registry_models_pb2.SetupDescriptor) -> SetupInfo | None:
+        """Convert proto SetupDescriptor to SetupInfo.
+
+        Args:
+            descriptor: Proto SetupDescriptor message.
+
+        Returns:
+            SetupInfo with mapped fields, or None if descriptor is empty.
+        """
+        if not descriptor.id:
+            return None
+        status_name = registry_enums_pb2.SetupStatus.Name(descriptor.status).removeprefix("SETUP_STATUS_")
+        visibility_name = registry_enums_pb2.Visibility.Name(descriptor.visibility).removeprefix("VISIBILITY_")
+        return SetupInfo(
+            setup_id=descriptor.id,
+            name=descriptor.name,
+            documentation=descriptor.documentation or None,
+            status=RegistrySetupStatus[status_name],
+            visibility=RegistryVisibility[visibility_name],
+            organization_id=descriptor.organization_id or None,
+            owner_id=descriptor.owner_id or None,
+            card_id=descriptor.card_id or None,
+            module_id=descriptor.module_id or None,
+            setup_version_id=descriptor.setup_version_id or None,
+            setup_version=descriptor.setup_version or None,
+            config=dict(descriptor.config) if descriptor.config else 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]
+
+    def get_setup(self, setup_id: str) -> SetupInfo | None:
+        """Get setup info.
+
+        Args:
+            setup_id: The setup identifier.
+
+        Returns:
+            SetupInfo if successful, None otherwise.
+
+        Raises:
+            RegistryServiceError: If gRPC call fails.
+        """
+        logger.debug("Getting setup", extra={"setup_id": setup_id})
+        with self.handle_grpc_errors("GetSetup", RegistryServiceError):
+            try:
+                response = self.exec_grpc_query(
+                    "GetSetup",
+                    registry_requests_pb2.GetSetupRequest(setup_id=setup_id),
+                )
+            except ServerError as e:
+                msg = f"Failed to get setup '{setup_id}': {e}"
+                logger.error(msg)
+                raise RegistryServiceError(msg) from e
+            return self._proto_to_setup_info(response)
+
+    async def deregister(self, module_id: str) -> bool:  # noqa: PLR6301
+        """Deregister a module from the registry.
+
+        Note: The registry protocol uses heartbeat expiration for deregistration.
+        When a module stops sending heartbeats, it becomes inactive. This method
+        logs the deregistration intent for observability.
+
+        Args:
+            module_id: The module identifier to deregister.
+
+        Returns:
+            True always (heartbeat expiration handles actual deregistration).
+        """
+        logger.info(
+            "Module deregistration initiated (will become inactive via heartbeat expiration)",
+            extra={"module_id": module_id},
+        )
+        return True

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,116 @@
-"""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,
+    SetupInfo,
+)
 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 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 get_by_id(self, module_id: str) -> None:
-        """Get services from the registry."""
+    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
+
+    @abstractmethod
+    def get_setup(self, setup_id: str) -> SetupInfo | None:
+        """Get setup info."""
+        raise NotImplementedError
+
+    @abstractmethod
+    async def deregister(self, module_id: str) -> bool:
+        """Deregister a module from the registry.
+
+        Args:
+            module_id: The module identifier to deregister.
+
+        Returns:
+            True if deregistration was successful, False otherwise.
+        """
         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

@@ -173,7 +173,7 @@ class DefaultSetup(SetupStrategy):
         return [
             value
             for value in self.setup_versions[setup_version_dict["setup_id"]].values()
-            if setup_version_dict["query_versions"] in value.version or setup_version_dict["name"] in value.name
+            if setup_version_dict["query_versions"] in value.version
         ]
 
     def update_setup_version(self, setup_version_dict: dict[str, Any]) -> bool:

digitalkin/services/setup/grpc_setup.py

@@ -5,7 +5,7 @@ from contextlib import contextmanager
 from typing import Any
 
 import grpc
-from digitalkin_proto.agentic_mesh_protocol.setup.v1 import (
+from agentic_mesh_protocol.setup.v1 import (
     setup_pb2,
     setup_service_pb2_grpc,
 )

digitalkin/services/storage/grpc_storage.py

@@ -1,6 +1,6 @@
 """This module implements the default storage strategy."""
 
-from digitalkin_proto.agentic_mesh_protocol.storage.v1 import data_pb2, storage_service_pb2_grpc
+from agentic_mesh_protocol.storage.v1 import data_pb2, storage_service_pb2_grpc
 from google.protobuf import json_format
 from google.protobuf.struct_pb2 import Struct
 from pydantic import BaseModel

digitalkin/services/user_profile/__init__.py

@@ -1 +1,12 @@
 """UserProfile service package."""
+
+from digitalkin.services.user_profile.default_user_profile import DefaultUserProfile
+from digitalkin.services.user_profile.grpc_user_profile import GrpcUserProfile
+from digitalkin.services.user_profile.user_profile_strategy import UserProfileServiceError, UserProfileStrategy
+
+__all__ = [
+    "DefaultUserProfile",
+    "GrpcUserProfile",
+    "UserProfileServiceError",
+    "UserProfileStrategy",
+]

digitalkin/services/user_profile/grpc_user_profile.py

@@ -2,7 +2,7 @@
 
 from typing import Any
 
-from digitalkin_proto.agentic_mesh_protocol.user_profile.v1 import (
+from agentic_mesh_protocol.user_profile.v1 import (
     user_profile_pb2,
     user_profile_service_pb2_grpc,
 )
@@ -49,7 +49,7 @@ class GrpcUserProfile(UserProfileStrategy, GrpcClientWrapper, GrpcErrorHandlerMi
             ServerError: If gRPC operation fails
         """
         with self.handle_grpc_errors("GetUserProfile", UserProfileServiceError):
-            # mission_id maps to the user context in the proto request
+            # mission_id typically contains user context
             request = user_profile_pb2.GetUserProfileRequest(mission_id=self.mission_id)
             response = self.exec_grpc_query("GetUserProfile", request)