diff --git a/packages/toolbox-adk/src/toolbox_adk/tool.py b/packages/toolbox-adk/src/toolbox_adk/tool.py index ae9eca015..504616186 100644 --- a/packages/toolbox-adk/src/toolbox_adk/tool.py +++ b/packages/toolbox-adk/src/toolbox_adk/tool.py @@ -395,3 +395,21 @@ def bind_params(self, bounded_params: Dict[str, Any]) -> "ToolboxTool": adk_token_getters=self._adk_token_getters, telemetry_attributes=self._telemetry_attributes, ) + + def bind_param(self, param_name: str, param_value: Any) -> "ToolboxTool": + """Allows runtime binding of a parameter, delegating to core tool.""" + return self.bind_params({param_name: param_value}) + + def bind_secure_params(self, bound_secure_params: Dict[str, Any]) -> "ToolboxTool": + """Allows runtime binding of secure parameters, delegating to core tool.""" + new_core_tool = self._core_tool.bind_secure_params(bound_secure_params) + return ToolboxTool( + core_tool=new_core_tool, + auth_config=self._auth_config, + adk_token_getters=self._adk_token_getters, + telemetry_attributes=self._telemetry_attributes, + ) + + def bind_secure_param(self, param_name: str, param_value: Any) -> "ToolboxTool": + """Allows runtime binding of a secure parameter, delegating to core tool.""" + return self.bind_secure_params({param_name: param_value}) diff --git a/packages/toolbox-adk/src/toolbox_adk/toolset.py b/packages/toolbox-adk/src/toolbox_adk/toolset.py index 956e54b7b..1f0267274 100644 --- a/packages/toolbox-adk/src/toolbox_adk/toolset.py +++ b/packages/toolbox-adk/src/toolbox_adk/toolset.py @@ -42,6 +42,7 @@ def __init__( Dict[str, Union[str, Callable[[], str], Callable[[], Awaitable[str]]]] ] = None, bound_params: Optional[Mapping[str, Union[Callable[[], Any], Any]]] = None, + secure_params: Optional[Mapping[str, Union[Callable[[], Any], Any]]] = None, auth_token_getters: Optional[ Mapping[ str, @@ -72,6 +73,7 @@ def __init__( credentials: Authentication configuration. additional_headers: Extra headers (static or dynamic). bound_params: Parameters to bind globally to loaded tools. + secure_params: Secure parameters to bind globally to loaded tools. auth_token_getters: Mapping of auth service names to token getters. telemetry_attributes: Telemetry attributes (model, user id, agent id) sent to the server with each tool invocation. Either a @@ -89,6 +91,7 @@ def __init__( self.__toolset_name = toolset_name self.__tool_names = tool_names self.__bound_params = bound_params + self.__secure_params = secure_params self.__auth_token_getters = auth_token_getters self.__telemetry_attributes = telemetry_attributes @@ -112,10 +115,14 @@ async def get_tools( tools = [] # 1. Load specific toolset if requested + load_kwargs: dict[str, Any] = {"bound_params": self.__bound_params or {}} + if self.__secure_params is not None: + load_kwargs["secure_params"] = self.__secure_params + if self.__toolset_name: core_tools = await self.client.load_toolset( self.__toolset_name, - bound_params=self.__bound_params or {}, + **load_kwargs, ) tools.extend(core_tools) @@ -124,7 +131,7 @@ async def get_tools( for name in self.__tool_names: core_tool = await self.client.load_tool( name, - bound_params=self.__bound_params or {}, + **load_kwargs, ) tools.append(core_tool) @@ -132,7 +139,7 @@ async def get_tools( if not self.__toolset_name and not self.__tool_names: core_tools = await self.client.load_toolset( None, - bound_params=self.__bound_params or {}, + **load_kwargs, ) tools.extend(core_tools) diff --git a/packages/toolbox-adk/tests/unit/test_tool.py b/packages/toolbox-adk/tests/unit/test_tool.py index b507c3498..8ad17da38 100644 --- a/packages/toolbox-adk/tests/unit/test_tool.py +++ b/packages/toolbox-adk/tests/unit/test_tool.py @@ -546,3 +546,52 @@ class EmptyTool: tool = ToolboxTool(core_tool) assert tool.name == "valid_tool" assert tool.description == "valid description" + + def test_bind_param(self): + core_tool = MagicMock() + core_tool.__name__ = "valid_tool" + core_tool.__doc__ = "valid doc" + new_core_tool = MagicMock() + new_core_tool.__name__ = "valid_tool" + new_core_tool.__doc__ = "valid doc" + core_tool.bind_params.return_value = new_core_tool + + tool = ToolboxTool(core_tool) + bound_tool = tool.bind_param("param_a", "val_a") + + core_tool.bind_params.assert_called_once_with({"param_a": "val_a"}) + assert isinstance(bound_tool, ToolboxTool) + assert bound_tool._core_tool is new_core_tool + + def test_bind_secure_param(self): + core_tool = MagicMock() + core_tool.__name__ = "valid_tool" + core_tool.__doc__ = "valid doc" + new_core_tool = MagicMock() + new_core_tool.__name__ = "valid_tool" + new_core_tool.__doc__ = "valid doc" + core_tool.bind_secure_params.return_value = new_core_tool + + tool = ToolboxTool(core_tool) + bound_tool = tool.bind_secure_param("api_key", "secret123") + + core_tool.bind_secure_params.assert_called_once_with({"api_key": "secret123"}) + assert isinstance(bound_tool, ToolboxTool) + assert bound_tool._core_tool is new_core_tool + + def test_bind_secure_params(self): + core_tool = MagicMock() + core_tool.__name__ = "valid_tool" + core_tool.__doc__ = "valid doc" + new_core_tool = MagicMock() + new_core_tool.__name__ = "valid_tool" + new_core_tool.__doc__ = "valid doc" + core_tool.bind_secure_params.return_value = new_core_tool + + tool = ToolboxTool(core_tool) + sec_dict = {"api_key": "secret123", "db_pass": "pass"} + bound_tool = tool.bind_secure_params(sec_dict) + + core_tool.bind_secure_params.assert_called_once_with(sec_dict) + assert isinstance(bound_tool, ToolboxTool) + assert bound_tool._core_tool is new_core_tool diff --git a/packages/toolbox-adk/tests/unit/test_toolset.py b/packages/toolbox-adk/tests/unit/test_toolset.py index 440a8b5c2..5f5308b07 100644 --- a/packages/toolbox-adk/tests/unit/test_toolset.py +++ b/packages/toolbox-adk/tests/unit/test_toolset.py @@ -171,3 +171,26 @@ def test_init_with_protocol(self, mock_client_cls): mock_client_cls.assert_called_once() call_kwargs = mock_client_cls.call_args[1] assert call_kwargs["protocol"] == Protocol.MCP + + @patch("toolbox_adk.toolset.ToolboxClient") + @pytest.mark.asyncio + async def test_get_tools_with_secure_params(self, mock_client_cls): + """Test that secure_params is passed to load_toolset.""" + mock_client = mock_client_cls.return_value + t1 = MagicMock() + t1.__name__ = "tool1" + t1.__doc__ = "doc1" + mock_client.load_toolset = AsyncMock(return_value=[t1]) + + sec_params = {"api_key": "secret123"} + toolset = ToolboxToolset( + "url", toolset_name="my_toolset", secure_params=sec_params + ) + tools = await toolset.get_tools() + + assert len(tools) == 1 + mock_client.load_toolset.assert_awaited_with( + "my_toolset", + bound_params={}, + secure_params=sec_params, + )