digitalkin 0.3.2.dev4__py3-none-any.whl → 0.3.2.dev6__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.dev4"
+    __version__ = "0.3.2.dev6"

digitalkin/grpc_servers/module_server.py

@@ -174,7 +174,7 @@ class ModuleServer(BaseServer):
             "Attempting to register module with registry",
             extra={
                 "module_id": module_id,
-                "address": self.server_config.address,
+                "address": self.server_config.host,
                 "port": self.server_config.port,
                 "version": version,
                 "registry_address": self.server_config.registry_address,
@@ -183,7 +183,7 @@ class ModuleServer(BaseServer):
 
         result = self.registry.register(
             module_id=module_id,
-            address=self.server_config.address,
+            address=self.server_config.host,
             port=self.server_config.port,
             version=version,
         )
@@ -193,7 +193,7 @@ class ModuleServer(BaseServer):
                 "Module registered successfully",
                 extra={
                     "module_id": result.module_id,
-                    "address": self.server_config.address,
+                    "address": self.server_config.host,
                     "port": self.server_config.port,
                     "registry_address": self.server_config.registry_address,
                 },

digitalkin/grpc_servers/utils/utility_schema_extender.py

@@ -67,7 +67,8 @@ class UtilitySchemaExtender:
         original_annotation = base_model.model_fields["root"].annotation
         original_types = cls._extract_union_types(original_annotation)
         extended_types = (*original_types, *cls._output_protocols)
-        extended_root = Annotated[extended_types, Field(discriminator="protocol")]  # type: ignore[valid-type]
+        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,
@@ -88,7 +89,8 @@ class UtilitySchemaExtender:
         original_annotation = base_model.model_fields["root"].annotation
         original_types = cls._extract_union_types(original_annotation)
         extended_types = (*original_types, *cls._input_protocols)
-        extended_root = Annotated[extended_types, Field(discriminator="protocol")]  # type: ignore[valid-type]
+        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,

digitalkin/models/module/module_context.py

@@ -6,6 +6,7 @@ from types import SimpleNamespace
 from typing import TYPE_CHECKING, Any
 from zoneinfo import ZoneInfo
 
+from digitalkin.models.module.module_helpers import ModuleHelpers
 from digitalkin.models.module.tool_cache import ToolCache
 from digitalkin.services.agent.agent_strategy import AgentStrategy
 from digitalkin.services.communication.communication_strategy import CommunicationStrategy
@@ -100,7 +101,7 @@ class ModuleContext:
     session: Session
     callbacks: SimpleNamespace
     metadata: SimpleNamespace
-    helpers: SimpleNamespace
+    helpers: ModuleHelpers
     state: SimpleNamespace = SimpleNamespace()
     tool_cache: ToolCache
 
@@ -152,7 +153,7 @@ class ModuleContext:
 
         self.metadata = SimpleNamespace(**metadata)
         self.session = Session(**session)
-        self.helpers = SimpleNamespace(**helpers)
+        self.helpers = ModuleHelpers(context=self, **helpers)
         self.callbacks = SimpleNamespace(**callbacks)
         self.tool_cache = tool_cache or ToolCache()
 

digitalkin/models/module/module_helpers.py

@@ -0,0 +1,189 @@
+"""Module helpers for inter-module communication."""
+
+from collections.abc import AsyncGenerator, Callable, Coroutine
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Any
+
+from digitalkin.logger import logger
+
+if TYPE_CHECKING:
+    from digitalkin.models.module.module_context import ModuleContext
+
+
+class ModuleHelpers(SimpleNamespace):
+    """Helpers for module-to-module communication.
+
+    Extends SimpleNamespace to allow dynamic attribute assignment
+    while providing built-in helper methods.
+    """
+
+    def __init__(self, context: "ModuleContext", **kwargs: dict[str, Any]) -> None:
+        """Initialize helpers with context reference.
+
+        Args:
+            context: ModuleContext providing access to services.
+            **kwargs: Additional attributes to set on the namespace.
+        """
+        super().__init__(**kwargs)
+        self._context = context
+
+    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._context.registry.discover_by_id(module_id)
+
+        logger.debug(
+            "Calling module by ID",
+            extra={
+                "module_id": module_id,
+                "address": module_info.address,
+                "port": module_info.port,
+            },
+        )
+
+        async for response in self._context.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:
+            module_id: Module identifier to look up in registry
+            llm_format: If True, return LLM-optimized schema format
+
+        Returns:
+            Dictionary containing schemas: {"input": ..., "output": ..., "setup": ..., "secret": ...}
+        """
+        module_info = self._context.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,
+            },
+        )
+
+        return await self._context.communication.get_module_schemas(
+            module_address=module_info.address,
+            module_port=module_info.port,
+            llm_format=llm_format,
+        )
+
+    async def create_openai_style_tool(self, module_id: 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:
+            module_id: Module ID to look up (checks cache first, then registry)
+
+        Returns:
+            OpenAI-style tool schema if found:
+            {
+                "type": "function",
+                "function": {
+                    "name": "...",
+                    "description": "...",
+                    "parameters": {...}  # Input JSON Schema
+                }
+            }
+            None if tool not found.
+        """
+        module_info = self._context.tool_cache.check_and_get(module_id, self._context.registry)
+        if not module_info:
+            return None
+
+        schemas = await self._context.communication.get_module_schemas(
+            module_address=module_info.address,
+            module_port=module_info.port,
+            llm_format=True,
+        )
+
+        return {
+            "type": "function",
+            "function": {
+                "name": module_info.name or module_info.module_id,
+                "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.
+            The function accepts **kwargs matching the tool's input schema
+            and yields dict responses.
+        """
+        module_info = self._context.tool_cache.check_and_get(module_id, self._context.registry)
+        if not module_info:
+            return None
+
+        # Capture references for closure
+        communication = self._context.communication
+        session = self._context.session
+        address = module_info.address
+        port = module_info.port
+
+        async def tool_function(**kwargs: Any) -> AsyncGenerator[dict, None]:  # noqa: ANN401
+            """Call remote tool module and yield responses.
+
+            Yields:
+                dict: Each response from the module until end_of_stream.
+            """
+            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

digitalkin/models/module/tool_cache.py

@@ -12,9 +12,9 @@ from typing import TYPE_CHECKING
 from pydantic import BaseModel, Field
 
 from digitalkin.logger import logger
+from digitalkin.models.services.registry import ModuleInfo
 
 if TYPE_CHECKING:
-    from digitalkin.models.services.registry import ModuleInfo
     from digitalkin.services.registry import RegistryStrategy
 
 
@@ -217,8 +217,6 @@ class ToolCache(BaseModel):
         Returns:
             ToolCache instance.
         """
-        from digitalkin.models.services.registry import ModuleInfo  # noqa: PLC0415
-
         cache = cls()
         for slug, entry_data in data.items():
             cache.entries[slug] = ToolCacheEntry(

digitalkin/services/communication/grpc_communication.py

@@ -187,6 +187,17 @@ class GrpcCommunication(CommunicationStrategy, GrpcClientWrapper):
                 # Convert protobuf Struct to dict
                 output_dict = json_format.MessageToDict(response.output)
 
+                # Check for end_of_stream signal
+                if output_dict.get("root", {}).get("protocol") == "end_of_stream":
+                    logger.debug(
+                        "End of stream received",
+                        extra={
+                            "module_address": module_address,
+                            "module_port": module_port,
+                        },
+                    )
+                    break
+
                 # Add job_id and success flag
                 response_dict = {
                     "success": response.success,

{digitalkin-0.3.2.dev4.dist-info → digitalkin-0.3.2.dev6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.3.2.dev4
+Version: 0.3.2.dev6
 Summary: SDK to build kin used in DigitalKin
 Author-email: "DigitalKin.ai" <contact@digitalkin.ai>
 License: Attribution-NonCommercial-ShareAlike 4.0 International

{digitalkin-0.3.2.dev4.dist-info → digitalkin-0.3.2.dev6.dist-info}/RECORD

@@ -7,7 +7,7 @@ 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=k7ywgivs-7wcYmxpZLVkbYfyQZo7NuQeNm6D2K0VbkI,195
+digitalkin/__version__.py,sha256=eog3-5zb_1ZBSzfiDIaGQRokNStum1iswtrpGBkQF5s,195
 digitalkin/logger.py,sha256=8ze_tjt2G6mDTuQcsf7-UTXWP3UHZ7LZVSs_iqF4rX4,4685
 digitalkin/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 digitalkin/core/__init__.py,sha256=FJRcJ-B1Viyn-38L8XpOpZ8KOnf1I7PCDOAmKXLQhqc,71
@@ -27,13 +27,13 @@ digitalkin/core/task_manager/task_executor.py,sha256=8xh5_1zuRAaGZIH_gWyNsA4T7YY
 digitalkin/core/task_manager/task_session.py,sha256=5jw21bT_SPXUzWE7tk6YG62EXqlRJcrSakFXDFDRy28,12730
 digitalkin/grpc_servers/__init__.py,sha256=ZIRMJ1Lcas8yQ106GCup6hn2UBOsx1sNk8ap0lpEDnY,72
 digitalkin/grpc_servers/_base_server.py,sha256=ZVeCDwI7w7fFbPTXPkeJb_SOuLfd2T7za3T4oCu2UWY,18680
-digitalkin/grpc_servers/module_server.py,sha256=7jO6Bz4pxno0AjRFercZMwgOFW7dutywynoCCht3dnQ,7858
+digitalkin/grpc_servers/module_server.py,sha256=Ec3izzV2YpdN8rGs_cX-iVulQ00FkLR5dBflHlQ8a6Y,7849
 digitalkin/grpc_servers/module_servicer.py,sha256=a5qVOd0beAwXlFAch9VfLNR8PdaYj_fF2O8rXcwInSU,21087
 digitalkin/grpc_servers/utils/__init__.py,sha256=ZnAIb_F8z4NhtPypqkdmzgRSzolKnJTk3oZx5GfWH5Y,38
 digitalkin/grpc_servers/utils/exceptions.py,sha256=LtaDtlqXCeT6iqApogs4pbtezotOVeg4fhnFzGBvFsY,692
 digitalkin/grpc_servers/utils/grpc_client_wrapper.py,sha256=ElGvp6evY5q-EBmDVyQZaDJktfShtMsptMmq16jfgxA,3285
 digitalkin/grpc_servers/utils/grpc_error_handler.py,sha256=0wPEU4713_ZlgIilaeXJV2bi90tHwYO1myDrSLeenKk,1848
-digitalkin/grpc_servers/utils/utility_schema_extender.py,sha256=XTHIzb7-o0nKgyDuspW8U4HQ0mZJd5azfBCtmPx2tgw,3557
+digitalkin/grpc_servers/utils/utility_schema_extender.py,sha256=qTnGrRcze2qe_W61FOfcXGWRILNctidm7z_1GKlCM1E,3719
 digitalkin/mixins/__init__.py,sha256=d6ljaoyJZJT9XxOrXZG5FVNvbLURb3_CZrkp4GPZWYM,590
 digitalkin/mixins/base_mixin.py,sha256=uLkg6MbDtVc9DysjdfNIGKahxQLnnjuL3DYpuyNLbk8,486
 digitalkin/mixins/callback_mixin.py,sha256=90nHm9-pbKT14GAy3CB3fsBtpYu5IH0woOQdNLM2e_Y,836
@@ -53,10 +53,11 @@ digitalkin/models/grpc_servers/types.py,sha256=rQ78s4nAet2jy-NIDj_PUWriT0kuGHr_w
 digitalkin/models/module/__init__.py,sha256=N55wan3rAUVPEGLIDjXoFM_-DYY_zxvbQOZHzNDfwoY,751
 digitalkin/models/module/base_types.py,sha256=oIylVNqo0idTFj4dRgCt7P19daNZ-AlvgCPpL9TJvto,1850
 digitalkin/models/module/module.py,sha256=k0W8vfJJFth8XdDzkHm32SyTuSf3h2qF0hSrxAfGF1s,956
-digitalkin/models/module/module_context.py,sha256=cnk5OpgevgqkmD_gU5KWSBzZ2vsNFUyN8rEzW0TNI-s,6354
+digitalkin/models/module/module_context.py,sha256=cjMP3cYPC4ZjTWXY3JhF_77PMT72ZglhHfq2I4JIl6I,6430
+digitalkin/models/module/module_helpers.py,sha256=n5lx25c5zRmuivcThC-HdumsjyWDv31xC6tI6NMLR7U,6438
 digitalkin/models/module/module_types.py,sha256=C9azCNBk76xMa-Mww8_6AiwQR8MLAsEyUOvBYxytovI,739
 digitalkin/models/module/setup_types.py,sha256=e9CG8x75Mkh7dDIQ66UWU99VcWmyk-JAqHCpUtcgyns,21323
-digitalkin/models/module/tool_cache.py,sha256=H8k2ue9m61oNHGMDgF9i3e4FxfFQoNW1iIZqoWyZNOc,7523
+digitalkin/models/module/tool_cache.py,sha256=WFoPFtHLs1Q3jdeQZpu-ZycnJwhocyYyeAo9k0JXj5o,7434
 digitalkin/models/module/tool_reference.py,sha256=A4EewKMlfhdI2hW7F3kgPc6_UinUrN-hhOweM0KzSa0,4302
 digitalkin/models/module/utility.py,sha256=gnbYfWpXGbomUI0fWf7T-Qm_VvT-LXDv1OuA9zObwVg,5589
 digitalkin/models/services/__init__.py,sha256=jhfVw6egq0OcHmos_fypH9XFehbHTBw09wluVFVFEyw,226
@@ -82,7 +83,7 @@ digitalkin/services/agent/default_agent.py,sha256=4N_E_eQxJGOx1KVUUg5jNOje-3ncMx
 digitalkin/services/communication/__init__.py,sha256=YbRjrg_BbNURgS56SlXSCzRBvyUELGgrIkuX19eA1f8,408
 digitalkin/services/communication/communication_strategy.py,sha256=u35O-al9LQxiORM9ZrGeHrBtT7CegtCsMwGdwvu2MlQ,2335
 digitalkin/services/communication/default_communication.py,sha256=zBExaitUxpuIyoWANdx5JHY9pHDRnSffTiHLvU5QqwY,3166
-digitalkin/services/communication/grpc_communication.py,sha256=0iaEQF4ihGJz2mWjxGjXVzzslEeI5CnbMJA5ddQlNAE,7581
+digitalkin/services/communication/grpc_communication.py,sha256=I78vlWlvn4ACTgOiUedgksApQvX2v1_FkitAGwSgCjs,8023
 digitalkin/services/cost/__init__.py,sha256=sD_a5LrnLluASOC5m5vgIqjaco-MzZJd6XhillIBHr0,400
 digitalkin/services/cost/cost_strategy.py,sha256=MpPX33P_S5b2by6F4zT-rcyeRuh2V4NYPZe05VpDOGQ,2649
 digitalkin/services/cost/default_cost.py,sha256=XE7kNFde8NmbulU9m1lc3mi-vHFkbaJf0XHUc0D4UHE,3945
@@ -121,7 +122,7 @@ digitalkin/utils/development_mode_action.py,sha256=2hznh0ajW_4ZTysfoc0Y49161f_PQ
 digitalkin/utils/dynamic_schema.py,sha256=5-B3dBGlCYYv6uRJkgudtc0ZpBOTYxl0yKedDGsteZQ,15184
 digitalkin/utils/llm_ready_schema.py,sha256=JjMug_lrQllqFoanaC091VgOqwAd-_YzcpqFlS7p778,2375
 digitalkin/utils/package_discover.py,sha256=sa6Zp5Kape1Zr4iYiNrnZxiHDnqM06ODk6yfWHom53w,13465
-digitalkin-0.3.2.dev4.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
+digitalkin-0.3.2.dev6.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
 modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 modules/cpu_intensive_module.py,sha256=GZlirQDZdYuXrI46sv1q4RNAHZjL4EptHVQTvgK9zz8,8363
 modules/dynamic_setup_module.py,sha256=tKvUWZdlYZkfAgKR0mLuFcLiFGKpVgpsz10LeJ6B2QI,11410
@@ -129,7 +130,7 @@ modules/minimal_llm_module.py,sha256=N9aIzZQI-miyH4AB4xTmGHpMvdSLnYyXNOD4Z3YFzis
 modules/text_transform_module.py,sha256=MfhI_Ki1U6qk379ne6oazNEu4PhO4R3cRezEcr0nGPw,7251
 services/filesystem_module.py,sha256=U4dgqtuDadaXz8PJ1d_uQ_1EPncBqudAQCLUICF9yL4,7421
 services/storage_module.py,sha256=Wz2MzLvqs2D_bnBBgtnujYcAKK2V2KFMk8K21RoepSE,6972
-digitalkin-0.3.2.dev4.dist-info/METADATA,sha256=9owo2Dtpnwbm4Rjkx9iXA_rKkzstywsbhe8hU9Mn8pY,29724
-digitalkin-0.3.2.dev4.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-digitalkin-0.3.2.dev4.dist-info/top_level.txt,sha256=gcjqlyrZuLjIyxrOIavCQM_olpr6ND5kPKkZd2j0xGo,40
-digitalkin-0.3.2.dev4.dist-info/RECORD,,
+digitalkin-0.3.2.dev6.dist-info/METADATA,sha256=AYxYAY30Ax3Zr1fmomZcVKwdjb9tSxiE8l0OSH7HfaY,29724
+digitalkin-0.3.2.dev6.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+digitalkin-0.3.2.dev6.dist-info/top_level.txt,sha256=gcjqlyrZuLjIyxrOIavCQM_olpr6ND5kPKkZd2j0xGo,40
+digitalkin-0.3.2.dev6.dist-info/RECORD,,