digitalkin 0.3.2.dev20__py3-none-any.whl → 0.3.2.dev22__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.dev20"
+    __version__ = "0.3.2.dev22"

digitalkin/grpc_servers/module_servicer.py

@@ -165,7 +165,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
         if isinstance(updated_setup_data, ModuleCodeModel):
             logger.error(
                 "Config setup failed",
-                extra={"job_id": job_id, "code": updated_setup_data.code, "message": updated_setup_data.message},
+                extra={"job_id": job_id, "code": updated_setup_data.code, "error_message": updated_setup_data.message},
             )
             context.set_code(grpc.StatusCode.INTERNAL)
             context.set_details(updated_setup_data.message or "Config setup failed")
@@ -178,7 +178,7 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
                 extra={
                     "job_id": job_id,
                     "code": updated_setup_data["code"],
-                    "message": updated_setup_data.get("message"),
+                    "error_message": updated_setup_data.get("message"),
                 },
             )
             context.set_code(grpc.StatusCode.INTERNAL)

digitalkin/models/module/__init__.py

@@ -14,8 +14,7 @@ from digitalkin.models.module.tool_cache import (
 )
 from digitalkin.models.module.tool_reference import (
     ToolReference,
-    ToolReferenceConfig,
-    ToolSelectionMode,
+    ToolReferenceInput,
 )
 from digitalkin.models.module.utility import (
     EndOfStreamOutput,
@@ -36,8 +35,7 @@ __all__ = [
     "ToolModuleInfo",
     "ToolParameter",
     "ToolReference",
-    "ToolReferenceConfig",
-    "ToolSelectionMode",
+    "ToolReferenceInput",
     "UtilityProtocol",
     "UtilityRegistry",
 ]

digitalkin/models/module/module_context.py

@@ -227,21 +227,19 @@ class ModuleContext:
             llm_format=llm_format,
         )
 
-    async def create_openai_style_tools(self, module_id: str) -> list[dict[str, Any]]:
+    def create_openai_style_tools(self, slug: str) -> list[dict[str, Any]]:
         """Create OpenAI-style function calling schemas for a tool module.
 
         Uses tool cache (fast path) with registry fallback. Returns one schema
         per ToolDefinition (protocol) in the module.
 
         Args:
-            module_id: Module ID to look up (checks cache first, then registry).
+            slug: Module ID to look up (checks cache first, then registry).
 
         Returns:
             List of OpenAI-style tool schemas, one per protocol. Empty if not found.
         """
-        tool_module_info = await self.tool_cache.get(
-            module_id, registry=self.registry, communication=self.communication
-        )
+        tool_module_info = self.tool_cache.get(slug)
         if not tool_module_info:
             return []
 
@@ -250,8 +248,8 @@ class ModuleContext:
                 "type": "function",
                 "function": {
                     "module_id": tool_module_info.module_id,
-                    "toolkit_name": tool_module_info.name or "undefined",
-                    "name": tool_def.name,
+                    "toolkit_name": tool_module_info.tool_name or "undefined",
+                    "name": tool_module_info.slug + "_" + tool_def.name,
                     "description": tool_def.description,
                     "parameters": ModuleContext._build_parameters_schema(tool_def.parameters),
                 },
@@ -277,9 +275,9 @@ class ModuleContext:
 
     def create_tool_functions(
         self,
-        module_id: str,
+        slug: str,
     ) -> list[tuple[ToolDefinition, Callable[..., AsyncGenerator[dict, None]]]]:
-        """Create tool functions for all protocols in a tool module.
+        """Create tool functions for all protocols in a tool setup.
 
         Returns an async generator per ToolDefinition that calls the remote tool
         module via gRPC with the protocol auto-injected.
@@ -288,12 +286,12 @@ class ModuleContext:
         in sync contexts like __init__ methods.
 
         Args:
-            module_id: Module ID to look up in cache.
+            slug: Setup ID to look up in cache.
 
         Returns:
             List of (ToolDefinition, async_generator_function) tuples. Empty if not found.
         """
-        tool_module_info = self.tool_cache.entries.get(module_id)
+        tool_module_info = self.tool_cache.entries.get(slug)
         if not tool_module_info:
             return []
 

digitalkin/models/module/setup_types.py

@@ -29,7 +29,6 @@ SetupModelT = TypeVar("SetupModelT", bound="SetupModel")
 class SetupModel(BaseModel, Generic[SetupModelT]):
     """Base setup model with dynamic schema and tool cache support."""
 
-    model_config = ConfigDict(extra="allow")
     _clean_model_cache: ClassVar[dict[tuple[type, bool, bool], type]] = {}
     resolved_tools: dict[str, ToolModuleInfo] = Field(
         default_factory=dict,
@@ -347,7 +346,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             resolved_tools: Cache of already resolved tools.
         """
         if isinstance(field_value, ToolReference):
-            await cls._resolve_single_tool_reference(
+            await cls._resolve_tool_reference(
                 field_name,
                 field_value,
                 registry,
@@ -377,7 +376,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             )
 
     @classmethod
-    async def _resolve_single_tool_reference(
+    async def _resolve_tool_reference(
         cls,
         field_name: str,
         tool_ref: ToolReference,
@@ -385,7 +384,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         communication: "CommunicationStrategy",
         resolved_tools: dict[str, ToolModuleInfo],
     ) -> None:
-        """Resolve a single ToolReference.
+        """Resolve a ToolReference (may contain multiple selected tools).
 
         Args:
             field_name: Name of the field for logging.
@@ -394,21 +393,23 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             communication: Communication service for module schemas.
             resolved_tools: Cache of already resolved tools.
         """
-        logger.info("Resolving ToolReference '%s' with setup_id='%s'", field_name, tool_ref.config.setup_id)
+        logger.info("Resolving ToolReference '%s' with %d selected tools", field_name, len(tool_ref.selected_tools))
 
-        slug = tool_ref.slug
-        if slug:
-            cached = resolved_tools.get(slug)
-            if cached:
-                tool_ref._cached_info = cached  # noqa: SLF001
-                logger.info("ToolReference '%s' resolved from cache -> %s", field_name, cached)
-                return
+        if not tool_ref.selected_tools:
+            logger.info("ToolReference '%s' has no selected tools, skipping", field_name)
+            return
+
+        tools_to_resolve = [tool for tool in tool_ref.selected_tools if tool.slug and tool.slug not in resolved_tools]
+
+        if not tools_to_resolve:
+            logger.info("All tools for '%s' already cached", field_name)
+            return
 
         try:
-            info = await tool_ref.resolve(registry, communication)
-            if info and info.setup_id:
+            infos = await tool_ref.resolve(registry, communication)
+            for info in infos:
                 resolved_tools[info.setup_id] = info
-            logger.info("Resolved ToolReference '%s' -> %s", field_name, tool_ref.tool_module_info)
+                logger.info("Resolved tool '%s' -> module_id=%s", info.setup_id, info.module_id)
         except Exception:
             logger.exception("Failed to resolve ToolReference '%s'", field_name)
 
@@ -430,7 +431,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         """
         for item in items:
             if isinstance(item, ToolReference):
-                await cls._resolve_single_tool_reference(
+                await cls._resolve_tool_reference(
                     "list_item",
                     item,
                     registry,
@@ -463,7 +464,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         """
         for item in mapping.values():
             if isinstance(item, ToolReference):
-                await cls._resolve_single_tool_reference(
+                await cls._resolve_tool_reference(
                     "dict_value",
                     item,
                     registry,
@@ -500,10 +501,10 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             if field_value is None:
                 continue
             if isinstance(field_value, ToolReference):
-                module_info = self.resolved_tools.get(field_value.slug or "") or field_value.tool_module_info
-                if module_info:
-                    self.resolved_tools[module_info.setup_id] = module_info
-                    cache.add(module_info.module_id, module_info)
+                for tool in field_value.selected_tools:
+                    tool_module_info = self.resolved_tools.get(tool.slug or "")
+                    if tool_module_info:
+                        cache.add(tool_module_info)
             elif isinstance(field_value, BaseModel):
                 self._build_tool_cache_recursive(field_value, cache)
             elif isinstance(field_value, list):
@@ -520,10 +521,10 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         """
         for item in items:
             if isinstance(item, ToolReference):
-                module_info = self.resolved_tools.get(item.slug or "") or item.tool_module_info
-                if module_info:
-                    self.resolved_tools[module_info.setup_id] = module_info
-                    cache.add(module_info.module_id, module_info)
+                for tool in item.selected_tools:
+                    tool_module_info = self.resolved_tools.get(tool.slug or "")
+                    if tool_module_info:
+                        cache.add(tool_module_info)
             elif isinstance(item, BaseModel):
                 self._build_tool_cache_recursive(item, cache)
 
@@ -536,9 +537,9 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         """
         for item in mapping.values():
             if isinstance(item, ToolReference):
-                module_info = self.resolved_tools.get(item.slug or "") or item.tool_module_info
-                if module_info:
-                    self.resolved_tools[module_info.setup_id] = module_info
-                    cache.add(module_info.module_id, module_info)
+                for tool in item.selected_tools:
+                    tool_module_info = self.resolved_tools.get(tool.slug or "")
+                    if tool_module_info:
+                        cache.add(tool_module_info)
             elif isinstance(item, BaseModel):
                 self._build_tool_cache_recursive(item, cache)

digitalkin/models/module/tool_cache.py

@@ -6,7 +6,16 @@ from pydantic import BaseModel, Field
 
 from digitalkin.logger import logger
 from digitalkin.models.services.registry import ModuleInfo
-from digitalkin.services.registry import RegistryStrategy
+
+
+class SelectedTool(BaseModel):
+    """Selected tool information."""
+
+    setup_id: str = ""
+    module_id: str = ""
+    slug: str = ""
+    name: str = ""
+
 
 if TYPE_CHECKING:
     from digitalkin.services.communication import CommunicationStrategy
@@ -51,8 +60,14 @@ class ToolDefinition(BaseModel):
 class ToolModuleInfo(ModuleInfo):
     """Module info for tool modules."""
 
-    tools: list[ToolDefinition]
-    setup_id: str
+    tools: list[ToolDefinition] = Field(default_factory=list)
+    tool_name: str = ""
+    setup_id: str = ""
+
+    @property
+    def slug(self) -> str:
+        """Module ID."""
+        return self.setup_id + "_" + self.tool_name
 
 
 class ToolCache(BaseModel):
@@ -60,53 +75,35 @@ class ToolCache(BaseModel):
 
     entries: dict[str, ToolModuleInfo] = Field(default_factory=dict)
 
-    def add(self, setup_id: str, tool_module_info: ToolModuleInfo) -> None:
+    def add(self, 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
+        self.entries[tool_module_info.slug] = tool_module_info
         logger.debug(
             "Tool cached",
-            extra={"setup_id": setup_id, "module_id": tool_module_info.module_id},
+            extra={
+                "slug": tool_module_info.slug,
+                "module_id": tool_module_info.module_id,
+                "setup_id": tool_module_info.setup_id,
+            },
         )
 
-    async def get(
+    def get(
         self,
-        setup_id: str,
-        *,
-        registry: RegistryStrategy | None = None,
-        communication: "CommunicationStrategy | None" = None,
+        slug: str,
     ) -> 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.
+            slug: Field name to look up.
 
         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
+        return self.entries.get(slug)
 
     def clear(self) -> None:
         """Clear all cache entries."""
@@ -123,7 +120,7 @@ class ToolCache(BaseModel):
 
 async def module_info_to_tool_module_info(
     module_info: ModuleInfo,
-    setup_id: str,
+    tool: SelectedTool,
     communication: "CommunicationStrategy",
     *,
     llm_format: bool = True,
@@ -135,7 +132,7 @@ async def module_info_to_tool_module_info(
 
     Args:
         module_info: Module info from registry.
-        setup_id: Setup ID from tool configuration.
+        tool: Selected tool information.
         communication: Communication strategy for gRPC calls.
         llm_format: Use LLM-friendly schema format.
 
@@ -160,11 +157,12 @@ async def module_info_to_tool_module_info(
         address=module_info.address,
         port=module_info.port,
         version=module_info.version,
-        name=module_info.name,
+        module_name=module_info.module_name,
         documentation=module_info.documentation,
         status=module_info.status,
         tools=tools,
-        setup_id=setup_id,
+        setup_id=tool.setup_id,
+        tool_name=tool.name,
     )
 
 

digitalkin/models/module/tool_reference.py

@@ -1,136 +1,103 @@
 """Tool reference types for module configuration."""
 
-from enum import Enum
+from typing import Annotated
 
-from pydantic import BaseModel, Field, PrivateAttr, model_validator
+from pydantic import BaseModel, BeforeValidator, Field
+from pydantic.json_schema import GetJsonSchemaHandler, JsonSchemaValue
+from pydantic_core import CoreSchema
 
-from digitalkin.models.module.tool_cache import ToolModuleInfo, module_info_to_tool_module_info
+from digitalkin.models.module.tool_cache import (
+    SelectedTool,
+    ToolModuleInfo,
+    module_info_to_tool_module_info,
+)
 from digitalkin.services.communication.communication_strategy import CommunicationStrategy
 from digitalkin.services.registry import RegistryStrategy
 
 
-class ToolSelectionMode(str, Enum):
-    """Tool selection mode."""
-
-    TAG = "tag"
-    FIXED = "fixed"
-    DISCOVERABLE = "discoverable"
-
-
-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="")
-
-    @model_validator(mode="after")
-    def validate_config(self) -> "ToolReferenceConfig":
-        """Validate required fields based on mode.
-
-        Returns:
-            Self if validation passes.
-
-        Raises:
-            ValueError: If required field is missing for the mode.
-        """
-        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"
-            raise ValueError(msg)
-        return self
-
-
 class ToolReference(BaseModel):
-    """Reference to a tool module, resolved via registry during config setup."""
-
-    config: ToolReferenceConfig
-    _cached_info: ToolModuleInfo | None = PrivateAttr(default=None)
+    """Tool selection configuration and reference.
+
+    The mode determines validation requirements and resolution behavior:
+    - FIXED: Requires setup_id, resolves to exact tool
+    - MODULE: Requires module_id, returns constraint for frontend selection
+    - TAG: Requires tag, returns constraint for frontend selection
+    - DISCOVERABLE: Optional module_id/tag constraints, returns constraint info
+    """
+
+    selected_tools: list[SelectedTool] = Field(default=[], description="Tools selected by the user.")
+    setup_ids: list[str] = Field(default=[], description="Setup IDs for the user to choose from.")
+    module_ids: list[str] = Field(default=[], description="Module IDs for the user to choose from.")
+    tags: list[str] = Field(default=[], description="Tags for the user to choose from.")
+    max_tools: int = Field(default=0, description="Maximum tools to select. 0 for unlimited.")
+
+    async def resolve(self, registry: RegistryStrategy, communication: CommunicationStrategy) -> list[ToolModuleInfo]:
+        """Resolve this reference using the registry.
 
-    @property
-    def slug(self) -> str:
-        """Cache key (same as module_id).
+        Args:
+            registry: Registry service for module discovery.
+            communication: Communication service for module schemas.
 
         Returns:
-            Module ID used as cache key.
+            List of ToolModuleInfo if resolved.
         """
-        return self.config.setup_id
-
-    @property
-    def module_id(self) -> str:
-        """Module identifier.
+        resolved: list[ToolModuleInfo] = []
+        for tool in self.selected_tools:
+            setup = registry.get_setup(tool.setup_id)
+            if setup and setup.module_id:
+                info = registry.discover_by_id(setup.module_id)
+                tool.slug = tool.setup_id
+                tool.module_id = setup.module_id
+                tool.name = setup.name
+                if info:
+                    resolved.append(await module_info_to_tool_module_info(info, tool, communication))
 
-        Returns:
-            Module ID or empty string if not set.
-        """
-        return self.config.module_id
+        return resolved
 
-    @property
-    def setup_id(self) -> str:
-        """Setup identifier.
 
-        Returns:
-            Setup ID or empty string if not set.
-        """
-        return self.config.setup_id
+def _convert_to_tool_reference(v: object) -> "ToolReference | object":
+    """Convert list of setup IDs to ToolReference.
 
-    @property
-    def tool_module_info(self) -> ToolModuleInfo | None:
-        """Resolved module information.
+    Args:
+        v: Input value, either a list of setup IDs or passthrough.
 
-        Returns:
-            ToolModuleInfo if resolved, None otherwise.
-        """
-        return self._cached_info
+    Returns:
+        ToolReference if input is list, otherwise original value.
+    """
+    if isinstance(v, list):
+        return ToolReference(selected_tools=[SelectedTool(setup_id=sid, slug=sid) for sid in v])
+    return v
 
-    @property
-    def is_resolved(self) -> bool:
-        """Whether this reference has been resolved.
 
-        Returns:
-            True if resolved, False otherwise.
-        """
-        return self._cached_info is not None
+class _ToolReferenceInputSchema:
+    """Custom JSON schema generator that wraps ToolReference in anyOf with array option."""
 
-    async def resolve(self, registry: RegistryStrategy, communication: CommunicationStrategy) -> ToolModuleInfo | None:
-        """Resolve this reference using the registry.
+    @staticmethod
+    def __get_pydantic_json_schema__(  # noqa: PLW3201
+        schema: CoreSchema,
+        handler: GetJsonSchemaHandler,
+    ) -> JsonSchemaValue:
+        """Generate JSON schema accepting both list[str] and ToolReference.
 
         Args:
-            registry: Registry service for module discovery.
-            communication: Communication service for module schemas.
+            schema: The core schema from Pydantic.
+            handler: Handler to generate JSON schema from core schema.
 
         Returns:
-            ToolModuleInfo if resolved, None for DISCOVERABLE mode or if not found.
+            JSON schema with anyOf accepting array or ToolReference.
         """
-        if self.config.mode == ToolSelectionMode.DISCOVERABLE:
-            return None
-
-        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, self.config.setup_id, communication)
-                    self._cached_info = tool_module_info
-                    return tool_module_info
-
-        if self.config.mode == ToolSelectionMode.TAG and self.config.tag:
-            results = registry.search(
-                name=self.config.tag,
-                module_type="tool",
-                organization_id=self.config.organization_id,
-            )
-            if results:
-                tool_module_info = await module_info_to_tool_module_info(
-                    results[0], self.config.setup_id, communication
-                )
-                self._cached_info = tool_module_info
-                self.config.module_id = tool_module_info.module_id
-                return tool_module_info
-
-        return None
+        json_schema = handler(schema)
+        return {
+            "anyOf": [
+                {"type": "array", "items": {"type": "string"}},
+                json_schema,
+            ]
+        }
+
+
+ToolReferenceInput = Annotated[
+    ToolReference,
+    BeforeValidator(_convert_to_tool_reference),
+    _ToolReferenceInputSchema,
+]
+"""Type alias for ToolReference fields that accept list[str] input from frontend."""

digitalkin/models/services/registry.py

@@ -26,12 +26,12 @@ class RegistryModuleType(str, Enum):
 class ModuleInfo(BaseModel):
     """Module information from registry."""
 
-    module_id: str
-    module_type: RegistryModuleType
-    address: str
-    port: int
-    version: str
-    name: str = ""
+    module_id: str = ""
+    module_type: RegistryModuleType = RegistryModuleType.UNSPECIFIED
+    address: str = ""
+    port: int = 0
+    version: str = ""
+    module_name: str = ""
     documentation: str | None = None
     status: RegistryModuleStatus | None = None
 

digitalkin/services/registry/grpc_registry.py

@@ -74,7 +74,7 @@ class GrpcRegistry(RegistryStrategy, GrpcClientWrapper, GrpcErrorHandlerMixin):
             address=descriptor.address,
             port=descriptor.port,
             version=descriptor.version,
-            name=descriptor.name,
+            module_name=descriptor.name,
             documentation=descriptor.documentation or None,
         )
 

digitalkin/utils/schema_splitter.py

@@ -17,24 +17,31 @@ class SchemaSplitter:
             Tuple of (jsonschema, uischema).
         """
         defs_ui: dict[str, dict[str, Any]] = {}
-        if "$defs" in combined_schema:
-            for def_name, def_value in combined_schema["$defs"].items():
+        schema_defs = combined_schema.get("$defs", {})
+        if schema_defs:
+            for def_name, def_value in schema_defs.items():
                 if isinstance(def_value, dict):
                     defs_ui[def_name] = {}
-                    cls._extract_ui_properties(def_value, defs_ui[def_name])
+                    cls._extract_ui_properties(def_value, defs_ui[def_name], schema_defs)
 
         json_schema: dict[str, Any] = {}
         ui_schema: dict[str, Any] = {}
-        cls._process_object(combined_schema, json_schema, ui_schema, defs_ui)
+        cls._process_object(combined_schema, json_schema, ui_schema, defs_ui, schema_defs)
         return json_schema, ui_schema
 
     @classmethod
-    def _extract_ui_properties(cls, source: dict[str, Any], ui_target: dict[str, Any]) -> None:  # noqa: C901
+    def _extract_ui_properties(
+        cls,
+        source: dict[str, Any],
+        ui_target: dict[str, Any],
+        defs: dict[str, Any] | None = None,
+    ) -> None:
         """Extract ui:* properties from source into ui_target recursively.
 
         Args:
             source: Source dict to extract from.
             ui_target: Target dict for ui properties.
+            defs: The $defs dictionary to resolve $ref references.
         """
         for key, value in source.items():
             if key.startswith("ui:"):
@@ -43,20 +50,26 @@ class SchemaSplitter:
                 for prop_name, prop_value in value.items():
                     if isinstance(prop_value, dict):
                         prop_ui: dict[str, Any] = {}
-                        cls._extract_ui_properties(prop_value, prop_ui)
+                        cls._extract_ui_properties(prop_value, prop_ui, defs)
                         if prop_ui:
                             ui_target[prop_name] = prop_ui
             elif key == "items" and isinstance(value, dict):
                 items_ui: dict[str, Any] = {}
-                cls._extract_ui_properties(value, items_ui)
+                cls._extract_ui_properties(value, items_ui, defs)
                 if items_ui:
                     ui_target["items"] = items_ui
             elif key == "allOf" and isinstance(value, list):
                 for item in value:
                     if isinstance(item, dict):
-                        cls._extract_ui_properties(item, ui_target)
+                        cls._extract_ui_properties(item, ui_target, defs)
             elif key in {"if", "then", "else"} and isinstance(value, dict):
-                cls._extract_ui_properties(value, ui_target)
+                cls._extract_ui_properties(value, ui_target, defs)
+            elif key == "$ref" and isinstance(value, str) and defs is not None:
+                # Resolve $ref and extract UI properties from the referenced definition
+                if value.startswith("#/$defs/"):
+                    def_name = value[8:]
+                    if def_name in defs:
+                        cls._extract_ui_properties(defs[def_name], ui_target, defs)
 
     @classmethod
     def _process_object(  # noqa: C901, PLR0912
@@ -65,6 +78,7 @@ class SchemaSplitter:
         json_target: dict[str, Any],
         ui_target: dict[str, Any],
         defs_ui: dict[str, dict[str, Any]],
+        schema_defs: dict[str, Any] | None = None,
     ) -> None:
         """Process an object, splitting json and ui properties.
 
@@ -73,6 +87,7 @@ class SchemaSplitter:
             json_target: Target dict for json schema.
             ui_target: Target dict for ui schema.
             defs_ui: Pre-extracted UI properties from $defs.
+            schema_defs: The original $defs dictionary to resolve $ref in _extract_ui_properties.
         """
         for key, value in source.items():
             if key.startswith("ui:"):
@@ -113,14 +128,14 @@ class SchemaSplitter:
                         cls._strip_ui_properties(item, item_json)
                         json_target["allOf"].append(item_json)
                         # Extract UI properties from allOf item
-                        cls._extract_ui_properties(item, ui_target)
+                        cls._extract_ui_properties(item, ui_target, schema_defs)
                     else:
                         json_target["allOf"].append(item)
             elif key in {"if", "then", "else"} and isinstance(value, dict):
                 json_target[key] = {}
                 cls._strip_ui_properties(value, json_target[key])
                 # Extract UI properties from conditional
-                cls._extract_ui_properties(value, ui_target)
+                cls._extract_ui_properties(value, ui_target, schema_defs)
             elif key == "hidden":
                 # Strip hidden key from json schema
                 continue

{digitalkin-0.3.2.dev20.dist-info → digitalkin-0.3.2.dev22.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.3.2.dev20
+Version: 0.3.2.dev22
 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.dev20.dist-info → digitalkin-0.3.2.dev22.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=dN3fPmuqWC0wGuJWF7p4H8rHGbT_gGfHvNrei8I6TU8,196
+digitalkin/__version__.py,sha256=dZgyJRHQFh_03su8ZElZY21agbb3n1gwA2QNWlIpunk,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
@@ -28,7 +28,7 @@ digitalkin/core/task_manager/task_session.py,sha256=5jw21bT_SPXUzWE7tk6YG62EXqlR
 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=Ec3izzV2YpdN8rGs_cX-iVulQ00FkLR5dBflHlQ8a6Y,7849
-digitalkin/grpc_servers/module_servicer.py,sha256=KsWXPwCQ0s2pygn0cq89KndJleAkjvkvLFs1Z2fWQIk,21705
+digitalkin/grpc_servers/module_servicer.py,sha256=yhS5zOxxOTh3aXNVZBbgomRtFa5k6Q7jm-PZh54Gus8,21717
 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=nGG8QdKnBH0UG9qbKrlPwIvcvPgW3osw7O3cImxisPE,3279
@@ -50,18 +50,18 @@ 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=e2a_AUmobkpyITQKvMkDaDxvb-GOMHhDF9fn0q5_EnQ,959
+digitalkin/models/module/__init__.py,sha256=XlFlLGmLUYPV4ZOgBxfGclq9wupGDXNUyWB4EB6iTXA,909
 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=QDdjZdhIJpvU_2Tn7kkJsZ1givB4dM1-ksopdF4VySw,12176
+digitalkin/models/module/module_context.py,sha256=kwaAobsA82Du6L0XH9Is9u6Qj1rEjhYqVk00FonMaw4,12087
 digitalkin/models/module/module_types.py,sha256=C9azCNBk76xMa-Mww8_6AiwQR8MLAsEyUOvBYxytovI,739
-digitalkin/models/module/setup_types.py,sha256=Xb_KZ5vKcpLqkbPdQfooqtI6TFp2cklvdXLUBa55c30,19346
-digitalkin/models/module/tool_cache.py,sha256=5e30A_GxT2W-w1LZFmVUqOxDjPcrZ8s_eW7p9impO64,7153
-digitalkin/models/module/tool_reference.py,sha256=eIWJrT6syyEaXAWRXIlWYTst-j0XuvtU_va9m3tj_KU,4470
+digitalkin/models/module/setup_types.py,sha256=8xEI6DvCYBHSOsuoeuZjtBL5rPsR_7DwWFxQnn2bo3c,19249
+digitalkin/models/module/tool_cache.py,sha256=xAMSyY73aduYGPz6C54g0YMAa2OnXPN6QRS-W4TK_js,6568
+digitalkin/models/module/tool_reference.py,sha256=rdkke5ydAgXcFfu34yxiHyaiRGcclnNGwPJK-01B2t8,3829
 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=mFehnPAVLGimodHquNrltXbH_aE0jEa-PxfyNm6J38E,1828
+digitalkin/models/services/registry.py,sha256=bCdVpsiu8bqKL-7P5r3OPmllZj4MAtpsfsdlsg5TAww,1887
 digitalkin/models/services/storage.py,sha256=wp7F-AvTsU46ujGPcguqM5kUKRZx4399D4EGAAJt2zs,1143
 digitalkin/modules/__init__.py,sha256=vTQk8DWopxQSJ17BjE5dNhq247Rou55iQLJdBxoPUmo,296
 digitalkin/modules/_base_module.py,sha256=0XC0aQAxlNfvz0KK9ut7K0JbZql3cZMU4aeg7ISEsD0,21971
@@ -97,7 +97,7 @@ digitalkin/services/identity/identity_strategy.py,sha256=skappBbds1_qa0Gr24FGrNX
 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=uBQbETfB3Gheteo2CWZ4l3Ho4QvSOJ0OCoRgzy7ZrVM,13079
+digitalkin/services/registry/grpc_registry.py,sha256=Myrt2WqmRQG1yLu2voUOLH4vp3Kwf_WitOBG3YwhtB4,13086
 digitalkin/services/registry/registry_models.py,sha256=DJEwMJg5_BewpgHDtY8xIGWj9jA9H07iYgHLCv81giY,331
 digitalkin/services/registry/registry_strategy.py,sha256=oxCm5mQcO1PSwMwOtEv8dwzQx_F0uGKgiWeVMeiV51s,2905
 digitalkin/services/setup/__init__.py,sha256=t6xcvEWqTbcRZstBFK9cESEqaZKvpW14VtYygxIqfYQ,65
@@ -122,10 +122,10 @@ digitalkin/utils/development_mode_action.py,sha256=2hznh0ajW_4ZTysfoc0Y49161f_PQ
 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=7KFiRUhFE9HsH_Z8p0POSbIhiJlgZEeSxL_7TDY0n1U,10374
-digitalkin-0.3.2.dev20.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
+digitalkin/utils/schema_splitter.py,sha256=6PeDrEBGEFyVEPgzGnwcCZuwxLobThmuPwJhGlRpU_0,11137
+digitalkin-0.3.2.dev22.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/archetype_with_tools_module.py,sha256=kJkVhAFWG0aDDqzupOXOnV3l8j3z5bEdWos_6Z9rUP8,7303
 modules/cpu_intensive_module.py,sha256=GZlirQDZdYuXrI46sv1q4RNAHZjL4EptHVQTvgK9zz8,8363
 modules/dynamic_setup_module.py,sha256=tKvUWZdlYZkfAgKR0mLuFcLiFGKpVgpsz10LeJ6B2QI,11410
 modules/minimal_llm_module.py,sha256=N9aIzZQI-miyH4AB4xTmGHpMvdSLnYyXNOD4Z3YFzis,11216
@@ -138,7 +138,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.dev20.dist-info/METADATA,sha256=EosqD8N82EhYPNyjWKp-_n2az8875Qp8iRfdZph5uB4,29725
-digitalkin-0.3.2.dev20.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
-digitalkin-0.3.2.dev20.dist-info/top_level.txt,sha256=AYVIesKrO0jnedQ-Muog9JBehG81WeTCNeOFoJgwsgE,51
-digitalkin-0.3.2.dev20.dist-info/RECORD,,
+digitalkin-0.3.2.dev22.dist-info/METADATA,sha256=6Zn0jPPta8ZRH7ohcui_18g4PsaoCBO6bSi59U4kg9o,29725
+digitalkin-0.3.2.dev22.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+digitalkin-0.3.2.dev22.dist-info/top_level.txt,sha256=AYVIesKrO0jnedQ-Muog9JBehG81WeTCNeOFoJgwsgE,51
+digitalkin-0.3.2.dev22.dist-info/RECORD,,

modules/archetype_with_tools_module.py

@@ -8,11 +8,7 @@ from pydantic import BaseModel, Field
 from digitalkin.models.grpc_servers.models import ClientConfig, SecurityMode, ServerMode
 from digitalkin.models.module.module_context import ModuleContext
 from digitalkin.models.module.setup_types import SetupModel
-from digitalkin.models.module.tool_reference import (
-    ToolReference,
-    ToolReferenceConfig,
-    ToolSelectionMode,
-)
+from digitalkin.models.module.tool_reference import ToolReference
 from digitalkin.modules._base_module import BaseModule  # noqa: PLC2701
 from digitalkin.services.services_models import ServicesStrategy
 
@@ -64,29 +60,21 @@ class ArchetypeSetup(SetupModel):
 
     search_tool: ToolReference = Field(
         default_factory=lambda: ToolReference(
-            config=ToolReferenceConfig(
-                mode=ToolSelectionMode.FIXED,
-                module_id="search-tool-v1",
-            )
+            module_ids=["search-tool-v1"],
         ),
         json_schema_extra={"config": True},
     )
 
     calculator_tool: ToolReference = Field(
         default_factory=lambda: ToolReference(
-            config=ToolReferenceConfig(
-                mode=ToolSelectionMode.TAG,
-                tag="math-calculator",
-            )
+            tags=["math-calculator"],
         ),
         json_schema_extra={"config": True},
     )
 
     dynamic_tool: ToolReference = Field(
         default_factory=lambda: ToolReference(
-            config=ToolReferenceConfig(
-                mode=ToolSelectionMode.DISCOVERABLE,
-            )
+            tags=["discoverable"],
         ),
         json_schema_extra={"config": True},
     )