digitalkin 0.3.2.dev2__py3-none-any.whl

digitalkin/services/communication/grpc_communication.py

@@ -0,0 +1,223 @@
+"""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)
+
+                # 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/__init__.py

@@ -0,0 +1,14 @@
+"""This module is responsible for handling the cost services."""
+
+from digitalkin.services.cost.cost_strategy import CostConfig, CostData, CostStrategy, CostType
+from digitalkin.services.cost.default_cost import DefaultCost
+from digitalkin.services.cost.grpc_cost import GrpcCost
+
+__all__ = [
+    "CostConfig",
+    "CostData",
+    "CostStrategy",
+    "CostType",
+    "DefaultCost",
+    "GrpcCost",
+]

digitalkin/services/cost/cost_strategy.py

@@ -0,0 +1,100 @@
+"""This module contains the abstract base class for cost strategies."""
+
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Literal
+
+from pydantic import BaseModel
+
+from digitalkin.services.base_strategy import BaseStrategy
+
+
+class CostType(Enum):
+    """Enum defining the types of costs that can be registered."""
+
+    OTHER = "OTHER"
+    TOKEN_INPUT = "TOKEN_INPUT"
+    TOKEN_OUTPUT = "TOKEN_OUTPUT"
+    API_CALL = "API_CALL"
+    STORAGE = "STORAGE"
+    TIME = "TIME"
+
+
+class CostConfig(BaseModel):
+    """Pydantic model that defines a cost configuration.
+
+    :param cost_name: Name of the cost (unique identifier in the service).
+    :param cost_type: The type/category of the cost.
+    :param description: A short description of the cost.
+    :param unit: The unit of measurement (e.g. token, call, MB).
+    :param rate: The cost per unit (e.g. dollars per token).
+    """
+
+    cost_name: str
+    cost_type: Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]
+    description: str | None = None
+    unit: str
+    rate: float
+
+
+class CostData(BaseModel):
+    """Data model for cost operations."""
+
+    cost: float
+    mission_id: str
+    name: str
+    cost_type: CostType
+    unit: str
+    rate: float
+    setup_version_id: str
+    quantity: float
+
+
+class CostServiceError(Exception):
+    """Custom exception for CostService errors."""
+
+
+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.
+
+        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
+        """
+        super().__init__(mission_id, setup_id, setup_version_id)
+        self.config = config
+
+    @abstractmethod
+    def add(
+        self,
+        name: str,
+        cost_config_name: str,
+        quantity: float,
+    ) -> None:
+        """Register a new cost."""
+
+    @abstractmethod
+    def get(
+        self,
+        name: str,
+    ) -> list[CostData]:
+        """Get a cost."""
+
+    @abstractmethod
+    def get_filtered(
+        self,
+        names: list[str] | None = None,
+        cost_types: list[Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]] | None = None,
+    ) -> list[CostData]:
+        """Get filtered costs."""

digitalkin/services/cost/default_cost.py

@@ -0,0 +1,114 @@
+"""Default cost."""
+
+from typing import Literal
+
+from digitalkin.logger import logger
+from digitalkin.services.cost.cost_strategy import (
+    CostConfig,
+    CostData,
+    CostServiceError,
+    CostStrategy,
+    CostType,
+)
+
+
+class DefaultCost(CostStrategy):
+    """Default cost strategy."""
+
+    def __init__(self, mission_id: str, setup_id: str, setup_version_id: str, config: dict[str, CostConfig]) -> None:
+        """Initialize the strategy.
+
+        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: The configuration dictionary for the cost
+        """
+        super().__init__(mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id, config=config)
+        self.db: dict[str, list[CostData]] = {}
+
+    def add(
+        self,
+        name: str,
+        cost_config_name: str,
+        quantity: float,
+    ) -> None:
+        """Create a new record in the cost database.
+
+        Args:
+            name: The name of the cost
+            cost_config_name: The name of the cost config
+            quantity: The quantity of the cost
+
+        Raises:
+            CostServiceError: If the cost data is invalid or if the cost already exists
+        """
+        cost_config = self.config.get(cost_config_name)
+        if cost_config is None:
+            msg = f"Cost config {cost_config_name} not found in the configuration."
+            logger.error(msg)
+            raise CostServiceError(msg)
+        cost_data = CostData.model_validate({
+            "name": name,
+            "cost": cost_config.rate * quantity,
+            "unit": cost_config.unit,
+            "cost_type": getattr(CostType, cost_config.cost_type),
+            "mission_id": self.mission_id,
+            "rate": cost_config.rate,
+            "quantity": quantity,
+            "setup_version_id": self.setup_version_id,
+        })
+        if cost_data.mission_id not in self.db:
+            self.db[cost_data.mission_id] = []
+        if cost_data.name in [cost.name for cost in self.db[cost_data.mission_id]]:
+            msg = f"Cost with name {cost_data.name} already exists in mission {cost_data.mission_id}"
+            logger.error(msg)
+            raise CostServiceError(msg)
+        self.db[cost_data.mission_id].append(cost_data)
+
+    def get(self, name: str) -> list[CostData]:
+        """Get a record from the database.
+
+        Args:
+            name: The name of the cost
+
+        Returns:
+            list[CostData]: The cost data
+
+        Raises:
+            CostServiceError: If the cost data is invalid or if the cost does not exist
+        """
+        if self.mission_id not in self.db:
+            msg = f"Mission {self.mission_id} not found in the database."
+            logger.warning(msg)
+            raise CostServiceError(msg)
+
+        return [cost for cost in self.db[self.mission_id] if cost.name == name] or []
+
+    def get_filtered(
+        self,
+        names: list[str] | None = None,
+        cost_types: list[Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]] | None = None,
+    ) -> list[CostData]:
+        """Get records from the database.
+
+        Args:
+            names: The names of the costs
+            cost_types: The types of the costs
+
+        Returns:
+            list[CostData]: The list of records
+
+        Raises:
+            CostServiceError: If the cost data is invalid or if the cost does not exist
+        """
+        if self.mission_id not in self.db:
+            msg = f"Mission {self.mission_id} not found in the database."
+            logger.warning(msg)
+            raise CostServiceError(msg)
+
+        return [
+            cost
+            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)
+        ]

digitalkin/services/cost/grpc_cost.py

@@ -0,0 +1,138 @@
+"""This module implements the gRPC Cost strategy."""
+
+from typing import Literal
+
+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
+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.services.cost.cost_strategy import (
+    CostConfig,
+    CostData,
+    CostServiceError,
+    CostStrategy,
+    CostType,
+)
+
+
+class GrpcCost(CostStrategy, GrpcClientWrapper, GrpcErrorHandlerMixin):
+    """This class implements the default Cost strategy."""
+
+    def __init__(
+        self,
+        mission_id: str,
+        setup_id: str,
+        setup_version_id: str,
+        config: dict[str, CostConfig],
+        client_config: ClientConfig,
+    ) -> None:
+        """Initialize the cost."""
+        super().__init__(mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id, config=config)
+        channel = self._init_channel(client_config)
+        self.stub = cost_service_pb2_grpc.CostServiceStub(channel)
+        logger.debug("Channel client 'Cost' initialized successfully")
+
+    def add(
+        self,
+        name: str,
+        cost_config_name: str,
+        quantity: float,
+    ) -> None:
+        """Create a new record in the cost database.
+
+        Args:
+            name: The name of the cost
+            cost_config_name: The name of the cost config
+            quantity: The quantity of the cost
+
+        Raises:
+            CostServiceError: If the cost config is invalid
+        """
+        with self.handle_grpc_errors("AddCost", CostServiceError):
+            cost_config = self.config.get(cost_config_name)
+            if cost_config is None:
+                msg = f"Cost config {cost_config_name} not found in the configuration."
+                logger.error(msg)
+                raise CostServiceError(msg)
+            valid_data = CostData.model_validate({
+                "name": name,
+                "cost": cost_config.rate * quantity,
+                "unit": cost_config.unit,
+                "cost_type": CostType[cost_config.cost_type],
+                "mission_id": self.mission_id,
+                "rate": cost_config.rate,
+                "quantity": quantity,
+                "setup_version_id": self.setup_version_id,
+            })
+            request = cost_pb2.AddCostRequest(
+                cost=valid_data.cost,
+                name=valid_data.name,
+                unit=valid_data.unit,
+                cost_type=valid_data.cost_type.name,
+                mission_id=valid_data.mission_id,
+                rate=valid_data.rate,
+                quantity=valid_data.quantity,
+                setup_version_id=valid_data.setup_version_id,
+            )
+            self.exec_grpc_query("AddCost", request)
+            logger.debug("Cost added with cost_dict: %s", valid_data.model_dump())
+
+    def get(self, name: str) -> list[CostData]:
+        """Get a record from the database.
+
+        Args:
+            name: The name of the cost
+
+        Returns:
+            CostData: The cost data
+        """
+        with self.handle_grpc_errors("GetCost", CostServiceError):
+            request = cost_pb2.GetCostRequest(name=name, mission_id=self.mission_id)
+            response: cost_pb2.GetCostResponse = self.exec_grpc_query("GetCost", request)
+            cost_data_list = [
+                json_format.MessageToDict(
+                    cost,
+                    preserving_proto_field_name=True,
+                    always_print_fields_with_no_presence=True,
+                )
+                for cost in response.costs
+            ]
+            logger.debug("Costs retrieved with cost_dict: %s", cost_data_list)
+            return [CostData.model_validate(cost_data) for cost_data in cost_data_list]
+
+    def get_filtered(
+        self,
+        names: list[str] | None = None,
+        cost_types: list[Literal["TOKEN_INPUT", "TOKEN_OUTPUT", "API_CALL", "STORAGE", "TIME", "OTHER"]] | None = None,
+    ) -> list[CostData]:
+        """Get a list of records from the database.
+
+        Args:
+            names: The names of the costs
+            cost_types: The types of the costs
+
+        Returns:
+            list[CostData]: The cost data
+        """
+        with self.handle_grpc_errors("GetCosts", CostServiceError):
+            request = cost_pb2.GetCostsRequest(
+                mission_id=self.mission_id,
+                filter=cost_pb2.CostFilter(
+                    names=names or [],
+                    cost_types=cost_types or [],
+                ),
+            )
+            response: cost_pb2.GetCostsResponse = self.exec_grpc_query("GetCosts", request)
+            cost_data_list = [
+                json_format.MessageToDict(
+                    cost,
+                    preserving_proto_field_name=True,
+                    always_print_fields_with_no_presence=True,
+                )
+                for cost in response.costs
+            ]
+            logger.debug("Filtered costs retrieved with cost_dict: %s", cost_data_list)
+            return [CostData.model_validate(cost_data) for cost_data in cost_data_list]

digitalkin/services/filesystem/__init__.py

@@ -0,0 +1,7 @@
+"""This module is responsible for handling the filesystem services."""
+
+from digitalkin.services.filesystem.default_filesystem import DefaultFilesystem
+from digitalkin.services.filesystem.filesystem_strategy import FilesystemStrategy
+from digitalkin.services.filesystem.grpc_filesystem import GrpcFilesystem
+
+__all__ = ["DefaultFilesystem", "FilesystemStrategy", "GrpcFilesystem"]