digitalkin 0.3.2.dev15__py3-none-any.whl → 0.3.2.dev17__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.dev15"
+    __version__ = "0.3.2.dev17"

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/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, 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,63 +227,113 @@ 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, tool_name: 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).
 
         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(
+            tool_name, 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
+        address = tool_module_info.address
+        port = tool_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, address, port, tool_def)
+            result.append((tool_def, fn))
+
+        return result
+
+    @staticmethod
+    def _create_single_tool_function(
+        communication: CommunicationStrategy,
+        session: Session,
+        address: str,
+        port: int,
+        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.
+            address: Module address.
+            port: Module 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,
@@ -294,7 +344,7 @@ class ModuleContext:
             ):
                 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/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/models/module/tool_cache.py

@@ -1,55 +1,107 @@
 """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]
+
 
 class ToolCache(BaseModel):
     """Registry cache storing resolved tool references by setup field name."""
 
-    entries: dict[str, ModuleInfo] = Field(default_factory=dict)
+    entries: dict[str, ToolModuleInfo] = Field(default_factory=dict)
 
-    def add(self, setup_tool_name: str, module_info: ModuleInfo) -> None:
+    def add(self, setup_tool_name: str, tool_module_info: ToolModuleInfo) -> None:
         """Add a tool to the cache.
 
         Args:
             setup_tool_name: Field name from SetupModel used as cache key.
-            module_info: Resolved module information.
+            tool_module_info: Resolved tool module information.
         """
-        self.entries[setup_tool_name] = module_info
+        self.entries[setup_tool_name] = tool_module_info
         logger.debug(
             "Tool cached",
-            extra={"setup_tool_name": setup_tool_name, "module_id": module_info.module_id},
+            extra={"setup_tool_name": setup_tool_name, "module_id": tool_module_info.module_id},
         )
 
-    def get(
+    async def get(
         self,
         setup_tool_name: str,
         *,
         registry: RegistryStrategy | None = None,
-    ) -> ModuleInfo | None:
+        communication: "CommunicationStrategy | None" = None,
+    ) -> ToolModuleInfo | None:
         """Get a tool from cache, optionally querying registry on miss.
 
         Args:
             setup_tool_name: Field name to look up.
             registry: Optional registry to query on cache miss.
+            communication: Optional communication strategy for schema fetching.
 
         Returns:
-            ModuleInfo if found, None otherwise.
+            ToolModuleInfo if found, None otherwise.
         """
         cached = self.entries.get(setup_tool_name)
         if cached:
             return cached
 
-        if registry:
+        if registry and communication:
             try:
-                info = registry.discover_by_id(setup_tool_name)
-                if info:
-                    self.add(setup_tool_name, info)
-                    return info
+                setup_info = registry.get_setup(setup_tool_name)
+                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, communication)
+                        self.add(setup_tool_name, tool_info)
+                        return tool_info
             except Exception:
                 logger.exception("Registry lookup failed", extra={"setup_tool_name": setup_tool_name})
 
@@ -66,3 +118,106 @@ class ToolCache(BaseModel):
             List of setup field names in cache.
         """
         return list(self.entries.keys())
+
+
+async def module_info_to_tool_module_info(
+    module_info: ModuleInfo,
+    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.
+        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,
+    )
+
+
+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

digitalkin/models/module/tool_reference.py

@@ -4,7 +4,8 @@ from enum import Enum
 
 from pydantic import BaseModel, Field, PrivateAttr, model_validator
 
-from digitalkin.models.services.registry import ModuleInfo
+from digitalkin.models.module.tool_cache import ToolModuleInfo, module_info_to_tool_module_info
+from digitalkin.services.communication.communication_strategy import CommunicationStrategy
 from digitalkin.services.registry import RegistryStrategy
 
 
@@ -20,6 +21,7 @@ class ToolReferenceConfig(BaseModel):
     """Tool selection configuration. The module_id serves as both identifier and cache key."""
 
     mode: ToolSelectionMode = Field(default=ToolSelectionMode.FIXED)
+    setup_id: str = Field(default="")
     module_id: str = Field(default="")
     tag: str = Field(default="")
     organization_id: str = Field(default="")
@@ -34,8 +36,8 @@ class ToolReferenceConfig(BaseModel):
         Raises:
             ValueError: If required field is missing for the mode.
         """
-        if self.mode == ToolSelectionMode.FIXED and not self.module_id:
-            msg = "module_id required when mode is FIXED"
+        if self.mode == ToolSelectionMode.FIXED and not self.setup_id:
+            msg = "setup_id required when mode is FIXED"
             raise ValueError(msg)
         if self.mode == ToolSelectionMode.TAG and not self.tag:
             msg = "tag required when mode is TAG"
@@ -47,7 +49,7 @@ class ToolReference(BaseModel):
     """Reference to a tool module, resolved via registry during config setup."""
 
     config: ToolReferenceConfig
-    _cached_info: ModuleInfo | None = PrivateAttr(default=None)
+    _cached_info: ToolModuleInfo | None = PrivateAttr(default=None)
 
     @property
     def slug(self) -> str:
@@ -56,7 +58,7 @@ class ToolReference(BaseModel):
         Returns:
             Module ID used as cache key.
         """
-        return self.config.module_id
+        return self.config.setup_id
 
     @property
     def module_id(self) -> str:
@@ -68,11 +70,20 @@ class ToolReference(BaseModel):
         return self.config.module_id
 
     @property
-    def module_info(self) -> ModuleInfo | None:
+    def setup_id(self) -> str:
+        """Setup identifier.
+
+        Returns:
+            Setup ID or empty string if not set.
+        """
+        return self.config.setup_id
+
+    @property
+    def tool_module_info(self) -> ToolModuleInfo | None:
         """Resolved module information.
 
         Returns:
-            ModuleInfo if resolved, None otherwise.
+            ToolModuleInfo if resolved, None otherwise.
         """
         return self._cached_info
 
@@ -85,23 +96,28 @@ class ToolReference(BaseModel):
         """
         return self._cached_info is not None
 
-    def resolve(self, registry: RegistryStrategy) -> ModuleInfo | None:
+    async def resolve(self, registry: RegistryStrategy, communication: CommunicationStrategy) -> ToolModuleInfo | None:
         """Resolve this reference using the registry.
 
         Args:
             registry: Registry service for module discovery.
+            communication: Communication service for module schemas.
 
         Returns:
-            ModuleInfo if resolved, None for DISCOVERABLE mode or if not found.
+            ToolModuleInfo if resolved, None for DISCOVERABLE mode or if not found.
         """
         if self.config.mode == ToolSelectionMode.DISCOVERABLE:
             return None
 
-        if self.config.mode == ToolSelectionMode.FIXED and self.config.module_id:
-            info = registry.discover_by_id(self.config.module_id)
-            if info:
-                self._cached_info = info
-            return info
+        if self.config.mode == ToolSelectionMode.FIXED and self.config.setup_id:
+            setup = registry.get_setup(self.config.setup_id)
+            if setup and setup.module_id:
+                self.config.module_id = setup.module_id
+                info = registry.discover_by_id(self.config.module_id)
+                if info:
+                    tool_module_info = await module_info_to_tool_module_info(info, communication)
+                    self._cached_info = tool_module_info
+                    return tool_module_info
 
         if self.config.mode == ToolSelectionMode.TAG and self.config.tag:
             results = registry.search(
@@ -110,8 +126,9 @@ class ToolReference(BaseModel):
                 organization_id=self.config.organization_id,
             )
             if results:
-                self._cached_info = results[0]
-                self.config.module_id = results[0].module_id
-                return results[0]
+                tool_module_info = await module_info_to_tool_module_info(results[0], communication)
+                self._cached_info = tool_module_info
+                self.config.module_id = tool_module_info.module_id
+                return tool_module_info
 
         return None

digitalkin/models/services/registry.py

@@ -1,6 +1,7 @@
 """Registry data models."""
 
 from enum import Enum
+from typing import Any
 
 from pydantic import BaseModel
 
@@ -33,3 +34,44 @@ class ModuleInfo(BaseModel):
     name: str = ""
     documentation: str | None = None
     status: RegistryModuleStatus | None = None
+
+
+class RegistrySetupStatus(str, Enum):
+    """Setup status in the registry."""
+
+    UNSPECIFIED = "unspecified"
+    DRAFT = "draft"
+    WAITING_FOR_APPROVAL = "waiting_for_approval"
+    READY = "ready"
+    PAUSED = "paused"
+    FAILED = "failed"
+    ARCHIVED = "archived"
+    NEEDS_CONFIGURATION = "needs_configuration"
+    CONFIGURATION_FAILED = "configuration_failed"
+    CONFIGURATION_SUCCEEDED = "configuration_succeeded"
+
+
+class RegistryVisibility(str, Enum):
+    """Visibility in the registry."""
+
+    UNSPECIFIED = "unspecified"
+    PUBLIC = "public"
+    PRIVATE = "private"
+    INTERNAL = "internal"
+
+
+class SetupInfo(BaseModel):
+    """Setup information from registry."""
+
+    setup_id: str
+    name: str
+    documentation: str | None = None
+    status: RegistrySetupStatus | None = None
+    visibility: RegistryVisibility | None = None
+    organization_id: str | None = None
+    owner_id: str | None = None
+    card_id: str | None = None
+    module_id: str | None = None
+    setup_version_id: str | None = None
+    setup_version: str | None = None
+    config: dict[str, Any] | None = None

digitalkin/modules/_base_module.py

@@ -518,11 +518,14 @@ class BaseModule(  # noqa: PLR0904
             config_setup_data: Setup data containing tool references.
         """
         logger.info("Starting tool resolution", extra=self.context.session.current_ids())
-        if self.context.registry is not None:
-            config_setup_data.resolve_tool_references(self.context.registry)
+        if self.context.registry is not None and self.context.communication is not None:
+            await config_setup_data.resolve_tool_references(self.context.registry, self.context.communication)
             logger.info("Tool references resolved", extra=self.context.session.current_ids())
         else:
-            logger.warning("No registry available, skipping tool resolution", extra=self.context.session.current_ids())
+            logger.warning(
+                "No registry or communication available, skipping tool resolution",
+                extra=self.context.session.current_ids(),
+            )
 
         tool_cache = config_setup_data.build_tool_cache()
         self.context.tool_cache = tool_cache

digitalkin/services/filesystem/grpc_filesystem.py

@@ -97,7 +97,7 @@ class GrpcFilesystem(FilesystemStrategy, GrpcClientWrapper, GrpcErrorHandlerMixi
             case "mission":
                 context_id = self.mission_id
         return filesystem_pb2.FileFilter(
-            **filters.model_dump(exclude={"file_types", "status"}),
+            **filters.model_dump(exclude={"file_types", "status", "context"}),
             file_types=[self._file_type_to_enum(file_type) for file_type in filters.file_types]
             if filters.file_types
             else None,

digitalkin/services/registry/grpc_registry.py

@@ -22,6 +22,9 @@ from digitalkin.models.services.registry import (
     ModuleInfo,
     RegistryModuleStatus,
     RegistryModuleType,
+    RegistrySetupStatus,
+    RegistryVisibility,
+    SetupInfo,
 )
 from digitalkin.services.registry.exceptions import (
     RegistryModuleNotFoundError,
@@ -75,6 +78,35 @@ class GrpcRegistry(RegistryStrategy, GrpcClientWrapper, GrpcErrorHandlerMixin):
             documentation=descriptor.documentation or None,
         )
 
+    @staticmethod
+    def _proto_to_setup_info(descriptor: registry_models_pb2.SetupDescriptor) -> SetupInfo | None:
+        """Convert proto SetupDescriptor to SetupInfo.
+
+        Args:
+            descriptor: Proto SetupDescriptor message.
+
+        Returns:
+            SetupInfo with mapped fields, or None if descriptor is empty.
+        """
+        if not descriptor.id:
+            return None
+        status_name = registry_enums_pb2.SetupStatus.Name(descriptor.status).removeprefix("SETUP_STATUS_")
+        visibility_name = registry_enums_pb2.Visibility.Name(descriptor.visibility).removeprefix("VISIBILITY_")
+        return SetupInfo(
+            setup_id=descriptor.id,
+            name=descriptor.name,
+            documentation=descriptor.documentation or None,
+            status=RegistrySetupStatus[status_name],
+            visibility=RegistryVisibility[visibility_name],
+            organization_id=descriptor.organization_id or None,
+            owner_id=descriptor.owner_id or None,
+            card_id=descriptor.card_id or None,
+            module_id=descriptor.module_id or None,
+            setup_version_id=descriptor.setup_version_id or None,
+            setup_version=descriptor.setup_version or None,
+            config=dict(descriptor.config) if descriptor.config else None,
+        )
+
     def discover_by_id(self, module_id: str) -> ModuleInfo:
         """Get module info by ID.
 
@@ -304,3 +336,28 @@ class GrpcRegistry(RegistryStrategy, GrpcClientWrapper, GrpcErrorHandlerMixin):
                 extra={"module_id": module_id, "status": status_name},
             )
             return RegistryModuleStatus[status_name]
+
+    def get_setup(self, setup_id: str) -> SetupInfo | None:
+        """Get setup info.
+
+        Args:
+            setup_id: The setup identifier.
+
+        Returns:
+            SetupInfo if successful, None otherwise.
+
+        Raises:
+            RegistryServiceError: If gRPC call fails.
+        """
+        logger.debug("Getting setup", extra={"setup_id": setup_id})
+        with self.handle_grpc_errors("GetSetup", RegistryServiceError):
+            try:
+                response = self.exec_grpc_query(
+                    "GetSetup",
+                    registry_requests_pb2.GetSetupRequest(setup_id=setup_id),
+                )
+            except ServerError as e:
+                msg = f"Failed to get setup '{setup_id}': {e}"
+                logger.error(msg)
+                raise RegistryServiceError(msg) from e
+            return self._proto_to_setup_info(response)

digitalkin/services/registry/registry_strategy.py

@@ -6,6 +6,7 @@ from typing import Any
 from digitalkin.models.services.registry import (
     ModuleInfo,
     RegistryModuleStatus,
+    SetupInfo,
 )
 from digitalkin.services.base_strategy import BaseStrategy
 from digitalkin.services.registry.registry_models import ModuleStatusInfo
@@ -96,3 +97,8 @@ class RegistryStrategy(BaseStrategy, ABC):
             RegistryModuleNotFoundError: If module not found.
         """
         raise NotImplementedError
+
+    @abstractmethod
+    def get_setup(self, setup_id: str) -> SetupInfo | None:
+        """Get setup info."""
+        raise NotImplementedError

digitalkin/utils/__init__.py

@@ -2,6 +2,7 @@
 
 from digitalkin.utils.dynamic_schema import (
     DEFAULT_TIMEOUT,
+    Dynamic,
     DynamicField,
     Fetcher,
     ResolveResult,
@@ -12,9 +13,6 @@ from digitalkin.utils.dynamic_schema import (
     resolve_safe,
 )
 
-# Alias for cleaner API: `Dynamic` is shorter than `DynamicField`
-Dynamic = DynamicField
-
 __all__ = [
     "DEFAULT_TIMEOUT",
     "Dynamic",

digitalkin/utils/dynamic_schema.py

@@ -136,6 +136,10 @@ class DynamicField:
         return hash(tuple(sorted(self.fetchers.keys())))
 
 
+# Alias for cleaner API: `Dynamic` is shorter than `DynamicField`
+Dynamic = DynamicField
+
+
 def get_dynamic_metadata(field_info: FieldInfo) -> DynamicField | None:
     """Extract DynamicField metadata from a FieldInfo's metadata list.
 

{digitalkin-0.3.2.dev15.dist-info → digitalkin-0.3.2.dev17.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.3.2.dev15
+Version: 0.3.2.dev17
 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.dev15.dist-info → digitalkin-0.3.2.dev17.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=UXz39XjMGBN8ihUTeUoa1hmsph_yYsSHufRaJHP4Fws,196
+digitalkin/__version__.py,sha256=Y3kvX8EPStPR7XivYOXc5tCfaRbBYw08EWY8LQ_zjgM,196
 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
@@ -50,21 +50,21 @@ digitalkin/models/core/task_monitor.py,sha256=CW-jydSgXMV464W0pqfar0HpgqlSxqdujm
 digitalkin/models/grpc_servers/__init__.py,sha256=0tA71nPSXgRrh9DoLvx-TSwZXdYIRUEItoadpTL1cTo,42
 digitalkin/models/grpc_servers/models.py,sha256=gRX94eL71a5mLIie-lCOwE7a0As_AuGduxPPzTHbAe4,13797
 digitalkin/models/grpc_servers/types.py,sha256=rQ78s4nAet2jy-NIDj_PUWriT0kuGHr_w6ELjmjgBao,539
-digitalkin/models/module/__init__.py,sha256=N55wan3rAUVPEGLIDjXoFM_-DYY_zxvbQOZHzNDfwoY,751
+digitalkin/models/module/__init__.py,sha256=e2a_AUmobkpyITQKvMkDaDxvb-GOMHhDF9fn0q5_EnQ,959
 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=qpjyMYgTCyQS0lW2RhTKZsoQKwFpv6l08FIM0sqC6DQ,10326
+digitalkin/models/module/module_context.py,sha256=u4yvvtYNwbSJSnV0omLasGN7wRZgaLS7trBBneeT2CI,12178
 digitalkin/models/module/module_types.py,sha256=C9azCNBk76xMa-Mww8_6AiwQR8MLAsEyUOvBYxytovI,739
-digitalkin/models/module/setup_types.py,sha256=XMKyDm-7n_Za9tDBv1AkgXUowlUj46SMxWANa_yg_LU,18311
-digitalkin/models/module/tool_cache.py,sha256=RP3JwASV8dFUZEKudyALbu0_tq81lstuEc4QHMQpmvM,2073
-digitalkin/models/module/tool_reference.py,sha256=6GhLQC4pqCV3b7Nmm0mr0TZC13RDHIUOm1J-pIdWsmo,3479
+digitalkin/models/module/setup_types.py,sha256=namyC-0iA4ryHqKCzjYsVYeFDduwMqYCYRmQSSH8414,19356
+digitalkin/models/module/tool_cache.py,sha256=bWZAjL2UMJIpQw0SeoOGa0S1flU6V8eKUYe09mluNog,7111
+digitalkin/models/module/tool_reference.py,sha256=DBOw22_ZtmOwooejl2J8QZ_gRazNkhGyeD-CvnMpHvE,4388
 digitalkin/models/module/utility.py,sha256=gnbYfWpXGbomUI0fWf7T-Qm_VvT-LXDv1OuA9zObwVg,5589
 digitalkin/models/services/__init__.py,sha256=jhfVw6egq0OcHmos_fypH9XFehbHTBw09wluVFVFEyw,226
 digitalkin/models/services/cost.py,sha256=9PXvd5RrIk9vCrRjcUGQ9ZyAokEbwLg4s0RfnE-aLP4,1616
-digitalkin/models/services/registry.py,sha256=hz_r03-K633XHu2fOb5HWsE59EPFusBBipy0Gz6OBvI,705
+digitalkin/models/services/registry.py,sha256=mFehnPAVLGimodHquNrltXbH_aE0jEa-PxfyNm6J38E,1828
 digitalkin/models/services/storage.py,sha256=wp7F-AvTsU46ujGPcguqM5kUKRZx4399D4EGAAJt2zs,1143
 digitalkin/modules/__init__.py,sha256=vTQk8DWopxQSJ17BjE5dNhq247Rou55iQLJdBxoPUmo,296
-digitalkin/modules/_base_module.py,sha256=4-fS376f9ti3xSgM946yRccHcEGttdo2QYKaxdqk1kY,22451
+digitalkin/modules/_base_module.py,sha256=vQX-vqGXafc9DkNJXzTic-1zJ3yrP1BC30fTH8yHB1k,22592
 digitalkin/modules/archetype_module.py,sha256=XC9tl1Yr6QlbPn_x0eov6UUZwQgwW--BYPPMYVJH_NU,505
 digitalkin/modules/tool_module.py,sha256=GBis7bKCkvWFCYLRvaS9oZVmLBBve1w8BhVnKOU2sCc,506
 digitalkin/modules/trigger_handler.py,sha256=qPNMi-8NHqscOxciHeaXtpwjXApT3YzjMF23zQAjaZY,1770
@@ -90,16 +90,16 @@ digitalkin/services/cost/grpc_cost.py,sha256=2xJ3baTNubCnawpbNftTuFVnPpDZPouSrUW
 digitalkin/services/filesystem/__init__.py,sha256=BhwMl_BUvM0d65fmglkp0SVwn3RfYiUOKJgIMnOCaGM,381
 digitalkin/services/filesystem/default_filesystem.py,sha256=WQbU-Bsi9r-28VqhKbrplce3otzjSKS-5iqKEpGWdQU,15117
 digitalkin/services/filesystem/filesystem_strategy.py,sha256=zibVLvX_IBQ-kgh-KYzHdszDeiHFPEAZszu_k99x1GQ,9487
-digitalkin/services/filesystem/grpc_filesystem.py,sha256=bdjAagsBedY45v41b-9tJh1UWM62vGN9teSVijHq5rU,12373
+digitalkin/services/filesystem/grpc_filesystem.py,sha256=7vUR0n8_1KmG35mXsSWgfFRIGsb6_m9A3NAOIHDfHyw,12384
 digitalkin/services/identity/__init__.py,sha256=InkeyLgFYYwItx8mePA8HpfacOMWZwwuc0G4pWtKq9s,270
 digitalkin/services/identity/default_identity.py,sha256=Y2auZHrGSZTIN5D8HyjLvLcNbYFM1CNUE23x7p5VIGw,386
 digitalkin/services/identity/identity_strategy.py,sha256=skappBbds1_qa0Gr24FGrNX1N0_OYhYT1Lh7dUaAirE,429
 digitalkin/services/registry/__init__.py,sha256=WPGQM3U-QvMXhsaOy9BN0kVMU3QkPFwAMT3lGmTR-Ko,835
 digitalkin/services/registry/default_registry.py,sha256=tOqw9Ve9w_BzhqrZmHuUl5Ps-J_KTEwYg3tu1gNIHmw,4258
 digitalkin/services/registry/exceptions.py,sha256=tAcVXioCzDqfBvxB_P0uQpaK_LDLrFb0KpymROuqs-8,1371
-digitalkin/services/registry/grpc_registry.py,sha256=2_He-I9I4SXdjYA9QYxFWQrr_xBao6LIYXsgiICkeHY,10822
+digitalkin/services/registry/grpc_registry.py,sha256=uBQbETfB3Gheteo2CWZ4l3Ho4QvSOJ0OCoRgzy7ZrVM,13079
 digitalkin/services/registry/registry_models.py,sha256=DJEwMJg5_BewpgHDtY8xIGWj9jA9H07iYgHLCv81giY,331
-digitalkin/services/registry/registry_strategy.py,sha256=N6oLxMwnYkM_XVW7WDxZoSfFafjUtsFDig3F8VQZj7Q,2745
+digitalkin/services/registry/registry_strategy.py,sha256=oxCm5mQcO1PSwMwOtEv8dwzQx_F0uGKgiWeVMeiV51s,2905
 digitalkin/services/setup/__init__.py,sha256=t6xcvEWqTbcRZstBFK9cESEqaZKvpW14VtYygxIqfYQ,65
 digitalkin/services/setup/default_setup.py,sha256=zTG9Nvw3Tqy9qctSEnjWE4Mlqol321AUaTlIm6Jl1k0,8195
 digitalkin/services/setup/grpc_setup.py,sha256=aEAHiG0WSzNSQtJjn7xAGQwj70b0HrFO8GBKFXRx3aI,13925
@@ -115,14 +115,14 @@ digitalkin/services/user_profile/__init__.py,sha256=RKEZCsgCHS7fmswhWgUoQd6vZ_1p
 digitalkin/services/user_profile/default_user_profile.py,sha256=46DH_VBCHKXJVyagVcc8kH5sLwRK54Fe_0ahqYJ1maA,1847
 digitalkin/services/user_profile/grpc_user_profile.py,sha256=xDiUC5Ceofa6QtGPmqJV3ik5j8HDHc1zxtpia49rlRw,2780
 digitalkin/services/user_profile/user_profile_strategy.py,sha256=CH8kT__1MUwA21k5djjmB5ZZ6pYg57OWbe_7owBCgwU,681
-digitalkin/utils/__init__.py,sha256=gDoEk5CP8tiBmIq-38aPhIMlnDUiqRrduuVXhpSPQnc,541
+digitalkin/utils/__init__.py,sha256=xusG_EoixTznt9ENCB8XkNbi67SZNt6y9hx5gsPWH_k,464
 digitalkin/utils/arg_parser.py,sha256=wzscRlE1Qp1gGl-lAJlkkwnbU1O2oezj6BwK_BZFBIk,3158
 digitalkin/utils/development_mode_action.py,sha256=2hznh0ajW_4ZTysfoc0Y49161f_PQPATRgNk8NAn1_o,1623
-digitalkin/utils/dynamic_schema.py,sha256=5-B3dBGlCYYv6uRJkgudtc0ZpBOTYxl0yKedDGsteZQ,15184
+digitalkin/utils/dynamic_schema.py,sha256=y5csxjuqVHjWDpnTUzxbcUuI_wou9-ibRVHQlBs_btY,15275
 digitalkin/utils/llm_ready_schema.py,sha256=JjMug_lrQllqFoanaC091VgOqwAd-_YzcpqFlS7p778,2375
 digitalkin/utils/package_discover.py,sha256=sa6Zp5Kape1Zr4iYiNrnZxiHDnqM06ODk6yfWHom53w,13465
 digitalkin/utils/schema_splitter.py,sha256=KMvYRHDHlwdhh_c6FJxkWvLStZo9Kbj-jd3pIGPZfxk,9317
-digitalkin-0.3.2.dev15.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
+digitalkin-0.3.2.dev17.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
 modules/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 modules/archetype_with_tools_module.py,sha256=PXTS6IXmC_OjxTmVrL_pYVI0MKwXjD5I1UJO_2xa10Q,7632
 modules/cpu_intensive_module.py,sha256=GZlirQDZdYuXrI46sv1q4RNAHZjL4EptHVQTvgK9zz8,8363
@@ -137,7 +137,7 @@ monitoring/digitalkin_observability/prometheus.py,sha256=gDmM9ySaVwPAe7Yg84pLxmE
 monitoring/tests/test_metrics.py,sha256=ugnYfAwqBPO6zA8z4afKTlyBWECTivacYSN-URQCn2E,5856
 services/filesystem_module.py,sha256=U4dgqtuDadaXz8PJ1d_uQ_1EPncBqudAQCLUICF9yL4,7421
 services/storage_module.py,sha256=Wz2MzLvqs2D_bnBBgtnujYcAKK2V2KFMk8K21RoepSE,6972
-digitalkin-0.3.2.dev15.dist-info/METADATA,sha256=XYJQ68LW34K4PFvn5wjmE-b1N8fa0ZI9WVOZWc_4cWs,29725
-digitalkin-0.3.2.dev15.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-digitalkin-0.3.2.dev15.dist-info/top_level.txt,sha256=AYVIesKrO0jnedQ-Muog9JBehG81WeTCNeOFoJgwsgE,51
-digitalkin-0.3.2.dev15.dist-info/RECORD,,
+digitalkin-0.3.2.dev17.dist-info/METADATA,sha256=rhqhGY-Zx3iq_pSrr0nRpsygx7UhYYsVLjUoVDzXMq8,29725
+digitalkin-0.3.2.dev17.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+digitalkin-0.3.2.dev17.dist-info/top_level.txt,sha256=AYVIesKrO0jnedQ-Muog9JBehG81WeTCNeOFoJgwsgE,51
+digitalkin-0.3.2.dev17.dist-info/RECORD,,