digitalkin 0.2.18__py3-none-any.whl → 0.2.20__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.18"
+    __version__ = "0.2.20"

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,8 +120,8 @@ 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.setup_id,
             setup_version.id,
         )
 
@@ -177,6 +178,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             input_data,
             setup_data,
             mission_id=request.mission_id,
+            setup_id=setup_data_class.current_setup_version.setup_id,
             setup_version_id=setup_data_class.current_setup_version.id,
         )
 
@@ -333,7 +335,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 +371,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 +407,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 +443,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 +479,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/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/models/module/module_types.py

@@ -1,43 +1,105 @@
 """Types for module models."""
 
-from typing import TypeVar
+from datetime import datetime, timezone
+from typing import Any, ClassVar, Generic, TypeVar, cast
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ConfigDict, Field, create_model
 
+from digitalkin.logger import logger
 
-class InputTrigger(BaseModel):
+
+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(InputModel):
-            root: InputTrigger
+        class MyInput(DataModel):
+            root: DataTrigger
             user_define_data: Any
 
         # Usage
-        my_input = MyInput(root=InputTrigger(protocol="message"))
+        my_input = MyInput(root=DataTrigger(protocol="message"))
         print(my_input.root.protocol)  # Output: message
     """
 
-    protocol: str
+    protocol: ClassVar[str]
+    created_at: str = datetime.now(tz=timezone.utc).isoformat()
+
+
+DataTriggerT = TypeVar("DataTriggerT", bound=DataTrigger)
 
 
-class InputModel(BaseModel):
+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(InputModel):
+        class ModuleInput(DataModel):
             root: FileInput | MessageInput
     """
 
-    root: InputTrigger
+    root: DataTriggerT
+    annotations: dict[str, str] = Field(
+        default={},
+        title="Annotations",
+        description="Additional metadata or annotations related to the output. ex {'role': 'user'}",
+    )
 
 
-ConfigSetupModelT = TypeVar("ConfigSetupModelT", bound=BaseModel | None)
-InputModelT = TypeVar("InputModelT", bound=InputModel)
-OutputModelT = TypeVar("OutputModelT", bound=BaseModel)
-SetupModelT = TypeVar("SetupModelT", bound=BaseModel)
+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/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,9 @@ 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]]
 
@@ -80,21 +76,38 @@ class BaseModule(  # noqa: PLR0904
     snapshot: SnapshotStrategy
     storage: StorageStrategy
 
+    # runtime params
+    job_id: str
+    mission_id: str
+    setup_id: str
+    setup_version_id: str
+    _status: ModuleStatus
+    _task: asyncio.Task | None
+
     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, self.setup_version_id)
+            service = self.services_config.init_strategy(
+                service_name,
+                self.mission_id,
+                self.setup_id,
+                self.setup_version_id,
+            )
             setattr(self, service_name, service)
 
     def __init__(
         self,
         job_id: str,
         mission_id: str,
+        setup_id: str,
         setup_version_id: str,
     ) -> None:
         """Initialize the module."""
         self.job_id: str = job_id
         self.mission_id: str = mission_id
+        # Setup reference needed for the overall Kin scope as the filesystem context
+        self.setup_id: str = setup_id
+        # SetupVersion reference needed for the precise Kin scope as the cost
         self.setup_version_id: str = setup_version_id
         self._status = ModuleStatus.CREATED
         self._task: asyncio.Task | None = None
@@ -172,20 +185,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 +213,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 +230,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 +245,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 +307,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 +381,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 +391,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 +399,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 +423,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 +439,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 +449,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/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/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:
@@ -88,6 +87,7 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT, ConfigSetupModel
         input_data: InputModelT,
         setup_data: SetupModelT,
         mission_id: str,
+        setup_id: str,
         setup_version_id: str,
     ) -> str:
         """Create and start a new job for the module's instance.
@@ -96,7 +96,8 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT, ConfigSetupModel
             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.
+            setup_version_id: The setup version ID associated with the module.
 
         Returns:
             str: The unique identifier (job ID) of the created job.
@@ -120,9 +121,9 @@ 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_id: str,
         setup_version_id: str,
     ) -> str:
         """Create and start a new module job.
@@ -132,9 +133,9 @@ 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.
+            setup_id: The setup ID.
+            setup_version_id: The setup version ID.
 
         Returns:
             str: The unique identifier (job ID) of the created job.