From bc771d9a3ac8b7c1f19990ef5d173b30c7305c49 Mon Sep 17 00:00:00 2001 From: Henry Su Date: Mon, 27 Jul 2026 14:05:58 -0500 Subject: [PATCH] fix(adk): support map additionalProperties and bind_param in ToolboxTool --- packages/toolbox-adk/src/toolbox_adk/tool.py | 26 ++++++-- packages/toolbox-adk/tests/unit/test_tool.py | 63 ++++++++++++++++++++ 2 files changed, 85 insertions(+), 4 deletions(-) diff --git a/packages/toolbox-adk/src/toolbox_adk/tool.py b/packages/toolbox-adk/src/toolbox_adk/tool.py index 0d94821b6..c697c1565 100644 --- a/packages/toolbox-adk/src/toolbox_adk/tool.py +++ b/packages/toolbox-adk/src/toolbox_adk/tool.py @@ -14,7 +14,7 @@ import inspect import logging -from typing import Any, Awaitable, Callable, Dict, Mapping, Optional +from typing import Any, Awaitable, Callable, Dict, Mapping, Optional, Union import toolbox_core from fastapi.openapi.models import OAuth2, OAuthFlowAuthorizationCode, OAuthFlows @@ -95,7 +95,7 @@ def _build_schema(self, param: Any) -> Schema: properties = {} required = [] schema_items = None - schema_additional_properties = None + schema_additional_properties: Optional[Union[Schema, bool]] = None if schema_type == Type.ARRAY: if hasattr(param, "items") and param.items: @@ -107,12 +107,19 @@ def _build_schema(self, param: Any) -> Schema: properties[k] = self._build_schema(v) if getattr(v, "required", False): required.append(k) + add_props = getattr(param, "additionalProperties", None) + if add_props is not None: + if isinstance(add_props, bool): + schema_additional_properties = add_props + elif hasattr(add_props, "type"): + schema_additional_properties = self._build_schema(add_props) return Schema( type=schema_type, description=getattr(param, "description", "") or "", properties=properties or None, required=required or None, items=schema_items, + additional_properties=schema_additional_properties, ) @override @@ -299,12 +306,23 @@ async def run_async( if reset_token: USER_TOKEN_CONTEXT_VAR.reset(reset_token) - def bind_params(self, bounded_params: Dict[str, Any]) -> "ToolboxTool": + def bind_params( + self, + bound_params: Optional[Dict[str, Any]] = None, + bounded_params: Optional[Dict[str, Any]] = None, + ) -> "ToolboxTool": """Allows runtime binding of parameters, delegating to core tool.""" - new_core_tool = self._core_tool.bind_params(bounded_params) + params_to_bind = ( + bound_params if bound_params is not None else bounded_params or {} + ) + new_core_tool = self._core_tool.bind_params(params_to_bind) # Return a new wrapper return ToolboxTool( core_tool=new_core_tool, auth_config=self._auth_config, adk_token_getters=self._adk_token_getters, ) + + def bind_param(self, param_name: str, param_value: Any) -> "ToolboxTool": + """Binds a single parameter to a value or callable.""" + return self.bind_params({param_name: param_value}) diff --git a/packages/toolbox-adk/tests/unit/test_tool.py b/packages/toolbox-adk/tests/unit/test_tool.py index 3b55290ea..ca46251c5 100644 --- a/packages/toolbox-adk/tests/unit/test_tool.py +++ b/packages/toolbox-adk/tests/unit/test_tool.py @@ -416,3 +416,66 @@ class EmptyTool: tool = ToolboxTool(core_tool) assert tool.name == "valid_tool" assert tool.description == "valid description" + + def test_get_declaration_additional_properties(self): + class MockParam: + def __init__(self, name, param_type, description, required): + self.name = name + self.type = param_type + self.description = description + self.required = required + + class MockAddProps: + def __init__(self, prop_type): + self.type = prop_type + + core_tool = MagicMock() + core_tool.__name__ = "mock_tool" + core_tool.__doc__ = "mock doc" + + map_param = MockParam("my_map", "object", "A typed map", True) + map_param.additionalProperties = MockAddProps("string") + + bool_map_param = MockParam("bool_map", "object", "A bool map", False) + bool_map_param.additionalProperties = True + + core_tool._params = [map_param, bool_map_param] + + tool = ToolboxTool(core_tool) + declaration = tool._get_declaration() + + parameters = declaration.parameters + assert parameters is not None + + # Verify typed map additionalProperties + map_schema = parameters.properties["my_map"] + assert map_schema.type == Type.OBJECT + assert map_schema.additional_properties is not None + assert map_schema.additional_properties.type == Type.STRING + + # Verify bool map additionalProperties + bool_map_schema = parameters.properties["bool_map"] + assert bool_map_schema.type == Type.OBJECT + assert bool_map_schema.additional_properties is True + + def test_bind_param_and_bind_params_keyword(self): + mock_core = MagicMock() + mock_core.__name__ = "mock" + mock_core.__doc__ = "mock" + + new_core_mock = MagicMock() + new_core_mock.__name__ = "bound_mock" + new_core_mock.__doc__ = "bound mock" + mock_core.bind_params.return_value = new_core_mock + + tool = ToolboxTool(mock_core) + + # Test bind_param (singular) + new_tool1 = tool.bind_param("a", 1) + assert isinstance(new_tool1, ToolboxTool) + mock_core.bind_params.assert_called_with({"a": 1}) + + # Test bind_params with bound_params keyword argument + new_tool2 = tool.bind_params(bound_params={"b": 2}) + assert isinstance(new_tool2, ToolboxTool) + mock_core.bind_params.assert_called_with({"b": 2})