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

digitalkin/modules/triggers/healthcheck_services_trigger.py

@@ -0,0 +1,63 @@
+"""Healthcheck services trigger - reports service health."""
+
+from typing import Any, ClassVar
+
+from digitalkin.mixins import BaseMixin
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.models.module.utility import (
+    HealthcheckServicesInput,
+    HealthcheckServicesOutput,
+    ServiceHealthStatus,
+)
+from digitalkin.modules.trigger_handler import TriggerHandler
+
+
+class HealthcheckServicesTrigger(TriggerHandler, BaseMixin):
+    """Handler for services healthcheck.
+
+    Reports the health status of all configured services (storage, cost, filesystem, etc.).
+    """
+
+    protocol: ClassVar[str] = "healthcheck_services"
+    input_format = HealthcheckServicesInput
+
+    def __init__(self, context: ModuleContext) -> None:
+        """Initialize the handler."""
+
+    async def handle(
+        self,
+        input_data: HealthcheckServicesInput,  # noqa: ARG002
+        setup_data: Any,  # noqa: ANN401, ARG002
+        context: ModuleContext,
+    ) -> None:
+        """Handle services healthcheck request.
+
+        Args:
+            input_data: The input trigger data (unused for healthcheck).
+            setup_data: The setup configuration (unused for healthcheck).
+            context: The module context.
+        """
+        service_names = ["storage", "cost", "filesystem", "registry", "user_profile"]
+        services_status: list[ServiceHealthStatus] = []
+
+        for name in service_names:
+            service = getattr(context, name, None)
+            if service is not None:
+                services_status.append(ServiceHealthStatus(name=name, status="healthy"))
+            else:
+                services_status.append(ServiceHealthStatus(name=name, status="unknown", message="Not configured"))
+
+        # Determine overall status
+        healthy_count = sum(1 for s in services_status if s.status == "healthy")
+        if healthy_count == len(services_status):
+            overall_status = "healthy"
+        elif healthy_count > 0:
+            overall_status = "degraded"
+        else:
+            overall_status = "unhealthy"
+
+        output = HealthcheckServicesOutput(
+            services=services_status,
+            overall_status=overall_status,  # type: ignore[arg-type]
+        )
+        await self.send_message(context, output)

digitalkin/modules/triggers/healthcheck_status_trigger.py

@@ -0,0 +1,52 @@
+"""Healthcheck status trigger - comprehensive module status."""
+
+import time
+from typing import Any, ClassVar
+
+from digitalkin.mixins import BaseMixin
+from digitalkin.models.module.module_context import ModuleContext
+from digitalkin.models.module.utility import (
+    HealthcheckStatusInput,
+    HealthcheckStatusOutput,
+)
+from digitalkin.modules.trigger_handler import TriggerHandler
+
+
+class HealthcheckStatusTrigger(TriggerHandler, BaseMixin):
+    """Handler for comprehensive status healthcheck.
+
+    Reports detailed module status including uptime, active jobs, and metadata.
+    """
+
+    protocol: ClassVar[str] = "healthcheck_status"
+    input_format = HealthcheckStatusInput
+    _start_time: ClassVar[float] = time.time()
+
+    def __init__(self, context: ModuleContext) -> None:
+        """Initialize the handler."""
+
+    async def handle(
+        self,
+        input_data: HealthcheckStatusInput,  # noqa: ARG002
+        setup_data: Any,  # noqa: ANN401, ARG002
+        context: ModuleContext,
+    ) -> None:
+        """Handle status healthcheck request.
+
+        Args:
+            input_data: The input trigger data (unused for healthcheck).
+            setup_data: The setup configuration (unused for healthcheck).
+            context: The module context.
+        """
+        output = HealthcheckStatusOutput(
+            module_name=context.session.setup_id,
+            module_status="RUNNING",
+            uptime_seconds=time.time() - self._start_time,
+            active_jobs=1,
+            metadata={
+                "job_id": context.session.job_id,
+                "mission_id": context.session.mission_id,
+                "setup_version_id": context.session.setup_version_id,
+            },
+        )
+        await self.send_message(context, output)

digitalkin/services/__init__.py

@@ -1,6 +1,7 @@
 """This package contains the abstract base class for all services."""
 
 from digitalkin.services.agent import AgentStrategy, DefaultAgent
+from digitalkin.services.communication import CommunicationStrategy, DefaultCommunication, GrpcCommunication
 from digitalkin.services.cost import CostStrategy, DefaultCost
 from digitalkin.services.filesystem import DefaultFilesystem, FilesystemStrategy
 from digitalkin.services.identity import DefaultIdentity, IdentityStrategy
@@ -10,8 +11,10 @@ from digitalkin.services.storage import DefaultStorage, StorageStrategy
 
 __all__ = [
     "AgentStrategy",
+    "CommunicationStrategy",
     "CostStrategy",
     "DefaultAgent",
+    "DefaultCommunication",
     "DefaultCost",
     "DefaultFilesystem",
     "DefaultIdentity",
@@ -19,6 +22,7 @@ __all__ = [
     "DefaultSnapshot",
     "DefaultStorage",
     "FilesystemStrategy",
+    "GrpcCommunication",
     "IdentityStrategy",
     "RegistryStrategy",
     "SnapshotStrategy",

digitalkin/services/communication/__init__.py

@@ -0,0 +1,7 @@
+"""Communication service for module-to-module interaction."""
+
+from digitalkin.services.communication.communication_strategy import CommunicationStrategy
+from digitalkin.services.communication.default_communication import DefaultCommunication
+from digitalkin.services.communication.grpc_communication import GrpcCommunication
+
+__all__ = ["CommunicationStrategy", "DefaultCommunication", "GrpcCommunication"]

digitalkin/services/communication/communication_strategy.py

@@ -0,0 +1,76 @@
+"""Abstract base class for communication strategies."""
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncGenerator, Awaitable, Callable
+
+from digitalkin.services.base_strategy import BaseStrategy
+
+
+class CommunicationStrategy(BaseStrategy, ABC):
+    """Abstract base class for module-to-module communication.
+
+    This service enables:
+    - Archetype → Tool communication
+    - Archetype → Archetype communication
+    - Tool → Tool communication
+    - Any module → Any module communication
+
+    The service wraps the Module Service protocol from agentic-mesh-protocol.
+    """
+
+    @abstractmethod
+    async def get_module_schemas(
+        self,
+        module_address: str,
+        module_port: int,
+        *,
+        llm_format: bool = False,
+    ) -> dict[str, dict]:
+        """Get module schemas (input/output/setup/secret).
+
+        Args:
+            module_address: Target module address
+            module_port: Target module port
+            llm_format: Return LLM-friendly format (simplified schema)
+
+        Returns:
+            Dictionary containing schemas:
+            {
+                "input": {...},
+                "output": {...},
+                "setup": {...},
+                "secret": {...}
+            }
+        """
+        raise NotImplementedError
+
+    @abstractmethod
+    async def call_module(
+        self,
+        module_address: str,
+        module_port: int,
+        input_data: dict,
+        setup_id: str,
+        mission_id: str,
+        callback: Callable[[dict], Awaitable[None]] | None = None,
+    ) -> AsyncGenerator[dict, None]:
+        """Call a module and stream responses.
+
+        Uses Module Service StartModule RPC to execute the module.
+        Streams responses as they are generated by the module.
+
+        Args:
+            module_address: Target module address
+            module_port: Target module port
+            input_data: Input data as dictionary
+            setup_id: Setup configuration ID
+            mission_id: Mission context ID
+            callback: Optional callback for each response
+
+        Yields:
+            Streaming responses from module as dictionaries
+        """
+        # Make this an actual async generator to satisfy type checkers
+        if False:  # pragma: no cover
+            yield {}
+        raise NotImplementedError

digitalkin/services/communication/default_communication.py

@@ -0,0 +1,101 @@
+"""Default communication implementation (local, for testing)."""
+
+from collections.abc import AsyncGenerator, Awaitable, Callable
+
+from digitalkin.logger import logger
+from digitalkin.services.communication.communication_strategy import CommunicationStrategy
+
+
+class DefaultCommunication(CommunicationStrategy):
+    """Default communication strategy (local implementation).
+
+    This implementation is primarily for testing and development.
+    For production, use GrpcCommunication to connect to remote modules.
+    """
+
+    def __init__(
+        self,
+        mission_id: str,
+        setup_id: str,
+        setup_version_id: str,
+    ) -> None:
+        """Initialize the default communication service.
+
+        Args:
+            mission_id: Mission identifier
+            setup_id: Setup identifier
+            setup_version_id: Setup version identifier
+        """
+        super().__init__(mission_id, setup_id, setup_version_id)
+        logger.debug("Initialized DefaultCommunication (local)")
+
+    async def get_module_schemas(  # noqa: PLR6301
+        self,
+        module_address: str,
+        module_port: int,
+        *,
+        llm_format: bool = False,
+    ) -> dict[str, dict]:
+        """Get module schemas (local implementation returns empty schemas).
+
+        Args:
+            module_address: Target module address
+            module_port: Target module port
+            llm_format: Return LLM-friendly format
+
+        Returns:
+            Empty schemas dictionary
+        """
+        logger.debug(
+            "DefaultCommunication.get_module_schemas called (returns empty)",
+            extra={
+                "module_address": module_address,
+                "module_port": module_port,
+                "llm_format": llm_format,
+            },
+        )
+        return {
+            "input": {},
+            "output": {},
+            "setup": {},
+            "secret": {},
+        }
+
+    async def call_module(  # noqa: PLR6301
+        self,
+        module_address: str,
+        module_port: int,
+        input_data: dict,  # noqa: ARG002
+        setup_id: str,
+        mission_id: str,
+        callback: Callable[[dict], Awaitable[None]] | None = None,
+    ) -> AsyncGenerator[dict, None]:
+        """Call module (local implementation yields empty response).
+
+        Args:
+            module_address: Target module address
+            module_port: Target module port
+            input_data: Input data
+            setup_id: Setup ID
+            mission_id: Mission ID
+            callback: Optional callback
+
+        Yields:
+            Empty response dictionary
+        """
+        logger.debug(
+            "DefaultCommunication.call_module called (returns empty)",
+            extra={
+                "module_address": module_address,
+                "module_port": module_port,
+                "setup_id": setup_id,
+                "mission_id": mission_id,
+            },
+        )
+
+        # Yield empty response
+        response = {"status": "error", "message": "Local communication not implemented"}
+        if callback:
+            await callback(response)
+        yield response
+        return  # Explicit return for async generator

digitalkin/services/communication/grpc_communication.py

@@ -0,0 +1,234 @@
+"""gRPC client implementation for Communication service."""
+
+import asyncio
+from collections.abc import AsyncGenerator, Awaitable, Callable
+
+from agentic_mesh_protocol.module.v1 import (
+    information_pb2,
+    lifecycle_pb2,
+    module_service_pb2_grpc,
+)
+from google.protobuf import json_format, struct_pb2
+
+from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
+from digitalkin.logger import logger
+from digitalkin.models.grpc_servers.models import ClientConfig
+from digitalkin.services.base_strategy import BaseStrategy
+from digitalkin.services.communication.communication_strategy import CommunicationStrategy
+
+
+class GrpcCommunication(CommunicationStrategy, GrpcClientWrapper):
+    """gRPC client for module-to-module communication.
+
+    This class provides methods to communicate with remote modules
+    using the Module Service gRPC protocol.
+    """
+
+    def __init__(
+        self,
+        mission_id: str,
+        setup_id: str,
+        setup_version_id: str,
+        client_config: ClientConfig,
+    ) -> None:
+        """Initialize the gRPC communication client.
+
+        Args:
+            mission_id: Mission identifier
+            setup_id: Setup identifier
+            setup_version_id: Setup version identifier
+            client_config: Client configuration for gRPC connection
+        """
+        BaseStrategy.__init__(self, mission_id, setup_id, setup_version_id)
+        self.client_config = client_config
+
+        logger.debug(
+            "Initialized GrpcCommunication",
+            extra={"security": client_config.security},
+        )
+
+    def _create_stub(self, module_address: str, module_port: int) -> module_service_pb2_grpc.ModuleServiceStub:
+        """Create a new stub for the target module.
+
+        Args:
+            module_address: Module host address
+            module_port: Module port
+
+        Returns:
+            ModuleServiceStub for the target module
+        """
+        logger.debug(
+            "Creating connection",
+            extra={"address": module_address, "port": module_port},
+        )
+
+        config = ClientConfig(
+            host=module_address,
+            port=module_port,
+            mode=self.client_config.mode,
+            security=self.client_config.security,
+            credentials=self.client_config.credentials,
+            channel_options=self.client_config.channel_options,
+        )
+
+        channel = self._init_channel(config)
+        return module_service_pb2_grpc.ModuleServiceStub(channel)
+
+    async def get_module_schemas(
+        self,
+        module_address: str,
+        module_port: int,
+        *,
+        llm_format: bool = False,
+    ) -> dict[str, dict]:
+        """Get module schemas via gRPC.
+
+        Args:
+            module_address: Target module address
+            module_port: Target module port
+            llm_format: Return LLM-friendly format
+
+        Returns:
+            Dictionary containing schemas
+        """
+        stub = self._create_stub(module_address, module_port)
+
+        # Create requests
+        input_request = information_pb2.GetModuleInputRequest(llm_format=llm_format)
+        output_request = information_pb2.GetModuleOutputRequest(llm_format=llm_format)
+        setup_request = information_pb2.GetModuleSetupRequest(llm_format=llm_format)
+        secret_request = information_pb2.GetModuleSecretRequest(llm_format=llm_format)
+
+        # Get all schemas in parallel
+        try:
+            input_response, output_response, setup_response, secret_response = await asyncio.gather(
+                asyncio.to_thread(stub.GetModuleInput, input_request),
+                asyncio.to_thread(stub.GetModuleOutput, output_request),
+                asyncio.to_thread(stub.GetModuleSetup, setup_request),
+                asyncio.to_thread(stub.GetModuleSecret, secret_request),
+            )
+
+            logger.debug(
+                "Retrieved module schemas",
+                extra={
+                    "module_address": module_address,
+                    "module_port": module_port,
+                    "llm_format": llm_format,
+                },
+            )
+
+            return {
+                "input": json_format.MessageToDict(input_response.input_schema),
+                "output": json_format.MessageToDict(output_response.output_schema),
+                "setup": json_format.MessageToDict(setup_response.setup_schema),
+                "secret": json_format.MessageToDict(secret_response.secret_schema),
+            }
+        except Exception:
+            logger.exception(
+                "Failed to get module schemas",
+                extra={
+                    "module_address": module_address,
+                    "module_port": module_port,
+                },
+            )
+            raise
+
+    async def call_module(
+        self,
+        module_address: str,
+        module_port: int,
+        input_data: dict,
+        setup_id: str,
+        mission_id: str,
+        callback: Callable[[dict], Awaitable[None]] | None = None,
+    ) -> AsyncGenerator[dict, None]:
+        """Call a module and stream responses via gRPC.
+
+        Args:
+            module_address: Target module address
+            module_port: Target module port
+            input_data: Input data as dictionary
+            setup_id: Setup configuration ID
+            mission_id: Mission context ID
+            callback: Optional callback for each response
+
+        Yields:
+            Streaming responses from module as dictionaries
+        """
+        stub = self._create_stub(module_address, module_port)
+
+        # Convert input data to protobuf Struct
+        input_struct = struct_pb2.Struct()
+        input_struct.update(input_data)
+
+        # Create request
+        request = lifecycle_pb2.StartModuleRequest(
+            input=input_struct,
+            setup_id=setup_id,
+            mission_id=mission_id,
+        )
+
+        logger.debug(
+            "Calling module",
+            extra={
+                "module_address": module_address,
+                "module_port": module_port,
+                "setup_id": setup_id,
+                "mission_id": mission_id,
+            },
+        )
+
+        try:
+            # Call StartModule with streaming response
+            response_stream = stub.StartModule(request)
+
+            # Stream responses
+            for response in response_stream:
+                # Convert protobuf Struct to dict
+                output_dict = json_format.MessageToDict(response.output)
+
+                # Check for end_of_stream signal
+                if output_dict.get("root", {}).get("protocol") == "end_of_stream":
+                    logger.debug(
+                        "End of stream received",
+                        extra={
+                            "module_address": module_address,
+                            "module_port": module_port,
+                        },
+                    )
+                    break
+
+                # Add job_id and success flag
+                response_dict = {
+                    "success": response.success,
+                    "job_id": response.job_id,
+                    "output": output_dict,
+                }
+
+                logger.debug(
+                    "Received module response",
+                    extra={
+                        "module_address": module_address,
+                        "module_port": module_port,
+                        "success": response.success,
+                        "job_id": response.job_id,
+                    },
+                )
+
+                # Call callback if provided
+                if callback:
+                    await callback(response_dict)
+
+                yield response_dict
+
+        except Exception:
+            logger.exception(
+                "Failed to call module",
+                extra={
+                    "module_address": module_address,
+                    "module_port": module_port,
+                    "setup_id": setup_id,
+                    "mission_id": mission_id,
+                },
+            )
+            raise

digitalkin/services/cost/grpc_cost.py

@@ -2,7 +2,7 @@
 
 from typing import Literal
 
-from digitalkin_proto.agentic_mesh_protocol.cost.v1 import cost_pb2, cost_service_pb2_grpc
+from agentic_mesh_protocol.cost.v1 import cost_pb2, cost_service_pb2_grpc
 from google.protobuf import json_format
 
 from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper

digitalkin/services/filesystem/grpc_filesystem.py

@@ -2,7 +2,7 @@
 
 from typing import Any, Literal
 
-from digitalkin_proto.agentic_mesh_protocol.filesystem.v1 import filesystem_pb2, filesystem_service_pb2_grpc
+from agentic_mesh_protocol.filesystem.v1 import filesystem_pb2, filesystem_service_pb2_grpc
 from google.protobuf import struct_pb2
 from google.protobuf.json_format import MessageToDict
 

digitalkin/services/registry/__init__.py

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