digitalkin 0.3.2.dev17__py3-none-any.whl → 0.3.2.dev19__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.dev17"
+    __version__ = "0.3.2.dev19"

digitalkin/grpc_servers/module_servicer.py

@@ -17,7 +17,7 @@ from digitalkin.core.job_manager.base_job_manager import BaseJobManager
 from digitalkin.grpc_servers.utils.exceptions import ServicerError
 from digitalkin.logger import logger
 from digitalkin.models.core.job_manager_models import JobManagerMode
-from digitalkin.models.module.module import ModuleStatus
+from digitalkin.models.module.module import ModuleCodeModel, ModuleStatus
 from digitalkin.modules._base_module import BaseModule
 from digitalkin.services.registry import GrpcRegistry, RegistryStrategy
 from digitalkin.services.services_models import ServicesMode
@@ -159,7 +159,32 @@ class ModuleServicer(module_service_pb2_grpc.ModuleServiceServicer, ArgParser):
             return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
 
         updated_setup_data = await self.job_manager.generate_config_setup_module_response(job_id)
-        logger.info("Setup updated", extra={"job_id": job_id})
+        logger.info("Setup response received", extra={"job_id": job_id})
+
+        # Check if response is an error
+        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},
+            )
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(updated_setup_data.message or "Config setup failed")
+            return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
+
+        if isinstance(updated_setup_data, dict) and "code" in updated_setup_data:
+            # ModuleCodeModel was serialized to dict
+            logger.error(
+                "Config setup failed",
+                extra={
+                    "job_id": job_id,
+                    "code": updated_setup_data["code"],
+                    "message": updated_setup_data.get("message"),
+                },
+            )
+            context.set_code(grpc.StatusCode.INTERNAL)
+            context.set_details(updated_setup_data.get("message") or "Config setup failed")
+            return lifecycle_pb2.ConfigSetupModuleResponse(success=False)
+
         logger.debug("Updated setup data", extra={"job_id": job_id, "setup_data": updated_setup_data})
         setup_version.content = json_format.ParseDict(
             updated_setup_data,

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, ToolDefinition, ToolParameter
+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,20 +227,20 @@ class ModuleContext:
             llm_format=llm_format,
         )
 
-    async def create_openai_style_tools(self, tool_name: str) -> list[dict[str, Any]]:
+    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. 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:
             List of OpenAI-style tool schemas, one per protocol. Empty if not found.
         """
         tool_module_info = await self.tool_cache.get(
-            tool_name, registry=self.registry, communication=self.communication
+            module_id, registry=self.registry, communication=self.communication
         )
         if not tool_module_info:
             return []
@@ -299,13 +299,11 @@ class ModuleContext:
 
         communication = self.communication
         session = self.session
-        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)
+            fn = ModuleContext._create_single_tool_function(communication, session, tool_module_info, tool_def)
             result.append((tool_def, fn))
 
         return result
@@ -314,8 +312,7 @@ class ModuleContext:
     def _create_single_tool_function(
         communication: CommunicationStrategy,
         session: Session,
-        address: str,
-        port: int,
+        tool_module_info: ToolModuleInfo,
         tool_def: ToolDefinition,
     ) -> Callable[..., AsyncGenerator[dict, None]]:
         """Create a single tool function for a specific protocol.
@@ -323,8 +320,7 @@ class ModuleContext:
         Args:
             communication: Communication strategy for gRPC calls.
             session: Current session with setup_id and mission_id.
-            address: Module address.
-            port: Module port.
+            tool_module_info: Tool module information containing address and port.
             tool_def: Tool definition with protocol name.
 
         Returns:
@@ -336,10 +332,10 @@ class ModuleContext:
             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

digitalkin/models/module/setup_types.py

@@ -29,65 +29,12 @@ 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]] = {}
-
-    def __init_subclass__(cls, **kwargs: Any) -> None:  # noqa: ANN401
-        """Inject hidden companion fields for ToolReference annotations.
-
-        Args:
-            **kwargs: Keyword arguments passed to parent.
-        """
-        super().__init_subclass__(**kwargs)
-        cls._inject_tool_cache_fields()
-
-    @classmethod
-    def _inject_tool_cache_fields(cls) -> None:
-        """Inject hidden companion fields for ToolReference annotations."""
-        annotations = getattr(cls, "__annotations__", {})
-        new_annotations: dict[str, Any] = {}
-
-        for field_name, annotation in annotations.items():
-            if cls._is_tool_reference_annotation(annotation):
-                cache_field_name = f"{field_name}_cache"
-                if cache_field_name not in annotations:
-                    # Check if it's a list type
-                    origin = get_origin(annotation)
-                    if origin is list:
-                        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] = ToolModuleInfo | None
-                        setattr(
-                            cls,
-                            cache_field_name,
-                            Field(default=None, json_schema_extra={"hidden": True}),
-                        )
-
-        if new_annotations:
-            cls.__annotations__ = {**annotations, **new_annotations}
-
-    @classmethod
-    def _is_tool_reference_annotation(cls, annotation: object) -> bool:
-        """Check if annotation is ToolReference or Optional[ToolReference].
-
-        Args:
-            annotation: Type annotation to check.
-
-        Returns:
-            True if annotation is or contains ToolReference.
-        """
-        origin = get_origin(annotation)
-        if origin is typing.Union or origin is types.UnionType:
-            return any(
-                arg is ToolReference or (isinstance(arg, type) and issubclass(arg, ToolReference))
-                for arg in get_args(annotation)
-                if arg is not type(None)
-            )
-        return annotation is ToolReference or (isinstance(annotation, type) and issubclass(annotation, ToolReference))
+    resolved_tools: dict[str, ToolModuleInfo] = Field(
+        default_factory=dict,
+        json_schema_extra={"hidden": True},
+    )
 
     @classmethod
     async def get_clean_model(
@@ -285,7 +232,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         if not has_changes:
             return model_cls
 
-        root_extra = cls.model_config.get("json_schema_extra", {})
+        root_extra = model_cls.model_config.get("json_schema_extra", {})
 
         return create_model(
             model_cls.__name__,
@@ -346,7 +293,12 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             communication: Communication service for module schemas.
         """
         logger.info("Starting resolve_tool_references")
-        await self._resolve_tool_references_recursive(self, registry, communication)
+        await self._resolve_tool_references_recursive(
+            self,
+            registry,
+            communication,
+            self.resolved_tools,
+        )
         logger.info("Finished resolve_tool_references")
 
     @classmethod
@@ -355,6 +307,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         model_instance: BaseModel,
         registry: "RegistryStrategy",
         communication: "CommunicationStrategy",
+        resolved_tools: dict[str, ToolModuleInfo],
     ) -> None:
         """Recursively resolve ToolReference fields in a model.
 
@@ -362,11 +315,18 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             model_instance: Model instance to process.
             registry: Registry service for resolution.
             communication: Communication service for module schemas.
+            resolved_tools: Cache of already resolved tools.
         """
         for field_name, field_value in model_instance.__dict__.items():
             if field_value is None:
                 continue
-            await cls._resolve_field_value(field_name, field_value, registry, communication)
+            await cls._resolve_field_value(
+                field_name,
+                field_value,
+                registry,
+                communication,
+                resolved_tools,
+            )
 
     @classmethod
     async def _resolve_field_value(
@@ -375,6 +335,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         field_value: "BaseModel | ToolReference | list | dict",
         registry: "RegistryStrategy",
         communication: "CommunicationStrategy",
+        resolved_tools: dict[str, ToolModuleInfo],
     ) -> None:
         """Resolve a single field value based on its type.
 
@@ -383,15 +344,37 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             field_value: Value to process.
             registry: Registry service for resolution.
             communication: Communication service for module schemas.
+            resolved_tools: Cache of already resolved tools.
         """
         if isinstance(field_value, ToolReference):
-            await cls._resolve_single_tool_reference(field_name, field_value, registry, communication)
+            await cls._resolve_single_tool_reference(
+                field_name,
+                field_value,
+                registry,
+                communication,
+                resolved_tools,
+            )
         elif isinstance(field_value, BaseModel):
-            await cls._resolve_tool_references_recursive(field_value, registry, communication)
+            await cls._resolve_tool_references_recursive(
+                field_value,
+                registry,
+                communication,
+                resolved_tools,
+            )
         elif isinstance(field_value, list):
-            await cls._resolve_list_items(field_value, registry, communication)
+            await cls._resolve_list_items(
+                field_value,
+                registry,
+                communication,
+                resolved_tools,
+            )
         elif isinstance(field_value, dict):
-            await cls._resolve_dict_values(field_value, registry, communication)
+            await cls._resolve_dict_values(
+                field_value,
+                registry,
+                communication,
+                resolved_tools,
+            )
 
     @classmethod
     async def _resolve_single_tool_reference(
@@ -400,6 +383,7 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
         tool_ref: ToolReference,
         registry: "RegistryStrategy",
         communication: "CommunicationStrategy",
+        resolved_tools: dict[str, ToolModuleInfo],
     ) -> None:
         """Resolve a single ToolReference.
 
@@ -408,17 +392,33 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             tool_ref: ToolReference to resolve.
             registry: Registry service for resolution.
             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)
+
+        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
+
         try:
-            await tool_ref.resolve(registry, communication)
+            info = await tool_ref.resolve(registry, communication)
+            if info and info.setup_id:
+                resolved_tools[info.setup_id] = info
             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
     async def _resolve_list_items(
-        cls, items: list, registry: "RegistryStrategy", communication: "CommunicationStrategy"
+        cls,
+        items: list,
+        registry: "RegistryStrategy",
+        communication: "CommunicationStrategy",
+        resolved_tools: dict[str, ToolModuleInfo],
     ) -> None:
         """Resolve ToolReference instances in a list.
 
@@ -426,16 +426,32 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             items: List of items to process.
             registry: Registry service for resolution.
             communication: Communication service for module schemas.
+            resolved_tools: Cache of already resolved tools.
         """
         for item in items:
             if isinstance(item, ToolReference):
-                await cls._resolve_single_tool_reference("list_item", item, registry, communication)
+                await cls._resolve_single_tool_reference(
+                    "list_item",
+                    item,
+                    registry,
+                    communication,
+                    resolved_tools,
+                )
             elif isinstance(item, BaseModel):
-                await cls._resolve_tool_references_recursive(item, registry, communication)
+                await cls._resolve_tool_references_recursive(
+                    item,
+                    registry,
+                    communication,
+                    resolved_tools,
+                )
 
     @classmethod
     async def _resolve_dict_values(
-        cls, mapping: dict, registry: "RegistryStrategy", communication: "CommunicationStrategy"
+        cls,
+        mapping: dict,
+        registry: "RegistryStrategy",
+        communication: "CommunicationStrategy",
+        resolved_tools: dict[str, ToolModuleInfo],
     ) -> None:
         """Resolve ToolReference instances in dict values.
 
@@ -443,63 +459,86 @@ class SetupModel(BaseModel, Generic[SetupModelT]):
             mapping: Dict to process.
             registry: Registry service for resolution.
             communication: Communication service for module schemas.
+            resolved_tools: Cache of already resolved tools.
         """
         for item in mapping.values():
             if isinstance(item, ToolReference):
-                await cls._resolve_single_tool_reference("dict_value", item, registry, communication)
+                await cls._resolve_single_tool_reference(
+                    "dict_value",
+                    item,
+                    registry,
+                    communication,
+                    resolved_tools,
+                )
             elif isinstance(item, BaseModel):
-                await cls._resolve_tool_references_recursive(item, registry, communication)
+                await cls._resolve_tool_references_recursive(
+                    item,
+                    registry,
+                    communication,
+                    resolved_tools,
+                )
 
     def build_tool_cache(self) -> ToolCache:
-        """Build tool cache from resolved ToolReferences, populating companion fields.
+        """Build tool cache from resolved ToolReferences.
 
         Returns:
             ToolCache with field names as keys and ToolModuleInfo as values.
         """
-        logger.info("Building tool cache")
         cache = ToolCache()
         self._build_tool_cache_recursive(self, cache)
         logger.info("Tool cache built: %d entries", len(cache.entries))
         return cache
 
-    def _build_tool_cache_recursive(self, model_instance: BaseModel, cache: ToolCache) -> None:  # noqa: C901
-        """Recursively build tool cache and populate companion fields.
+    def _build_tool_cache_recursive(self, model_instance: BaseModel, cache: ToolCache) -> None:
+        """Recursively build tool cache from ToolReferences.
 
         Args:
             model_instance: Model instance to process.
             cache: ToolCache to populate.
         """
-        for field_name, field_value in model_instance.__dict__.items():
+        for field_value in model_instance.__dict__.values():
             if field_value is None:
                 continue
             if isinstance(field_value, ToolReference):
-                cache_field_name = f"{field_name}_cache"
-
-                cached_info = getattr(model_instance, cache_field_name, None)
-                module_info = field_value.tool_module_info or cached_info
+                module_info = self.resolved_tools.get(field_value.slug or "") or field_value.tool_module_info
                 if module_info:
-                    if not cached_info:
-                        setattr(model_instance, cache_field_name, module_info)
+                    self.resolved_tools[module_info.setup_id] = module_info
                     cache.add(module_info.module_id, module_info)
-                    logger.debug("Added tool to cache: %s", module_info.module_id)
             elif isinstance(field_value, BaseModel):
                 self._build_tool_cache_recursive(field_value, cache)
             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[ToolModuleInfo] = []
-
-                for idx, item in enumerate(field_value):
-                    if isinstance(item, ToolReference):
-                        # Use resolved info or fallback to cached
-                        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)
-                            logger.debug("Added tool to cache: %s", module_info.module_id)
-                    elif isinstance(item, BaseModel):
-                        self._build_tool_cache_recursive(item, cache)
-
-                # Update companion field with resolved infos
-                if resolved_infos:
-                    setattr(model_instance, cache_field_name, resolved_infos)
+                self._process_list_items(field_value, cache)
+            elif isinstance(field_value, dict):
+                self._process_dict_values(field_value, cache)
+
+    def _process_list_items(self, items: list, cache: ToolCache) -> None:
+        """Process list items for ToolReferences.
+
+        Args:
+            items: List to process.
+            cache: ToolCache to populate.
+        """
+        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)
+            elif isinstance(item, BaseModel):
+                self._build_tool_cache_recursive(item, cache)
+
+    def _process_dict_values(self, mapping: dict, cache: ToolCache) -> None:
+        """Process dict values for ToolReferences.
+
+        Args:
+            mapping: Dict to process.
+            cache: ToolCache to populate.
+        """
+        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)
+            elif isinstance(item, BaseModel):
+                self._build_tool_cache_recursive(item, cache)

digitalkin/models/module/tool_cache.py

@@ -52,6 +52,7 @@ class ToolModuleInfo(ModuleInfo):
     """Module info for tool modules."""
 
     tools: list[ToolDefinition]
+    setup_id: str
 
 
 class ToolCache(BaseModel):
@@ -59,22 +60,22 @@ class ToolCache(BaseModel):
 
     entries: dict[str, ToolModuleInfo] = Field(default_factory=dict)
 
-    def add(self, setup_tool_name: str, tool_module_info: ToolModuleInfo) -> None:
+    def add(self, setup_id: str, tool_module_info: ToolModuleInfo) -> None:
         """Add a tool to the cache.
 
         Args:
-            setup_tool_name: Field name from SetupModel used as cache key.
+            setup_id: Field name from SetupModel used as cache key.
             tool_module_info: Resolved tool module information.
         """
-        self.entries[setup_tool_name] = tool_module_info
+        self.entries[setup_id] = tool_module_info
         logger.debug(
             "Tool cached",
-            extra={"setup_tool_name": setup_tool_name, "module_id": tool_module_info.module_id},
+            extra={"setup_id": setup_id, "module_id": tool_module_info.module_id},
         )
 
     async def get(
         self,
-        setup_tool_name: str,
+        setup_id: str,
         *,
         registry: RegistryStrategy | None = None,
         communication: "CommunicationStrategy | None" = None,
@@ -82,28 +83,28 @@ class ToolCache(BaseModel):
         """Get a tool from cache, optionally querying registry on miss.
 
         Args:
-            setup_tool_name: Field name to look up.
+            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_tool_name)
+        cached = self.entries.get(setup_id)
         if cached:
             return cached
 
         if registry and communication:
             try:
-                setup_info = registry.get_setup(setup_tool_name)
+                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, communication)
-                        self.add(setup_tool_name, tool_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_tool_name": setup_tool_name})
+                logger.exception("Registry lookup failed", extra={"setup_id": setup_id})
 
         return None
 
@@ -122,6 +123,7 @@ class ToolCache(BaseModel):
 
 async def module_info_to_tool_module_info(
     module_info: ModuleInfo,
+    setup_id: str,
     communication: "CommunicationStrategy",
     *,
     llm_format: bool = True,
@@ -133,6 +135,7 @@ async def module_info_to_tool_module_info(
 
     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.
 
@@ -161,6 +164,7 @@ async def module_info_to_tool_module_info(
         documentation=module_info.documentation,
         status=module_info.status,
         tools=tools,
+        setup_id=setup_id,
     )
 
 

digitalkin/models/module/tool_reference.py

@@ -115,7 +115,7 @@ class ToolReference(BaseModel):
                 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)
+                    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
 
@@ -126,7 +126,9 @@ class ToolReference(BaseModel):
                 organization_id=self.config.organization_id,
             )
             if results:
-                tool_module_info = await module_info_to_tool_module_info(results[0], communication)
+                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

digitalkin/modules/_base_module.py

@@ -556,21 +556,7 @@ class BaseModule(  # noqa: PLR0904
             await self._resolve_tools(config_setup_data)
             updated_config = await self.run_config_setup(self.context, config_setup_data)
 
-            # Build wrapper: original structure with updated content
-            wrapper = config_setup_data.model_dump()
-            wrapper["content"] = updated_config.model_dump()
-
-            # Debug logging
-            content = wrapper.get("content", {})
-            logger.info(
-                "Config setup wrapper: keys=%s, content_keys=%s, tools_cache=%s",
-                list(wrapper.keys()),
-                list(content.keys()) if isinstance(content, dict) else "N/A",
-                content.get("tools_cache") if isinstance(content, dict) else "N/A",
-                extra=self.context.session.current_ids(),
-            )
-
-            setup_model = await self.create_setup_model(wrapper)
+            setup_model = await self.create_setup_model(updated_config.model_dump())
             await callback(setup_model)
             self._status = ModuleStatus.STOPPING
         except Exception:

digitalkin/utils/__init__.py

@@ -1,5 +1,12 @@
 """General utils folder."""
 
+from digitalkin.utils.conditional_schema import (
+    Conditional,
+    ConditionalField,
+    ConditionalSchemaMixin,
+    get_conditional_metadata,
+    has_conditional,
+)
 from digitalkin.utils.dynamic_schema import (
     DEFAULT_TIMEOUT,
     Dynamic,
@@ -14,13 +21,20 @@ from digitalkin.utils.dynamic_schema import (
 )
 
 __all__ = [
+    # Dynamic schema
     "DEFAULT_TIMEOUT",
+    # Conditional schema
+    "Conditional",
+    "ConditionalField",
+    "ConditionalSchemaMixin",
     "Dynamic",
     "DynamicField",
     "Fetcher",
     "ResolveResult",
+    "get_conditional_metadata",
     "get_dynamic_metadata",
     "get_fetchers",
+    "has_conditional",
     "has_dynamic",
     "resolve",
     "resolve_safe",

digitalkin/utils/conditional_schema.py

@@ -0,0 +1,260 @@
+"""Conditional field visibility for react-jsonschema-form.
+
+This module provides a clean way to mark fields as conditional using Annotated metadata,
+generating JSON Schema with if/then clauses for react-jsonschema-form.
+
+Example:
+    from typing import Annotated, Literal
+    from pydantic import BaseModel, Field
+    from digitalkin.utils import Conditional, ConditionalSchemaMixin
+
+    class Tools(ConditionalSchemaMixin, BaseModel):
+        web_search_enabled: bool = Field(...)
+
+        web_search_engine: Annotated[
+            Literal["duckduckgo", "tavily"],
+            Conditional(trigger="web_search_enabled", show_when=True),
+        ] = Field(...)
+
+See Also:
+    - Documentation: docs/api/conditional_schema.md
+    - Tests: tests/utils/test_conditional_schema.py
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, ClassVar
+
+from pydantic import BaseModel
+
+if TYPE_CHECKING:
+    from pydantic.annotated_handlers import GetJsonSchemaHandler
+    from pydantic.fields import FieldInfo
+    from pydantic.json_schema import JsonSchemaValue
+    from pydantic_core.core_schema import CoreSchema
+
+
+@dataclass
+class ConditionalField:
+    """Metadata for conditional field visibility.
+
+    Use with typing.Annotated to mark fields that should only appear
+    when a trigger field has a specific value.
+
+    Args:
+        trigger: Name of the field that controls visibility.
+        show_when: Value(s) that trigger field must have to show this field.
+            Can be a boolean, string, or list of strings for multiple values.
+        required_when_shown: Whether field is required when visible. Defaults to True.
+
+    Example:
+        # Boolean condition
+        web_search_engine: Annotated[
+            str,
+            Conditional(trigger="web_search_enabled", show_when=True),
+        ] = Field(...)
+
+        # Enum condition
+        advanced_option: Annotated[
+            str,
+            Conditional(trigger="mode", show_when="advanced"),
+        ] = Field(...)
+
+        # Multiple values condition
+        shared_feature: Annotated[
+            bool,
+            Conditional(trigger="mode", show_when=["standard", "advanced"]),
+        ] = Field(...)
+    """
+
+    trigger: str
+    show_when: bool | str | list[str]
+    required_when_shown: bool = True
+
+    def __post_init__(self) -> None:
+        """Normalize single-item lists to scalar values."""
+        if isinstance(self.show_when, list) and len(self.show_when) == 1:
+            self.show_when = self.show_when[0]
+
+
+# Short alias for cleaner API
+Conditional = ConditionalField
+
+
+def get_conditional_metadata(field_info: FieldInfo) -> ConditionalField | None:
+    """Extract ConditionalField from field metadata.
+
+    Args:
+        field_info: The Pydantic FieldInfo object to inspect.
+
+    Returns:
+        The ConditionalField metadata instance if found, None otherwise.
+    """
+    for meta in field_info.metadata:
+        if isinstance(meta, ConditionalField):
+            return meta
+    return None
+
+
+def has_conditional(field_info: FieldInfo) -> bool:
+    """Check if field has ConditionalField metadata.
+
+    Args:
+        field_info: The Pydantic FieldInfo object to check.
+
+    Returns:
+        True if the field has ConditionalField metadata, False otherwise.
+    """
+    return get_conditional_metadata(field_info) is not None
+
+
+def _collect_conditions(
+    model_fields: dict[str, FieldInfo],
+    props: dict[str, Any],
+) -> tuple[dict[tuple[str, Any], list[tuple[str, bool]]], set[str]]:
+    """Collect conditional fields grouped by trigger and show_when value.
+
+    Args:
+        model_fields: The model's field definitions.
+        props: The schema properties dict.
+
+    Returns:
+        Tuple of (conditions dict, fields to remove set).
+    """
+    conditions: dict[tuple[str, Any], list[tuple[str, bool]]] = {}
+    fields_to_remove: set[str] = set()
+
+    for field_name, field_info in model_fields.items():
+        cond = get_conditional_metadata(field_info)
+        if cond is None or field_name not in props:
+            continue
+
+        show_key = tuple(cond.show_when) if isinstance(cond.show_when, list) else cond.show_when
+        key = (cond.trigger, show_key)
+
+        if key not in conditions:
+            conditions[key] = []
+        conditions[key].append((field_name, cond.required_when_shown))
+        fields_to_remove.add(field_name)
+
+    return conditions, fields_to_remove
+
+
+def _build_if_clause(trigger: str, *, show_when: bool | str | tuple[str, ...]) -> dict[str, Any]:
+    """Build the if clause for a conditional.
+
+    Args:
+        trigger: The trigger field name.
+        show_when: The value(s) that trigger visibility.
+
+    Returns:
+        The if clause dict.
+    """
+    if isinstance(show_when, tuple):
+        return {"properties": {trigger: {"enum": list(show_when)}}, "required": [trigger]}
+    return {"properties": {trigger: {"const": show_when}}, "required": [trigger]}
+
+
+def _resolve_field_schema(
+    field_schema: dict[str, Any],
+    handler: GetJsonSchemaHandler,
+) -> dict[str, Any]:
+    """Resolve $ref in field schema if present.
+
+    Args:
+        field_schema: The field's schema dict.
+        handler: The JSON schema handler for resolving refs.
+
+    Returns:
+        The resolved schema dict.
+    """
+    if "$ref" not in field_schema:
+        return field_schema
+
+    resolved = handler.resolve_ref_schema(field_schema)
+    extra = {k: v for k, v in field_schema.items() if k != "$ref"}
+    return {**resolved, **extra}
+
+
+class ConditionalSchemaMixin(BaseModel):
+    """Mixin for automatic conditional field processing in JSON schema.
+
+    Inherit from this mixin to automatically generate JSON Schema with
+    if/then clauses for fields marked with ConditionalField metadata.
+
+    The mixin processes Annotated fields with Conditional metadata and:
+    1. Removes conditional fields from main properties
+    2. Adds them to allOf with if/then clauses
+    3. Groups multiple fields with the same condition together
+
+    Example:
+        class Config(ConditionalSchemaMixin, BaseModel):
+            mode: Literal["basic", "advanced"] = Field(...)
+
+            advanced_option: Annotated[
+                str,
+                Conditional(trigger="mode", show_when="advanced"),
+            ] = Field(...)
+
+        # Generates schema with:
+        # {
+        #     "properties": {"mode": {...}},
+        #     "allOf": [{
+        #         "if": {"properties": {"mode": {"const": "advanced"}}},
+        #         "then": {"properties": {"advanced_option": {...}}}
+        #     }]
+        # }
+    """
+
+    model_fields: ClassVar[dict[str, FieldInfo]]  # type: ignore[misc]
+
+    @classmethod
+    def __get_pydantic_json_schema__(  # noqa: PLW3201
+        cls,
+        core_schema: CoreSchema,
+        handler: GetJsonSchemaHandler,
+    ) -> JsonSchemaValue:
+        """Generate JSON schema with conditional field handling.
+
+        Args:
+            core_schema: The Pydantic core schema.
+            handler: The JSON schema handler for resolving refs.
+
+        Returns:
+            The JSON schema with if/then clauses for conditional fields.
+        """
+        schema = handler(core_schema)
+        props = schema.get("properties", {})
+        if not props:
+            return schema
+
+        conditions, fields_to_remove = _collect_conditions(cls.model_fields, props)
+        if not conditions:
+            return schema
+
+        all_of = schema.setdefault("allOf", [])
+
+        for (trigger, show_when), field_list in conditions.items():
+            then_props: dict[str, Any] = {}
+            then_required: list[str] = []
+
+            for field_name, required in field_list:
+                then_props[field_name] = _resolve_field_schema(props[field_name], handler)
+                if required:
+                    then_required.append(field_name)
+
+            if_clause = _build_if_clause(trigger, show_when=show_when)
+            then_clause: dict[str, Any] = {"properties": then_props}
+            if then_required:
+                then_clause["required"] = then_required
+
+            all_of.append({"if": if_clause, "then": then_clause})
+
+        for field_name in fields_to_remove:
+            del props[field_name]
+
+        if "required" in schema:
+            schema["required"] = [r for r in schema["required"] if r not in fields_to_remove]
+
+        return schema

digitalkin/utils/schema_splitter.py

@@ -55,6 +55,8 @@ class SchemaSplitter:
                 for item in value:
                     if isinstance(item, dict):
                         cls._extract_ui_properties(item, ui_target)
+            elif key in {"if", "then", "else"} and isinstance(value, dict):
+                cls._extract_ui_properties(value, ui_target)
 
     @classmethod
     def _process_object(  # noqa: C901, PLR0912
@@ -107,11 +109,15 @@ class SchemaSplitter:
                         item_json: dict[str, Any] = {}
                         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)
                     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)
             else:
                 json_target[key] = value
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: digitalkin
-Version: 0.3.2.dev17
+Version: 0.3.2.dev19
 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.dev17.dist-info → digitalkin-0.3.2.dev19.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=Y3kvX8EPStPR7XivYOXc5tCfaRbBYw08EWY8LQ_zjgM,196
+digitalkin/__version__.py,sha256=udtzuXsaIao_ZAPzuFlc1q_JtSznP6hI13v4Zl8o1Zo,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=7GQOyAPYMxHVaJGplgDNiVoKr1oaAIL-zdZpyDpznTA,20530
+digitalkin/grpc_servers/module_servicer.py,sha256=KsWXPwCQ0s2pygn0cq89KndJleAkjvkvLFs1Z2fWQIk,21705
 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
@@ -53,18 +53,18 @@ digitalkin/models/grpc_servers/types.py,sha256=rQ78s4nAet2jy-NIDj_PUWriT0kuGHr_w
 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=u4yvvtYNwbSJSnV0omLasGN7wRZgaLS7trBBneeT2CI,12178
+digitalkin/models/module/module_context.py,sha256=QDdjZdhIJpvU_2Tn7kkJsZ1givB4dM1-ksopdF4VySw,12176
 digitalkin/models/module/module_types.py,sha256=C9azCNBk76xMa-Mww8_6AiwQR8MLAsEyUOvBYxytovI,739
-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/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/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/storage.py,sha256=wp7F-AvTsU46ujGPcguqM5kUKRZx4399D4EGAAJt2zs,1143
 digitalkin/modules/__init__.py,sha256=vTQk8DWopxQSJ17BjE5dNhq247Rou55iQLJdBxoPUmo,296
-digitalkin/modules/_base_module.py,sha256=vQX-vqGXafc9DkNJXzTic-1zJ3yrP1BC30fTH8yHB1k,22592
+digitalkin/modules/_base_module.py,sha256=0XC0aQAxlNfvz0KK9ut7K0JbZql3cZMU4aeg7ISEsD0,21971
 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
@@ -115,14 +115,15 @@ 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=xusG_EoixTznt9ENCB8XkNbi67SZNt6y9hx5gsPWH_k,464
+digitalkin/utils/__init__.py,sha256=_RXXALotIr_JToSQdhwDi3eizghKz9KrQtvC8Um80H8,808
 digitalkin/utils/arg_parser.py,sha256=wzscRlE1Qp1gGl-lAJlkkwnbU1O2oezj6BwK_BZFBIk,3158
+digitalkin/utils/conditional_schema.py,sha256=ZyGvY-DWkM_d0OuLElUpqnvUeDgZ2hXX3dFnKXbgL5k,8373
 digitalkin/utils/development_mode_action.py,sha256=2hznh0ajW_4ZTysfoc0Y49161f_PQPATRgNk8NAn1_o,1623
 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.dev17.dist-info/licenses/LICENSE,sha256=Ies4HFv2r2hzDRakJYxk3Y60uDFLiG-orIgeTpstnIo,20327
+digitalkin/utils/schema_splitter.py,sha256=9PHC-bvEDQudyYZNgXyjFtp7EJlmw4C_gPCJ-JmGDk0,9704
+digitalkin-0.3.2.dev19.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 +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.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,,
+digitalkin-0.3.2.dev19.dist-info/METADATA,sha256=HBJG5Mpck63wVwh73EffNv54xKZpapsRcSgOsu2izKs,29725
+digitalkin-0.3.2.dev19.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+digitalkin-0.3.2.dev19.dist-info/top_level.txt,sha256=AYVIesKrO0jnedQ-Muog9JBehG81WeTCNeOFoJgwsgE,51
+digitalkin-0.3.2.dev19.dist-info/RECORD,,

{digitalkin-0.3.2.dev17.dist-info → digitalkin-0.3.2.dev19.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.9.0)
+Generator: setuptools (80.10.1)
 Root-Is-Purelib: true
 Tag: py3-none-any