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

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}")

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.