digitalkin 0.3.1.dev1__py3-none-any.whl → 0.3.2a2__py3-none-any.whl

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

digitalkin/services/cost/grpc_cost.py

@@ -2,13 +2,14 @@
 
 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
 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.cost import AmountLimit, QuantityLimit
 from digitalkin.services.cost.cost_strategy import (
     CostConfig,
     CostData,
@@ -30,11 +31,49 @@ class GrpcCost(CostStrategy, GrpcClientWrapper, GrpcErrorHandlerMixin):
         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)
+        super().__init__(mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id)
+        self.config = config
+        self._limits: dict[str, QuantityLimit | AmountLimit] = {}
+        self._accumulated: dict[str, float] = {}
         channel = self._init_channel(client_config)
         self.stub = cost_service_pb2_grpc.CostServiceStub(channel)
         logger.debug("Channel client 'Cost' initialized successfully")
 
+    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,
         name: str,
@@ -136,3 +175,60 @@ class GrpcCost(CostStrategy, GrpcClientWrapper, GrpcErrorHandlerMixin):
             ]
             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]
+
+    def get_cost_config(self) -> list[CostConfig]:
+        """Get cost configuration from the database.
+
+        Returns:
+            List of CostConfig objects from the database.
+        """
+        with self.handle_grpc_errors("GetCostConfig", CostServiceError):
+            request = cost_pb2.GetCostConfigRequest(setup_version_id=self.setup_version_id)
+            response: cost_pb2.GetCostConfigResponse = self.exec_grpc_query("GetCostConfig", request)
+            config_list = []
+            for config in response.configs:
+                config_dict = json_format.MessageToDict(
+                    config,
+                    preserving_proto_field_name=True,
+                    always_print_fields_with_no_presence=True,
+                )
+                # Map proto field names to CostConfig field names
+                config_list.append(
+                    CostConfig(
+                        cost_name=config_dict.get("name", ""),
+                        cost_type=config_dict.get("cost_type", "OTHER"),
+                        description=config_dict.get("description"),
+                        unit=config_dict.get("unit", ""),
+                        rate=config_dict.get("rate", 0.0),
+                    )
+                )
+            logger.debug("Cost configs retrieved: %s", config_list)
+            return config_list
+
+    def set_cost_config(self, configs: list[CostConfig]) -> bool:
+        """Store cost configuration in the database.
+
+        Args:
+            configs: List of CostConfig objects to store.
+
+        Returns:
+            True if successfully stored.
+        """
+        with self.handle_grpc_errors("SetCostConfig", CostServiceError):
+            proto_configs = [
+                cost_pb2.CostConfig(
+                    name=config.cost_name,
+                    cost_type=config.cost_type,
+                    description=config.description or "",
+                    unit=config.unit,
+                    rate=config.rate,
+                )
+                for config in configs
+            ]
+            request = cost_pb2.SetCostConfigRequest(
+                setup_version_id=self.setup_version_id,
+                configs=proto_configs,
+            )
+            response: cost_pb2.SetCostConfigResponse = self.exec_grpc_query("SetCostConfig", request)
+            logger.debug("Cost configs stored, success: %s", response.success)
+            return response.success

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
 
@@ -90,12 +90,19 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper, GrpcErrorHandlerMixi
         Returns:
             filesystem_pb2.FileFilter: The converted FileFilter proto message
         """
+        context_id = "unknown"
+        match filters.context:
+            case "setup":
+                context_id = self.setup_id
+            case "mission":
+                context_id = self.mission_id
         return filesystem_pb2.FileFilter(
-            **filters.model_dump(exclude={"file_types", "status"}),
+            **filters.model_dump(exclude={"file_types", "status", "context"}),
             file_types=[self._file_type_to_enum(file_type) for file_type in filters.file_types]
             if filters.file_types
             else None,
             status=self._file_status_to_enum(filters.status) if filters.status else None,
+            context=context_id,
         )
 
     def __init__(

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",
+]

digitalkin/services/registry/default_registry.py

@@ -1,10 +1,162 @@
-"""Default registry."""
+"""Default registry implementation."""
 
+from typing import ClassVar
+
+from digitalkin.models.services.registry import (
+    ModuleInfo,
+    RegistryModuleStatus,
+    RegistryModuleType,
+)
+from digitalkin.services.registry.exceptions import RegistryModuleNotFoundError
+from digitalkin.services.registry.registry_models import ModuleStatusInfo
 from digitalkin.services.registry.registry_strategy import RegistryStrategy
 
 
 class DefaultRegistry(RegistryStrategy):
-    """Default registry strategy."""
+    """Default registry strategy using in-memory storage."""
+
+    _modules: ClassVar[dict[str, ModuleInfo]] = {}
+
+    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.
+        """
+        if module_id not in self._modules:
+            raise RegistryModuleNotFoundError(module_id)
+        return self._modules[module_id]
+
+    def search(
+        self,
+        name: str | None = None,
+        module_type: str | None = None,
+        organization_id: str | None = None,  # noqa: ARG002
+    ) -> list[ModuleInfo]:
+        """Search for modules by criteria.
+
+        Args:
+            name: Filter by name (partial match).
+            module_type: Filter by type (archetype, tool).
+            organization_id: Filter by organization (not used in local storage).
+
+        Returns:
+            List of matching modules.
+        """
+        results = list(self._modules.values())
+
+        if name:
+            results = [m for m in results if name in m.name]
+
+        if module_type:
+            results = [m for m in results if m.module_type == module_type]
+
+        return results
+
+    def get_status(self, module_id: str) -> ModuleStatusInfo:
+        """Get module status.
+
+        Args:
+            module_id: The module identifier.
+
+        Returns:
+            ModuleStatusInfo with current status.
+
+        Raises:
+            RegistryModuleNotFoundError: If module not found.
+        """
+        if module_id not in self._modules:
+            raise RegistryModuleNotFoundError(module_id)
+
+        module = self._modules[module_id]
+        return ModuleStatusInfo(
+            module_id=module_id,
+            status=module.status or RegistryModuleStatus.UNSPECIFIED,
+        )
+
+    def register(
+        self,
+        module_id: str,
+        address: str,
+        port: int,
+        version: str,
+    ) -> ModuleInfo | None:
+        """Register a module with the registry.
+
+        Note: Updates existing module or creates new one in local storage.
+
+        Args:
+            module_id: Unique module identifier.
+            address: Network address.
+            port: Network port.
+            version: Module version.
+
+        Returns:
+            ModuleInfo if successful, None otherwise.
+        """
+        existing = self._modules.get(module_id)
+        self._modules[module_id] = ModuleInfo(
+            module_id=module_id,
+            module_type=existing.module_type if existing else RegistryModuleType.UNSPECIFIED,
+            address=address,
+            port=port,
+            version=version,
+            name=existing.name if existing else module_id,
+            status=RegistryModuleStatus.ACTIVE,
+        )
+        return self._modules[module_id]
+
+    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.
+        """
+        if module_id not in self._modules:
+            raise RegistryModuleNotFoundError(module_id)
+
+        module = self._modules[module_id]
+        # Update status to ACTIVE on heartbeat
+        self._modules[module_id] = ModuleInfo(
+            module_id=module.module_id,
+            module_type=module.module_type,
+            address=module.address,
+            port=module.port,
+            version=module.version,
+            name=module.name,
+            status=RegistryModuleStatus.ACTIVE,
+        )
+        return RegistryModuleStatus.ACTIVE
+
+    async def deregister(self, module_id: str) -> bool:
+        """Deregister a module from the registry.
+
+        Args:
+            module_id: The module identifier to deregister.
+
+        Returns:
+            True if module was removed, False if not found.
+        """
+        if module_id in self._modules:
+            del self._modules[module_id]
+            return True
+        return False
+
+    def get_setup(self, setup_id: str) -> None:
+        """Get setup info (not supported in default registry).
 
-    def get_by_id(self, module_id: str) -> None:
-        """Get services from the registry."""
+        Args:
+            setup_id: The setup identifier.
+        """

digitalkin/services/registry/exceptions.py

@@ -0,0 +1,47 @@
+"""Registry-specific exceptions.
+
+This module contains custom exceptions for registry service operations.
+"""
+
+
+class RegistryServiceError(Exception):
+    """Base exception for registry service errors."""
+
+
+class RegistryModuleNotFoundError(RegistryServiceError):
+    """Raised when a module is not found in the registry."""
+
+    def __init__(self, module_id: str) -> None:
+        """Initialize the exception.
+
+        Args:
+            module_id: The ID of the module that was not found.
+        """
+        self.module_id = module_id
+        super().__init__(f"Module '{module_id}' not found in registry")
+
+
+class ModuleAlreadyExistsError(RegistryServiceError):
+    """Raised when attempting to register an already-registered module."""
+
+    def __init__(self, module_id: str) -> None:
+        """Initialize the exception.
+
+        Args:
+            module_id: The ID of the module that already exists.
+        """
+        self.module_id = module_id
+        super().__init__(f"Module '{module_id}' already registered")
+
+
+class InvalidStatusError(RegistryServiceError):
+    """Raised when an invalid status is provided."""
+
+    def __init__(self, status: int) -> None:
+        """Initialize the exception.
+
+        Args:
+            status: The invalid status value.
+        """
+        self.status = status
+        super().__init__(f"Invalid module status: {status}")