digitalkin 0.3.2.dev7__py3-none-any.whl → 0.3.2.dev10__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.dev7"
+    __version__ = "0.3.2.dev10"

digitalkin/grpc_servers/module_servicer.py

@@ -130,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(
@@ -147,15 +145,6 @@ 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,

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.

digitalkin/grpc_servers/utils/utility_schema_extender.py

@@ -3,6 +3,7 @@
 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
@@ -50,7 +51,7 @@ class UtilitySchemaExtender:
             inner_args = get_args(annotation)
             if inner_args:
                 return cls._extract_union_types(inner_args[0])
-        if get_origin(annotation) is Union:
+        if get_origin(annotation) is Union or isinstance(annotation, types.UnionType):
             return get_args(annotation)
         return (annotation,)
 

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/module_context.py

@@ -1,12 +1,13 @@
 """Define the module context used in the triggers."""
 
 import os
+from collections.abc import AsyncGenerator, Callable, Coroutine
 from datetime import tzinfo
 from types import SimpleNamespace
-from typing import TYPE_CHECKING, Any
+from typing import Any
 from zoneinfo import ZoneInfo
 
-from digitalkin.models.module.module_helpers import ModuleHelpers
+from digitalkin.logger import logger
 from digitalkin.models.module.tool_cache import ToolCache
 from digitalkin.services.agent.agent_strategy import AgentStrategy
 from digitalkin.services.communication.communication_strategy import CommunicationStrategy
@@ -18,9 +19,6 @@ from digitalkin.services.snapshot.snapshot_strategy import SnapshotStrategy
 from digitalkin.services.storage.storage_strategy import StorageStrategy
 from digitalkin.services.user_profile.user_profile_strategy import UserProfileStrategy
 
-if TYPE_CHECKING:
-    from digitalkin.models.services.registry import ModuleInfo
-
 
 class Session(SimpleNamespace):
     """Session data container with mandatory setup_id and mission_id."""
@@ -101,7 +99,7 @@ class ModuleContext:
     session: Session
     callbacks: SimpleNamespace
     metadata: SimpleNamespace
-    helpers: ModuleHelpers
+    helpers: SimpleNamespace
     state: SimpleNamespace = SimpleNamespace()
     tool_cache: ToolCache
 
@@ -140,7 +138,6 @@ class ModuleContext:
             callbacks: Functions allowing user to agent interaction.
             tool_cache: ToolCache with pre-resolved tool references from setup.
         """
-        # Core services
         self.agent = agent
         self.communication = communication
         self.cost = cost
@@ -153,35 +150,151 @@ class ModuleContext:
 
         self.metadata = SimpleNamespace(**metadata)
         self.session = Session(**session)
-        self.helpers = ModuleHelpers(context=self, **helpers)
+        self.helpers = SimpleNamespace(**helpers)
         self.callbacks = SimpleNamespace(**callbacks)
         self.tool_cache = tool_cache or ToolCache()
 
-    def get_tool(self, slug: str) -> "ModuleInfo | None":
-        """Get resolved tool info by slug.
+    async def call_module_by_id(
+        self,
+        module_id: str,
+        input_data: dict,
+        setup_id: str,
+        mission_id: str,
+        callback: Callable[[dict], Coroutine[Any, Any, None]] | None = None,
+    ) -> AsyncGenerator[dict, None]:
+        """Call a module by ID, discovering address/port from registry.
+
+        Args:
+            module_id: Module identifier to look up in registry.
+            input_data: Input data as dictionary.
+            setup_id: Setup configuration ID.
+            mission_id: Mission context ID.
+            callback: Optional callback for each response.
+
+        Yields:
+            Streaming responses from module as dictionaries.
+        """
+        module_info = self.registry.discover_by_id(module_id)
 
-        Fast lookup from the pre-populated tool cache.
+        logger.debug(
+            "Calling module by ID",
+            extra={
+                "module_id": module_id,
+                "address": module_info.address,
+                "port": module_info.port,
+            },
+        )
+
+        async for response in self.communication.call_module(
+            module_address=module_info.address,
+            module_port=module_info.port,
+            input_data=input_data,
+            setup_id=setup_id,
+            mission_id=mission_id,
+            callback=callback,
+        ):
+            yield response
+
+    async def get_module_schemas_by_id(
+        self,
+        module_id: str,
+        *,
+        llm_format: bool = False,
+    ) -> dict[str, dict]:
+        """Get module schemas by ID, discovering address/port from registry.
 
         Args:
-            slug: The tool slug to look up.
+            module_id: Module identifier to look up in registry.
+            llm_format: If True, return LLM-optimized schema format.
 
         Returns:
-            ModuleInfo if found and valid, None otherwise.
+            Dictionary containing schemas: {"input": ..., "output": ..., "setup": ..., "secret": ...}
         """
-        return self.tool_cache.get(slug)
+        module_info = self.registry.discover_by_id(module_id)
+
+        logger.debug(
+            "Getting module schemas by ID",
+            extra={
+                "module_id": module_id,
+                "address": module_info.address,
+                "port": module_info.port,
+            },
+        )
 
-    def check_and_get_tool(self, slug: str) -> "ModuleInfo | None":
-        """Check cache first, then query registry if not found.
+        return await self.communication.get_module_schemas(
+            module_address=module_info.address,
+            module_port=module_info.port,
+            llm_format=llm_format,
+        )
 
-        This is the primary method for LLMs to discover tools. It:
-        1. Checks the pre-populated cache (fast path)
-        2. If not in cache, queries the registry
-        3. If found via registry, caches the result
+    async def create_openai_style_tool(self, tool_name: str) -> dict[str, Any] | None:
+        """Create OpenAI-style function calling schema for a tool.
+
+        Uses tool cache (fast path) with registry fallback. Fetches the tool's
+        input schema and wraps it in OpenAI function calling format.
 
         Args:
-            slug: The tool slug to look up.
+            tool_name: Module ID to look up (checks cache first, then registry).
 
         Returns:
-            ModuleInfo if found, None otherwise.
+            OpenAI-style tool schema if found, None otherwise.
         """
-        return self.tool_cache.check_and_get(slug, self.registry)
+        module_info = self.tool_cache.get(tool_name, registry=self.registry)
+        if not module_info:
+            return None
+
+        schemas = await self.communication.get_module_schemas(
+            module_address=module_info.address,
+            module_port=module_info.port,
+            llm_format=True,
+        )
+
+        return {
+            "type": "function",
+            "function": {
+                "module_id": module_info.module_id,
+                "name": module_info.name or "undefined",
+                "description": module_info.documentation or "",
+                "parameters": schemas["input"],
+            },
+        }
+
+    def create_tool_function(
+        self,
+        module_id: str,
+    ) -> Callable[..., AsyncGenerator[dict, None]] | None:
+        """Create async generator function for a tool.
+
+        Returns an async generator that calls the remote tool module via gRPC
+        and yields each response as it arrives until end_of_stream or gRPC ends.
+
+        Args:
+            module_id: Module ID to look up (checks cache first, then registry).
+
+        Returns:
+            Async generator function if tool found, None otherwise.
+        """
+        module_info = self.tool_cache.get(module_id, registry=self.registry)
+        if not module_info:
+            return None
+
+        communication = self.communication
+        session = self.session
+        address = module_info.address
+        port = module_info.port
+
+        async def tool_function(**kwargs: Any) -> AsyncGenerator[dict, None]:  # noqa: ANN401
+            wrapped_input = {"root": kwargs}
+            async for response in communication.call_module(
+                module_address=address,
+                module_port=port,
+                input_data=wrapped_input,
+                setup_id=session.setup_id,
+                mission_id=session.mission_id,
+            ):
+                yield response
+
+        tool_function.__name__ = module_info.name or module_info.module_id
+        tool_function.__doc__ = module_info.documentation or ""
+
+        return tool_function