digitalkin 0.3.2.dev16__tar.gz → 0.3.2.dev18__tar.gz

{digitalkin-0.3.2.dev16 → digitalkin-0.3.2.dev18}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.3.2.dev16
+Version: 0.3.2.dev18
 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.dev16 → digitalkin-0.3.2.dev18}/pyproject.toml

@@ -12,7 +12,7 @@
 
     keywords = [ "digitalkin", "kin", "agent", "gprc", "sdk" ]
 
-    version = "0.3.2.dev16"
+    version = "0.3.2.dev18"
     classifiers = [
         "Development Status :: 3 - Alpha",
         "Intended Audience :: Developers",

{digitalkin-0.3.2.dev16 → digitalkin-0.3.2.dev18}/src/digitalkin/__version__.py

@@ -5,4 +5,4 @@ from importlib.metadata import PackageNotFoundError, version
 try:
     __version__ = version("digitalkin")
 except PackageNotFoundError:
-    __version__ = "0.3.2.dev16"
+    __version__ = "0.3.2.dev18"

{digitalkin-0.3.2.dev16 → digitalkin-0.3.2.dev18}/src/digitalkin/models/module/__init__.py

@@ -6,6 +6,12 @@ from digitalkin.models.module.module_types import (
     DataTrigger,
     SetupModel,
 )
+from digitalkin.models.module.tool_cache import (
+    ToolCache,
+    ToolDefinition,
+    ToolModuleInfo,
+    ToolParameter,
+)
 from digitalkin.models.module.tool_reference import (
     ToolReference,
     ToolReferenceConfig,
@@ -25,6 +31,10 @@ __all__ = [
     "ModuleContext",
     "ModuleStartInfoOutput",
     "SetupModel",
+    "ToolCache",
+    "ToolDefinition",
+    "ToolModuleInfo",
+    "ToolParameter",
     "ToolReference",
     "ToolReferenceConfig",
     "ToolSelectionMode",

{digitalkin-0.3.2.dev16 → digitalkin-0.3.2.dev18}/src/digitalkin/models/module/module_context.py

@@ -8,7 +8,7 @@ from typing import Any
 from zoneinfo import ZoneInfo
 
 from digitalkin.logger import logger
-from digitalkin.models.module.tool_cache import ToolCache
+from digitalkin.models.module.tool_cache import ToolCache, ToolDefinition, ToolModuleInfo, ToolParameter
 from digitalkin.services.agent.agent_strategy import AgentStrategy
 from digitalkin.services.communication.communication_strategy import CommunicationStrategy
 from digitalkin.services.cost.cost_strategy import CostStrategy
@@ -227,74 +227,120 @@ class ModuleContext:
             llm_format=llm_format,
         )
 
-    async def create_openai_style_tool(self, tool_name: str) -> dict[str, Any] | None:
-        """Create OpenAI-style function calling schema for a tool.
+    async def create_openai_style_tools(self, module_id: str) -> list[dict[str, Any]]:
+        """Create OpenAI-style function calling schemas for a tool module.
 
-        Uses tool cache (fast path) with registry fallback. Fetches the tool's
-        input schema and wraps it in OpenAI function calling format.
+        Uses tool cache (fast path) with registry fallback. Returns one schema
+        per ToolDefinition (protocol) in the module.
 
         Args:
-            tool_name: Module ID to look up (checks cache first, then registry).
+            module_id: Module ID to look up (checks cache first, then registry).
 
         Returns:
-            OpenAI-style tool schema if found, None otherwise.
+            List of OpenAI-style tool schemas, one per protocol. Empty if not found.
         """
-        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,
+        tool_module_info = await self.tool_cache.get(
+            module_id, registry=self.registry, communication=self.communication
         )
+        if not tool_module_info:
+            return []
+
+        return [
+            {
+                "type": "function",
+                "function": {
+                    "module_id": tool_module_info.module_id,
+                    "toolkit_name": tool_module_info.name or "undefined",
+                    "name": tool_def.name,
+                    "description": tool_def.description,
+                    "parameters": ModuleContext._build_parameters_schema(tool_def.parameters),
+                },
+            }
+            for tool_def in tool_module_info.tools
+        ]
+
+    @staticmethod
+    def _build_parameters_schema(params: list[ToolParameter]) -> dict[str, Any]:
+        """Convert ToolParameter list to JSON Schema.
+
+        Args:
+            params: List of tool parameters.
 
+        Returns:
+            JSON Schema object with properties and required fields.
+        """
         return {
-            "type": "function",
-            "function": {
-                "module_id": module_info.module_id,
-                "name": module_info.name or "undefined",
-                "description": module_info.documentation or "",
-                "parameters": schemas["input"],
-            },
+            "type": "object",
+            "properties": {p.name: {"type": p.type, "description": p.description or ""} for p in params},
+            "required": [p.name for p in params if p.required],
         }
 
-    def create_tool_function(
+    def create_tool_functions(
         self,
         module_id: str,
-    ) -> Callable[..., AsyncGenerator[dict, None]] | None:
-        """Create async generator function for a tool.
+    ) -> list[tuple[ToolDefinition, Callable[..., AsyncGenerator[dict, None]]]]:
+        """Create tool functions for all protocols in a tool module.
+
+        Returns an async generator per ToolDefinition that calls the remote tool
+        module via gRPC with the protocol auto-injected.
 
-        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.
+        This method only uses the tool cache (no registry fallback). Use this
+        in sync contexts like __init__ methods.
 
         Args:
-            module_id: Module ID to look up (checks cache first, then registry).
+            module_id: Module ID to look up in cache.
 
         Returns:
-            Async generator function if tool found, None otherwise.
+            List of (ToolDefinition, async_generator_function) tuples. Empty if not found.
         """
-        module_info = self.tool_cache.get(module_id, registry=self.registry)
-        if not module_info:
-            return None
+        tool_module_info = self.tool_cache.entries.get(module_id)
+        if not tool_module_info:
+            return []
 
         communication = self.communication
         session = self.session
-        address = module_info.address
-        port = module_info.port
+
+        result = []
+        for tool_def in tool_module_info.tools:
+            # Capture tool_def in closure via separate method
+            fn = ModuleContext._create_single_tool_function(communication, session, tool_module_info, tool_def)
+            result.append((tool_def, fn))
+
+        return result
+
+    @staticmethod
+    def _create_single_tool_function(
+        communication: CommunicationStrategy,
+        session: Session,
+        tool_module_info: ToolModuleInfo,
+        tool_def: ToolDefinition,
+    ) -> Callable[..., AsyncGenerator[dict, None]]:
+        """Create a single tool function for a specific protocol.
+
+        Args:
+            communication: Communication strategy for gRPC calls.
+            session: Current session with setup_id and mission_id.
+            tool_module_info: Tool module information containing address and port.
+            tool_def: Tool definition with protocol name.
+
+        Returns:
+            Async generator function that calls the module with protocol injected.
+        """
+        protocol = tool_def.name
 
         async def tool_function(**kwargs: Any) -> AsyncGenerator[dict, None]:  # noqa: ANN401
+            kwargs["protocol"] = protocol
             wrapped_input = {"root": kwargs}
             async for response in communication.call_module(
-                module_address=address,
-                module_port=port,
+                module_address=tool_module_info.address,
+                module_port=tool_module_info.port,
                 input_data=wrapped_input,
-                setup_id=session.setup_id,
+                setup_id=tool_module_info.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 ""
+        tool_function.__name__ = tool_def.name
+        tool_function.__doc__ = tool_def.description
 
         return tool_function

{digitalkin-0.3.2.dev16 → digitalkin-0.3.2.dev18}/src/digitalkin/models/module/setup_types.py

@@ -8,9 +8,8 @@ from typing import TYPE_CHECKING, Any, ClassVar, Generic, TypeVar, cast, get_arg
 from pydantic import BaseModel, ConfigDict, Field, create_model
 
 from digitalkin.logger import logger
-from digitalkin.models.module.tool_cache import ToolCache
+from digitalkin.models.module.tool_cache import ToolCache, ToolModuleInfo
 from digitalkin.models.module.tool_reference import ToolReference
-from digitalkin.models.services.registry import ModuleInfo
 from digitalkin.utils.dynamic_schema import (
     DynamicField,
     get_fetchers,
@@ -21,6 +20,7 @@ from digitalkin.utils.dynamic_schema import (
 if TYPE_CHECKING:
     from pydantic.fields import FieldInfo
 
+    from digitalkin.services.communication import CommunicationStrategy
     from digitalkin.services.registry import RegistryStrategy
 
 SetupModelT = TypeVar("SetupModelT", bound="SetupModel")
@@ -53,14 +53,14 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
                     # Check if it's a list type
                     origin = get_origin(annotation)
                     if origin is list:
-                        new_annotations[cache_field_name] = list[ModuleInfo]
+                        new_annotations[cache_field_name] = list[ToolModuleInfo]
                         setattr(
                             cls,
                             cache_field_name,
                             Field(default_factory=list, json_schema_extra={"hidden": True}),
                         )
                     else:
-                        new_annotations[cache_field_name] = ModuleInfo | None
+                        new_annotations[cache_field_name] = ToolModuleInfo | None
                         setattr(
                             cls,
                             cache_field_name,
@@ -336,39 +336,45 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
 
         return new_field_info
 
-    def resolve_tool_references(self, registry: "RegistryStrategy") -> None:
+    async def resolve_tool_references(
+        self, registry: "RegistryStrategy", communication: "CommunicationStrategy"
+    ) -> None:
         """Resolve all ToolReference fields recursively.
 
         Args:
             registry: Registry service for module discovery.
+            communication: Communication service for module schemas.
         """
         logger.info("Starting resolve_tool_references")
-        self._resolve_tool_references_recursive(self, registry)
+        await self._resolve_tool_references_recursive(self, registry, communication)
         logger.info("Finished resolve_tool_references")
 
     @classmethod
-    def _resolve_tool_references_recursive(
+    async def _resolve_tool_references_recursive(
         cls,
         model_instance: BaseModel,
         registry: "RegistryStrategy",
+        communication: "CommunicationStrategy",
     ) -> None:
         """Recursively resolve ToolReference fields in a model.
 
         Args:
             model_instance: Model instance to process.
             registry: Registry service for resolution.
+            communication: Communication service for module schemas.
         """
         for field_name, field_value in model_instance.__dict__.items():
             if field_value is None:
                 continue
-            cls._resolve_field_value(field_name, field_value, registry)
+            await cls._resolve_field_value(field_name, field_value, registry, communication)
 
     @classmethod
-    def _resolve_field_value(
+    async def _resolve_field_value(
         cls,
         field_name: str,
         field_value: "BaseModel | ToolReference | list | dict",
         registry: "RegistryStrategy",
+        communication: "CommunicationStrategy",
     ) -> None:
         """Resolve a single field value based on its type.
 
@@ -376,22 +382,24 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             field_name: Name of the field.
             field_value: Value to process.
             registry: Registry service for resolution.
+            communication: Communication service for module schemas.
         """
         if isinstance(field_value, ToolReference):
-            cls._resolve_single_tool_reference(field_name, field_value, registry)
+            await cls._resolve_single_tool_reference(field_name, field_value, registry, communication)
         elif isinstance(field_value, BaseModel):
-            cls._resolve_tool_references_recursive(field_value, registry)
+            await cls._resolve_tool_references_recursive(field_value, registry, communication)
         elif isinstance(field_value, list):
-            cls._resolve_list_items(field_value, registry)
+            await cls._resolve_list_items(field_value, registry, communication)
         elif isinstance(field_value, dict):
-            cls._resolve_dict_values(field_value, registry)
+            await cls._resolve_dict_values(field_value, registry, communication)
 
     @classmethod
-    def _resolve_single_tool_reference(
+    async def _resolve_single_tool_reference(
         cls,
         field_name: str,
         tool_ref: ToolReference,
         registry: "RegistryStrategy",
+        communication: "CommunicationStrategy",
     ) -> None:
         """Resolve a single ToolReference.
 
@@ -399,47 +407,54 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             field_name: Name of the field for logging.
             tool_ref: ToolReference to resolve.
             registry: Registry service for resolution.
+            communication: Communication service for module schemas.
         """
-        logger.info("Resolving ToolReference '%s' with module_id='%s'", field_name, tool_ref.config.module_id)
+        logger.info("Resolving ToolReference '%s' with setup_id='%s'", field_name, tool_ref.config.setup_id)
         try:
-            tool_ref.resolve(registry)
-            logger.info("Resolved ToolReference '%s' -> %s", field_name, tool_ref.module_info)
+            await tool_ref.resolve(registry, communication)
+            logger.info("Resolved ToolReference '%s' -> %s", field_name, tool_ref.tool_module_info)
         except Exception:
             logger.exception("Failed to resolve ToolReference '%s'", field_name)
 
     @classmethod
-    def _resolve_list_items(cls, items: list, registry: "RegistryStrategy") -> None:
+    async def _resolve_list_items(
+        cls, items: list, registry: "RegistryStrategy", communication: "CommunicationStrategy"
+    ) -> None:
         """Resolve ToolReference instances in a list.
 
         Args:
             items: List of items to process.
             registry: Registry service for resolution.
+            communication: Communication service for module schemas.
         """
         for item in items:
             if isinstance(item, ToolReference):
-                cls._resolve_single_tool_reference("list_item", item, registry)
+                await cls._resolve_single_tool_reference("list_item", item, registry, communication)
             elif isinstance(item, BaseModel):
-                cls._resolve_tool_references_recursive(item, registry)
+                await cls._resolve_tool_references_recursive(item, registry, communication)
 
     @classmethod
-    def _resolve_dict_values(cls, mapping: dict, registry: "RegistryStrategy") -> None:
+    async def _resolve_dict_values(
+        cls, mapping: dict, registry: "RegistryStrategy", communication: "CommunicationStrategy"
+    ) -> None:
         """Resolve ToolReference instances in dict values.
 
         Args:
             mapping: Dict to process.
             registry: Registry service for resolution.
+            communication: Communication service for module schemas.
         """
         for item in mapping.values():
             if isinstance(item, ToolReference):
-                cls._resolve_single_tool_reference("dict_value", item, registry)
+                await cls._resolve_single_tool_reference("dict_value", item, registry, communication)
             elif isinstance(item, BaseModel):
-                cls._resolve_tool_references_recursive(item, registry)
+                await cls._resolve_tool_references_recursive(item, registry, communication)
 
     def build_tool_cache(self) -> ToolCache:
         """Build tool cache from resolved ToolReferences, populating companion fields.
 
         Returns:
-            ToolCache with field names as keys and ModuleInfo as values.
+            ToolCache with field names as keys and ToolModuleInfo as values.
         """
         logger.info("Building tool cache")
         cache = ToolCache()
@@ -461,7 +476,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
                 cache_field_name = f"{field_name}_cache"
 
                 cached_info = getattr(model_instance, cache_field_name, None)
-                module_info = field_value.module_info or cached_info
+                module_info = field_value.tool_module_info or cached_info
                 if module_info:
                     if not cached_info:
                         setattr(model_instance, cache_field_name, module_info)
@@ -472,12 +487,12 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             elif isinstance(field_value, list):
                 cache_field_name = f"{field_name}_cache"
                 cached_infos = getattr(model_instance, cache_field_name, None) or []
-                resolved_infos: list[ModuleInfo] = []
+                resolved_infos: list[ToolModuleInfo] = []
 
                 for idx, item in enumerate(field_value):
                     if isinstance(item, ToolReference):
                         # Use resolved info or fallback to cached
-                        module_info = item.module_info or (cached_infos[idx] if idx < len(cached_infos) else None)
+                        module_info = item.tool_module_info or (cached_infos[idx] if idx < len(cached_infos) else None)
                         if module_info:
                             resolved_infos.append(module_info)
                             cache.add(module_info.module_id, module_info)

digitalkin-0.3.2.dev18/src/digitalkin/models/module/tool_cache.py

@@ -0,0 +1,227 @@
+"""Tool cache for resolved tool references."""
+
+from typing import TYPE_CHECKING, Any
+
+from pydantic import BaseModel, Field
+
+from digitalkin.logger import logger
+from digitalkin.models.services.registry import ModuleInfo
+from digitalkin.services.registry import RegistryStrategy
+
+if TYPE_CHECKING:
+    from digitalkin.services.communication import CommunicationStrategy
+
+
+class ToolParameter(BaseModel):
+    """Definition of a single tool parameter.
+
+    Attributes:
+        name: Parameter name.
+        type: JSON Schema type (string, integer, number, boolean, array, object).
+        description: Parameter description for the LLM.
+        required: Whether this parameter is required.
+        enum: Optional list of allowed values.
+        items: Optional schema for array item types.
+        properties: Optional schema for object properties.
+    """
+
+    name: str
+    type: str
+    description: str
+    required: bool = True
+    enum: list[str] | None = None
+    items: dict[str, Any] | None = None
+    properties: dict[str, Any] | None = None
+
+
+class ToolDefinition(BaseModel):
+    """Complete definition of an LLM tool with grouped parameters.
+
+    Attributes:
+        name: Tool name (from protocol const or trigger class name).
+        description: Tool description (from trigger docstring).
+        parameters: List of parameter definitions.
+    """
+
+    name: str
+    description: str
+    parameters: list[ToolParameter] = Field(default_factory=list)
+
+
+class ToolModuleInfo(ModuleInfo):
+    """Module info for tool modules."""
+
+    tools: list[ToolDefinition]
+    setup_id: str
+
+
+class ToolCache(BaseModel):
+    """Registry cache storing resolved tool references by setup field name."""
+
+    entries: dict[str, ToolModuleInfo] = Field(default_factory=dict)
+
+    def add(self, setup_id: str, tool_module_info: ToolModuleInfo) -> None:
+        """Add a tool to the cache.
+
+        Args:
+            setup_id: Field name from SetupModel used as cache key.
+            tool_module_info: Resolved tool module information.
+        """
+        self.entries[setup_id] = tool_module_info
+        logger.debug(
+            "Tool cached",
+            extra={"setup_id": setup_id, "module_id": tool_module_info.module_id},
+        )
+
+    async def get(
+        self,
+        setup_id: str,
+        *,
+        registry: RegistryStrategy | None = None,
+        communication: "CommunicationStrategy | None" = None,
+    ) -> ToolModuleInfo | None:
+        """Get a tool from cache, optionally querying registry on miss.
+
+        Args:
+            setup_id: Field name to look up.
+            registry: Optional registry to query on cache miss.
+            communication: Optional communication strategy for schema fetching.
+
+        Returns:
+            ToolModuleInfo if found, None otherwise.
+        """
+        cached = self.entries.get(setup_id)
+        if cached:
+            return cached
+
+        if registry and communication:
+            try:
+                setup_info = registry.get_setup(setup_id)
+                if setup_info and setup_info.module_id:
+                    info = registry.discover_by_id(setup_info.module_id)
+                    if info:
+                        tool_info = await module_info_to_tool_module_info(info, setup_id, communication)
+                        self.add(setup_id, tool_info)
+                        return tool_info
+            except Exception:
+                logger.exception("Registry lookup failed", extra={"setup_id": setup_id})
+
+        return None
+
+    def clear(self) -> None:
+        """Clear all cache entries."""
+        self.entries.clear()
+
+    def list_tools(self) -> list[str]:
+        """List all cached tool names.
+
+        Returns:
+            List of setup field names in cache.
+        """
+        return list(self.entries.keys())
+
+
+async def module_info_to_tool_module_info(
+    module_info: ModuleInfo,
+    setup_id: str,
+    communication: "CommunicationStrategy",
+    *,
+    llm_format: bool = True,
+) -> ToolModuleInfo:
+    """Convert ModuleInfo to ToolModuleInfo by fetching schemas via gRPC.
+
+    Fetches the module's input schema and extracts tool definitions from
+    the discriminated union structure.
+
+    Args:
+        module_info: Module info from registry.
+        setup_id: Setup ID from tool configuration.
+        communication: Communication strategy for gRPC calls.
+        llm_format: Use LLM-friendly schema format.
+
+    Returns:
+        ToolModuleInfo with tools extracted from input schema.
+    """
+    schemas = await communication.get_module_schemas(
+        module_info.address,
+        module_info.port,
+        llm_format=llm_format,
+    )
+
+    input_schema = schemas.get("input", {})
+    if llm_format:
+        input_schema = input_schema.get("json_schema", input_schema)
+
+    tools = _extract_tools_from_schema(input_schema)
+
+    return ToolModuleInfo(
+        module_id=module_info.module_id,
+        module_type=module_info.module_type,
+        address=module_info.address,
+        port=module_info.port,
+        version=module_info.version,
+        name=module_info.name,
+        documentation=module_info.documentation,
+        status=module_info.status,
+        tools=tools,
+        setup_id=setup_id,
+    )
+
+
+def _extract_tools_from_schema(schema: dict[str, Any]) -> list[ToolDefinition]:
+    """Extract tool definitions from a discriminated union input schema.
+
+    Args:
+        schema: JSON schema with $defs containing protocol-based types.
+
+    Returns:
+        List of ToolDefinition with parameters grouped by trigger.
+    """
+    tools: list[ToolDefinition] = []
+    defs = schema.get("$defs", {})
+
+    # Skip SDK utility protocols
+    utility_protocols = {"HealthcheckPingInput", "HealthcheckServicesInput", "HealthcheckStatusInput"}
+
+    for def_name, def_schema in defs.items():
+        if def_name in utility_protocols:
+            continue
+
+        properties = def_schema.get("properties", {})
+        protocol_prop = properties.get("protocol", {})
+
+        # Skip if no protocol const (not a tool input type)
+        if "const" not in protocol_prop:
+            continue
+
+        # Extract tool-level info from trigger
+        tool_name = protocol_prop.get("const", def_name)
+        tool_description = def_schema.get("description", "")
+
+        required_fields = set(def_schema.get("required", []))
+        parameters: list[ToolParameter] = []
+
+        for prop_name, prop_info in properties.items():
+            if prop_name in {"protocol", "created_at"}:
+                continue
+
+            param = ToolParameter(
+                name=prop_name,
+                type=prop_info.get("type", "string"),
+                description=prop_info.get("description", ""),
+                required=prop_name in required_fields,
+                enum=prop_info.get("enum"),
+                items=prop_info.get("items"),
+                properties=prop_info.get("properties"),
+            )
+            parameters.append(param)
+
+        tools.append(
+            ToolDefinition(
+                name=tool_name,
+                description=tool_description,
+                parameters=parameters,
+            )
+        )
+
+    return tools