digitalkin 0.3.1.dev1__py3-none-any.whl → 0.3.2a2__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,87 @@
+"""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 cleanup(self) -> None:
+        """Clean up communication resources.
+
+        This method should release any held resources such as
+        gRPC channels, connection pools, etc.
+        """
+        raise NotImplementedError
+
+    @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/cost).
+
+        Args:
+            module_address: Target module address
+            module_port: Target module port
+            llm_format: Return LLM-friendly format (simplified schema).
+                Note: cost always returns actual data regardless of this flag.
+
+        Returns:
+            Dictionary containing schemas:
+            {
+                "input": {...},
+                "output": {...},
+                "setup": {...},
+                "secret": {...},
+                "cost": {...}
+            }
+        """
+        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,104 @@
+"""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
+
+    async def cleanup(self) -> None:
+        """No-op for local communication."""

digitalkin/services/communication/grpc_communication.py

@@ -0,0 +1,264 @@
+"""gRPC client implementation for Communication service."""
+
+import asyncio
+from collections.abc import AsyncGenerator, Awaitable, Callable
+
+import grpc
+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
+        self._channel_pool: dict[tuple[str, int], grpc.Channel] = {}
+
+        logger.debug(
+            "Initialized GrpcCommunication",
+            extra={"security": client_config.security},
+        )
+
+    def _get_or_create_channel(self, module_address: str, module_port: int) -> grpc.Channel:
+        """Get existing channel or create new one for the target module.
+
+        Args:
+            module_address: Module host address
+            module_port: Module port
+
+        Returns:
+            gRPC channel for the target module
+        """
+        key = (module_address, module_port)
+        if key not in self._channel_pool:
+            logger.debug(
+                "Creating new channel",
+                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,
+            )
+            self._channel_pool[key] = self._init_channel(config)
+        return self._channel_pool[key]
+
+    def close_all_channels(self) -> None:
+        """Close all pooled gRPC channels."""
+        for channel in self._channel_pool.values():
+            channel.close()
+        self._channel_pool.clear()
+
+    async def cleanup(self) -> None:
+        """Clean up all gRPC channels."""
+        self.close_all_channels()
+
+    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
+        """
+        channel = self._get_or_create_channel(module_address, module_port)
+        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: input, output, setup, secret, cost
+        """
+        stub = self._create_stub(module_address, module_port)
+
+        # Create requests
+        # Note: cost always uses llm_format=False to get actual config data (rates, units)
+        # No LLM are allowed to set costs
+        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)
+        cost_request = information_pb2.GetModuleCostRequest(llm_format=False)
+
+        # Get all schemas in parallel
+        try:
+            input_response, output_response, setup_response, secret_response, cost_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),
+                asyncio.to_thread(stub.GetModuleCost, cost_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),
+                "cost": json_format.MessageToDict(cost_response.cost_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/cost_strategy.py

@@ -6,6 +6,7 @@ from typing import Literal
 
 from pydantic import BaseModel
 
+from digitalkin.models.services.cost import AmountLimit, QuantityLimit
 from digitalkin.services.base_strategy import BaseStrategy
 
 
@@ -57,23 +58,25 @@ class CostServiceError(Exception):
 class CostStrategy(BaseStrategy, ABC):
     """Abstract base class for cost strategies."""
 
-    def __init__(
-        self,
-        mission_id: str,
-        setup_id: str,
-        setup_version_id: str,
-        config: dict[str, CostConfig],
-    ) -> None:
-        """Initialize the strategy.
+    @abstractmethod
+    def set_limits(self, limits: list[QuantityLimit | AmountLimit]) -> None:
+        """Set cost limits for this session.
+
+        Args:
+            limits: List of CostLimit objects to enforce.
+        """
+
+    @abstractmethod
+    def check_limit(self, cost_config_name: str, quantity: float) -> bool:
+        """Check if adding this cost would exceed any limits.
 
         Args:
-            mission_id: The ID of the mission this strategy is associated with
-            setup_id: The ID of the setup
-            setup_version_id: The ID of the setup version this strategy is associated with
-            config: Configuration dictionary for the strategy
+            cost_config_name: Name of the cost config.
+            quantity: Quantity to add.
+
+        Returns:
+            True if within limits, False if would exceed.
         """
-        super().__init__(mission_id, setup_id, setup_version_id)
-        self.config = config
 
     @abstractmethod
     def add(
@@ -98,3 +101,22 @@ class CostStrategy(BaseStrategy, ABC):
         cost_types: list[Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]] | None = None,
     ) -> list[CostData]:
         """Get filtered costs."""
+
+    @abstractmethod
+    def get_cost_config(self) -> list[CostConfig]:
+        """Get cost configuration for the current setup version.
+
+        Returns:
+            List of CostConfig objects from the database.
+        """
+
+    @abstractmethod
+    def set_cost_config(self, configs: list[CostConfig]) -> bool:
+        """Store cost configuration for the current setup version.
+
+        Args:
+            configs: List of CostConfig objects to store.
+
+        Returns:
+            True if successfully stored.
+        """