digitalkin 0.2.10__py3-none-any.whl → 0.2.12__py3-none-any.whl

digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.2.10"
+    __version__ = "0.2.12"

digitalkin/grpc_servers/module_servicer.py

@@ -96,6 +96,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
             input_data,
             setup_data,
             mission_id=request.mission_id,
+            setup_version_id=setup_data_class.current_setup_version.id,
             callback=self.add_to_queue,
         )
         job_id, module = result

digitalkin/modules/_base_module.py

@@ -48,17 +48,19 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
     def _init_strategies(self) -> None:
         """Initialize the services configuration."""
         for service_name in self.services_config.valid_strategy_names():
-            service = self.services_config.init_strategy(service_name, self.mission_id)
+            service = self.services_config.init_strategy(service_name, self.mission_id, self.setup_version_id)
             setattr(self, service_name, service)
 
     def __init__(
         self,
         job_id: str,
         mission_id: str,
+        setup_version_id: str,
     ) -> None:
         """Initialize the module."""
         self.job_id: str = job_id
         self.mission_id: str = mission_id
+        self.setup_version_id: str = setup_version_id
         self._status = ModuleStatus.CREATED
         self._task: asyncio.Task | None = None
         # Initialize services configuration

digitalkin/modules/job_manager.py

@@ -79,6 +79,7 @@ class JobManager(ArgParser):
         input_data: InputModelT,
         setup_data: SetupModelT,
         mission_id: str,
+        setup_version_id: str,
         callback: Callable[[str, OutputModelT], Coroutine[Any, Any, None]],
     ) -> tuple[str, BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT]]:  # type: ignore
         """Start new module job in background (asyncio).
@@ -94,7 +95,7 @@ class JobManager(ArgParser):
         job_id = str(uuid.uuid4())
         """TODO: check uniqueness of the job_id"""
         # Création et démarrage du module
-        module = self.module_class(job_id, mission_id=mission_id)
+        module = self.module_class(job_id, mission_id=mission_id, setup_version_id=setup_version_id)
         self.modules[job_id] = module
         try:
             await module.start(input_data, setup_data, await JobManager._job_specific_callback(callback, job_id))

digitalkin/services/base_strategy.py

@@ -9,10 +9,12 @@ class BaseStrategy(ABC):
     This class defines the interface for all strategies.
     """
 
-    def __init__(self, mission_id: str) -> None:
+    def __init__(self, mission_id: str, setup_version_id: str) -> None:
         """Initialize the strategy.
 
         Args:
             mission_id: The ID of the mission this strategy is associated with
+            setup_version_id: The ID of the setup version this strategy is associated with
         """
         self.mission_id: str = mission_id
+        self.setup_version_id: str = setup_version_id

digitalkin/services/cost/cost_strategy.py

@@ -1,8 +1,8 @@
 """This module contains the abstract base class for cost strategies."""
 
 from abc import ABC, abstractmethod
-from enum import Enum, auto
-from typing import Any
+from enum import Enum
+from typing import Literal
 
 from pydantic import BaseModel
 
@@ -10,36 +10,84 @@ from digitalkin.services.base_strategy import BaseStrategy
 
 
 class CostType(Enum):
-    """."""
+    """Enum defining the types of costs that can be registered."""
 
-    OTHER = auto()
-    TOKEN_INPUT = auto()
-    TOKEN_OUTPUT = auto()
-    API_CALL = auto()
-    STORAGE = auto()
-    TIME = auto()
+    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
-    type: CostType
+    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_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_version_id: The ID of the setup version this strategy is associated with
+            config: Configuration dictionary for the strategy
+        """
+        super().__init__(mission_id, setup_version_id)
+        self.config = config
+
     @abstractmethod
-    def add_cost(self, cost_dict: dict[str, Any]) -> str:
+    def add(
+        self,
+        name: str,
+        cost_config_name: str,
+        quantity: float,
+    ) -> None:
         """Register a new cost."""
 
-    def __post_init__(self, *args, **kwargs) -> None:  # noqa: ANN002, ANN003
-        """Allow post init configuration."""
-
     @abstractmethod
-    def get(self, cost_dict: dict[str, Any]) -> list[CostData]:
+    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

@@ -1,9 +1,9 @@
 """Default cost."""
 
 import logging
-from typing import Any
+from typing import Literal
 
-from digitalkin.services.cost.cost_strategy import CostData, CostStrategy
+from digitalkin.services.cost.cost_strategy import CostConfig, CostData, CostServiceError, CostStrategy, CostType
 
 logger = logging.getLogger(__name__)
 
@@ -11,20 +11,99 @@ logger = logging.getLogger(__name__)
 class DefaultCost(CostStrategy):
     """Default cost strategy."""
 
-    def add_cost(self, cost_dict: dict[str, Any]) -> str:  # noqa: PLR6301
+    def __init__(self, mission_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_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_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:
-            str: The ID of the new record
+            list[CostData]: The cost data
+
+        Raises:
+            CostServiceError: If the cost data is invalid or if the cost does not exist
         """
-        logger.info("Cost added with cost_dict: %s", cost_dict)
-        return ""
+        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(self, cost_dict: dict[str, Any]) -> list[CostData]:  # noqa: PLR6301
+    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
         """
-        logger.info("Costs querried with cost_dict: %s", cost_dict)
-        return []
+        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

@@ -1,15 +1,17 @@
 """This module implements the default Cost strategy."""
 
 import logging
-from typing import Any
+from collections.abc import Generator
+from contextlib import contextmanager
+from typing import Any, Literal
 
 from digitalkin_proto.digitalkin.cost.v1 import cost_pb2, cost_service_pb2_grpc
 from google.protobuf import json_format
-from pydantic import ValidationError
 
+from digitalkin.grpc_servers.utils.exceptions import ServerError
 from digitalkin.grpc_servers.utils.grpc_client_wrapper import GrpcClientWrapper
-from digitalkin.grpc_servers.utils.models import ServerConfig
-from digitalkin.services.cost.cost_strategy import CostData, CostStrategy
+from digitalkin.grpc_servers.utils.models import ClientConfig
+from digitalkin.services.cost.cost_strategy import CostConfig, CostData, CostServiceError, CostStrategy, CostType
 
 logger = logging.getLogger(__name__)
 
@@ -17,65 +19,144 @@ logger = logging.getLogger(__name__)
 class GrpcCost(CostStrategy, GrpcClientWrapper):
     """This class implements the default Cost strategy."""
 
-    def _get_costs_by_name(self, cost_dict: dict[str, Any]) -> list[CostData]:
-        request = cost_pb2.GetCostsByNameRequest(
-            mission_id=cost_dict["mission_id"],
-            name=cost_dict["name"],
-        )
-        response = self.exec_grpc_query("GetCostsByName", request)
-        return [CostData(**json_format.MessageToDict(cost)) for cost in response.costs]
-
-    def _get_costs_by_mission(self, cost_dict: dict[str, Any]) -> list[CostData]:
-        request = cost_pb2.GetCostsByMissionRequest(mission_id=cost_dict["mission_id"])
-        response = self.exec_grpc_query("GetCostsByMission", request)
-        return [CostData(**json_format.MessageToDict(cost)) for cost in response.costs]
-
-    def _get_costs_by_type(self, cost_dict: dict[str, Any]) -> list[CostData]:
-        request = cost_pb2.GetCostsByTypeRequest(
-            mission_id=cost_dict["mission_id"],
-            type=cost_dict["type"],
-        )
-        response = self.exec_grpc_query("GetCostsBytype", request)
-        return [CostData(**json_format.MessageToDict(cost)) for cost in response.costs]
-
-    def __init__(self, mission_id: str, config: ServerConfig) -> None:
+    @staticmethod
+    @contextmanager
+    def _handle_grpc_errors(operation: str) -> Generator[Any, Any, Any]:
+        """Context manager for consistent gRPC error handling.
+
+        Yields:
+            Allow error handling in context.
+
+        Args:
+            operation: Description of the operation being performed.
+
+        Raises:
+            ValueError: Error with the model validation.
+            ServerError: from gRPC Client.
+            CostServiceError: Unexpected error.
+        """
+        try:
+            yield
+        except CostServiceError as e:
+            msg = f"CostServiceError in {operation}: {e}"
+            logger.exception(msg)
+            raise CostServiceError(msg) from e
+        except ServerError as e:
+            msg = f"gRPC {operation} failed: {e}"
+            logger.exception(msg)
+            raise ServerError(msg) from e
+        except Exception as e:
+            msg = f"Unexpected error in {operation}"
+            logger.exception(msg)
+            raise CostServiceError(msg) from e
+
+    def __init__(
+        self, mission_id: str, setup_version_id: str, config: dict[str, CostConfig], client_config: ClientConfig
+    ) -> None:
         """Initialize the cost."""
-        super().__init__(mission_id)
-        channel = self._init_channel(config)
+        super().__init__(mission_id=mission_id, setup_version_id=setup_version_id, config=config)
+        channel = self._init_channel(client_config)
         self.stub = cost_service_pb2_grpc.CostServiceStub(channel)
         logger.info("Channel client 'Cost' initialized succesfully")
 
-    def add_cost(self, cost_dict: dict[str, Any]) -> str:
+    def add(
+        self,
+        name: str,
+        cost_config_name: str,
+        quantity: float,
+    ) -> None:
         """Create a new record in the cost database.
 
-        Required arguments:
-            data: Object representation of CostData
+        Args:
+            name: The name of the cost
+            cost_config_name: The name of the cost config
+            quantity: The quantity of the cost
 
-        Returns:
-            str: The ID of the new record
+        Raises:
+            CostServiceError: If the cost config is invalid
         """
-        try:
-            valid_data = CostData.model_validate(cost_dict["data"])
-        except ValidationError:
-            logger.exception("Validation failed for model StorageData")
-            return ""
-        except KeyError:
-            logger.exception("Missing mandatory 'data' in dict.")
-            return ""
-
-        request = cost_pb2.AddCostRequest(**valid_data.model_dump())
-        return self.exec_grpc_query("AddCost", request)
+        with self._handle_grpc_errors("AddCost"):
+            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
 
-    def get(self, cost_dict: dict[str, Any]) -> list[CostData]:
-        """Get records from the database.
+        Returns:
+            CostData: The cost data
+        """
+        with self._handle_grpc_errors("GetCost"):
+            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 list of records
+            list[CostData]: The cost data
         """
-        if "mission_id" not in cost_dict:
-            return []
-        if "name" in cost_dict:
-            return self._get_costs_by_name(cost_dict)
-        if "type" in cost_dict:
-            return self._get_costs_by_type(cost_dict)
-        return self._get_costs_by_mission(cost_dict)
+        with self._handle_grpc_errors("GetCosts"):
+            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/default_filesystem.py

@@ -18,14 +18,15 @@ logger = logging.getLogger(__name__)
 class DefaultFilesystem(FilesystemStrategy):
     """Default state filesystem strategy."""
 
-    def __init__(self, mission_id: str, config: dict[str, str]) -> None:
+    def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, str]) -> None:
         """Initialize the default filesystem strategy.
 
         Args:
             mission_id: The ID of the mission this strategy is associated with
+            setup_version_id: The ID of the setup version this strategy is associated with
             config: A dictionary mapping names to Pydantic model classes
         """
-        super().__init__(mission_id, config)
+        super().__init__(mission_id, setup_version_id, config)
         self.temp_root: str = self.config.get("temp_root", "") or tempfile.gettempdir()
         os.makedirs(self.temp_root, exist_ok=True)
         self.db: dict[str, FilesystemData] = {}

digitalkin/services/filesystem/filesystem_strategy.py

@@ -35,14 +35,15 @@ class FilesystemData(BaseModel):
 class FilesystemStrategy(BaseStrategy, ABC):
     """Abstract base class for filesystem strategies."""
 
-    def __init__(self, mission_id: str, config: dict[str, str]) -> None:
+    def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, str]) -> None:
         """Initialize the strategy.
 
         Args:
             mission_id: The ID of the mission this strategy is associated with
+            setup_version_id: The ID of the setup version this strategy is associated with
             config: configuration dictionary for the filesystem strategy
         """
-        super().__init__(mission_id)
+        super().__init__(mission_id, setup_version_id)
         self.config: dict[str, str] = config
 
     @abstractmethod

digitalkin/services/filesystem/grpc_filesystem.py

@@ -24,26 +24,6 @@ logger = logging.getLogger(__name__)
 class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
     """Default state filesystem strategy."""
 
-    def __init__(
-        self,
-        mission_id: str,
-        config: dict[str, str],
-        client_config: ClientConfig,
-        **kwargs,  # noqa: ANN003, ARG002
-    ) -> None:
-        """Initialize the default filesystem strategy.
-
-        Args:
-            mission_id: The ID of the mission this strategy is associated with
-            config: A dictionary mapping names to Pydantic model classes
-            client_config: The client configuration object
-            kwargs: other optional arguments to pass to the parent class constructor
-        """
-        super().__init__(mission_id, config)
-        channel = self._init_channel(client_config)
-        self.stub = filesystem_service_pb2_grpc.FilesystemServiceStub(channel)
-        logger.info("Channel client 'Filesystem' initialized succesfully")
-
     @staticmethod
     @contextmanager
     def _handle_grpc_errors(operation: str) -> Generator[Any, Any, Any]:
@@ -56,7 +36,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
             operation: Description of the operation being performed.
 
         Raises:
-            ValueError: Error wiht the model validation.
+            ValueError: Error with the model validation.
             ServerError: from gRPC Client.
             FilesystemServiceError: Filesystem service internal.
         """
@@ -71,6 +51,28 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
             logger.exception(msg)
             raise FilesystemServiceError(msg) from e
 
+    def __init__(
+        self,
+        mission_id: str,
+        setup_version_id: str,
+        config: dict[str, str],
+        client_config: ClientConfig,
+        **kwargs,  # noqa: ANN003, ARG002
+    ) -> None:
+        """Initialize the default filesystem strategy.
+
+        Args:
+            mission_id: The ID of the mission this strategy is associated with
+            setup_version_id: The ID of the setup version this strategy is associated with
+            config: A dictionary mapping names to Pydantic model classes
+            client_config: The server configuration object
+            kwargs: other optional arguments to pass to the parent class constructor
+        """
+        super().__init__(mission_id, setup_version_id, config)
+        channel = self._init_channel(client_config)
+        self.stub = filesystem_service_pb2_grpc.FilesystemServiceStub(channel)
+        logger.info("Channel client 'Filesystem' initialized succesfully")
+
     def upload(self, content: bytes, name: str, file_type: FileType) -> FilesystemData:
         """Create a new file in the file system.
 

digitalkin/services/services_config.py

@@ -5,7 +5,7 @@ from typing import Any, ClassVar
 from pydantic import BaseModel, Field, PrivateAttr
 
 from digitalkin.services.agent import AgentStrategy, DefaultAgent
-from digitalkin.services.cost import CostStrategy, DefaultCost
+from digitalkin.services.cost import CostStrategy, DefaultCost, GrpcCost
 from digitalkin.services.filesystem import DefaultFilesystem, FilesystemStrategy, GrpcFilesystem
 from digitalkin.services.identity import DefaultIdentity, IdentityStrategy
 from digitalkin.services.registry import DefaultRegistry, RegistryStrategy
@@ -30,7 +30,7 @@ class ServicesConfig(BaseModel):
     )
     _config_storage: dict[str, Any | None] = PrivateAttr(default_factory=dict)
     _cost: ServicesStrategy[CostStrategy] = PrivateAttr(
-        default_factory=lambda: ServicesStrategy(local=DefaultCost, remote=DefaultCost)
+        default_factory=lambda: ServicesStrategy(local=DefaultCost, remote=GrpcCost)
     )
     _config_cost: dict[str, Any | None] = PrivateAttr(default_factory=dict)
     _snapshot: ServicesStrategy[SnapshotStrategy] = PrivateAttr(
@@ -111,12 +111,13 @@ class ServicesConfig(BaseModel):
         """
         return getattr(self, f"_config_{name}", {})
 
-    def init_strategy(self, name: str, mission_id: str) -> ServicesStrategy:
+    def init_strategy(self, name: str, mission_id: str, setup_version_id: str) -> ServicesStrategy:
         """Initialize a specific strategy.
 
         Args:
             name: The name of the strategy to initialize
             mission_id: The ID of the mission this strategy is associated with
+            setup_version_id: The setup version ID for the strategy
 
         Returns:
             The initialized strategy instance
@@ -129,8 +130,8 @@ class ServicesConfig(BaseModel):
             msg = f"Strategy {name} not found in ServicesConfig."
             raise ValueError(msg)
 
-        # Instantiate the strategy with the mission ID and configuration
-        return strategy_type(mission_id, **self.get_strategy_config(name) or {})
+        # Instantiate the strategy with the mission ID, setup version ID, and configuration
+        return strategy_type(mission_id, setup_version_id, **self.get_strategy_config(name) or {})
 
     @property
     def storage(self) -> type[StorageStrategy]:

digitalkin/services/storage/default_storage.py

@@ -209,12 +209,13 @@ class DefaultStorage(StorageStrategy):
     def __init__(
         self,
         mission_id: str,
+        setup_version_id: str,
         config: dict[str, type[BaseModel]],
         storage_file_path: str = "local_storage",
         **kwargs,  # noqa: ANN003, ARG002
     ) -> None:
         """Initialize the storage."""
-        super().__init__(mission_id=mission_id, config=config)
+        super().__init__(mission_id=mission_id, setup_version_id=setup_version_id, config=config)
         self.storage_file_path = f"{self.mission_id}_{storage_file_path}.json"
         self.storage_file = Path(self.storage_file_path)
         self.storage = self._load_from_file()

digitalkin/services/storage/grpc_storage.py

@@ -37,7 +37,7 @@ class GrpcStorage(StorageStrategy, GrpcClientWrapper):
         dtype = DataType[raw["data_type"]]
         payload = raw.get("data", {})
 
-        validated = self._validate_data(rid, payload)
+        validated = self._validate_data(coll, payload)
         return StorageRecord(
             mission_id=mission,
             collection=coll,
@@ -185,12 +185,13 @@ class GrpcStorage(StorageStrategy, GrpcClientWrapper):
     def __init__(
         self,
         mission_id: str,
+        setup_version_id: str,
         config: dict[str, type[BaseModel]],
         client_config: ClientConfig,
         **kwargs,  # noqa: ANN003, ARG002
     ) -> None:
         """Initialize the storage."""
-        super().__init__(mission_id=mission_id, config=config)
+        super().__init__(mission_id=mission_id, setup_version_id=setup_version_id, config=config)
 
         channel = self._init_channel(client_config)
         self.stub = storage_service_pb2_grpc.StorageServiceStub(channel)

digitalkin/services/storage/storage_strategy.py

@@ -39,11 +39,11 @@ class StorageRecord(BaseModel):
 class StorageStrategy(BaseStrategy, ABC):
     """Define CRUD + list/remove-collection against a collection/record store."""
 
-    def _validate_data(self, record_id: str, data: dict[str, Any]) -> BaseModel:
+    def _validate_data(self, collection: str, data: dict[str, Any]) -> BaseModel:
         """Validate data against the model schema for the given key.
 
         Args:
-            record_id: The unique ID for the record
+            collection: The unique name for the record type
             data: The data to validate
 
         Returns:
@@ -52,15 +52,15 @@ class StorageStrategy(BaseStrategy, ABC):
         Raises:
             ValueError: If the key has no associated model or validation fails
         """
-        model_cls = self.config.get(record_id)
+        model_cls = self.config.get(collection)
         if not model_cls:
-            msg = f"No schema registered for collection '{record_id}'"
+            msg = f"No schema registered for collection '{collection}'"
             raise ValueError(msg)
 
         try:
             return model_cls.model_validate(data)
         except Exception as e:
-            msg = f"Validation failed for '{record_id}': {e!s}"
+            msg = f"Validation failed for '{collection}': {e!s}"
             raise ValueError(msg) from e
 
     def _create_storage_record(
@@ -163,14 +163,15 @@ class StorageStrategy(BaseStrategy, ABC):
             True if the deletion was successful, False otherwise
         """
 
-    def __init__(self, mission_id: str, config: dict[str, type[BaseModel]]) -> None:
+    def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, type[BaseModel]]) -> None:
         """Initialize the storage strategy.
 
         Args:
             mission_id: The ID of the mission this strategy is associated with
+            setup_version_id: The ID of the setup version
             config: A dictionary mapping names to Pydantic model classes
         """
-        super().__init__(mission_id)
+        super().__init__(mission_id, setup_version_id)
         # Schema configuration mapping keys to model classes
         self.config: dict[str, type[BaseModel]] = config
 
@@ -200,7 +201,7 @@ class StorageStrategy(BaseStrategy, ABC):
             raise ValueError(msg)
         record_id = record_id or uuid4().hex
         data_type_enum = DataType[data_type]
-        validated_data = self._validate_data(record_id, {**data, "mission_id": self.mission_id})
+        validated_data = self._validate_data(collection, {**data, "mission_id": self.mission_id})
         record = self._create_storage_record(collection, record_id, validated_data, data_type_enum)
         return self._store(record)
 
@@ -227,7 +228,7 @@ class StorageStrategy(BaseStrategy, ABC):
         Returns:
             StorageRecord: The modified record
         """
-        validated_data = self._validate_data(record_id, data)
+        validated_data = self._validate_data(collection, data)
         return self._update(collection, record_id, validated_data)
 
     def remove(self, collection: str, record_id: str) -> bool:

{digitalkin-0.2.10.dist-info → digitalkin-0.2.12.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.2.10
+Version: 0.2.12
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International
@@ -452,7 +452,7 @@ Classifier: License :: Other/Proprietary License
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: digitalkin-proto>=0.1.8
+Requires-Dist: digitalkin-proto>=0.1.10
 Requires-Dist: grpcio-health-checking>=1.71.0
 Requires-Dist: grpcio-reflection>=1.71.0
 Requires-Dist: grpcio-status>=1.71.0

{digitalkin-0.2.10.dist-info → digitalkin-0.2.12.dist-info}/RECORD

@@ -7,13 +7,13 @@ base_server/mock/__init__.py,sha256=YZFT-F1l_TpvJYuIPX-7kTeE1CfOjhx9YmNRXVoi-jQ,
 base_server/mock/mock_pb2.py,sha256=sETakcS3PAAm4E-hTCV1jIVaQTPEAIoVVHupB8Z_k7Y,1843
 base_server/mock/mock_pb2_grpc.py,sha256=BbOT70H6q3laKgkHfOx1QdfmCS_HxCY4wCOX84YAdG4,3180
 digitalkin/__init__.py,sha256=7LLBAba0th-3SGqcpqFO-lopWdUkVLKzLZiMtB-mW3M,162
-digitalkin/__version__.py,sha256=_lmtRGNN0s_rEpZfYljDoaivZSEvHnwiOHxMyt0abjo,191
+digitalkin/__version__.py,sha256=NhqSkQ1MfEwCjG1RWIYDONpYEZDeyyeAwIaGaE1q-Tk,191
 digitalkin/logger.py,sha256=9cDgyJV2QXXT8F--xRODFlZyDgjuTTXNdpCU3GdqCsk,382
 digitalkin/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 digitalkin/grpc_servers/__init__.py,sha256=0cJBlwipSmFdXkyH3T0i6OJ1WpAtNsZgYX7JaSnkbtg,804
 digitalkin/grpc_servers/_base_server.py,sha256=ec4xmgAuOMVg45a63O_PEa2T7mI4tJ6boxcXauFyZ5g,18649
 digitalkin/grpc_servers/module_server.py,sha256=0jC5MN-H4MGA9O9J87gnlM8EHcbrkb0chuCm8yZh7gI,10232
-digitalkin/grpc_servers/module_servicer.py,sha256=gm1dhmHhBGqY2xlhAp5i15O6waOA1mDHbFJxsT2RUT8,13662
+digitalkin/grpc_servers/module_servicer.py,sha256=mbXbfCXjRZHl8ueaYIedWMqw0OneoBU60mtbHV5yAe4,13734
 digitalkin/grpc_servers/registry_server.py,sha256=PmWaH4Xmg5Sj7NtFVLBNTOzkOfqo7dw_qyVBnaW4jy4,2238
 digitalkin/grpc_servers/registry_servicer.py,sha256=mCAjNhdMq5DozZMEPsJK__DIxePEYxSWV-gAq-Xctk4,16469
 digitalkin/grpc_servers/utils/exceptions.py,sha256=I00OM8p8up20He4dU1fiHsvdLj1DymjR_UmoeUm2MSA,785
@@ -29,26 +29,26 @@ digitalkin/models/services/__init__.py,sha256=HsW7MUGFPvH7Ri28WN4BHHBfEQk5dzU_9F
 digitalkin/models/services/cost.py,sha256=QTEuFD6xz62nob0z4ksE-INJWcZ-iFiuNW5mvXhpFes,1599
 digitalkin/models/services/storage.py,sha256=cYTVIriGKiprF9OerhSxmc_jM6fUTVwmeon1yQCinkE,143
 digitalkin/modules/__init__.py,sha256=ppYARmhvdVi55ofC0QZerIempSlcJYDeCXhcl4qXObw,278
-digitalkin/modules/_base_module.py,sha256=WU2oUW4-FjFtieuyRfXv3R3P4N1xT78zzMWrSTdhmAg,9090
+digitalkin/modules/_base_module.py,sha256=jzsB5d51Sp6-K_e9PxJBVLAv0TeVRo7WhYOKM2SvZ_4,9198
 digitalkin/modules/archetype_module.py,sha256=T2Ehj7EpAC2MO9WQbJv39hqRw7rh3exhVZTEL3JPM8U,421
-digitalkin/modules/job_manager.py,sha256=QHcrm3F99I9tlNzXefkyhkkwPITf_H1bEb_obMDbUic,6177
+digitalkin/modules/job_manager.py,sha256=EB53_HZJZOkOR3MXidN8ZiJte8_kOJZ3hnOFVfm5v64,6243
 digitalkin/modules/tool_module.py,sha256=86g0M1wHZ1ReIc7AkKfyjnlGN2QYJBGxrEQpKVlyrZI,421
 digitalkin/modules/trigger_module.py,sha256=kVoI4Gdkw7WWUP5T6hSCNqw5FxibTxL6Tpq9KP7gg78,379
 digitalkin/services/__init__.py,sha256=LqGk_5DJy8Bzz62ajIq9jCeYNKQUIgtSCpafZk15FLc,910
-digitalkin/services/base_strategy.py,sha256=U7J09LkNvTaj4bkX-gF86iU_FcUeHK6AKp4zyUPx6eI,470
-digitalkin/services/services_config.py,sha256=6mNYtw07ZNZbk2b4e3xKEpYfozD_NUZyUY3iiein_10,7256
+digitalkin/services/base_strategy.py,sha256=QAQnJw1BbqcYMSzwlFyhHP5juBH2WKrZzWxqDr_sDHI,638
+digitalkin/services/services_config.py,sha256=4hc7-rHgSigoS3SuV0V9FReD2Dz7XoMXcD6iMBP2KKM,7391
 digitalkin/services/services_models.py,sha256=5zXkWcfKnXGwQi9sN4OAL3XrgqOcmsTl8ai5Mi4RPsw,1668
 digitalkin/services/agent/__init__.py,sha256=vJc8JN0pdtA8ecypLBeHrwAUIW6H2C8NyW-dk24rTpk,244
 digitalkin/services/agent/agent_strategy.py,sha256=42Q9RciHX6tg3CgDQkbrlIx4h_TX0WIuSpLmCjitVmA,492
 digitalkin/services/agent/default_agent.py,sha256=4N_E_eQxJGOx1KVUUg5jNOje-3ncMxF3ePB-uDuGrJc,345
 digitalkin/services/cost/__init__.py,sha256=Wi9ZB4LSXFsUYgkX-V1UJQkVXYDNDpp8q2dXccR2uRM,303
-digitalkin/services/cost/cost_strategy.py,sha256=KkVLT8X2EY58ul-27Gj6lFx7dsk1SZgg-UunBC0D8Ks,986
-digitalkin/services/cost/default_cost.py,sha256=LaTOBIdGOwh4-9vDlnefm8eKDvSWglJOS7BFLOT_G9Q,811
-digitalkin/services/cost/grpc_cost.py,sha256=k0Df_41dwrowBjtz11eqYQyWCMIjY-wEPEytbo2j4gE,3215
+digitalkin/services/cost/cost_strategy.py,sha256=VhHeqi9WnE1yoDBBVp5qmqwIt5tTZHU6_Z_jld8CVeE,2535
+digitalkin/services/cost/default_cost.py,sha256=Be_OZiFBSD4XeM--Cd1B8J2GYMhdoX5z9xAszExmKDg,3860
+digitalkin/services/cost/grpc_cost.py,sha256=nGJD6vj8QSPYXTNX9-TJfqeWGHH35EsDZaEISKcznBs,6213
 digitalkin/services/filesystem/__init__.py,sha256=BhwMl_BUvM0d65fmglkp0SVwn3RfYiUOKJgIMnOCaGM,381
-digitalkin/services/filesystem/default_filesystem.py,sha256=tTOsBt0MSjexKgWm657y8BLppMYJtmJIVdqzsSRRaYk,7264
-digitalkin/services/filesystem/filesystem_strategy.py,sha256=VH93jkijFtyVA6XaH_rsfdSyeEpkqea54NAvHFLJDjM,2216
-digitalkin/services/filesystem/grpc_filesystem.py,sha256=KifdA_1evYA19rib5ZR5WMIrm9_jzr3PF5Sq3s2uI3Q,8180
+digitalkin/services/filesystem/default_filesystem.py,sha256=PNRmLiBOkvAdo84WNZRNLYFXhSTPcJLPtfvY9vRBGPk,7396
+digitalkin/services/filesystem/filesystem_strategy.py,sha256=g5duep7Trc_NWq9WNtCrnheVxfMUvX6vJyx4W0hwwiY,2348
+digitalkin/services/filesystem/grpc_filesystem.py,sha256=9wZRnS6S6sssyRJEDtGDunietfxo1bJjJB2prRY155w,8320
 digitalkin/services/identity/__init__.py,sha256=InkeyLgFYYwItx8mePA8HpfacOMWZwwuc0G4pWtKq9s,270
 digitalkin/services/identity/default_identity.py,sha256=Y2auZHrGSZTIN5D8HyjLvLcNbYFM1CNUE23x7p5VIGw,386
 digitalkin/services/identity/identity_strategy.py,sha256=skappBbds1_qa0Gr24FGrNX1N0_OYhYT1Lh7dUaAirE,429
@@ -62,18 +62,18 @@ digitalkin/services/snapshot/__init__.py,sha256=Uzlnzo0CYlSpVsdiI37hW7xQk8hu3YA1
 digitalkin/services/snapshot/default_snapshot.py,sha256=Mb8QwWRsHh9I_tN0ln_ZiFa1QCZxOVWmuVLemQOTWpc,1058
 digitalkin/services/snapshot/snapshot_strategy.py,sha256=B1TU3V_k9A-OdqBkdyc41-ihnrW5Btcwd1KyQdHT46A,898
 digitalkin/services/storage/__init__.py,sha256=T-ocYLLphudkQgzvG47jBOm5GQsRFRIGA88y7Ur4akg,341
-digitalkin/services/storage/default_storage.py,sha256=qzLPrND92NR9hB7Ok6BF3Yxot14Efa_CHIvVte6kLsU,7817
-digitalkin/services/storage/grpc_storage.py,sha256=Vp1fGVGdL7Fu0EyWA_wq-1aHElMdP5knzPum22QtBF8,6936
-digitalkin/services/storage/storage_strategy.py,sha256=PPRFWmZTy2HEaWMGwiGzcx5NIhiGWaX3h0_NOb3a4Oo,8621
+digitalkin/services/storage/default_storage.py,sha256=AWcayUtdSI1ns4loOqEYkIFiXKPV_0lclt6n94UgqSU,7883
+digitalkin/services/storage/grpc_storage.py,sha256=Or94uk9KqjhLvz9IeJcuSXrNBOLczgZ1GhftQAMmn14,7003
+digitalkin/services/storage/storage_strategy.py,sha256=vZWPk49AyJs7RwhSfUwDiTG1A7C4Ccd8y-VVPslyf0w,8734
 digitalkin/utils/__init__.py,sha256=sJnY-ZUgsjMfojAjONC1VN14mhgIDnzyOlGkw21rRnM,28
 digitalkin/utils/arg_parser.py,sha256=3YyI6oZhhrlTmPTrzlwpQzbCNWDFAT3pggcLxNtJoc0,4388
 digitalkin/utils/llm_ready_schema.py,sha256=JjMug_lrQllqFoanaC091VgOqwAd-_YzcpqFlS7p778,2375
-digitalkin-0.2.10.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
+digitalkin-0.2.12.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
 modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 modules/minimal_llm_module.py,sha256=W-E3OrRbAsRJ6hvSeTU8pzmacdJC_PbcWfDapRv5A1A,5617
-modules/storage_module.py,sha256=528tfWyRw3q5h0lUlP9Or8E7m3AnnyXezXgwisXd8BI,6399
+modules/storage_module.py,sha256=U9zxEgNORR5SuunWOZvbVivzXOgMGljY-0h791bztVw,6388
 modules/text_transform_module.py,sha256=1KaA7abwxltKKtbmiW1rkkIK3BTYFPegUq54px0LOQs,7277
-digitalkin-0.2.10.dist-info/METADATA,sha256=u46XIpQfozHrIaBpU0t5w6R80y964eJorhuxLEyvt_A,29126
-digitalkin-0.2.10.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-digitalkin-0.2.10.dist-info/top_level.txt,sha256=5_5e35inSM5YfWNZE21p5wGBojiVtQQML_WzbEk4BRU,31
-digitalkin-0.2.10.dist-info/RECORD,,
+digitalkin-0.2.12.dist-info/METADATA,sha256=QYAoblLIWAUNEkl70DJzK9D0hO8EIvCLqMaMhGsu2ZU,29127
+digitalkin-0.2.12.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+digitalkin-0.2.12.dist-info/top_level.txt,sha256=5_5e35inSM5YfWNZE21p5wGBojiVtQQML_WzbEk4BRU,31
+digitalkin-0.2.12.dist-info/RECORD,,

modules/storage_module.py

@@ -62,7 +62,7 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
 
     # Define services_config_params with default values
     services_config_strategies = {}
-    services_config_params = {"storage": {"config": {"example_outputs": ExampleOutput}}, "filesystem": {"config": {}}}
+    services_config_params = {"storage": {"config": {"example": ExampleOutput}}, "filesystem": {"config": {}}}
 
     def __init__(self, job_id: str, mission_id: str) -> None:
         """Initialize the example module.
@@ -172,7 +172,7 @@ async def test_module() -> None:
 def test_storage_directly() -> None:
     """Test the storage service directly."""
     # Initialize storage service
-    storage = ServicesConfig().storage(mission_id="test-mission", config={"test_table": ExampleStorage})
+    storage = ServicesConfig().storage(mission_id="test-mission", config={"example": ExampleStorage})
 
     # Create a test record
     storage.store("example", "test_table", {"test_key": "test_value"}, "OUTPUT")