digitalkin 0.2.13__py3-none-any.whl → 0.2.14__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.13"
+    __version__ = "0.2.14"

digitalkin/grpc_servers/module_servicer.py

@@ -81,6 +81,62 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         )
         self.setup = GrpcSetup() if self.args.services_mode == ServicesMode.REMOTE else DefaultSetup()
 
+    async def ConfigSetupModule(  # noqa: N802
+        self,
+        request: lifecycle_pb2.ConfigSetupModuleRequest,
+        context: grpc.aio.ServicerContext,
+    ) -> lifecycle_pb2.ConfigSetupModuleResponse:
+        """Configure the module setup.
+
+        Args:
+            request: The configuration request.
+            context: The gRPC context.
+
+        Returns:
+            A response indicating success or failure.
+
+        Raises:
+            ServicerError: if the setup data is not returned or job creation fails.
+        """
+        logger.info("ConfigSetupVersion called for module: '%s'", self.module_class.__name__)
+        # Process the module input
+        # TODO: Secret should be used here as well
+        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)
+        )
+
+        if not setup_version_data:
+            msg = "No setup data returned."
+            raise ServicerError(msg)
+
+        if not config_setup_data:
+            msg = "No config setup data returned."
+            raise ServicerError(msg)
+
+        # 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,
+        )
+
+        if job_id is None:
+            context.set_code(grpc.StatusCode.NOT_FOUND)
+            context.set_details("Failed to create module instance")
+            return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
+
+        updated_setup_data = await self.job_manager.generate_config_setup_module_response(job_id)
+        logger.warning(f"Updated setup data: {updated_setup_data=}")
+        setup_version.content = json_format.ParseDict(
+            updated_setup_data,
+            struct_pb2.Struct(),
+            ignore_unknown_fields=True,
+        )
+        return lifecycle_pb2.ConfigSetupModuleResponse(success=True, setup_version=setup_version)
+
     async def StartModule(  # noqa: N802
         self,
         request: lifecycle_pb2.StartModuleRequest,
@@ -116,7 +172,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         setup_data = self.module_class.create_setup_model(setup_data_class.current_setup_version.content)
 
         # create a task to run the module in background
-        job_id = await self.job_manager.create_job(
+        job_id = await self.job_manager.create_module_instance_job(
             input_data,
             setup_data,
             mission_id=request.mission_id,
@@ -391,3 +447,39 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             success=True,
             secret_schema=secret_format_struct,
         )
+
+    async def GetConfigSetupModule(  # noqa: N802
+        self,
+        request: information_pb2.GetConfigSetupModuleRequest,
+        context: grpc.ServicerContext,
+    ) -> information_pb2.GetConfigSetupModuleResponse:
+        """Get information about the module's setup and configuration.
+
+        Args:
+            request: The get module setup request.
+            context: The gRPC context.
+
+        Returns:
+            A response with the module's setup information.
+        """
+        logger.debug("GetConfigSetupModule called for module: '%s'", self.module_class.__name__)
+
+        # Get setup schema if available
+        try:
+            # Convert schema to proto format
+            config_setup_schema_proto = self.module_class.get_config_setup_format(llm_format=request.llm_format)
+            config_setup_format_struct = json_format.Parse(
+                text=config_setup_schema_proto,
+                message=struct_pb2.Struct(),  # pylint: disable=no-member
+                ignore_unknown_fields=True,
+            )
+        except NotImplementedError as e:
+            logger.warning(e)
+            context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+            context.set_details(e)
+            return information_pb2.GetConfigSetupModuleResponse()
+
+        return information_pb2.GetConfigSetupModuleResponse(
+            success=True,
+            config_setup_schema=config_setup_format_struct,
+        )

digitalkin/grpc_servers/utils/exceptions.py

@@ -31,3 +31,7 @@ class ReflectionError(ServerError):
 
 class HealthCheckError(ServerError):
     """Error related to gRPC health check service."""
+
+
+class OptionalFeatureNotImplementedError(NotImplementedError):
+    """Raised when an optional feature is not implemented, but was requested."""

digitalkin/models/module/__init__.py

@@ -2,10 +2,11 @@
 
 from digitalkin.models.module.module import Module, ModuleStatus
 from digitalkin.models.module.module_types import (
+    ConfigSetupModelT,
     InputModelT,
     OutputModelT,
     SecretModelT,
     SetupModelT,
 )
 
-__all__ = ["InputModelT", "Module", "ModuleStatus", "OutputModelT", "SecretModelT", "SetupModelT"]
+__all__ = ["ConfigSetupModelT", "InputModelT", "Module", "ModuleStatus", "OutputModelT", "SecretModelT", "SetupModelT"]

digitalkin/models/module/module_types.py

@@ -4,6 +4,7 @@ from typing import TypeVar
 
 from pydantic import BaseModel
 
+ConfigSetupModelT = TypeVar("ConfigSetupModelT", bound=BaseModel | None)
 InputModelT = TypeVar("InputModelT", bound=BaseModel)
 OutputModelT = TypeVar("OutputModelT", bound=BaseModel)
 SetupModelT = TypeVar("SetupModelT", bound=BaseModel)

digitalkin/modules/_base_module.py

@@ -9,8 +9,16 @@ 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 InputModelT, ModuleStatus, OutputModelT, SecretModelT, SetupModelT
+from digitalkin.models.module import (
+    ConfigSetupModelT,
+    InputModelT,
+    ModuleStatus,
+    OutputModelT,
+    SecretModelT,
+    SetupModelT,
+)
 from digitalkin.services.agent.agent_strategy import AgentStrategy
 from digitalkin.services.cost.cost_strategy import CostStrategy
 from digitalkin.services.filesystem.filesystem_strategy import FilesystemStrategy
@@ -30,11 +38,22 @@ class ModuleErrorModel(BaseModel):
     short_description: str
 
 
-class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretModelT]):
+class BaseModule(
+    ABC,
+    Generic[
+        InputModelT,
+        OutputModelT,
+        SetupModelT,
+        SecretModelT,
+        ConfigSetupModelT,
+    ],
+):
     """BaseModule is the abstract base for all modules in the DigitalKin SDK."""
 
     name: str
     description: str
+
+    config_setup_format: type[ConfigSetupModelT]
     input_format: type[InputModelT]
     output_format: type[OutputModelT]
     setup_format: type[SetupModelT]
@@ -136,6 +155,23 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         msg = "'%s' class does not define an 'output_format'."
         raise NotImplementedError(msg)
 
+    @classmethod
+    def get_config_setup_format(cls, *, llm_format: bool) -> str:
+        """Gets the JSON schema of the config setup format model.
+
+        Raises:
+            OptionalFeatureNotImplementedError: If the `config_setup_format` is not defined.
+
+        Returns:
+            The JSON schema of the config setup format as a string.
+        """
+        if cls.config_setup_format is not None:
+            if llm_format:
+                return json.dumps(llm_ready_schema(cls.config_setup_format), indent=2)
+            return json.dumps(cls.config_setup_format.model_json_schema(), indent=2)
+        msg = "'%s' class does not define an 'config_setup_format'."
+        raise OptionalFeatureNotImplementedError(msg)
+
     @classmethod
     def get_setup_format(cls, *, llm_format: bool) -> str:
         """Gets the JSON schema of the setup format model.
@@ -153,6 +189,18 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         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:
+        """Create the setup model from the setup data.
+
+        Args:
+            config_setup_data: The setup data to create the model from.
+
+        Returns:
+            The setup model.
+        """
+        return cls.config_setup_format(**config_setup_data)
+
     @classmethod
     def create_input_model(cls, input_data: dict[str, Any]) -> InputModelT:
         """Create the input model from the input data.
@@ -201,6 +249,21 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         """
         return cls.output_format(**output_data)
 
+    @abstractmethod
+    async def run_config_setup(
+        self,
+        config_setup_data: ConfigSetupModelT,
+        setup_data: SetupModelT,
+        callback: Callable,
+    ) -> None:
+        """Run config setup the module.
+
+        Raises:
+            OptionalFeatureNotImplementedError: If the config setup feature is not implemented.
+        """
+        msg = f"'{self}' class does not define an optional 'run_config_setup' attribute."
+        raise OptionalFeatureNotImplementedError(msg)
+
     @abstractmethod
     async def initialize(self, setup_data: SetupModelT) -> None:
         """Initialize the module."""
@@ -302,3 +365,18 @@ class BaseModule(ABC, Generic[InputModelT, OutputModelT, SetupModelT, SecretMode
         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]],
+    ) -> 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)
+        except Exception:
+            self._status = ModuleStatus.FAILED
+            logger.exception("Error during module lifecyle")

digitalkin/modules/archetype_module.py

@@ -3,8 +3,18 @@
 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
 
 
-class ArchetypeModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT], ABC):
+class ArchetypeModule(
+    BaseModule[
+        InputModelT,
+        OutputModelT,
+        SetupModelT,
+        SecretModelT,
+        ConfigSetupModelT,
+    ],
+    ABC,
+):
     """ArchetypeModule extends BaseModule to implement specific module types."""

digitalkin/modules/job_manager/base_job_manager.py

@@ -7,12 +7,13 @@ 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]):
+class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT, ConfigSetupModelT]):
     """Abstract base class for managing background module jobs."""
 
     async def _start(self) -> None:
@@ -82,14 +83,14 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT]):
         """
 
     @abc.abstractmethod
-    async def create_job(
+    async def create_module_instance_job(
         self,
         input_data: InputModelT,
         setup_data: SetupModelT,
         mission_id: str,
         setup_version_id: str,
     ) -> str:
-        """Create and start a new job for the module.
+        """Create and start a new job for the module's instance.
 
         Args:
             input_data: The input data required to start the job.
@@ -101,6 +102,47 @@ class BaseJobManager(abc.ABC, Generic[InputModelT, SetupModelT]):
             str: The unique identifier (job ID) of the created job.
         """
 
+    @abc.abstractmethod
+    async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT:
+        """Generate a stream consumer for a module's output data.
+
+        This method creates an asynchronous generator that streams output data
+        from a specific module job. If the module does not exist, it generates
+        an error message.
+
+        Args:
+            job_id: The unique identifier of the job.
+
+        Returns:
+            SetupModelT: the SetupModelT object fully processed.
+        """
+
+    @abc.abstractmethod
+    async def create_config_setup_instance_job(
+        self,
+        config_setup_data: ConfigSetupModelT,
+        setup_data: SetupModelT,
+        mission_id: str,
+        setup_version_id: str,
+    ) -> str:
+        """Create and start a new module job.
+
+        This method initializes a new module job, assigns it a unique job ID,
+        and starts it in the background.
+
+        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.
+
+        Returns:
+            str: The unique identifier (job ID) of the created job.
+
+        Raises:
+            Exception: If the module fails to start.
+        """
+
     @abc.abstractmethod
     async def stop_module(self, job_id: str) -> bool:
         """Stop a running module job.

digitalkin/modules/job_manager/single_job_manager.py

@@ -10,13 +10,13 @@ import grpc
 
 from digitalkin.logger import logger
 from digitalkin.models import ModuleStatus
-from digitalkin.models.module import InputModelT, OutputModelT, SetupModelT
+from digitalkin.models.module import ConfigSetupModelT, InputModelT, OutputModelT, SetupModelT
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.modules.job_manager.base_job_manager import BaseJobManager
 from digitalkin.services.services_models import ServicesMode
 
 
-class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
+class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigSetupModelT]):
     """Manages a single instance of a module job.
 
     This class ensures that only one instance of a module job is active at a time.
@@ -44,6 +44,73 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
         self.modules: dict[str, BaseModule] = {}
         self.queues: dict[str, asyncio.Queue] = {}
 
+    async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT:
+        """Generate a stream consumer for a module's output data.
+
+        This method creates an asynchronous generator that streams output data
+        from a specific module job. If the module does not exist, it generates
+        an error message.
+
+        Args:
+            job_id: The unique identifier of the job.
+
+        Returns:
+            SetupModelT: the SetupModelT object fully processed.
+        """
+        module = self.modules.get(job_id, None)
+        logger.debug("Module %s found: %s", job_id, module)
+
+        try:
+            return await self.queues[job_id].get()
+        finally:
+            logger.info(f"{job_id=}: {self.queues[job_id].empty()}")
+            del self.queues[job_id]
+
+    async def create_config_setup_instance_job(
+        self,
+        config_setup_data: ConfigSetupModelT,
+        setup_data: SetupModelT,
+        mission_id: str,
+        setup_version_id: str,
+    ) -> str:
+        """Create and start a new module setup configuration job.
+
+        This method initializes a new module job, assigns it a unique job ID,
+        and starts the config setup it in the background.
+
+        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.
+
+        Returns:
+            str: The unique identifier (job ID) of the created job.
+
+        Raises:
+            Exception: If the module fails to start.
+        """
+        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)
+        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)
+        except Exception:
+            # Remove the module from the manager in case of an error.
+            del self.modules[job_id]
+            logger.exception("Failed to start module %s: %s", job_id)
+            raise
+        else:
+            return job_id
+
     async def add_to_queue(self, job_id: str, output_data: OutputModelT) -> None:  # type: ignore
         """Add output data to the queue for a specific job.
 
@@ -106,7 +173,7 @@ class SingleJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
 
         yield _stream()
 
-    async def create_job(
+    async def create_module_instance_job(
         self,
         input_data: InputModelT,
         setup_data: SetupModelT,

digitalkin/modules/job_manager/taskiq_broker.py

@@ -130,7 +130,7 @@ async def send_message_to_stream(job_id: str, output_data: OutputModelT) -> None
 
 
 @TASKIQ_BROKER.task
-async def run_task(
+async def run_start_module(
     mission_id: str,
     setup_version_id: str,
     module_class: type[BaseModule],
@@ -171,3 +171,44 @@ async def run_task(
         # TODO: should define a BaseModel for stream code / error
         done_callback=lambda _: asyncio.create_task(callback(StreamCodeModel(code="__END_OF_STREAM__"))),
     )
+
+
+@TASKIQ_BROKER.task
+async def run_config_module(
+    mission_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_version_id: The setup ID associated with the module.
+        module_class: type[BaseModule],
+        services_mode: ServicesMode,
+        config_setup_data: dict,
+        setup_data: dict,
+        context: Allow TaskIQ context access
+    """
+    logger.warning("%s", services_mode)
+    services_config = ServicesConfig(
+        services_config_strategies=module_class.services_config_strategies,
+        services_config_params=module_class.services_config_params,
+        mode=services_mode,
+    )
+    setattr(module_class, "services_config", services_config)
+    logger.warning("%s | %s", services_config, module_class.services_config)
+
+    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)
+
+    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 InputModelT, SetupModelT
+from digitalkin.models.module import ConfigSetupModelT, 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]):
+class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT, ConfigSetupModelT]):
     """Taskiq job manager for running modules in Taskiq tasks."""
 
     services_mode: ServicesMode
@@ -95,6 +95,92 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
         with contextlib.suppress(asyncio.CancelledError):
             await self.stream_consumer_task
 
+    def __init__(
+        self,
+        module_class: type[BaseModule],
+        services_mode: ServicesMode,
+    ) -> None:
+        """Initialize the Taskiq job manager."""
+        super().__init__(module_class, services_mode)
+
+        logger.warning("TaskiqJobManager initialized with app: %s", TASKIQ_BROKER)
+        self.services_mode = services_mode
+        self.job_queues: dict[str, asyncio.Queue] = {}
+        self.max_queue_size = 1000
+
+    async def generate_config_setup_module_response(self, job_id: str) -> SetupModelT:
+        """Generate a stream consumer for a module's output data.
+
+        This method creates an asynchronous generator that streams output data
+        from a specific module job. If the module does not exist, it generates
+        an error message.
+
+        Args:
+            job_id: The unique identifier of the job.
+
+        Returns:
+            SetupModelT: the SetupModelT object fully processed.
+        """
+        queue: asyncio.Queue = asyncio.Queue(maxsize=self.max_queue_size)
+        self.job_queues[job_id] = queue
+
+        try:
+            item = await queue.get()
+            queue.task_done()
+            return item
+        finally:
+            logger.info(f"generate_config_setup_module_response: {job_id=}: {self.job_queues[job_id].empty()}")
+            self.job_queues.pop(job_id, None)
+
+    async def create_config_setup_instance_job(
+        self,
+        config_setup_data: ConfigSetupModelT,
+        setup_data: SetupModelT,
+        mission_id: str,
+        setup_version_id: str,
+    ) -> str:
+        """Create and start a new module setup configuration job.
+
+        This method initializes a new module job, assigns it a unique job ID,
+        and starts the config setup it in the background.
+
+        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.
+
+        Returns:
+            str: The unique identifier (job ID) of the created job.
+
+        Raises:
+            TypeError: If the function is called with bad data type.
+            ValueError: If the module fails to start.
+        """
+        task = TASKIQ_BROKER.find_task("digitalkin.modules.job_manager.taskiq_broker:run_config_module")
+
+        if task is None:
+            msg = "Task not found"
+            raise ValueError(msg)
+
+        if config_setup_data is None:
+            msg = "config_setup_data must be a valid model with model_dump method"
+            raise TypeError(msg)
+
+        running_task: AsyncTaskiqTask[Any] = await task.kiq(
+            mission_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
+        result = await running_task.wait_result(timeout=10)
+        logger.info("Job %s with data %s", job_id, result)
+        return job_id
+
     @asynccontextmanager  # type: ignore
     async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:  # type: ignore
         """Generate a stream consumer for the RStream stream.
@@ -132,20 +218,7 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
         finally:
             self.job_queues.pop(job_id, None)
 
-    def __init__(
-        self,
-        module_class: type[BaseModule],
-        services_mode: ServicesMode,
-    ) -> None:
-        """Initialize the Taskiq job manager."""
-        super().__init__(module_class, services_mode)
-
-        logger.warning("TaskiqJobManager initialized with app: %s", TASKIQ_BROKER)
-        self.services_mode = services_mode
-        self.job_queues: dict[str, asyncio.Queue] = {}
-        self.max_queue_size = 1000
-
-    async def create_job(
+    async def create_module_instance_job(
         self,
         input_data: InputModelT,
         setup_data: SetupModelT,
@@ -166,7 +239,7 @@ class TaskiqJobManager(BaseJobManager, Generic[InputModelT, SetupModelT]):
         Raises:
             ValueError: If the task is not found.
         """
-        task = TASKIQ_BROKER.find_task("digitalkin.modules.job_manager.taskiq_broker:run_task")
+        task = TASKIQ_BROKER.find_task("digitalkin.modules.job_manager.taskiq_broker:run_start_module")
 
         if task is None:
             msg = "Task not found"

digitalkin/modules/tool_module.py

@@ -3,8 +3,9 @@
 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], ABC):
+class ToolModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT,ConfigSetupModelT,], ABC):
     """ToolModule extends BaseModule to implement specific module types."""

digitalkin/modules/trigger_module.py

@@ -2,8 +2,10 @@
 
 from abc import ABC
 
+from digitalkin.models.module.module_types import ConfigSetupModelT
 from digitalkin.modules._base_module import BaseModule, InputModelT, OutputModelT, SecretModelT, SetupModelT
 
 
-class TriggerModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT], ABC):
+class TriggerModule(BaseModule[InputModelT, OutputModelT, SetupModelT, SecretModelT,
+        ConfigSetupModelT,], ABC):
     """TriggerModule extends BaseModule to implement specific module types."""

{digitalkin-0.2.13.dist-info → digitalkin-0.2.14.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.2.13
+Version: 0.2.14
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International
@@ -452,36 +452,36 @@ Classifier: License :: Other/Proprietary License
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: digitalkin-proto>=0.1.10
+Requires-Dist: digitalkin-proto>=0.1.15
 Requires-Dist: grpcio-health-checking>=1.71.0
 Requires-Dist: grpcio-reflection>=1.71.0
 Requires-Dist: grpcio-status>=1.71.0
-Requires-Dist: pydantic>=2.11.4
+Requires-Dist: pydantic>=2.11.5
 Provides-Extra: dev
-Requires-Dist: typos>=1.32.0; extra == "dev"
-Requires-Dist: ruff>=0.11.9; extra == "dev"
-Requires-Dist: mypy>=1.15.0; extra == "dev"
-Requires-Dist: pyright>=1.1.400; 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: 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>=44.0.3; extra == "dev"
-Requires-Dist: taskiq[reload]>=0.11.17; extra == "dev"
+Requires-Dist: cryptography>=45.0.4; extra == "dev"
 Provides-Extra: examples
 Requires-Dist: openai>=1.75.0; extra == "examples"
 Provides-Extra: tests
-Requires-Dist: freezegun>=1.5.1; extra == "tests"
+Requires-Dist: freezegun>=1.5.2; 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.3.4; extra == "tests"
-Requires-Dist: pytest-asyncio>=0.26.0; extra == "tests"
+Requires-Dist: pytest>=8.4.0; extra == "tests"
+Requires-Dist: pytest-asyncio>=1.0.0; extra == "tests"
 Requires-Dist: pytest-cov>=6.1.0; extra == "tests"
 Provides-Extra: taskiq
-Requires-Dist: rstream>=0.20.9; 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.8; extra == "taskiq"
+Requires-Dist: taskiq-redis>=1.0.9; extra == "taskiq"
+Requires-Dist: taskiq[reload]>=0.11.17; extra == "taskiq"
 Dynamic: license-file
 
 # DigitalKin Python SDK

{digitalkin-0.2.13.dist-info → digitalkin-0.2.14.dist-info}/RECORD

@@ -7,37 +7,37 @@ 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=FYm4fUSZE0epsu7qiHq0IkVIvAs5ZzTEDWDQWZeryGs,191
+digitalkin/__version__.py,sha256=yRvTSc47o9ijGrVlx1OEwjguXmUPtnbjcgpjhitwyMw,191
 digitalkin/logger.py,sha256=cFbIAZHOFx3nddOssRNYLXyqUPzR4CgDR_c-5wmB-og,1685
 digitalkin/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 digitalkin/grpc_servers/__init__.py,sha256=0cJBlwipSmFdXkyH3T0i6OJ1WpAtNsZgYX7JaSnkbtg,804
 digitalkin/grpc_servers/_base_server.py,sha256=NXnnZPjJqUDWoumhEbb7EOEWB7d8XYwpQs-l97NTe4k,18647
 digitalkin/grpc_servers/module_server.py,sha256=tc4n8fu7bcMkwCI3CKkAKR4cZqyP_JTEW0lGtqsc5hs,10283
-digitalkin/grpc_servers/module_servicer.py,sha256=I6S8qDhKOwDD1tofoUhAfhN8VH6A7hAM5UAlwu91_Is,14747
+digitalkin/grpc_servers/module_servicer.py,sha256=DZ1c_3XAXt-m5I6-CGqmrQ8TAFrYCUQj8R5rqMDKeQc,18383
 digitalkin/grpc_servers/registry_server.py,sha256=StY18DKYoPKQIU1SIzgito6D4_QA1aMVddZ8O2WGlHY,2223
 digitalkin/grpc_servers/registry_servicer.py,sha256=dqsKGHZ0LnaIvGt4ipaAuigd37sbJBndT4MAT029GsY,16471
-digitalkin/grpc_servers/utils/exceptions.py,sha256=I00OM8p8up20He4dU1fiHsvdLj1DymjR_UmoeUm2MSA,785
+digitalkin/grpc_servers/utils/exceptions.py,sha256=SyOgvjggaUECYmSiqy8KJLHwHVt5IClSTxslHM-IZzI,931
 digitalkin/grpc_servers/utils/factory.py,sha256=jm6rFjiqmtSv7BIHNAOxsG9xXtSvWpx9TfzSQiX97MQ,5899
 digitalkin/grpc_servers/utils/grpc_client_wrapper.py,sha256=GXRqCTCufwWd8sVMoClnl-JpeBjgmdlSCD0a3YMU9D0,2594
 digitalkin/grpc_servers/utils/models.py,sha256=80F5oHiv8MOqMoDZJBXmJSRoVYJRyhaVcijQ2LinAig,8428
 digitalkin/grpc_servers/utils/types.py,sha256=rQ78s4nAet2jy-NIDj_PUWriT0kuGHr_w6ELjmjgBao,539
 digitalkin/models/__init__.py,sha256=hDHtUfswaNh8wo4NZaBItg9JqC0uNSRqXArNWSrGynY,163
-digitalkin/models/module/__init__.py,sha256=ihvRlemJuFvU4armZIL-Vq-zaJx9UrjDoJAVPbEG8jw,345
+digitalkin/models/module/__init__.py,sha256=iQJgq3OVvxoY6gkNCVdVz5bSGgLCy0TMY5H1FFxs-0Q,389
 digitalkin/models/module/module.py,sha256=NWOwCeD_NH0NL89AF_NJ0SwE7G1n-HTcRalodLCDmpM,749
-digitalkin/models/module/module_types.py,sha256=SxwzfDDnkyuQe50glVrxN6VPuJCkbB313FRyVZugIMw,312
+digitalkin/models/module/module_types.py,sha256=rlT6prjBN-h1sLKstSwjxWXiyFgtxEn1mjOSLgfI6uA,385
 digitalkin/models/services/__init__.py,sha256=HsW7MUGFPvH7Ri28WN4BHHBfEQk5dzU_9FOWAc-0lSE,56
 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=EwfCCIy3N1eBdBUj6KAdx0WzNRFJsyQDVCdA8Lwvtn4,10463
-digitalkin/modules/archetype_module.py,sha256=T2Ehj7EpAC2MO9WQbJv39hqRw7rh3exhVZTEL3JPM8U,421
-digitalkin/modules/tool_module.py,sha256=86g0M1wHZ1ReIc7AkKfyjnlGN2QYJBGxrEQpKVlyrZI,421
-digitalkin/modules/trigger_module.py,sha256=kVoI4Gdkw7WWUP5T6hSCNqw5FxibTxL6Tpq9KP7gg78,379
-digitalkin/modules/job_manager/base_job_manager.py,sha256=I3bduyGDjCJ6RsbwQvLp_NA6ucP8MJJVoIKpEbaP9ss,4706
+digitalkin/modules/_base_module.py,sha256=wFnGGN97RNu3TJMdgMktOKdYPWmJL9WnwR120bPRXQg,12922
+digitalkin/modules/archetype_module.py,sha256=7rleVOkZfFR7TQeaq6oZXLBuRWixKsUJMOr9HfDV0T8,566
+digitalkin/modules/tool_module.py,sha256=7eWem70dEJiyoCtkGIJP4Ak2G8MzSruOjrByQMqW9-I,508
+digitalkin/modules/trigger_module.py,sha256=FNiJmG3nTG1nDqKC1XdJI8PmFNZNPs1_QrH6RX7ev2Y,475
+digitalkin/modules/job_manager/base_job_manager.py,sha256=Ku_MT7Cuf1R_qJU7SV-yHkzUhhMtmDOa8t5DyetvMxw,6177
 digitalkin/modules/job_manager/job_manager_models.py,sha256=onHy-DfInLZQveniMIWIKwTKSQjojz500JvHB54x93c,1129
-digitalkin/modules/job_manager/single_job_manager.py,sha256=Dekcnn8wb4epqGrWFAZ86YvvII553Hw0xjipgAINl4A,7745
-digitalkin/modules/job_manager/taskiq_broker.py,sha256=OnHZy-hAtVDU5FU4W9folZZKCrDrUk16VfHlfoECuiQ,5930
-digitalkin/modules/job_manager/taskiq_job_manager.py,sha256=ODqC7Snf-Flo0ugotqicrTgjFu4wgll6bUu7CM3l7ZE,7459
+digitalkin/modules/job_manager/single_job_manager.py,sha256=DkQV3-E9XqgmuC6hoYKolTX2iv9lUJ8BSCde3zk1OU8,10264
+digitalkin/modules/job_manager/taskiq_broker.py,sha256=sMXIgheOBshMiBUWUKph_5rxBHkip4rKusrrxFoa7uY,7379
+digitalkin/modules/job_manager/taskiq_job_manager.py,sha256=evcE3o51DPWzW2y4ri6DFQZNZbFJXyvFb65SibExJH8,10152
 digitalkin/services/__init__.py,sha256=LqGk_5DJy8Bzz62ajIq9jCeYNKQUIgtSCpafZk15FLc,910
 digitalkin/services/base_strategy.py,sha256=QAQnJw1BbqcYMSzwlFyhHP5juBH2WKrZzWxqDr_sDHI,638
 digitalkin/services/services_config.py,sha256=4hc7-rHgSigoS3SuV0V9FReD2Dz7XoMXcD6iMBP2KKM,7391
@@ -74,13 +74,13 @@ digitalkin/utils/__init__.py,sha256=sJnY-ZUgsjMfojAjONC1VN14mhgIDnzyOlGkw21rRnM,
 digitalkin/utils/arg_parser.py,sha256=nvjI1pKDY1HfS0oGcMQPtdTQcggXLtpxXMbnMxNEKRU,3109
 digitalkin/utils/development_mode_action.py,sha256=TqRuAF_A7bDD4twRB4PnZcRoNeaiAnEdxM5kvy4aoaA,1511
 digitalkin/utils/llm_ready_schema.py,sha256=JjMug_lrQllqFoanaC091VgOqwAd-_YzcpqFlS7p778,2375
-digitalkin-0.2.13.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
+digitalkin-0.2.14.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
 modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-modules/cpu_intensive_module.py,sha256=LUfgd1qFJ0HfsWMapSXUivgf5KGDpANKPORbcQBUjc4,8089
-modules/minimal_llm_module.py,sha256=6cygPliUb4mroOrUo5YirZfQQ1IeOT3xUHzPpDC_0Tc,10045
+modules/cpu_intensive_module.py,sha256=Nd_dHcnOKwhhJOLdqlfGklrxZkABV-iG-2g8wTsPUpg,8389
+modules/minimal_llm_module.py,sha256=Ijld__ZnhzfLwpXD1XVkLZ7jyKZKyOFZczOpiPttJZc,11216
 modules/storage_module.py,sha256=KJA75oxXDUxZrLfoY5KxeyI8Gm1b3DxFuFj7LGOkeaQ,6389
 modules/text_transform_module.py,sha256=Db_WbJi7_gWyoNvQEVhHQf5VurDsU-jd9LMAalG0oZY,7277
-digitalkin-0.2.13.dist-info/METADATA,sha256=CUpDclUxsamv7icPsWN-ImippAec3_EQZUGm1cQGHFQ,30577
-digitalkin-0.2.13.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-digitalkin-0.2.13.dist-info/top_level.txt,sha256=5_5e35inSM5YfWNZE21p5wGBojiVtQQML_WzbEk4BRU,31
-digitalkin-0.2.13.dist-info/RECORD,,
+digitalkin-0.2.14.dist-info/METADATA,sha256=2srdCp7b3ieTRigCaUMqFJc2c3Q0tQq1LLBZgSoBUVw,30580
+digitalkin-0.2.14.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+digitalkin-0.2.14.dist-info/top_level.txt,sha256=5_5e35inSM5YfWNZE21p5wGBojiVtQQML_WzbEk4BRU,31
+digitalkin-0.2.14.dist-info/RECORD,,

{digitalkin-0.2.13.dist-info → digitalkin-0.2.14.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.7.1)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

modules/cpu_intensive_module.py

@@ -127,6 +127,16 @@ class CPUOutput(BaseModel):
     )
 
 
+class CPUConfigSetup(BaseModel):
+    """Config Setup model definining data that will be pre-computed for each setup and module instance."""
+
+    files: list[str] = Field(
+        ...,
+        title="Files to embed",
+        description="List of files to embed in the setup lifecycle.",
+    )
+
+
 class CPUSetup(BaseModel):
     """Setup model defining module configuration parameters."""
 
@@ -175,7 +185,7 @@ client_config = ClientConfig(
 )
 
 
-class CPUIntensiveModule(BaseModule[CPUInput, CPUOutput, CPUSetup, CPUToolSecret]):
+class CPUIntensiveModule(BaseModule[CPUInput, CPUOutput, CPUSetup, CPUToolSecret, None]):
     """A CPU endpoint tool module module."""
 
     name = "CPUIntensiveModule"

modules/minimal_llm_module.py

@@ -11,7 +11,6 @@ from pydantic import BaseModel, Field
 from digitalkin.grpc_servers.utils.models import ClientConfig, SecurityMode, ServerMode
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_models import ServicesStrategy
-from digitalkin.services.setup.setup_strategy import SetupData
 
 # Configure logging with clear formatting
 logging.basicConfig(
@@ -154,6 +153,16 @@ class OpenAISetup(BaseModel):
     )
 
 
+class OpenAIConfigSetup(BaseModel):
+    """Setup model defining module configuration parameters."""
+
+    rag_files: list[bytes] = Field(
+        ...,
+        title="RAG Files",
+        description="Files used for retrieval-augmented generation (RAG) with the OpenAI module.",
+    )
+
+
 class OpenAIToolSecret(BaseModel):
     """Secret model defining module configuration parameters."""
 
@@ -167,13 +176,22 @@ client_config = ClientConfig(
 )
 
 
-class OpenAIToolModule(BaseModule[OpenAIInput, OpenAIOutput, OpenAISetup, OpenAIToolSecret]):
+class OpenAIToolModule(
+    BaseModule[
+        OpenAIInput,
+        OpenAIOutput,
+        OpenAISetup,
+        OpenAIToolSecret,
+        OpenAIConfigSetup,
+    ]
+):
     """A openAI endpoint tool module module."""
 
     name = "OpenAIToolModule"
     description = "A module that interacts with OpenAI API to process text"
 
     # Define the schema formats for the module
+    config_setup_format = OpenAIConfigSetup
     input_format = OpenAIInput
     output_format = OpenAIOutput
     setup_format = OpenAISetup
@@ -205,7 +223,27 @@ class OpenAIToolModule(BaseModule[OpenAIInput, OpenAIOutput, OpenAISetup, OpenAI
         },
     }
 
-    async def initialize(self, setup_data: SetupData) -> None:
+    async def run_config_setup(
+        self,
+        config_setup_data: OpenAIConfigSetup,
+        setup_data: OpenAISetup,
+        callback: Callable,
+    ) -> None:
+        """Configure the module with additional setup data.
+
+        Args:
+            config_setup_data: Additional configuration content.
+            setup_data: Initial setup data for the module.
+            callback: Function to send output data back to the client.
+        """
+        logger.info("Configuring OpenAIToolModule with additional setup data. %s", config_setup_data)
+
+        # Here you can process config_content and update setup_data as needed
+        # For now, we just return the original setup_data
+        setup_data.developer_prompt = "| + |".join(f.decode("utf-8") for f in config_setup_data.rag_files)
+        await callback(setup_data)
+
+    async def initialize(self, setup_data: OpenAISetup) -> None:
         """Initialize the module capabilities.
 
         This method is called when the module is loaded by the server.