digitalkin 0.3.1.dev2__py3-none-any.whl → 0.3.2.dev14__py3-none-any.whl

digitalkin/grpc_servers/module_servicer.py

@@ -5,7 +5,7 @@ from collections.abc import AsyncGenerator
 from typing import Any
 
 import grpc
-from digitalkin_proto.agentic_mesh_protocol.module.v1 import (
+from agentic_mesh_protocol.module.v1 import (
     information_pb2,
     lifecycle_pb2,
     module_service_pb2_grpc,
@@ -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,
@@ -108,8 +130,6 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
                 "mission_id": request.mission_id,
             },
         )
-        # 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 = await self.module_class.create_setup_model(
@@ -139,8 +159,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(),
@@ -219,16 +239,16 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
                         yield lifecycle_pb2.StartModuleResponse(success=False, job_id=job_id)
                         break
 
-                    if message.get("code", None) is not None and message.get("code") == "__END_OF_STREAM__":
+                    logger.info("Yielding message from job %s: %s", job_id, message)
+                    proto = json_format.ParseDict(message, struct_pb2.Struct(), ignore_unknown_fields=True)
+                    yield lifecycle_pb2.StartModuleResponse(success=True, output=proto, job_id=job_id)
+
+                    if message.get("root", {}).get("protocol") == "end_of_stream":
                         logger.info(
-                            "End of stream via __END_OF_STREAM__",
+                            "End of stream signal received",
                             extra={"job_id": job_id, "mission_id": request.mission_id},
                         )
                         break
-
-                    logger.info("Yielding message from job %s: %s", job_id, message)
-                    proto = json_format.ParseDict(message, struct_pb2.Struct(), ignore_unknown_fields=True)
-                    yield lifecycle_pb2.StartModuleResponse(success=True, output=proto, job_id=job_id)
         finally:
             await self.job_manager.wait_for_completion(job_id)
             await self.job_manager.clean_session(job_id, mission_id=request.mission_id)
@@ -249,17 +269,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/grpc_servers/utils/grpc_client_wrapper.py

@@ -43,9 +43,9 @@ class GrpcClientWrapper:
                 private_key=private_key,
             )
 
-            return grpc.secure_channel(config.address, channel_credentials, options=config.channel_options)
+            return grpc.secure_channel(config.address, channel_credentials, options=config.grpc_options)
         # Insecure channel
-        return grpc.insecure_channel(config.address, options=config.channel_options)
+        return grpc.insecure_channel(config.address, options=config.grpc_options)
 
     def exec_grpc_query(self, query_endpoint: str, request: Any) -> Any:  # noqa: ANN401
         """Execute a gRPC query with from the query's rpc endpoint name.
@@ -58,15 +58,31 @@ class GrpcClientWrapper:
             corresponding gRPC reponse.
 
         Raises:
-            ServerError: gRPC error catching
+            ServerError: gRPC error catching with status code and details
         """
+        service_name = getattr(self, "service_name", "unknown")
         try:
-            # Call the register method
-            logger.debug("send request to %s", query_endpoint, extra={"request": request})
+            logger.debug(
+                "Sending gRPC request to %s",
+                query_endpoint,
+                extra={"request": str(request), "service": service_name},
+            )
             response = getattr(self.stub, query_endpoint)(request)
-            logger.debug("receive response from request to %s", query_endpoint, extra={"response": response})
+            logger.debug(
+                "Received gRPC response from %s",
+                query_endpoint,
+                extra={"response": str(response), "service": service_name},
+            )
         except grpc.RpcError as e:
-            logger.exception("RPC error during %s", query_endpoint, extra={"error": e.details()})
-            raise ServerError
+            status_code = e.code().name if hasattr(e, "code") else "UNKNOWN"
+            details = e.details() if hasattr(e, "details") else str(e)
+            msg = f"[{status_code}] {details}"
+            logger.error(
+                "gRPC %s failed: %s",
+                query_endpoint,
+                msg,
+                extra={"service": service_name},
+            )
+            raise ServerError(msg) from e
         else:
             return response

digitalkin/grpc_servers/utils/utility_schema_extender.py

@@ -0,0 +1,100 @@
+"""Utility schema extender for gRPC API responses.
+
+This module extends module schemas with SDK utility protocols for API responses.
+"""
+
+import types
+from typing import Annotated, Union, get_args, get_origin
+
+from pydantic import Field, create_model
+
+from digitalkin.models.module.module_types import DataModel
+from digitalkin.models.module.utility import (
+    EndOfStreamOutput,
+    HealthcheckPingInput,
+    HealthcheckPingOutput,
+    HealthcheckServicesInput,
+    HealthcheckServicesOutput,
+    HealthcheckStatusInput,
+    HealthcheckStatusOutput,
+)
+
+
+class UtilitySchemaExtender:
+    """Extends module schemas with SDK utility protocols for API responses.
+
+    This class provides methods to create extended Pydantic models that include
+    both user-defined protocols and SDK utility protocols in their schemas.
+    """
+
+    _output_protocols = (
+        EndOfStreamOutput,
+        HealthcheckPingOutput,
+        HealthcheckServicesOutput,
+        HealthcheckStatusOutput,
+    )
+
+    _input_protocols = (
+        HealthcheckPingInput,
+        HealthcheckServicesInput,
+        HealthcheckStatusInput,
+    )
+
+    @classmethod
+    def _extract_union_types(cls, annotation: type) -> tuple:
+        """Extract individual types from a Union or Annotated[Union, ...] annotation.
+
+        Returns:
+            A tuple of individual types contained in the Union.
+        """
+        if get_origin(annotation) is Annotated:
+            inner_args = get_args(annotation)
+            if inner_args:
+                return cls._extract_union_types(inner_args[0])
+        if get_origin(annotation) is Union or isinstance(annotation, types.UnionType):
+            return get_args(annotation)
+        return (annotation,)
+
+    @classmethod
+    def create_extended_output_model(cls, base_model: type[DataModel]) -> type[DataModel]:
+        """Create an extended output model that includes utility output protocols.
+
+        Args:
+            base_model: The module's output_format class (a DataModel subclass).
+
+        Returns:
+            A new DataModel subclass with root typed as Union[original_types, utility_types].
+        """
+        original_annotation = base_model.model_fields["root"].annotation
+        original_types = cls._extract_union_types(original_annotation)
+        extended_types = (*original_types, *cls._output_protocols)
+        union_type = Union[extended_types]  # type: ignore[valid-type] # noqa: UP007
+        extended_root = Annotated[union_type, Field(discriminator="protocol")]  # type: ignore[valid-type]
+        return create_model(
+            f"{base_model.__name__}Utilities",
+            __base__=DataModel,
+            root=(extended_root, ...),
+            annotations=(dict[str, str], Field(default={})),
+        )
+
+    @classmethod
+    def create_extended_input_model(cls, base_model: type[DataModel]) -> type[DataModel]:
+        """Create an extended input model that includes utility input protocols.
+
+        Args:
+            base_model: The module's input_format class (a DataModel subclass).
+
+        Returns:
+            A new DataModel subclass with root typed as Union[original_types, utility_types].
+        """
+        original_annotation = base_model.model_fields["root"].annotation
+        original_types = cls._extract_union_types(original_annotation)
+        extended_types = (*original_types, *cls._input_protocols)
+        union_type = Union[extended_types]  # type: ignore[valid-type] # noqa: UP007
+        extended_root = Annotated[union_type, Field(discriminator="protocol")]  # type: ignore[valid-type]
+        return create_model(
+            f"{base_model.__name__}Utilities",
+            __base__=DataModel,
+            root=(extended_root, ...),
+            annotations=(dict[str, str], Field(default={})),
+        )

digitalkin/models/__init__.py

@@ -1,6 +1,6 @@
 """This package contains the models for DigitalKin."""
 
-from digitalkin.models.module import Module, ModuleStatus
+from digitalkin.models.module.module import Module, ModuleStatus
 
 __all__ = [
     "Module",

digitalkin/models/core/job_manager_models.py

@@ -2,17 +2,9 @@
 
 from enum import Enum
 
-from pydantic import BaseModel
-
 from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 
 
-class StreamCodeModel(BaseModel):
-    """Typed error/code model."""
-
-    code: str
-
-
 class JobManagerMode(Enum):
     """Job manager mode."""
 

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/grpc_servers/models.py

@@ -65,6 +65,42 @@ class ServerCredentials(BaseModel):
         return v
 
 
+class RetryPolicy(BaseModel):
+    """gRPC retry policy configuration for resilient connections.
+
+    Attributes:
+        max_attempts: Maximum retry attempts including the original call
+        initial_backoff: Initial backoff duration (e.g., "0.1s")
+        max_backoff: Maximum backoff duration (e.g., "10s")
+        backoff_multiplier: Multiplier for exponential backoff
+        retryable_status_codes: gRPC status codes that trigger retry
+    """
+
+    max_attempts: int = Field(default=5, ge=1, le=10, description="Maximum retry attempts including the original call")
+    initial_backoff: str = Field(default="0.1s", description="Initial backoff duration (e.g., '0.1s')")
+    max_backoff: str = Field(default="10s", description="Maximum backoff duration (e.g., '10s')")
+    backoff_multiplier: float = Field(default=2.0, ge=1.0, description="Multiplier for exponential backoff")
+    retryable_status_codes: list[str] = Field(
+        default_factory=lambda: ["UNAVAILABLE", "RESOURCE_EXHAUSTED"],
+        description="gRPC status codes that trigger retry",
+    )
+
+    model_config = {"extra": "forbid", "frozen": True}
+
+    def to_service_config_json(self) -> str:
+        """Serialize to gRPC service config JSON string.
+
+        Returns:
+            JSON string for grpc.service_config channel option.
+        """
+        codes = "[" + ",".join(f'"{c}"' for c in self.retryable_status_codes) + "]"
+        return (
+            f'{{"methodConfig":[{{"name":[{{}}],"retryPolicy":{{"maxAttempts":{self.max_attempts},'
+            f'"initialBackoff":"{self.initial_backoff}","maxBackoff":"{self.max_backoff}",'
+            f'"backoffMultiplier":{self.backoff_multiplier},"retryableStatusCodes":{codes}}}}}]}}'
+        )
+
+
 class ClientCredentials(BaseModel):
     """Model for client credentials in secure mode.
 
@@ -170,15 +206,47 @@ class ClientConfig(ChannelConfig):
         security: Security mode (secure/insecure)
         credentials: Client credentials for secure mode
         channel_options: Additional channel options
+        retry_policy: Retry policy for failed RPCs
     """
 
     credentials: ClientCredentials | None = Field(None, description="Client credentials for secure mode")
+    retry_policy: RetryPolicy = Field(default_factory=lambda: RetryPolicy(), description="Retry policy for failed RPCs")  # noqa: PLW0108
     channel_options: list[tuple[str, Any]] = Field(
         default_factory=lambda: [
-            ("grpc.max_receive_message_length", 100 * 1024 * 1024),  # 100MB
-            ("grpc.max_send_message_length", 100 * 1024 * 1024),  # 100MB
+            ("grpc.max_receive_message_length", 100 * 1024 * 1024),
+            ("grpc.max_send_message_length", 100 * 1024 * 1024),
+            # === DNS Re-resolution (Critical for Container Environments) ===
+            # Minimum milliseconds between DNS re-resolution attempts (500 ms)
+            # When connection fails, gRPC will re-query DNS after this interval
+            # Solves: Container restarts with new IPs causing "No route to host"
+            ("grpc.dns_min_time_between_resolutions_ms", 500),
+            # Initial delay before first reconnection attempt (1 second)
+            ("grpc.initial_reconnect_backoff_ms", 1000),
+            # Maximum delay between reconnection attempts (10 seconds)
+            # Prevents overwhelming the network during extended outages
+            ("grpc.max_reconnect_backoff_ms", 10000),
+            # Minimum delay between reconnection attempts (500ms)
+            # Ensures rapid recovery for brief network glitches
+            ("grpc.min_reconnect_backoff_ms", 500),
+            # === Keepalive Settings (Detect Dead Connections) ===
+            # Send keepalive ping every 60 seconds when connection is idle
+            # Proactively detects dead connections before RPC calls fail
+            ("grpc.keepalive_time_ms", 60000),
+            # Wait 20 seconds for keepalive response before declaring connection dead
+            # Triggers reconnection (with DNS re-resolution) if pong not received
+            ("grpc.keepalive_timeout_ms", 20000),
+            # Send keepalive pings even when no RPCs are in flight
+            # Essential for long-lived connections that may sit idle
+            ("grpc.keepalive_permit_without_calls", True),
+            # Minimum interval between HTTP/2 pings (30 seconds)
+            # Must be >= server's grpc.http2.min_ping_interval_without_data_ms (10s)
+            ("grpc.http2.min_time_between_pings_ms", 30000),
+            # === Retry Configuration ===
+            # Enable automatic retry for failed RPCs (1 = enabled)
+            # Works with retryable status codes: UNAVAILABLE, RESOURCE_EXHAUSTED
+            ("grpc.enable_retries", 1),
         ],
-        description="Additional channel options",
+        description="Resilient gRPC channel options with DNS re-resolution, keepalive, and retries",
     )
 
     @field_validator("credentials")
@@ -204,6 +272,15 @@ class ClientConfig(ChannelConfig):
             raise ConfigurationError(msg)
         return v
 
+    @property
+    def grpc_options(self) -> list[tuple[str, Any]]:
+        """Get channel options with retry policy service config.
+
+        Returns:
+            Full list of gRPC channel options.
+        """
+        return [*self.channel_options, ("grpc.service_config", self.retry_policy.to_service_config_json())]
+
 
 class ServerConfig(ChannelConfig):
     """Base configuration for gRPC servers.
@@ -223,10 +300,18 @@ class ServerConfig(ChannelConfig):
     credentials: ServerCredentials | None = Field(None, description="Server credentials for secure mode")
     server_options: list[tuple[str, Any]] = Field(
         default_factory=lambda: [
-            ("grpc.max_receive_message_length", 100 * 1024 * 1024),  # 100MB
-            ("grpc.max_send_message_length", 100 * 1024 * 1024),  # 100MB
+            ("grpc.max_receive_message_length", 100 * 1024 * 1024),
+            ("grpc.max_send_message_length", 100 * 1024 * 1024),
+            # === Keepalive Permission (Required for Client Keepalive) ===
+            # Allow clients to send keepalive pings without active RPCs
+            # Without this, server rejects client keepalives with GOAWAY
+            ("grpc.keepalive_permit_without_calls", True),
+            # Minimum interval server allows between client pings (10 seconds)
+            # Prevents "too_many_pings" GOAWAY errors
+            # Must match or be less than client's http2.min_time_between_pings_ms
+            ("grpc.http2.min_ping_interval_without_data_ms", 10000),
         ],
-        description="Additional server options",
+        description="gRPC server options with keepalive support",
     )
     enable_reflection: bool = Field(default=True, description="Enable reflection for the server")
     enable_health_check: bool = Field(default=True, description="Enable health check service")

digitalkin/models/module/__init__.py

@@ -1,28 +1,33 @@
 """This module contains the models for the modules."""
 
-from digitalkin.models.module.module import Module, ModuleStatus
 from digitalkin.models.module.module_context import ModuleContext
 from digitalkin.models.module.module_types import (
     DataModel,
     DataTrigger,
-    DataTriggerT,
-    InputModelT,
-    OutputModelT,
-    SecretModelT,
     SetupModel,
-    SetupModelT,
+)
+from digitalkin.models.module.tool_reference import (
+    ToolReference,
+    ToolReferenceConfig,
+    ToolSelectionMode,
+)
+from digitalkin.models.module.utility import (
+    EndOfStreamOutput,
+    ModuleStartInfoOutput,
+    UtilityProtocol,
+    UtilityRegistry,
 )
 
 __all__ = [
     "DataModel",
     "DataTrigger",
-    "DataTriggerT",
-    "InputModelT",
-    "Module",
+    "EndOfStreamOutput",
     "ModuleContext",
-    "ModuleStatus",
-    "OutputModelT",
-    "SecretModelT",
+    "ModuleStartInfoOutput",
     "SetupModel",
-    "SetupModelT",
+    "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")