digitalkin 0.3.2.dev2__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,
+    ModuleStatusInfo,
+    RegistryModuleStatus,
+    RegistryModuleType,
+)
+from digitalkin.services.registry.exceptions import (
+    RegistryModuleNotFoundError,
+    RegistryServiceError,
+)
+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,43 @@
+"""Registry data models.
+
+This module contains Pydantic models for registry service data structures.
+"""
+
+from enum import IntEnum
+
+from pydantic import BaseModel
+
+
+class RegistryModuleStatus(IntEnum):
+    """Module status in the registry.
+
+    Maps to proto ModuleStatus enum values.
+    """
+
+    UNKNOWN = 0
+    READY = 1
+    ACTIVE = 2
+    OFFLINE = 3
+
+
+class ModuleInfo(BaseModel):
+    """Complete module information from registry.
+
+    Maps to proto ModuleDescriptor message.
+    """
+
+    module_id: str
+    module_type: str
+    address: str
+    port: int
+    version: str
+    name: str = ""
+    documentation: str | None = None
+    status: RegistryModuleStatus | None = None
+
+
+class ModuleStatusInfo(BaseModel):
+    """Module status response."""
+
+    module_id: str
+    status: RegistryModuleStatus

digitalkin/services/registry/registry_strategy.py

@@ -0,0 +1,98 @@
+"""Abstract base class for registry strategies."""
+
+from abc import ABC, abstractmethod
+from typing import Any
+
+from digitalkin.models.services.registry import (
+    ModuleInfo,
+    ModuleStatusInfo,
+    RegistryModuleStatus,
+)
+from digitalkin.services.base_strategy import BaseStrategy
+
+
+class RegistryStrategy(BaseStrategy, ABC):
+    """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 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

@@ -0,0 +1,200 @@
+"""Service Provider definitions."""
+
+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, 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):
+    """Service class describing the available services in a Module.
+
+    This class manages the strategy implementations for various services,
+    allowing them to be switched between local and remote modes.
+    """
+
+    # Mode setting for all strategies
+    mode: ServicesMode = Field(default=ServicesMode.LOCAL, description="The mode of the services (local or remote)")
+
+    # Strategy definitions with proper type annotations
+    _storage: ServicesStrategy[StorageStrategy] = PrivateAttr(
+        default_factory=lambda: ServicesStrategy(local=DefaultStorage, remote=GrpcStorage)
+    )
+    _config_storage: dict[str, Any | None] = PrivateAttr(default_factory=dict)
+    _cost: ServicesStrategy[CostStrategy] = PrivateAttr(
+        default_factory=lambda: ServicesStrategy(local=DefaultCost, remote=GrpcCost)
+    )
+    _config_cost: dict[str, Any | None] = PrivateAttr(default_factory=dict)
+    _snapshot: ServicesStrategy[SnapshotStrategy] = PrivateAttr(
+        default_factory=lambda: ServicesStrategy(local=DefaultSnapshot, remote=DefaultSnapshot)
+    )
+    _config_snapshot: dict[str, Any | None] = PrivateAttr(default_factory=dict)
+    _registry: ServicesStrategy[RegistryStrategy] = PrivateAttr(
+        default_factory=lambda: ServicesStrategy(local=DefaultRegistry, remote=GrpcRegistry)
+    )
+    _config_registry: dict[str, Any | None] = PrivateAttr(default_factory=dict)
+    _filesystem: ServicesStrategy[FilesystemStrategy] = PrivateAttr(
+        default_factory=lambda: ServicesStrategy(local=DefaultFilesystem, remote=GrpcFilesystem)
+    )
+    _config_filesystem: dict[str, Any | None] = PrivateAttr(default_factory=dict)
+    _agent: ServicesStrategy[AgentStrategy] = PrivateAttr(
+        default_factory=lambda: ServicesStrategy(local=DefaultAgent, remote=DefaultAgent)
+    )
+    _config_agent: dict[str, Any | None] = PrivateAttr(default_factory=dict)
+    _identity: ServicesStrategy[IdentityStrategy] = PrivateAttr(
+        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]] = {
+        "storage",
+        "cost",
+        "snapshot",
+        "registry",
+        "filesystem",
+        "agent",
+        "identity",
+        "communication",
+        "user_profile",
+    }
+
+    def __init__(
+        self,
+        services_config_strategies: dict[str, ServicesStrategy | None] = {},
+        services_config_params: dict[str, dict[str, Any | None] | None] = {},
+        mode: ServicesMode = ServicesMode.LOCAL,
+        **kwargs: dict[str, Any],
+    ) -> None:
+        """Initialize the service configuration with optional strategy overrides.
+
+        Args:
+            services_config_strategies: Dictionary mapping service names to strategy implementations
+            services_config_params: Dictionary mapping service names to configuration parameters
+            mode: The mode of the services (local or remote)
+            **kwargs: Additional keyword arguments passed to the parent class constructor
+        """
+        super().__init__(**kwargs)
+        self.mode = mode
+        # Apply any strategy overrides
+        if services_config_strategies:
+            for name, strategy in services_config_strategies.items():
+                if strategy is not None and name in self._valid_strategy_names:
+                    setattr(self, f"_{name}", strategy)
+
+        for name in self.valid_strategy_names():
+            setattr(self, f"_config_{name}", services_config_params.get(name, {}))
+
+    @classmethod
+    def valid_strategy_names(cls) -> set[str]:
+        """Get the list of valid strategy names.
+
+        Returns:
+            The set of valid strategy names.
+        """
+        return cls._valid_strategy_names
+
+    def get_strategy_config(self, name: str) -> dict[str, Any]:
+        """Get the configuration for a specific strategy.
+
+        Args:
+            name: The name of the strategy to retrieve the configuration for
+
+        Returns:
+            The configuration for the specified strategy, or None if not found
+        """
+        return getattr(self, f"_config_{name}", {})
+
+    def init_strategy(self, name: str, mission_id: str, setup_id: str, setup_version_id: str) -> ServicesStrategy:
+        """Initialize a specific strategy.
+
+        Args:
+            name: The name of the strategy to initialize
+            mission_id: The ID of the mission this strategy is associated with
+            setup_id: The setup ID for the strategy
+            setup_version_id: The setup version ID for the strategy
+
+        Returns:
+            The initialized strategy instance
+
+        Raises:
+            ValueError: If the strategy is not found
+        """
+        strategy_type = getattr(self, name, None)
+        if strategy_type is None:
+            msg = f"Strategy {name} not found in ServicesConfig."
+            raise ValueError(msg)
+
+        # Instantiate the strategy with the mission ID, setup version ID, and configuration
+        return strategy_type(mission_id, setup_id, setup_version_id, **self.get_strategy_config(name) or {})
+
+    @property
+    def storage(self) -> type[StorageStrategy]:
+        """Get the storage service strategy class based on the current mode."""
+        return self._storage[self.mode.value]
+
+    @property
+    def cost(self) -> type[CostStrategy]:
+        """Get the cost service strategy class based on the current mode."""
+        return self._cost[self.mode.value]
+
+    @property
+    def snapshot(self) -> type[SnapshotStrategy]:
+        """Get the snapshot service strategy class based on the current mode."""
+        return self._snapshot[self.mode.value]
+
+    @property
+    def registry(self) -> type[RegistryStrategy]:
+        """Get the registry service strategy class based on the current mode."""
+        return self._registry[self.mode.value]
+
+    @property
+    def filesystem(self) -> type[FilesystemStrategy]:
+        """Get the filesystem service strategy class based on the current mode."""
+        return self._filesystem[self.mode.value]
+
+    @property
+    def agent(self) -> type[AgentStrategy]:
+        """Get the agent service strategy class based on the current mode."""
+        return self._agent[self.mode.value]
+
+    @property
+    def identity(self) -> type[IdentityStrategy]:
+        """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.
+
+        Parameters:
+            mode: The new mode to use for all strategies
+        """
+        self.mode = mode