digitalkin 0.2.18__tar.gz → 0.2.19__tar.gz

{digitalkin-0.2.18 → digitalkin-0.2.19}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.2.18
+Version: 0.2.19
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International
@@ -458,30 +458,30 @@ Requires-Dist: grpcio-reflection>=1.71.0
 Requires-Dist: grpcio-status>=1.71.0
 Requires-Dist: pydantic>=2.11.5
 Provides-Extra: dev
-Requires-Dist: typos>=1.33.1; extra == "dev"
-Requires-Dist: ruff>=0.11.13; extra == "dev"
-Requires-Dist: mypy>=1.16.0; extra == "dev"
-Requires-Dist: pyright>=1.1.401; extra == "dev"
+Requires-Dist: typos>=1.34.0; extra == "dev"
+Requires-Dist: ruff>=0.12.5; extra == "dev"
+Requires-Dist: mypy>=1.17.0; extra == "dev"
+Requires-Dist: pyright>=1.1.403; extra == "dev"
 Requires-Dist: pre-commit>=4.2.0; extra == "dev"
 Requires-Dist: bump2version>=1.0.1; extra == "dev"
 Requires-Dist: build>=1.2.2; extra == "dev"
 Requires-Dist: twine>=6.1.0; extra == "dev"
-Requires-Dist: cryptography>=45.0.4; extra == "dev"
+Requires-Dist: cryptography>=45.0.5; extra == "dev"
 Provides-Extra: examples
 Requires-Dist: openai>=1.75.0; extra == "examples"
 Provides-Extra: tests
-Requires-Dist: freezegun>=1.5.2; extra == "tests"
+Requires-Dist: freezegun>=1.5.3; extra == "tests"
 Requires-Dist: hdrhistogram>=0.10.3; extra == "tests"
 Requires-Dist: grpcio-testing>=1.71.0; extra == "tests"
 Requires-Dist: psutil>=7.0.0; extra == "tests"
 Requires-Dist: pytest>=8.4.0; extra == "tests"
-Requires-Dist: pytest-asyncio>=1.0.0; extra == "tests"
+Requires-Dist: pytest-asyncio>=1.1.0; extra == "tests"
 Requires-Dist: pytest-cov>=6.1.0; extra == "tests"
 Provides-Extra: taskiq
 Requires-Dist: rstream>=0.30.0; extra == "taskiq"
-Requires-Dist: taskiq-aio-pika>=0.4.2; extra == "taskiq"
-Requires-Dist: taskiq-redis>=1.0.9; extra == "taskiq"
-Requires-Dist: taskiq[reload]>=0.11.17; extra == "taskiq"
+Requires-Dist: taskiq-aio-pika>=0.4.3; extra == "taskiq"
+Requires-Dist: taskiq-redis>=1.1.0; extra == "taskiq"
+Requires-Dist: taskiq[reload]>=0.11.18; extra == "taskiq"
 Dynamic: license-file
 
 # DigitalKin Python SDK

{digitalkin-0.2.18 → digitalkin-0.2.19}/pyproject.toml

@@ -12,7 +12,7 @@
 
     keywords = [ "digitalkin", "kin", "agent", "gprc", "sdk" ]
     # Version of the package automatically updated by bump2version (that is why it is separated)
-    version = "0.2.18"
+    version = "0.2.19"
 
     classifiers = [
         "Development Status :: 3 - Alpha",
@@ -34,35 +34,35 @@
         "grpcio-reflection>=1.71.0",
         "grpcio-status>=1.71.0",
         "pydantic>=2.11.5",
-]
+    ]
 
     [project.optional-dependencies]
         dev = [
-            "typos>=1.33.1",
-            "ruff>=0.11.13",
-            "mypy>=1.16.0",
-            "pyright>=1.1.401",
+            "typos>=1.34.0",
+            "ruff>=0.12.5",
+            "mypy>=1.17.0",
+            "pyright>=1.1.403",
             "pre-commit>=4.2.0",
             "bump2version>=1.0.1",
             "build>=1.2.2",
             "twine>=6.1.0",
-            "cryptography>=45.0.4",
+            "cryptography>=45.0.5",
         ]
         examples = [ "openai>=1.75.0" ]
         tests = [
-            "freezegun>=1.5.2",
+            "freezegun>=1.5.3",
             "hdrhistogram>=0.10.3",
             "grpcio-testing>=1.71.0",
             "psutil>=7.0.0",
             "pytest>=8.4.0",
-            "pytest-asyncio>=1.0.0",
+            "pytest-asyncio>=1.1.0",
             "pytest-cov>=6.1.0",
         ]
         taskiq = [
             "rstream>=0.30.0",
-            "taskiq-aio-pika>=0.4.2",
-            "taskiq-redis>=1.0.9",
-            "taskiq[reload]>=0.11.17",
+            "taskiq-aio-pika>=0.4.3",
+            "taskiq-redis>=1.1.0",
+            "taskiq[reload]>=0.11.18",
         ]
 
     [project.urls]

{digitalkin-0.2.18 → digitalkin-0.2.19}/src/digitalkin/__version__.py

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

{digitalkin-0.2.18 → digitalkin-0.2.19}/src/digitalkin/grpc_servers/module_servicer.py

@@ -105,7 +105,8 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         setup_version = request.setup_version
         config_setup_data = self.module_class.create_config_setup_model(json_format.MessageToDict(request.content))
         setup_version_data = self.module_class.create_setup_model(
-            json_format.MessageToDict(request.setup_version.content)
+            json_format.MessageToDict(request.setup_version.content),
+            config_fields=True,
         )
 
         if not setup_version_data:
@@ -119,7 +120,6 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         # create a task to run the module in background
         job_id = await self.job_manager.create_config_setup_instance_job(
             config_setup_data,
-            setup_version_data,
             request.mission_id,
             setup_version.id,
         )
@@ -333,7 +333,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         except NotImplementedError as e:
             logger.warning(e)
             context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-            context.set_details(e)
+            context.set_details(str(e))
             return information_pb2.GetModuleInputResponse()
 
         return information_pb2.GetModuleInputResponse(
@@ -369,7 +369,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         except NotImplementedError as e:
             logger.warning(e)
             context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-            context.set_details(e)
+            context.set_details(str(e))
             return information_pb2.GetModuleOutputResponse()
 
         return information_pb2.GetModuleOutputResponse(
@@ -405,7 +405,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         except NotImplementedError as e:
             logger.warning(e)
             context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-            context.set_details(e)
+            context.set_details(str(e))
             return information_pb2.GetModuleSetupResponse()
 
         return information_pb2.GetModuleSetupResponse(
@@ -441,7 +441,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         except NotImplementedError as e:
             logger.warning(e)
             context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-            context.set_details(e)
+            context.set_details(str(e))
             return information_pb2.GetModuleSecretResponse()
 
         return information_pb2.GetModuleSecretResponse(
@@ -477,7 +477,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         except NotImplementedError as e:
             logger.warning(e)
             context.set_code(grpc.StatusCode.UNIMPLEMENTED)
-            context.set_details(e)
+            context.set_details(str(e))
             return information_pb2.GetConfigSetupModuleResponse()
 
         return information_pb2.GetConfigSetupModuleResponse(

{digitalkin-0.2.18 → digitalkin-0.2.19}/src/digitalkin/models/module/__init__.py

@@ -3,24 +3,24 @@
 from digitalkin.models.module.module import Module, ModuleStatus
 from digitalkin.models.module.module_context import ModuleContext
 from digitalkin.models.module.module_types import (
-    ConfigSetupModelT,
-    InputModel,
+    DataModel,
+    DataTrigger,
     InputModelT,
-    InputTrigger,
     OutputModelT,
     SecretModelT,
+    SetupModel,
     SetupModelT,
 )
 
 __all__ = [
-    "ConfigSetupModelT",
-    "InputModel",
+    "DataModel",
+    "DataTrigger",
     "InputModelT",
-    "InputTrigger",
     "Module",
     "ModuleContext",
     "ModuleStatus",
     "OutputModelT",
     "SecretModelT",
+    "SetupModel",
     "SetupModelT",
 ]

digitalkin-0.2.19/src/digitalkin/models/module/module_types.py

@@ -0,0 +1,105 @@
+"""Types for module models."""
+
+from datetime import datetime, timezone
+from typing import Any, ClassVar, Generic, TypeVar, cast
+
+from pydantic import BaseModel, ConfigDict, Field, create_model
+
+from digitalkin.logger import logger
+
+
+class DataTrigger(BaseModel):
+    """Defines the root input model exposing the protocol.
+
+    The mandatory protocol is important to define the module beahvior following the user or agent input.
+
+    Example:
+        class MyInput(DataModel):
+            root: DataTrigger
+            user_define_data: Any
+
+        # Usage
+        my_input = MyInput(root=DataTrigger(protocol="message"))
+        print(my_input.root.protocol)  # Output: message
+    """
+
+    protocol: ClassVar[str]
+    created_at: str = datetime.now(tz=timezone.utc).isoformat()
+
+
+DataTriggerT = TypeVar("DataTriggerT", bound=DataTrigger)
+
+
+class DataModel(BaseModel, Generic[DataTriggerT]):
+    """Base definition of input model showing mandatory root fields.
+
+    The Model define the Module Input, usually referring to multiple input type defined by an union.
+
+    Example:
+        class ModuleInput(DataModel):
+            root: FileInput | MessageInput
+    """
+
+    root: DataTriggerT
+    annotations: dict[str, str] = Field(
+        default={},
+        title="Annotations",
+        description="Additional metadata or annotations related to the output. ex {'role': 'user'}",
+    )
+
+
+InputModelT = TypeVar("InputModelT", bound=DataModel)
+OutputModelT = TypeVar("OutputModelT", bound=DataModel)
+SecretModelT = TypeVar("SecretModelT", bound=BaseModel)
+SetupModelT = TypeVar("SetupModelT", bound="SetupModel")
+
+
+class SetupModel(BaseModel):
+    """Base definition of setup model showing mandatory root fields.
+
+    Optionally, the setup model can define a config option in json_schema_extra to be used to initialize the Kin.
+
+    Example:
+        class MySetup(SetupModel):
+            name: str = Field()
+            number: int = Field(..., json_schema_extra={"config": True})
+    """
+
+    @classmethod
+    def get_clean_model(cls, *, config_fields: bool, hidden_fields: bool) -> type[SetupModelT]:  # type: ignore
+        """Dynamically builds and returns a new BaseModel subclass.
+
+        containing only those fields where json_schema_extra["config"] == True.
+
+        Returns:
+            Type[BaseModel]: A new BaseModel subclass with the filtered fields.
+
+        Raises:
+            ValueError: If both config_fields and hidden_fields are set to True.
+        """
+        clean_fields: dict[str, Any] = {}
+        for name, field_info in cls.model_fields.items():
+            extra = getattr(field_info, "json_schema_extra", {}) or {}
+            is_config = bool(extra.get("config", False))
+            is_hidden = bool(extra.get("hidden", False))
+
+            # Skip config unless explicitly included
+            if is_config and not config_fields:
+                logger.debug("Skipping '%s' (config-only)", name)
+                continue
+
+            # Skip hidden unless explicitly included
+            if is_hidden and not hidden_fields:
+                logger.debug("Skipping '%s' (hidden-only)", name)
+                continue
+
+            clean_fields[name] = (field_info.annotation, field_info)
+
+        # Dynamically create a model e.g. "SetupModel"
+        m = create_model(
+            f"{cls.__name__}",
+            __base__=BaseModel,
+            __config__=ConfigDict(arbitrary_types_allowed=True),
+            **clean_fields,
+        )
+        return cast("type[SetupModelT]", m)

{digitalkin-0.2.18 → digitalkin-0.2.19}/src/digitalkin/modules/_base_module.py

@@ -9,10 +9,8 @@ from typing import Any, ClassVar, Generic
 
 from pydantic import BaseModel
 
-from digitalkin.grpc_servers.utils.exceptions import OptionalFeatureNotImplementedError
 from digitalkin.logger import logger
 from digitalkin.models.module import (
-    ConfigSetupModelT,
     InputModelT,
     ModuleStatus,
     OutputModelT,
@@ -33,11 +31,11 @@ from digitalkin.utils.llm_ready_schema import llm_ready_schema
 from digitalkin.utils.package_discover import ModuleDiscoverer
 
 
-class ModuleErrorModel(BaseModel):
+class ModuleCodeModel(BaseModel):
     """typed error/code model."""
 
     code: str
-    exception: str
+    message: str
     short_description: str
 
 
@@ -48,7 +46,6 @@ class BaseModule(  # noqa: PLR0904
         OutputModelT,
         SetupModelT,
         SecretModelT,
-        ConfigSetupModelT,
     ],
 ):
     """BaseModule is the abstract base for all modules in the DigitalKin SDK."""
@@ -56,10 +53,10 @@ class BaseModule(  # noqa: PLR0904
     name: str
     description: str
 
-    config_setup_format: type[ConfigSetupModelT]
+    setup_format: type[SetupModelT]
+
     input_format: type[InputModelT]
     output_format: type[OutputModelT]
-    setup_format: type[SetupModelT]
     secret_format: type[SecretModelT]
     metadata: ClassVar[dict[str, Any]]
 
@@ -172,20 +169,22 @@ class BaseModule(  # noqa: PLR0904
     def get_config_setup_format(cls, *, llm_format: bool) -> str:
         """Gets the JSON schema of the config setup format model.
 
+        The config setup format is used only to initialize the module with configuration data.
+        The setup format is used to initialize an run the module with setup data.
+
         Raises:
-            OptionalFeatureNotImplementedError: If the `config_setup_format` is not defined.
+            NotImplementedError: If the `setup_format` is not defined.
 
         Returns:
             The JSON schema of the config setup format as a string.
         """
-        config_setup_format = getattr(cls, "config_setup_format", None)
-
-        if config_setup_format is not None:
+        if cls.setup_format is not None:
+            setup_format = cls.setup_format.get_clean_model(config_fields=True, hidden_fields=False)
             if llm_format:
-                return json.dumps(llm_ready_schema(config_setup_format), indent=2)
-            return json.dumps(config_setup_format.model_json_schema(), indent=2)
+                return json.dumps(llm_ready_schema(setup_format), indent=2)
+            return json.dumps(setup_format.model_json_schema(), indent=2)
         msg = "'%s' class does not define an 'config_setup_format'."
-        raise OptionalFeatureNotImplementedError(msg)
+        raise NotImplementedError(msg)
 
     @classmethod
     def get_setup_format(cls, *, llm_format: bool) -> str:
@@ -198,14 +197,15 @@ class BaseModule(  # noqa: PLR0904
             The JSON schema of the setup format as a string.
         """
         if cls.setup_format is not None:
+            setup_format = cls.setup_format.get_clean_model(config_fields=False, hidden_fields=True)
             if llm_format:
-                return json.dumps(llm_ready_schema(cls.setup_format), indent=2)
-            return json.dumps(cls.setup_format.model_json_schema(), indent=2)
+                return json.dumps(llm_ready_schema(setup_format), indent=2)
+            return json.dumps(setup_format.model_json_schema(), indent=2)
         msg = "'%s' class does not define an 'setup_format'."
         raise NotImplementedError(msg)
 
     @classmethod
-    def create_config_setup_model(cls, config_setup_data: dict[str, Any]) -> ConfigSetupModelT:
+    def create_config_setup_model(cls, config_setup_data: dict[str, Any]) -> SetupModelT:
         """Create the setup model from the setup data.
 
         Args:
@@ -214,7 +214,7 @@ class BaseModule(  # noqa: PLR0904
         Returns:
             The setup model.
         """
-        return cls.config_setup_format(**config_setup_data)
+        return cls.setup_format(**config_setup_data)
 
     @classmethod
     def create_input_model(cls, input_data: dict[str, Any]) -> InputModelT:
@@ -229,16 +229,17 @@ class BaseModule(  # noqa: PLR0904
         return cls.input_format(**input_data)
 
     @classmethod
-    def create_setup_model(cls, setup_data: dict[str, Any]) -> SetupModelT:
+    def create_setup_model(cls, setup_data: dict[str, Any], *, config_fields: bool = False) -> SetupModelT:
         """Create the setup model from the setup data.
 
         Args:
             setup_data: The setup data to create the model from.
+            config_fields: If True, include only fields with json_schema_extra["config"] == True.
 
         Returns:
             The setup model.
         """
-        return cls.setup_format(**setup_data)
+        return cls.setup_format.get_clean_model(config_fields=config_fields, hidden_fields=True)(**setup_data)
 
     @classmethod
     def create_secret_model(cls, secret_data: dict[str, Any]) -> SecretModelT:
@@ -290,20 +291,20 @@ class BaseModule(  # noqa: PLR0904
         """
         return cls.triggers_discoverer.register_trigger(handler_cls)
 
-    @abstractmethod
-    async def run_config_setup(
-        self,
-        config_setup_data: ConfigSetupModelT,
-        setup_data: SetupModelT,
-        callback: Callable,
-    ) -> None:
+    async def run_config_setup(self, config_setup_data: SetupModelT) -> SetupModelT:  # noqa: PLR6301
         """Run config setup the module.
 
-        Raises:
-            OptionalFeatureNotImplementedError: If the config setup feature is not implemented.
+        The config setup is used to initialize the setup with configuration data.
+        This method is typically used to set up the module with necessary configuration before running it,
+        especially for processing data like files.
+        The function needs to save the setup in the storage.
+        The module will be initialize with the setup and not the config setup.
+        This method is optional, the config setup and setup can be the same.
+
+        Returns:
+            The updated setup model after running the config setup.
         """
-        msg = f"'{self}' class does not define an optional 'run_config_setup' attribute."
-        raise OptionalFeatureNotImplementedError(msg)
+        return config_setup_data
 
     @abstractmethod
     async def initialize(self, setup_data: SetupModelT) -> None:
@@ -364,9 +365,9 @@ class BaseModule(  # noqa: PLR0904
             asyncio.CancelledError: If the module is cancelled
         """
         try:
-            logger.warning("Starting module %s", self.name)
+            logger.info("Starting module %s", self.name)
             await self.run(input_data, setup_data, callback)
-            logger.warning("Module %s finished", self.name)
+            logger.info("Module %s finished", self.name)
         except asyncio.CancelledError:
             self._status = ModuleStatus.CANCELLED
             logger.error(f"Module {self.name} cancelled")
@@ -374,7 +375,7 @@ class BaseModule(  # noqa: PLR0904
             self._status = ModuleStatus.FAILED
             logger.exception("Error inside module %s", self.name)
         else:
-            self._status = ModuleStatus.STOPPED
+            self._status = ModuleStatus.STOPPING
         finally:
             await self.stop()
 
@@ -382,22 +383,22 @@ class BaseModule(  # noqa: PLR0904
         self,
         input_data: InputModelT,
         setup_data: SetupModelT,
-        callback: Callable[[OutputModelT | ModuleErrorModel], Coroutine[Any, Any, None]],
+        callback: Callable[[OutputModelT | ModuleCodeModel], Coroutine[Any, Any, None]],
         done_callback: Callable | None = None,
     ) -> None:
         """Start the module."""
         try:
-            logger.info("Inititalize module")
+            logger.debug("Inititalize module")
             await self.initialize(setup_data=setup_data)
         except Exception as e:
             self._status = ModuleStatus.FAILED
             short_description = "Error initializing module"
             logger.exception("%s: %s", short_description, e)
             await callback(
-                ModuleErrorModel(
+                ModuleCodeModel(
                     code=str(self._status),
                     short_description=short_description,
-                    exception=str(e),
+                    message=str(e),
                 )
             )
             if done_callback is not None:
@@ -406,9 +407,9 @@ class BaseModule(  # noqa: PLR0904
             return
 
         try:
-            logger.info("Init the discod input handlers.")
+            logger.debug("Init the discovered input handlers.")
             self.triggers_discoverer.init_handlers(self.context)
-            logger.info("Run lifecycle")
+            logger.debug("Run lifecycle")
             self._status = ModuleStatus.RUNNING
             self._task = asyncio.create_task(
                 self._run_lifecycle(input_data, setup_data, callback),
@@ -422,7 +423,8 @@ class BaseModule(  # noqa: PLR0904
 
     async def stop(self) -> None:
         """Stop the module."""
-        if self._status != ModuleStatus.RUNNING:
+        logger.info("Stopping module %s with status %s", self.name, self._status)
+        if self._status not in {ModuleStatus.RUNNING, ModuleStatus.STOPPING}:
             return
 
         try:
@@ -431,22 +433,29 @@ class BaseModule(  # noqa: PLR0904
                 self._task.cancel()
                 with contextlib.suppress(asyncio.CancelledError):
                     await self._task
+            logger.debug("Module %s stopped", self.name)
             await self.cleanup()
+            self._status = ModuleStatus.STOPPED
+            logger.debug("Module %s cleaned", self.name)
         except Exception:
             self._status = ModuleStatus.FAILED
             logger.exception("Error stopping module")
 
     async def start_config_setup(
         self,
-        config_setup_data: ConfigSetupModelT,
-        setup_data: SetupModelT,
-        callback: Callable[[OutputModelT | ModuleErrorModel], Coroutine[Any, Any, None]],
+        config_setup_data: SetupModelT,
+        callback: Callable[[SetupModelT | ModuleCodeModel], Coroutine[Any, Any, None]],
     ) -> None:
         """Start the module."""
         try:
             logger.info("Run Config Setup lifecycle")
             self._status = ModuleStatus.RUNNING
-            await self.run_config_setup(config_setup_data, setup_data, callback)
+            content = await self.run_config_setup(config_setup_data)
+
+            wrapper = config_setup_data.model_dump()
+            wrapper["content"] = content.model_dump()
+            await callback(self.create_setup_model(wrapper))
+            self._status = ModuleStatus.STOPPING
         except Exception:
             self._status = ModuleStatus.FAILED
             logger.exception("Error during module lifecyle")

{digitalkin-0.2.18 → digitalkin-0.2.19}/src/digitalkin/modules/archetype_module.py

@@ -3,7 +3,6 @@
 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
 
 
@@ -13,7 +12,6 @@ class ArchetypeModule(
         OutputModelT,
         SetupModelT,
         SecretModelT,
-        ConfigSetupModelT,
     ],
     ABC,
 ):

{digitalkin-0.2.18 → digitalkin-0.2.19}/src/digitalkin/modules/job_manager/base_job_manager.py

@@ -7,13 +7,12 @@ from typing import Any, Generic
 
 from digitalkin.models import ModuleStatus
 from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
-from digitalkin.models.module.module_types import ConfigSetupModelT
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_config import ServicesConfig
 from digitalkin.services.services_models import ServicesMode
 
 
-class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT, ConfigSetupModelT]):
+class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT]):
     """Abstract base class for managing background module jobs."""
 
     async def _start(self) -> None:
@@ -120,8 +119,7 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT, ConfigSetupModel
     @abc.abstractmethod
     async def create_config_setup_instance_job(
         self,
-        config_setup_data: ConfigSetupModelT,
-        setup_data: SetupModelT,
+        config_setup_data: SetupModelT,
         mission_id: str,
         setup_version_id: str,
     ) -> str:
@@ -132,7 +130,6 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT, ConfigSetupModel
 
         Args:
             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_version_id: The setup ID.
 

{digitalkin-0.2.18 → digitalkin-0.2.19}/src/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,8 +69,7 @@ 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_version_id: str,
     ) -> str:
@@ -99,7 +99,6 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
         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 +165,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]
@@ -202,12 +200,14 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
         module = self.module_class(job_id, mission_id=mission_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-0.2.18 → digitalkin-0.2.19}/src/digitalkin/modules/job_manager/taskiq_broker.py

@@ -180,7 +180,6 @@ async def run_config_module(
     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.
@@ -209,6 +208,5 @@ async def run_config_module(
 
     await module.start_config_setup(
         module_class.create_config_setup_model(config_setup_data),
-        module_class.create_setup_model(setup_data),
         callback,
     )

{digitalkin-0.2.18 → digitalkin-0.2.19}/src/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,8 +134,7 @@ 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_version_id: str,
     ) -> str:
@@ -173,7 +172,6 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigS
             self.module_class,
             self.services_mode,
             config_setup_data.model_dump(),  # type: ignore
-            setup_data.model_dump(),
         )
 
         job_id = running_task.task_id

{digitalkin-0.2.18 → digitalkin-0.2.19}/src/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-0.2.18 → digitalkin-0.2.19}/src/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,7 +193,7 @@ 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
 

{digitalkin-0.2.18 → digitalkin-0.2.19}/src/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:

{digitalkin-0.2.18 → digitalkin-0.2.19}/src/digitalkin.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.2.18
+Version: 0.2.19
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International
@@ -458,30 +458,30 @@ Requires-Dist: grpcio-reflection>=1.71.0
 Requires-Dist: grpcio-status>=1.71.0
 Requires-Dist: pydantic>=2.11.5
 Provides-Extra: dev
-Requires-Dist: typos>=1.33.1; extra == "dev"
-Requires-Dist: ruff>=0.11.13; extra == "dev"
-Requires-Dist: mypy>=1.16.0; extra == "dev"
-Requires-Dist: pyright>=1.1.401; extra == "dev"
+Requires-Dist: typos>=1.34.0; extra == "dev"
+Requires-Dist: ruff>=0.12.5; extra == "dev"
+Requires-Dist: mypy>=1.17.0; extra == "dev"
+Requires-Dist: pyright>=1.1.403; extra == "dev"
 Requires-Dist: pre-commit>=4.2.0; extra == "dev"
 Requires-Dist: bump2version>=1.0.1; extra == "dev"
 Requires-Dist: build>=1.2.2; extra == "dev"
 Requires-Dist: twine>=6.1.0; extra == "dev"
-Requires-Dist: cryptography>=45.0.4; extra == "dev"
+Requires-Dist: cryptography>=45.0.5; extra == "dev"
 Provides-Extra: examples
 Requires-Dist: openai>=1.75.0; extra == "examples"
 Provides-Extra: tests
-Requires-Dist: freezegun>=1.5.2; extra == "tests"
+Requires-Dist: freezegun>=1.5.3; extra == "tests"
 Requires-Dist: hdrhistogram>=0.10.3; extra == "tests"
 Requires-Dist: grpcio-testing>=1.71.0; extra == "tests"
 Requires-Dist: psutil>=7.0.0; extra == "tests"
 Requires-Dist: pytest>=8.4.0; extra == "tests"
-Requires-Dist: pytest-asyncio>=1.0.0; extra == "tests"
+Requires-Dist: pytest-asyncio>=1.1.0; extra == "tests"
 Requires-Dist: pytest-cov>=6.1.0; extra == "tests"
 Provides-Extra: taskiq
 Requires-Dist: rstream>=0.30.0; extra == "taskiq"
-Requires-Dist: taskiq-aio-pika>=0.4.2; extra == "taskiq"
-Requires-Dist: taskiq-redis>=1.0.9; extra == "taskiq"
-Requires-Dist: taskiq[reload]>=0.11.17; extra == "taskiq"
+Requires-Dist: taskiq-aio-pika>=0.4.3; extra == "taskiq"
+Requires-Dist: taskiq-redis>=1.1.0; extra == "taskiq"
+Requires-Dist: taskiq[reload]>=0.11.18; extra == "taskiq"
 Dynamic: license-file
 
 # DigitalKin Python SDK

{digitalkin-0.2.18 → digitalkin-0.2.19}/src/digitalkin.egg-info/requires.txt

@@ -5,30 +5,30 @@ grpcio-status>=1.71.0
 pydantic>=2.11.5
 
 [dev]
-typos>=1.33.1
-ruff>=0.11.13
-mypy>=1.16.0
-pyright>=1.1.401
+typos>=1.34.0
+ruff>=0.12.5
+mypy>=1.17.0
+pyright>=1.1.403
 pre-commit>=4.2.0
 bump2version>=1.0.1
 build>=1.2.2
 twine>=6.1.0
-cryptography>=45.0.4
+cryptography>=45.0.5
 
 [examples]
 openai>=1.75.0
 
 [taskiq]
 rstream>=0.30.0
-taskiq-aio-pika>=0.4.2
-taskiq-redis>=1.0.9
-taskiq[reload]>=0.11.17
+taskiq-aio-pika>=0.4.3
+taskiq-redis>=1.1.0
+taskiq[reload]>=0.11.18
 
 [tests]
-freezegun>=1.5.2
+freezegun>=1.5.3
 hdrhistogram>=0.10.3
 grpcio-testing>=1.71.0
 psutil>=7.0.0
 pytest>=8.4.0
-pytest-asyncio>=1.0.0
+pytest-asyncio>=1.1.0
 pytest-cov>=6.1.0

digitalkin-0.2.18/src/digitalkin/models/module/module_types.py

@@ -1,43 +0,0 @@
-"""Types for module models."""
-
-from typing import TypeVar
-
-from pydantic import BaseModel
-
-
-class InputTrigger(BaseModel):
-    """Defines the root input model exposing the protocol.
-
-    The mandatory protocol is important to define the module beahvior following the user or agent input.
-
-    Example:
-        class MyInput(InputModel):
-            root: InputTrigger
-            user_define_data: Any
-
-        # Usage
-        my_input = MyInput(root=InputTrigger(protocol="message"))
-        print(my_input.root.protocol)  # Output: message
-    """
-
-    protocol: str
-
-
-class InputModel(BaseModel):
-    """Base definition of input model showing mandatory root fields.
-
-    The Model define the Module Input, usually referring to multiple input type defined by an union.
-
-    Example:
-        class ModuleInput(InputModel):
-            root: FileInput | MessageInput
-    """
-
-    root: InputTrigger
-
-
-ConfigSetupModelT = TypeVar("ConfigSetupModelT", bound=BaseModel | None)
-InputModelT = TypeVar("InputModelT", bound=InputModel)
-OutputModelT = TypeVar("OutputModelT", bound=BaseModel)
-SetupModelT = TypeVar("SetupModelT", bound=BaseModel)
-SecretModelT = TypeVar("SecretModelT", bound=BaseModel)