digitalkin 0.2.18__py3-none-any.whl → 0.2.20__py3-none-any.whl

digitalkin/modules/job_manager/single_job_manager.py

@@ -10,13 +10,14 @@ import grpc
 
 from digitalkin.logger import logger
 from digitalkin.models import ModuleStatus
-from digitalkin.models.module import ConfigSetupModelT, InputModelT, OutputModelT, SetupModelT
+from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
+from digitalkin.modules.job_manager.job_manager_models import StreamCodeModel
 from digitalkin.services.services_models import ServicesMode
 
 
-class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigSetupModelT]):
+class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
     """Manages a single instance of a module job.
 
     This class ensures that only one instance of a module job is active at a time.
@@ -68,9 +69,9 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
 
     async def create_config_setup_instance_job(
         self,
-        config_setup_data: ConfigSetupModelT,
-        setup_data: SetupModelT,
+        config_setup_data: SetupModelT,
         mission_id: str,
+        setup_id: str,
         setup_version_id: str,
     ) -> str:
         """Create and start a new module setup configuration job.
@@ -82,6 +83,7 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
             config_setup_data: The input data required to start the job.
             setup_data: The setup configuration for the module.
             mission_id: The mission ID associated with the job.
+            setup_id: The setup ID associated with the module.
             setup_version_id: The setup ID.
 
         Returns:
@@ -92,14 +94,13 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
         """
         job_id = str(uuid.uuid4())
         # TODO: Ensure the job_id is unique.
-        module = self.module_class(job_id, mission_id=mission_id, setup_version_id=setup_version_id)
+        module = self.module_class(job_id, mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id)
         self.modules[job_id] = module
         self.queues[job_id] = asyncio.Queue()
 
         try:
             await module.start_config_setup(
                 config_setup_data,
-                setup_data,
                 await self.job_specific_callback(self.add_to_queue, job_id),
             )
             logger.debug("Module %s (%s) started successfully", job_id, module.name)
@@ -166,7 +167,6 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
                 ):
                     logger.info(f"{job_id=}: {module.status=}")
                     yield await self.queues[job_id].get()
-                logger.info(f"{job_id=}: {module.status=} | {self.queues[job_id].empty()}")
 
             finally:
                 del self.queues[job_id]
@@ -178,6 +178,7 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
         input_data: InputModelT,
         setup_data: SetupModelT,
         mission_id: str,
+        setup_id: str,
         setup_version_id: str,
     ) -> str:
         """Create and start a new module job.
@@ -189,7 +190,8 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
             input_data: The input data required to start the job.
             setup_data: The setup configuration for the module.
             mission_id: The mission ID associated with the job.
-            setup_version_id: The setup ID associated with the module.
+            setup_id: The setup ID associated with the module.
+            setup_version_id: The setup Version ID associated with the module.
 
         Returns:
             str: The unique identifier (job ID) of the created job.
@@ -199,15 +201,22 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
         """
         job_id = str(uuid.uuid4())
         # TODO: Ensure the job_id is unique.
-        module = self.module_class(job_id, mission_id=mission_id, setup_version_id=setup_version_id)
+        module = self.module_class(
+            job_id,
+            mission_id=mission_id,
+            setup_id=setup_id,
+            setup_version_id=setup_version_id,
+        )
         self.modules[job_id] = module
         self.queues[job_id] = asyncio.Queue()
+        callback = await self.job_specific_callback(self.add_to_queue, job_id)
 
         try:
             await module.start(
                 input_data,
                 setup_data,
-                await self.job_specific_callback(self.add_to_queue, job_id),
+                callback,
+                done_callback=lambda _: asyncio.create_task(callback(StreamCodeModel(code="__END_OF_STREAM__"))),
             )
             logger.debug("Module %s (%s) started successfully", job_id, module.name)
         except Exception:

digitalkin/modules/job_manager/taskiq_broker.py

@@ -132,6 +132,7 @@ async def send_message_to_stream(job_id: str, output_data: OutputModelT) -> None
 @TASKIQ_BROKER.task
 async def run_start_module(
     mission_id: str,
+    setup_id: str,
     setup_version_id: str,
     module_class: type[BaseModule],
     services_mode: ServicesMode,
@@ -143,6 +144,7 @@ async def run_start_module(
 
     Args:
         mission_id: str,
+        setup_id: The setup ID associated with the module.
         setup_version_id: The setup ID associated with the module.
         module_class: type[BaseModule],
         services_mode: ServicesMode,
@@ -161,7 +163,7 @@ async def run_start_module(
 
     job_id = context.message.task_id
     callback = await BaseJobManager.job_specific_callback(send_message_to_stream, job_id)
-    module = module_class(job_id, mission_id=mission_id, setup_version_id=setup_version_id)
+    module = module_class(job_id, mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id)
 
     await module.start(
         input_data,
@@ -176,17 +178,18 @@ async def run_start_module(
 @TASKIQ_BROKER.task
 async def run_config_module(
     mission_id: str,
+    setup_id: str,
     setup_version_id: str,
     module_class: type[BaseModule],
     services_mode: ServicesMode,
     config_setup_data: dict,
-    setup_data: dict,
     context: Context = TaskiqDepends(),
 ) -> None:
     """TaskIQ task allowing a module to compute in the background asynchronously.
 
     Args:
         mission_id: str,
+        setup_id: The setup ID associated with the module.
         setup_version_id: The setup ID associated with the module.
         module_class: type[BaseModule],
         services_mode: ServicesMode,
@@ -205,10 +208,9 @@ async def run_config_module(
 
     job_id = context.message.task_id
     callback = await BaseJobManager.job_specific_callback(send_message_to_stream, job_id)
-    module = module_class(job_id, mission_id=mission_id, setup_version_id=setup_version_id)
+    module = module_class(job_id, mission_id=mission_id, setup_id=setup_id, setup_version_id=setup_version_id)
 
     await module.start_config_setup(
         module_class.create_config_setup_model(config_setup_data),
-        module_class.create_setup_model(setup_data),
         callback,
     )

digitalkin/modules/job_manager/taskiq_job_manager.py

@@ -18,7 +18,7 @@ from typing import TYPE_CHECKING, Any, Generic
 from rstream import Consumer, ConsumerOffsetSpecification, MessageContext, OffsetType
 
 from digitalkin.logger import logger
-from digitalkin.models.module import ConfigSetupModelT, InputModelT, SetupModelT
+from digitalkin.models.module import InputModelT, SetupModelT
 from digitalkin.models.module.module import ModuleStatus
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
@@ -29,7 +29,7 @@ if TYPE_CHECKING:
     from taskiq.task import AsyncTaskiqTask
 
 
-class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigSetupModelT]):
+class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
     """Taskiq job manager for running modules in Taskiq tasks."""
 
     services_mode: ServicesMode
@@ -134,9 +134,9 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
 
     async def create_config_setup_instance_job(
         self,
-        config_setup_data: ConfigSetupModelT,
-        setup_data: SetupModelT,
+        config_setup_data: SetupModelT,
         mission_id: str,
+        setup_id: str,
         setup_version_id: str,
     ) -> str:
         """Create and start a new module setup configuration job.
@@ -148,6 +148,7 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
             config_setup_data: The input data required to start the job.
             setup_data: The setup configuration for the module.
             mission_id: The mission ID associated with the job.
+            setup_id: The setup ID associated with the module.
             setup_version_id: The setup ID.
 
         Returns:
@@ -169,11 +170,11 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
 
         running_task: AsyncTaskiqTask[Any] = await task.kiq(
             mission_id,
+            setup_id,
             setup_version_id,
             self.module_class,
             self.services_mode,
             config_setup_data.model_dump(),  # type: ignore
-            setup_data.model_dump(),
         )
 
         job_id = running_task.task_id
@@ -223,6 +224,7 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
         input_data: InputModelT,
         setup_data: SetupModelT,
         mission_id: str,
+        setup_id: str,
         setup_version_id: str,
     ) -> str:
         """Launches the module_task in Taskiq, returns the Taskiq task id as job_id.
@@ -231,6 +233,7 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
             input_data: Input data for the module
             setup_data: Setup data for the module
             mission_id: Mission ID for the module
+            setup_id: The setup ID associated with the module.
             setup_version_id: The setup ID associated with the module.
 
         Returns:
@@ -247,6 +250,7 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
 
         running_task: AsyncTaskiqTask[Any] = await task.kiq(
             mission_id,
+            setup_id,
             setup_version_id,
             self.module_class,
             self.services_mode,

digitalkin/modules/tool_module.py

@@ -3,9 +3,8 @@
 from abc import ABC
 
 from digitalkin.models.module import InputModelT, OutputModelT, SecretModelT, SetupModelT
-from digitalkin.models.module.module_types import ConfigSetupModelT
 from digitalkin.modules._base_module import BaseModule # type: ignore
 
 
-class ToolModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT,ConfigSetupModelT,], ABC):
+class ToolModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT,], ABC):
     """ToolModule extends BaseModule to implement specific module types."""

digitalkin/services/base_strategy.py

@@ -9,12 +9,14 @@ class BaseStrategy(ABC):
     This class defines the interface for all strategies.
     """
 
-    def __init__(self, mission_id: str, setup_version_id: str) -> None:
+    def __init__(self, mission_id: str, setup_id: str, setup_version_id: str) -> None:
         """Initialize the strategy.
 
         Args:
             mission_id: The ID of the mission this strategy is associated with
+            setup_id: The ID of the setup 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_id: str = setup_id
         self.setup_version_id: str = setup_version_id

digitalkin/services/cost/cost_strategy.py

@@ -57,15 +57,22 @@ class CostServiceError(Exception):
 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:
+    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_version_id)
+        super().__init__(mission_id, setup_id, setup_version_id)
         self.config = config
 
     @abstractmethod

digitalkin/services/cost/default_cost.py

@@ -15,15 +15,16 @@ from digitalkin.services.cost.cost_strategy import (
 class DefaultCost(CostStrategy):
     """Default cost strategy."""
 
-    def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, CostConfig]) -> None:
+    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_version_id=setup_version_id, config=config)
+        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(

digitalkin/services/cost/grpc_cost.py

@@ -57,12 +57,13 @@ class GrpcCost(CostStrategy, GrpcClientWrapper):
     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_version_id=setup_version_id, config=config)
+        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 succesfully")

digitalkin/services/filesystem/default_filesystem.py

@@ -25,14 +25,15 @@ class DefaultFilesystem(FilesystemStrategy):
     Files are stored in a temporary directory with proper metadata tracking.
     """
 
-    def __init__(self, mission_id: str, setup_version_id: str) -> None:
+    def __init__(self, mission_id: str, setup_id: str, setup_version_id: str) -> None:
         """Initialize the default filesystem 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
         """
-        super().__init__(mission_id, setup_version_id)
+        super().__init__(mission_id, setup_id, setup_version_id)
         self.temp_root: str = tempfile.mkdtemp()
         os.makedirs(self.temp_root, exist_ok=True)
         self.db: dict[str, FilesystemRecord] = {}
@@ -118,7 +119,7 @@ class DefaultFilesystem(FilesystemStrategy):
         for file in files:
             try:
                 # Check if file with same name exists in the context
-                context_dir = self._get_context_temp_dir(self.mission_id)
+                context_dir = self._get_context_temp_dir(self.setup_id)
                 file_path = os.path.join(context_dir, file.name)
                 if os.path.exists(file_path) and not file.replace_if_exists:
                     msg = f"File with name {file.name} already exists."
@@ -129,7 +130,7 @@ class DefaultFilesystem(FilesystemStrategy):
                 storage_uri = str(Path(file_path).resolve())
                 file_data = FilesystemRecord(
                     id=str(uuid.uuid4()),
-                    context=self.mission_id,
+                    context=self.setup_id,
                     name=file.name,
                     file_type=file.file_type,
                     content_type=file.content_type or "application/octet-stream",
@@ -138,14 +139,13 @@ class DefaultFilesystem(FilesystemStrategy):
                     metadata=file.metadata,
                     storage_uri=storage_uri,
                     file_url=storage_uri,
-                    status=file.status if hasattr(file, "status") and file.status else "ACTIVE",
+                    status="ACTIVE",
                 )
 
                 self.db[file_data.id] = file_data
                 uploaded_files.append(file_data)
                 total_uploaded += 1
                 logger.debug("Uploaded file %s", file_data)
-
             except Exception as e:  # noqa: PERF203
                 logger.exception("Error uploading file %s: %s", file.name, e)
                 total_failed += 1
@@ -199,6 +199,8 @@ class DefaultFilesystem(FilesystemStrategy):
             end_idx = start_idx + list_size
             paginated_files = filtered_files[start_idx:end_idx]
 
+            logger.critical(f"{filters=} | {paginated_files=}")
+
             if include_content:
                 for file in paginated_files:
                     file.content = Path(file.storage_uri).read_bytes()
@@ -306,7 +308,7 @@ class DefaultFilesystem(FilesystemStrategy):
             raise FilesystemServiceError(msg)
 
         try:
-            context_dir = self._get_context_temp_dir(self.mission_id)
+            context_dir = self._get_context_temp_dir(self.setup_id)
             file_path = os.path.join(context_dir, file_id)
             existing_file = self.db[file_id]
 

digitalkin/services/filesystem/filesystem_strategy.py

@@ -90,15 +90,22 @@ class FilesystemStrategy(BaseStrategy, ABC):
     filtering, and pagination.
     """
 
-    def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, Any] | None = None) -> None:
+    def __init__(
+        self,
+        mission_id: str,
+        setup_id: str,
+        setup_version_id: str,
+        config: dict[str, Any] | None = None,
+    ) -> 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 for the filesystem strategy
         """
-        super().__init__(mission_id, setup_version_id)
+        super().__init__(mission_id, setup_id, setup_version_id)
         self.config = config
 
     @abstractmethod
@@ -125,7 +132,7 @@ class FilesystemStrategy(BaseStrategy, ABC):
         file_id: str,
         *,
         include_content: bool = False,
-    ) -> tuple[FilesystemRecord, bytes | None]:
+    ) -> FilesystemRecord:
         """Get a specific file by ID or name.
 
         This method fetches detailed information about a single file,

digitalkin/services/filesystem/grpc_filesystem.py

@@ -95,7 +95,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
             filesystem_pb2.FileFilter: The converted FileFilter proto message
         """
         return filesystem_pb2.FileFilter(
-            context=self.mission_id,
+            context=self.setup_id,
             **filters.model_dump(exclude={"file_types", "status"}),
             file_types=[self._file_type_to_enum(file_type) for file_type in filters.file_types]
             if filters.file_types
@@ -114,7 +114,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
         """
         return FilesystemRecord(
             id=file.file_id,
-            context=self.mission_id,
+            context=self.setup_id,
             name=file.name,
             file_type=filesystem_pb2.FileType.Name(file.file_type),
             content_type=file.content_type,
@@ -130,6 +130,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
     def __init__(
         self,
         mission_id: str,
+        setup_id: str,
         setup_version_id: str,
         client_config: ClientConfig,
         config: dict[str, Any] | None = None,
@@ -138,11 +139,12 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
 
         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
             client_config: Configuration for the gRPC client connection
             config: Configuration for the filesystem strategy
         """
-        super().__init__(mission_id, setup_version_id, config)
+        super().__init__(mission_id, setup_id, setup_version_id, config)
         self.service_name = "FilesystemService"
         channel = self._init_channel(client_config)
         self.stub = filesystem_service_pb2_grpc.FilesystemServiceStub(channel)
@@ -170,7 +172,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
                     metadata_struct.update(file.metadata)
                 upload_files.append(
                     filesystem_pb2.UploadFileData(
-                        context=self.mission_id,
+                        context=self.setup_id,
                         name=file.name,
                         file_type=self._file_type_to_enum(file.file_type),
                         content_type=file.content_type or "application/octet-stream",
@@ -206,7 +208,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
         """
         with GrpcFilesystem._handle_grpc_errors("GetFile"):
             request = filesystem_pb2.GetFileRequest(
-                context=self.mission_id,
+                context=self.setup_id,
                 file_id=file_id,
                 include_content=include_content,
             )
@@ -254,7 +256,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
         """
         with GrpcFilesystem._handle_grpc_errors("UpdateFile"):
             request = filesystem_pb2.UpdateFileRequest(
-                context=self.mission_id,
+                context=self.setup_id,
                 file_id=file_id,
                 content=content,
                 file_type=self._file_type_to_enum(file_type) if file_type else None,
@@ -288,7 +290,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
         """
         with GrpcFilesystem._handle_grpc_errors("DeleteFiles"):
             request = filesystem_pb2.DeleteFilesRequest(
-                context=self.mission_id,
+                context=self.setup_id,
                 filters=self._filter_to_proto(filters),
                 permanent=permanent,
                 force=force,
@@ -320,7 +322,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper):
         """
         with GrpcFilesystem._handle_grpc_errors("GetFiles"):
             request = filesystem_pb2.GetFilesRequest(
-                context=self.mission_id,
+                context=self.setup_id,
                 filters=self._filter_to_proto(filters),
                 include_content=include_content,
                 list_size=list_size,

digitalkin/services/services_config.py

@@ -111,12 +111,13 @@ class ServicesConfig(BaseModel):
         """
         return getattr(self, f"_config_{name}", {})
 
-    def init_strategy(self, name: str, mission_id: str, setup_version_id: str) -> ServicesStrategy:
+    def init_strategy(self, name: str, mission_id: str, setup_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_id: The setup ID for the strategy
             setup_version_id: The setup version ID for the strategy
 
         Returns:
@@ -131,7 +132,7 @@ class ServicesConfig(BaseModel):
             raise ValueError(msg)
 
         # 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 {})
+        return strategy_type(mission_id, setup_id, setup_version_id, **self.get_strategy_config(name) or {})
 
     @property
     def storage(self) -> type[StorageStrategy]:

digitalkin/services/storage/default_storage.py

@@ -215,13 +215,14 @@ class DefaultStorage(StorageStrategy):
     def __init__(
         self,
         mission_id: str,
+        setup_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, setup_version_id=setup_version_id, config=config)
+        super().__init__(mission_id=mission_id, setup_id=setup_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

@@ -97,7 +97,7 @@ class GrpcStorage(StorageStrategy, GrpcClientWrapper):
             resp = self.exec_grpc_query("ReadRecord", req)
             return self._build_record_from_proto(resp.stored_data)
         except Exception:
-            logger.exception("gRPC ReadRecord failed for %s:%s", collection, record_id)
+            logger.warning("gRPC ReadRecord failed for %s:%s", collection, record_id)
             return None
 
     def _update(
@@ -149,7 +149,7 @@ class GrpcStorage(StorageStrategy, GrpcClientWrapper):
             )
             self.exec_grpc_query("RemoveRecord", req)
         except Exception:
-            logger.exception(
+            logger.warning(
                 "gRPC RemoveRecord failed for %s:%s",
                 collection,
                 record_id,
@@ -174,7 +174,7 @@ class GrpcStorage(StorageStrategy, GrpcClientWrapper):
             resp = self.exec_grpc_query("ListRecords", req)
             return [self._build_record_from_proto(r) for r in resp.records]
         except Exception:
-            logger.exception("gRPC ListRecords failed for %s", collection)
+            logger.warning("gRPC ListRecords failed for %s", collection)
             return []
 
     def _remove_collection(self, collection: str) -> bool:
@@ -193,20 +193,21 @@ class GrpcStorage(StorageStrategy, GrpcClientWrapper):
             )
             self.exec_grpc_query("RemoveCollection", req)
         except Exception:
-            logger.exception("gRPC RemoveCollection failed for %s", collection)
+            logger.warning("gRPC RemoveCollection failed for %s", collection)
             return False
         return True
 
     def __init__(
         self,
         mission_id: str,
+        setup_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, setup_version_id=setup_version_id, config=config)
+        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 = storage_service_pb2_grpc.StorageServiceStub(channel)

digitalkin/services/storage/storage_strategy.py

@@ -163,15 +163,22 @@ class StorageStrategy(BaseStrategy, ABC):
             True if the deletion was successful, False otherwise
         """
 
-    def __init__(self, mission_id: str, setup_version_id: str, config: dict[str, type[BaseModel]]) -> None:
+    def __init__(
+        self,
+        mission_id: str,
+        setup_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_id: The ID of the setup
             setup_version_id: The ID of the setup version
             config: A dictionary mapping names to Pydantic model classes
         """
-        super().__init__(mission_id, setup_version_id)
+        super().__init__(mission_id, setup_id, setup_version_id)
         # Schema configuration mapping keys to model classes
         self.config: dict[str, type[BaseModel]] = config
 

digitalkin/utils/package_discover.py

@@ -10,7 +10,7 @@ from pathlib import Path
 from typing import ClassVar
 
 from digitalkin.models.module.module_context import ModuleContext
-from digitalkin.models.module.module_types import InputTrigger
+from digitalkin.models.module.module_types import DataTrigger
 from digitalkin.modules.trigger_handler import TriggerHandler
 
 logger = logging.getLogger(__name__)
@@ -329,7 +329,7 @@ class ModuleDiscoverer:
         for protocol, handlers_cls in self._trigger_handlers_cls.items():
             self.trigger_handlers[protocol] = tuple(handler_cls(context) for handler_cls in set(handlers_cls))
 
-    def get_trigger(self, protocol: str, input_instance: InputTrigger) -> TriggerHandler:
+    def get_trigger(self, protocol: str, input_instance: DataTrigger) -> TriggerHandler:
         """Retrieve a trigger handler instance based on the provided protocol and input instance type.
 
         Args: