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

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.
+        """

digitalkin/services/cost/default_cost.py

@@ -3,6 +3,7 @@
 from typing import Literal
 
 from digitalkin.logger import logger
+from digitalkin.models.services.cost import AmountLimit, QuantityLimit
 from digitalkin.services.cost.cost_strategy import (
     CostConfig,
     CostData,
@@ -24,8 +25,46 @@ class DefaultCost(CostStrategy):
             setup_version_id: The ID of the setup version this strategy is associated with
             config: The configuration dictionary for the cost
         """
-        super().__init__(mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id, config=config)
+        super().__init__(mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id)
+        self.config = config
         self.db: dict[str, list[CostData]] = {}
+        self._limits: dict[str, QuantityLimit | AmountLimit] = {}
+        self._accumulated: dict[str, float] = {}
+
+    def set_limits(self, limits: list[QuantityLimit | AmountLimit]) -> None:
+        """Set cost limits for this session.
+
+        Args:
+            limits: List of CostLimit objects to enforce.
+        """
+        self._limits = {limit.name: limit for limit in limits}
+        self._accumulated = {}
+
+    def check_limit(self, cost_config_name: str, quantity: float) -> bool:
+        """Check if adding this cost would exceed any limits.
+
+        Args:
+            cost_config_name: Name of the cost config.
+            quantity: Quantity to add.
+
+        Returns:
+            True if within limits, False if would exceed.
+        """
+        limit = self._limits.get(cost_config_name)
+        if limit is None:
+            return True
+
+        cost_config = self.config.get(cost_config_name)
+        if cost_config is None:
+            return True
+
+        if limit.limit_type == "quantity":
+            current = self._accumulated.get(f"{cost_config_name}_quantity", 0)
+            return current + quantity <= limit.max_value
+
+        current = self._accumulated.get(f"{cost_config_name}_amount", 0)
+        projected = cost_config.rate * quantity
+        return current + projected <= limit.max_value
 
     def add(
         self,
@@ -112,3 +151,24 @@ class DefaultCost(CostStrategy):
             for cost in self.db[self.mission_id]
             if (names and cost.name in names) or (cost_types and cost.cost_type in cost_types)
         ]
+
+    def get_cost_config(self) -> list[CostConfig]:
+        """Get cost configuration from in-memory config.
+
+        Returns:
+            List of CostConfig objects from the config dictionary.
+        """
+        return list(self.config.values())
+
+    def set_cost_config(self, configs: list[CostConfig]) -> bool:
+        """Store cost configuration in memory.
+
+        Args:
+            configs: List of CostConfig objects to store.
+
+        Returns:
+            True if successfully stored.
+        """
+        self.config = {config.cost_name: config for config in configs}
+        logger.debug("Cost configs stored in memory: %s", self.config)
+        return True