digitalkin 0.2.3__py3-none-any.whl → 0.2.4__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.3"
+    __version__ = "0.2.4"

digitalkin/grpc_servers/module_servicer.py

@@ -51,7 +51,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         self.job_manager = JobManager(module_class)
         self.setup = GrpcSetup() if self.job_manager.args.services_mode == ServicesMode.REMOTE else DefaultSetup()
 
-    async def add_to_queue(self, job_id: str, output_data: OutputModelT) -> None:
+    async def add_to_queue(self, job_id: str, output_data: OutputModelT) -> None:  # type: ignore
         """Callback used to add the output data to the queue of messages."""
         logger.info("JOB: %s added an output_data: %s", job_id, output_data)
         await self.queue.put({job_id: output_data})
@@ -76,11 +76,11 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
         logger.info("StartModule called for module: '%s'", self.module_class.__name__)
         # Process the module input
         # TODO: Check failure of input data format
-        input_data = self.module_class.input_format(**dict(request.input.items()))
+        input_data = self.module_class.create_input_model(dict(request.input.items()))
         setup_data_class = self.setup.get_setup(
             setup_dict={
                 "setup_id": request.setup_id,
-                "mission_id": "missions:test_demo",
+                "mission_id": request.mission_id,
             }
         )
 
@@ -88,31 +88,44 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer):
             msg = "No setup data returned."
             raise ServicerError(msg)
         # TODO: Check failure of setup data format
-        setup_data = self.module_class.setup_format(**setup_data_class.current_setup_version.content)
+        setup_data = self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
 
         # setup_id should be use to request a precise setup from the module
         # Create a job for this execution
-        job_id, module = await self.job_manager.create_job(
+        result: tuple[str, BaseModule] = await self.job_manager.create_job(
             input_data,
             setup_data,
+            mission_id=request.mission_id,
             callback=self.add_to_queue,
         )
+        job_id, module = result
 
         while module.status == ModuleStatus.RUNNING or not self.queue.empty():
-            output_data = await self.queue.get()
+            output_data: dict = await self.queue.get()
+            if job_id not in output_data or output_data[job_id] not in self.job_manager.modules:
+                message = f"Job {job_id} not found"
+                logger.warning(message)
+                context.set_code(grpc.StatusCode.NOT_FOUND)
+                context.set_details(message)
+                yield lifecycle_pb2.StartModuleResponse(success=False)
+                return
+
             if output_data[job_id].get("error", None) is not None:
                 context.set_code(output_data[job_id]["error"]["code"])
                 context.set_details(output_data[job_id]["error"]["error_message"])
                 yield lifecycle_pb2.StartModuleResponse(success=False)
                 return
-            else:
-                # TODO: add a check for the job_id / error handling
-                output_proto = {key: str(value) for key, value in output_data[job_id].items()}
-                yield lifecycle_pb2.StartModuleResponse(
-                    success=True,
-                    output=output_proto,
-                    job_id=job_id,
-                )
+
+            output_proto = json_format.ParseDict(
+                output_data[job_id],
+                struct_pb2.Struct(),
+                ignore_unknown_fields=True,
+            )
+            yield lifecycle_pb2.StartModuleResponse(
+                success=True,
+                output=output_proto,
+                job_id=job_id,
+            )
 
     async def StopModule(  # noqa: N802
         self,

digitalkin/modules/_base_module.py

@@ -140,6 +140,54 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         msg = "'%s' class does not define an 'setup_format'."
         raise NotImplementedError(msg)
 
+    @classmethod
+    def create_input_model(cls, input_data: dict[str, Any]) -> InputModelT:
+        """Create the input model from the input data.
+
+        Args:
+            input_data: The input data to create the model from.
+
+        Returns:
+            The input model.
+        """
+        return cls.input_format(**input_data)
+
+    @classmethod
+    def create_setup_model(cls, setup_data: dict[str, Any]) -> SetupModelT:
+        """Create the setup model from the setup data.
+
+        Args:
+            setup_data: The setup data to create the model from.
+
+        Returns:
+            The setup model.
+        """
+        return cls.setup_format(**setup_data)
+
+    @classmethod
+    def create_secret_model(cls, secret_data: dict[str, Any]) -> SecretModelT:
+        """Create the secret model from the secret data.
+
+        Args:
+            secret_data: The secret data to create the model from.
+
+        Returns:
+            The secret model.
+        """
+        return cls.secret_format(**secret_data)
+
+    @classmethod
+    def create_output_model(cls, output_data: dict[str, Any]) -> OutputModelT:
+        """Create the output model from the output data.
+
+        Args:
+            output_data: The output data to create the model from.
+
+        Returns:
+            The output model.
+        """
+        return cls.output_format(**output_data)
+
     @abstractmethod
     async def initialize(self, setup_data: SetupModelT) -> None:
         """Initialize the module."""

digitalkin/modules/job_manager.py

@@ -8,7 +8,7 @@ from typing import Any
 
 from digitalkin.logger import logger
 from digitalkin.models import ModuleStatus
-from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
+from digitalkin.models.module import InputModelT, OutputModelT, SecretModelT, SetupModelT
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
@@ -78,8 +78,9 @@ class JobManager(ArgParser):
         self,
         input_data: InputModelT,
         setup_data: SetupModelT,
+        mission_id: str,
         callback: Callable[[str, OutputModelT], Coroutine[Any, Any, None]],
-    ) -> tuple[str, BaseModule]:
+    ) -> tuple[str, BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT]]:  # type: ignore
         """Start new module job in background (asyncio).
 
         Args:
@@ -91,7 +92,6 @@ class JobManager(ArgParser):
             str: job_id of the module entity
         """
         job_id = str(uuid.uuid4())
-        mission_id = "missions:test_demo"
         """TODO: check uniqueness of the job_id"""
         # Création et démarrage du module
         module = self.module_class(job_id, mission_id=mission_id)

{digitalkin-0.2.3.dist-info → digitalkin-0.2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.2.3
+Version: 0.2.4
 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.6
+Requires-Dist: digitalkin-proto>=0.1.7
 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.3.dist-info → digitalkin-0.2.4.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=HxifwzaIVOH4OARz9iIyKg-4JvPFh73sakEoHGNycsA,190
+digitalkin/__version__.py,sha256=mR_bGYOQA3eO6IjqeeDDhhBptzMqVTzT6k-jTNiddLY,190
 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=23Q6F9mSO9bCu2MmsLqM9WeviIzwfhDiBFrVspy0PCI,9065
-digitalkin/grpc_servers/module_servicer.py,sha256=togOyvogAYkBv353S-9X_pP5yI8MvxupOIBSvq39jA4,11816
+digitalkin/grpc_servers/module_servicer.py,sha256=FMzB2E1dQZck34vGWf9OXNXEIfu_KhTdZ_Kn79XEowc,12306
 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,9 +29,9 @@ 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=f4AEZFyl0meT_V6esnEFyUtuN7ZXvhHgIYa1ZOMe7tM,7651
+digitalkin/modules/_base_module.py,sha256=iVPRlooAGIW4_3lN9loUPe1q6uf62-vSsEE5THvn8eU,9007
 digitalkin/modules/archetype_module.py,sha256=T2Ehj7EpAC2MO9WQbJv39hqRw7rh3exhVZTEL3JPM8U,421
-digitalkin/modules/job_manager.py,sha256=q48ZJwWmOZ4xJKV2QApI6x_Nfgviej-AoIV8BJQ9B-o,6110
+digitalkin/modules/job_manager.py,sha256=QHcrm3F99I9tlNzXefkyhkkwPITf_H1bEb_obMDbUic,6177
 digitalkin/modules/tool_module.py,sha256=86g0M1wHZ1ReIc7AkKfyjnlGN2QYJBGxrEQpKVlyrZI,421
 digitalkin/modules/trigger_module.py,sha256=kVoI4Gdkw7WWUP5T6hSCNqw5FxibTxL6Tpq9KP7gg78,379
 digitalkin/services/__init__.py,sha256=LqGk_5DJy8Bzz62ajIq9jCeYNKQUIgtSCpafZk15FLc,910
@@ -67,12 +67,12 @@ digitalkin/services/storage/grpc_storage.py,sha256=BqWLK9w_03BEjlJCVbdrlbMvh2szG
 digitalkin/services/storage/storage_strategy.py,sha256=vGo4aYkEp_GZV11m7vd-xY_Z3gVa5K0gMTzbj2Au_3o,6600
 digitalkin/utils/__init__.py,sha256=sJnY-ZUgsjMfojAjONC1VN14mhgIDnzyOlGkw21rRnM,28
 digitalkin/utils/arg_parser.py,sha256=3YyI6oZhhrlTmPTrzlwpQzbCNWDFAT3pggcLxNtJoc0,4388
-digitalkin-0.2.3.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
+digitalkin-0.2.4.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
 modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-modules/minimal_llm_module.py,sha256=JrRPTx7DCYrbmnfkE98xWBU_6x5lr7i8MFP16ATKkfo,5531
-modules/storage_module.py,sha256=HktIy4j8-0w7ugMA7HqcNZymZIGruhchsWKXje9qR4o,6264
+modules/minimal_llm_module.py,sha256=i9KTbv8w1ekzRYWDiRsxD6xLCoZ3uDjD_hWgfY2PRys,5637
+modules/storage_module.py,sha256=bu52lW4RFcWB8VFDhrpBFfCaTSkVL6so3zrkfW4LO9E,6270
 modules/text_transform_module.py,sha256=fAC6r_Ujca1Tz1qdWL4hTPZFn3gFWIVNj5-rytQMObE,7191
-digitalkin-0.2.3.dist-info/METADATA,sha256=ckrSIFHPHRwPsf0W5wLoFue25_OSH4phzfRGiCIPoLk,29095
-digitalkin-0.2.3.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-digitalkin-0.2.3.dist-info/top_level.txt,sha256=5_5e35inSM5YfWNZE21p5wGBojiVtQQML_WzbEk4BRU,31
-digitalkin-0.2.3.dist-info/RECORD,,
+digitalkin-0.2.4.dist-info/METADATA,sha256=h_9hJnMdGTJZwkr0b5IZwSEO5l70sbtJmd5PNZhY0VI,29095
+digitalkin-0.2.4.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+digitalkin-0.2.4.dist-info/top_level.txt,sha256=5_5e35inSM5YfWNZE21p5wGBojiVtQQML_WzbEk4BRU,31
+digitalkin-0.2.4.dist-info/RECORD,,

modules/minimal_llm_module.py

@@ -82,7 +82,11 @@ class OpenAIToolModule(BaseModule[OpenAIToolInput, OpenAIToolOutput, OpenAIToolS
         "storage": {
             "config": {"setups": OpenAIToolSetup},
             "server_config": server_config,
-        }
+        },
+        "filesystem": {
+            "config": {},
+            "server_config": server_config,
+        },
     }
 
     async def initialize(self, setup_data: SetupData) -> None:

modules/storage_module.py

@@ -12,7 +12,7 @@ from digitalkin.models.module import ModuleStatus
 from digitalkin.modules.archetype_module import ArchetypeModule
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
-from digitalkin.services.storage.storage_strategy import DataType, StorageRecord
+from digitalkin.services.storage.storage_strategy import StorageRecord
 
 
 class ExampleInput(BaseModel):
@@ -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}}}
+    services_config_params = {"storage": {"config": {"example_outputs": ExampleOutput}}, "filesystem": {"config": {}}}
 
     def __init__(self, job_id: str, mission_id: str) -> None:
         """Initialize the example module.
@@ -119,7 +119,7 @@ class ExampleModule(ArchetypeModule[ExampleInput, ExampleOutput, ExampleSetup, E
         )
 
         # Store the output data in storage
-        storage_id = self.storage.store("example_outputs", output_data.model_dump(), data_type=DataType.OUTPUT)
+        storage_id = self.storage.store("example_outputs", output_data.model_dump(), data_type="OUTPUT")
 
         logger.info("Stored output data with ID: %s", storage_id)
 
@@ -170,7 +170,7 @@ def test_storage_directly() -> None:
     storage = ServicesConfig().storage(mission_id="test-mission", config={"test_table": ExampleStorage})
 
     # Create a test record
-    storage.store("test_table", {"test_key": "test_value"}, DataType.OUTPUT)
+    storage.store("test_table", {"test_key": "test_value"}, "OUTPUT")
 
     # Retrieve the record
     retrieved = storage.read("test_table")