digitalkin 0.3.2.dev2__py3-none-any.whl → 0.3.2.dev4__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.3.2.dev2"
+    __version__ = "0.3.2.dev4"

digitalkin/core/job_manager/single_job_manager.py

@@ -15,8 +15,8 @@ from digitalkin.core.task_manager.local_task_manager import LocalTaskManager
 from digitalkin.core.task_manager.task_session import TaskSession
 from digitalkin.logger import logger
 from digitalkin.models.core.task_monitor import TaskStatus
+from digitalkin.models.module.base_types import InputModelT, OutputModelT, SetupModelT
 from digitalkin.models.module.module import ModuleCodeModel
-from digitalkin.models.module.module_types import InputModelT, OutputModelT, SetupModelT
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.services_models import ServicesMode
 
@@ -86,7 +86,10 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
                 message=f"Module {job_id} did not respond within 30 seconds",
             )
         finally:
-            logger.info(f"{job_id=}: {session.queue.empty()}")
+            logger.debug(
+                "Config setup response retrieved",
+                extra={"job_id": job_id, "queue_empty": session.queue.empty()},
+            )
 
     async def create_config_setup_instance_job(
         self,
@@ -126,7 +129,7 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         except Exception:
             # Remove the module from the manager in case of an error.
             del self.tasks_sessions[job_id]
-            logger.exception("Failed to start module %s: %s", job_id)
+            logger.exception("Failed to start module", extra={"job_id": job_id})
             raise
         else:
             return job_id
@@ -140,7 +143,8 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
             job_id: The unique identifier of the job.
             output_data: The output data produced by the job.
         """
-        await self.tasks_sessions[job_id].queue.put(output_data.model_dump())
+        session = self.tasks_sessions[job_id]
+        await session.queue.put(output_data.model_dump())
 
     @asynccontextmanager  # type: ignore
     async def generate_stream_consumer(self, job_id: str) -> AsyncIterator[AsyncGenerator[dict[str, Any], None]]:  # type: ignore
@@ -259,6 +263,18 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         logger.info("Managed task started: '%s'", job_id, extra={"task_id": job_id})
         return job_id
 
+    async def clean_session(self, task_id: str, mission_id: str) -> bool:
+        """Clean a task's session.
+
+        Args:
+            task_id: Unique identifier for the task.
+            mission_id: Mission identifier.
+
+        Returns:
+            bool: True if the task was successfully cleaned, False otherwise.
+        """
+        return await self._task_manager.clean_session(task_id, mission_id)
+
     async def stop_module(self, job_id: str) -> bool:
         """Stop a running module job.
 
@@ -271,20 +287,23 @@ class SingleJobManager(BaseJobManager[InputModelT, OutputModelT, SetupModelT]):
         Raises:
             Exception: If an error occurs while stopping the module.
         """
-        logger.info(f"STOP required for {job_id=}")
+        logger.info("Stop module requested", extra={"job_id": job_id})
 
         async with self._lock:
             session = self.tasks_sessions.get(job_id)
 
             if not session:
-                logger.warning(f"session with id: {job_id} not found")
+                logger.warning("Session not found", extra={"job_id": job_id})
                 return False
             try:
                 await session.module.stop()
                 await self.cancel_task(job_id, session.mission_id)
-                logger.debug(f"session {job_id} ({session.module.name}) stopped successfully")
-            except Exception as e:
-                logger.error(f"Error while stopping module {job_id}: {e}")
+                logger.debug(
+                    "Module stopped successfully",
+                    extra={"job_id": job_id, "mission_id": session.mission_id},
+                )
+            except Exception:
+                logger.exception("Error stopping module", extra={"job_id": job_id})
                 raise
             else:
                 return True

digitalkin/core/task_manager/task_session.py

@@ -84,9 +84,12 @@ class TaskSession:
         self._heartbeat_interval = heartbeat_interval
 
         logger.info(
-            "TaskContext initialized for task: '%s'",
-            task_id,
-            extra={"task_id": task_id, "mission_id": mission_id, "heartbeat_interval": heartbeat_interval},
+            "TaskSession initialized",
+            extra={
+                "task_id": task_id,
+                "mission_id": mission_id,
+                "heartbeat_interval": str(heartbeat_interval),
+            },
         )
 
     @property
@@ -99,6 +102,21 @@ class TaskSession:
         """Task paused status."""
         return self._paused.is_set()
 
+    @property
+    def setup_id(self) -> str:
+        """Get setup_id from module context."""
+        return self.module.context.session.setup_id
+
+    @property
+    def setup_version_id(self) -> str:
+        """Get setup_version_id from module context."""
+        return self.module.context.session.setup_version_id
+
+    @property
+    def session_ids(self) -> dict[str, str]:
+        """Get all session IDs from module context for structured logging."""
+        return self.module.context.session.current_ids()
+
     async def send_heartbeat(self) -> bool:
         """Rate-limited heartbeat with connection resilience.
 
@@ -108,6 +126,8 @@ class TaskSession:
         heartbeat = HeartbeatMessage(
             task_id=self.task_id,
             mission_id=self.mission_id,
+            setup_id=self.setup_id,
+            setup_version_id=self.setup_version_id,
             timestamp=datetime.datetime.now(datetime.timezone.utc),
         )
 
@@ -120,23 +140,17 @@ class TaskSession:
                     return True
             except Exception as e:
                 logger.error(
-                    "Heartbeat exception for task: '%s'",
-                    self.task_id,
-                    extra={"task_id": self.task_id, "error": str(e)},
+                    "Heartbeat exception",
+                    extra={**self.session_ids, "error": str(e)},
                     exc_info=True,
                 )
-            logger.error(
-                "Initial heartbeat failed for task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id},
-            )
+            logger.error("Initial heartbeat failed", extra=self.session_ids)
             return False
 
         if (heartbeat.timestamp - self._last_heartbeat) < self._heartbeat_interval:
             logger.debug(
-                "Heartbeat skipped due to rate limiting for task: '%s' | delta=%s",
-                self.task_id,
-                heartbeat.timestamp - self._last_heartbeat,
+                "Heartbeat skipped due to rate limiting",
+                extra={**self.session_ids, "delta": str(heartbeat.timestamp - self._last_heartbeat)},
             )
             return True
 
@@ -147,39 +161,24 @@ class TaskSession:
                 return True
         except Exception as e:
             logger.error(
-                "Heartbeat exception for task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id, "error": str(e)},
+                "Heartbeat exception",
+                extra={**self.session_ids, "error": str(e)},
                 exc_info=True,
             )
-        logger.warning(
-            "Heartbeat failed for task: '%s'",
-            self.task_id,
-            extra={"task_id": self.task_id},
-        )
+        logger.warning("Heartbeat failed", extra=self.session_ids)
         return False
 
     async def generate_heartbeats(self) -> None:
         """Periodic heartbeat generator with cancellation support."""
-        logger.debug(
-            "Heartbeat generator started for task: '%s'",
-            self.task_id,
-            extra={"task_id": self.task_id, "mission_id": self.mission_id},
-        )
+        logger.debug("Heartbeat generator started", extra=self.session_ids)
         while not self.cancelled:
             logger.debug(
-                "Heartbeat tick for task: '%s', cancelled=%s",
-                self.task_id,
-                self.cancelled,
-                extra={"task_id": self.task_id, "mission_id": self.mission_id},
+                "Heartbeat tick",
+                extra={**self.session_ids, "cancelled": self.cancelled},
             )
             success = await self.send_heartbeat()
             if not success:
-                logger.error(
-                    "Heartbeat failed, cancelling task: '%s'",
-                    self.task_id,
-                    extra={"task_id": self.task_id, "mission_id": self.mission_id},
-                )
+                logger.error("Heartbeat failed, cancelling task", extra=self.session_ids)
                 await self._handle_cancel(CancellationReason.HEARTBEAT_FAILURE)
                 break
             await asyncio.sleep(self._heartbeat_interval.total_seconds())
@@ -187,11 +186,7 @@ class TaskSession:
     async def wait_if_paused(self) -> None:
         """Block execution if task is paused."""
         if self._paused.is_set():
-            logger.info(
-                "Task paused, waiting for resume: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id},
-            )
+            logger.info("Task paused, waiting for resume", extra=self.session_ids)
             await self._paused.wait()
 
     async def listen_signals(self) -> None:  # noqa: C901
@@ -200,18 +195,14 @@ class TaskSession:
         Raises:
             CancelledError: Asyncio when task cancelling
         """
-        logger.info(
-            "Signal listener started for task: '%s'",
-            self.task_id,
-            extra={"task_id": self.task_id},
-        )
+        logger.info("Signal listener started", extra=self.session_ids)
         if self.signal_record_id is None:
             self.signal_record_id = (await self.db.select_by_task_id("tasks", self.task_id)).get("id")
 
         live_id, live_signals = await self.db.start_live("tasks")
         try:
             async for signal in live_signals:
-                logger.debug("Signal received for task '%s': %s", self.task_id, signal)
+                logger.debug("Signal received", extra={**self.session_ids, "signal": signal})
                 if self.cancelled:
                     break
 
@@ -228,26 +219,17 @@ class TaskSession:
                     await self._handle_status_request()
 
         except asyncio.CancelledError:
-            logger.debug(
-                "Signal listener cancelled for task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id},
-            )
+            logger.debug("Signal listener cancelled", extra=self.session_ids)
             raise
         except Exception as e:
             logger.error(
-                "Signal listener fatal error for task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id, "error": str(e)},
+                "Signal listener fatal error",
+                extra={**self.session_ids, "error": str(e)},
                 exc_info=True,
             )
         finally:
             await self.db.stop_live(live_id)
-            logger.info(
-                "Signal listener stopped for task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id},
-            )
+            logger.info("Signal listener stopped", extra=self.session_ids)
 
     async def _handle_cancel(self, reason: CancellationReason = CancellationReason.UNKNOWN) -> None:
         """Idempotent cancellation with acknowledgment and reason tracking.
@@ -257,13 +239,9 @@ class TaskSession:
         """
         if self.is_cancelled.is_set():
             logger.debug(
-                "Cancel ignored - task already cancelled: '%s' (existing reason: %s, new reason: %s)",
-                self.task_id,
-                self.cancellation_reason.value,
-                reason.value,
+                "Cancel ignored - already cancelled",
                 extra={
-                    "task_id": self.task_id,
-                    "mission_id": self.mission_id,
+                    **self.session_ids,
                     "existing_reason": self.cancellation_reason.value,
                     "new_reason": reason.value,
                 },
@@ -277,25 +255,13 @@ class TaskSession:
         # Log with appropriate level based on reason
         if reason in {CancellationReason.SUCCESS_CLEANUP, CancellationReason.FAILURE_CLEANUP}:
             logger.debug(
-                "Task cancelled (cleanup): '%s', reason: %s",
-                self.task_id,
-                reason.value,
-                extra={
-                    "task_id": self.task_id,
-                    "mission_id": self.mission_id,
-                    "cancellation_reason": reason.value,
-                },
+                "Task cancelled (cleanup)",
+                extra={**self.session_ids, "cancellation_reason": reason.value},
             )
         else:
             logger.info(
-                "Task cancelled: '%s', reason: %s",
-                self.task_id,
-                reason.value,
-                extra={
-                    "task_id": self.task_id,
-                    "mission_id": self.mission_id,
-                    "cancellation_reason": reason.value,
-                },
+                "Task cancelled",
+                extra={**self.session_ids, "cancellation_reason": reason.value},
             )
 
         # Resume if paused so cancellation can proceed
@@ -308,6 +274,8 @@ class TaskSession:
             SignalMessage(
                 task_id=self.task_id,
                 mission_id=self.mission_id,
+                setup_id=self.setup_id,
+                setup_version_id=self.setup_version_id,
                 action=SignalType.ACK_CANCEL,
                 status=self.status,
             ).model_dump(),
@@ -316,11 +284,7 @@ class TaskSession:
     async def _handle_pause(self) -> None:
         """Pause task execution."""
         if not self._paused.is_set():
-            logger.info(
-                "Pausing task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id},
-            )
+            logger.info("Task paused", extra=self.session_ids)
             self._paused.set()
 
         await self.db.update(
@@ -329,6 +293,8 @@ class TaskSession:
             SignalMessage(
                 task_id=self.task_id,
                 mission_id=self.mission_id,
+                setup_id=self.setup_id,
+                setup_version_id=self.setup_version_id,
                 action=SignalType.ACK_PAUSE,
                 status=self.status,
             ).model_dump(),
@@ -337,11 +303,7 @@ class TaskSession:
     async def _handle_resume(self) -> None:
         """Resume paused task."""
         if self._paused.is_set():
-            logger.info(
-                "Resuming task: '%s'",
-                self.task_id,
-                extra={"task_id": self.task_id},
-            )
+            logger.info("Task resumed", extra=self.session_ids)
             self._paused.clear()
 
         await self.db.update(
@@ -350,6 +312,8 @@ class TaskSession:
             SignalMessage(
                 task_id=self.task_id,
                 mission_id=self.mission_id,
+                setup_id=self.setup_id,
+                setup_version_id=self.setup_version_id,
                 action=SignalType.ACK_RESUME,
                 status=self.status,
             ).model_dump(),
@@ -361,18 +325,16 @@ class TaskSession:
             "tasks",
             self.signal_record_id,  # type: ignore
             SignalMessage(
-                mission_id=self.mission_id,
                 task_id=self.task_id,
+                mission_id=self.mission_id,
+                setup_id=self.setup_id,
+                setup_version_id=self.setup_version_id,
                 status=self.status,
                 action=SignalType.ACK_STATUS,
             ).model_dump(),
         )
 
-        logger.debug(
-            "Status report sent for task: '%s'",
-            self.task_id,
-            extra={"task_id": self.task_id},
-        )
+        logger.debug("Status report sent", extra=self.session_ids)
 
     async def cleanup(self) -> None:
         """Clean up task session resources.

digitalkin/grpc_servers/module_servicer.py

@@ -19,6 +19,7 @@ from digitalkin.logger import logger
 from digitalkin.models.core.job_manager_models import JobManagerMode
 from digitalkin.models.module.module import ModuleStatus
 from digitalkin.modules._base_module import BaseModule
+from digitalkin.services.registry import GrpcRegistry, RegistryStrategy
 from digitalkin.services.services_models import ServicesMode
 from digitalkin.services.setup.default_setup import DefaultSetup
 from digitalkin.services.setup.grpc_setup import GrpcSetup
@@ -40,6 +41,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
     args: Namespace
     setup: SetupStrategy
     job_manager: BaseJobManager
+    _registry_cache: RegistryStrategy | None = None
 
     def _add_parser_args(self, parser: ArgumentParser) -> None:
         super()._add_parser_args(parser)
@@ -82,6 +84,26 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         )
         self.setup = GrpcSetup() if self.args.services_mode == ServicesMode.REMOTE else DefaultSetup()
 
+    def _get_registry(self) -> RegistryStrategy | None:
+        """Get a cached registry instance if configured.
+
+        Returns:
+            Cached GrpcRegistry instance if registry config exists, None otherwise.
+        """
+        if self._registry_cache is not None:
+            return self._registry_cache
+
+        registry_config = self.module_class.services_config_params.get("registry")
+        if not registry_config:
+            return None
+
+        client_config = registry_config.get("client_config")
+        if not client_config:
+            return None
+
+        self._registry_cache = GrpcRegistry("", "", "", client_config)
+        return self._registry_cache
+
     async def ConfigSetupModule(  # noqa: N802
         self,
         request: lifecycle_pb2.ConfigSetupModuleRequest,
@@ -125,6 +147,15 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             msg = "No config setup data returned."
             raise ServicerError(msg)
 
+        # Resolve tool references in config_setup_data if registry is configured
+        # This also builds the tool_cache for LLM access during execution
+        registry = self._get_registry()
+        if registry:
+            if hasattr(config_setup_data, "resolve_tool_references"):
+                config_setup_data.resolve_tool_references(registry)
+            if hasattr(config_setup_data, "build_tool_cache"):
+                config_setup_data.build_tool_cache()
+
         # create a task to run the module in background
         job_id = await self.job_manager.create_config_setup_instance_job(
             config_setup_data,
@@ -139,8 +170,8 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
 
         updated_setup_data = await self.job_manager.generate_config_setup_module_response(job_id)
-        logger.info("Setup updated")
-        logger.debug(f"Updated setup data: {updated_setup_data=}")
+        logger.info("Setup updated", extra={"job_id": job_id})
+        logger.debug("Updated setup data", extra={"job_id": job_id, "setup_data": updated_setup_data})
         setup_version.content = json_format.ParseDict(
             updated_setup_data,
             struct_pb2.Struct(),
@@ -249,17 +280,19 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         Returns:
             A response indicating success or failure.
         """
-        logger.debug("StopModule called for module: '%s'", self.module_class.__name__)
+        logger.debug(
+            "StopModule called",
+            extra={"module_class": self.module_class.__name__, "job_id": request.job_id},
+        )
 
         response: bool = await self.job_manager.stop_module(request.job_id)
         if not response:
-            message = f"Job {request.job_id} not found"
-            logger.warning(message)
+            logger.warning("Job not found for stop request", extra={"job_id": request.job_id})
             context.set_code(grpc.StatusCode.NOT_FOUND)
-            context.set_details(message)
+            context.set_details(f"Job {request.job_id} not found")
             return lifecycle_pb2.StopModuleResponse(success=False)
 
-        logger.debug("Job %s stopped successfully", request.job_id, extra={"job_id": request.job_id})
+        logger.debug("Job stopped successfully", extra={"job_id": request.job_id})
         return lifecycle_pb2.StopModuleResponse(success=True)
 
     async def GetModuleStatus(  # noqa: N802

digitalkin/models/core/task_monitor.py

@@ -55,6 +55,8 @@ class SignalMessage(BaseModel):
 
     task_id: str = Field(..., description="Unique identifier for the task")
     mission_id: str = Field(..., description="Identifier for the mission")
+    setup_id: str = Field(default="", description="Identifier for the setup")
+    setup_version_id: str = Field(default="", description="Identifier for the setup version")
     status: TaskStatus = Field(..., description="Current status of the task")
     action: SignalType = Field(..., description="Type of signal action")
     timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))
@@ -67,4 +69,6 @@ class HeartbeatMessage(BaseModel):
 
     task_id: str = Field(..., description="Unique identifier for the task")
     mission_id: str = Field(..., description="Identifier for the mission")
+    setup_id: str = Field(default="", description="Identifier for the setup")
+    setup_version_id: str = Field(default="", description="Identifier for the setup version")
     timestamp: datetime = Field(default_factory=lambda: datetime.now(timezone.utc))

digitalkin/models/module/__init__.py

@@ -6,20 +6,28 @@ from digitalkin.models.module.module_types import (
     DataTrigger,
     SetupModel,
 )
+from digitalkin.models.module.tool_reference import (
+    ToolReference,
+    ToolReferenceConfig,
+    ToolSelectionMode,
+)
 from digitalkin.models.module.utility import (
     EndOfStreamOutput,
+    ModuleStartInfoOutput,
     UtilityProtocol,
     UtilityRegistry,
 )
 
 __all__ = [
-    # Core types (used by all SDK users)
     "DataModel",
     "DataTrigger",
-    # Utility (commonly used)
     "EndOfStreamOutput",
     "ModuleContext",
+    "ModuleStartInfoOutput",
     "SetupModel",
+    "ToolReference",
+    "ToolReferenceConfig",
+    "ToolSelectionMode",
     "UtilityProtocol",
     "UtilityRegistry",
 ]

digitalkin/models/module/base_types.py

@@ -0,0 +1,61 @@
+"""Base types for module models."""
+
+from __future__ import annotations
+
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar
+
+from pydantic import BaseModel, Field
+
+if TYPE_CHECKING:
+    from digitalkin.models.module.setup_types import SetupModel
+
+
+class DataTrigger(BaseModel):
+    """Defines the root input/output model exposing the protocol.
+
+    The mandatory protocol is important to define the module beahvior following the user or agent input/output.
+
+    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 = Field(
+        default_factory=lambda: datetime.now(tz=timezone.utc).isoformat(),
+        title="Created At",
+        description="Timestamp when the payload was created.",
+    )
+
+
+DataTriggerT = TypeVar("DataTriggerT", bound=DataTrigger)
+
+
+class DataModel(BaseModel, Generic[DataTriggerT]):
+    """Base definition of input/output model showing mandatory root fields.
+
+    The Model define the Module Input/output, usually referring to multiple input/output 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")