diff --git a/packages/toolbox-langchain/src/toolbox_langchain/async_client.py b/packages/toolbox-langchain/src/toolbox_langchain/async_client.py index 98b586a11..f1de43306 100644 --- a/packages/toolbox-langchain/src/toolbox_langchain/async_client.py +++ b/packages/toolbox-langchain/src/toolbox_langchain/async_client.py @@ -64,6 +64,7 @@ async def aload_tool( auth_headers: Optional[dict[str, Callable[[], str]]] = None, bound_params: dict[str, Union[Any, Callable[[], Any]]] = {}, telemetry_attributes: Optional[TelemetryAttributes] = None, + secure_params: dict[str, Union[Any, Callable[[], Any]]] = {}, ) -> AsyncToolboxTool: """ Loads the tool with the given tool name from the Toolbox service. @@ -78,6 +79,8 @@ async def aload_tool( bound values. telemetry_attributes: Optional telemetry attributes (model, user id, agent id) sent to the server with every tool invocation. + secure_params: An optional mapping of secure parameter names to their + bound values. Returns: A tool loaded from the Toolbox. @@ -108,10 +111,15 @@ async def aload_tool( ) auth_token_getters = auth_headers + kwargs: dict[str, Any] = {} + if secure_params: + kwargs["secure_params"] = secure_params + core_tool = await self.__core_client.load_tool( name=tool_name, auth_token_getters=auth_token_getters, bound_params=bound_params, + **kwargs, ) if telemetry_attributes is not None: core_tool = core_tool.add_telemetry_attributes(telemetry_attributes) @@ -126,6 +134,7 @@ async def aload_toolset( bound_params: dict[str, Union[Any, Callable[[], Any]]] = {}, strict: bool = False, telemetry_attributes: Optional[TelemetryAttributes] = None, + secure_params: dict[str, Union[Any, Callable[[], Any]]] = {}, ) -> list[AsyncToolboxTool]: """ Loads tools from the Toolbox service, optionally filtered by toolset @@ -141,12 +150,14 @@ async def aload_toolset( bound_params: An optional mapping of parameter names to their bound values. strict: If True, raises an error if *any* loaded tool instance fails - to utilize all of the given parameters or auth tokens. (if any + to utilize all of the given parameters, auth tokens, or secure parameters (if any provided). If False (default), raises an error only if a - user-provided parameter or auth token cannot be applied to *any* + user-provided parameter, auth token, or secure parameter cannot be applied to *any* loaded tool across the set. telemetry_attributes: Optional telemetry attributes (model, user id, agent id) sent to the server with every tool invocation. + secure_params: An optional mapping of secure parameter names to their + bound values. Returns: A list of all tools loaded from the Toolbox. @@ -177,11 +188,16 @@ async def aload_toolset( ) auth_token_getters = auth_headers + kwargs: dict[str, Any] = {} + if secure_params: + kwargs["secure_params"] = secure_params + core_tools = await self.__core_client.load_toolset( name=toolset_name, auth_token_getters=auth_token_getters, bound_params=bound_params, strict=strict, + **kwargs, ) tools = [] @@ -198,6 +214,7 @@ def load_tool( auth_tokens: Optional[dict[str, Callable[[], str]]] = None, auth_headers: Optional[dict[str, Callable[[], str]]] = None, bound_params: dict[str, Union[Any, Callable[[], Any]]] = {}, + secure_params: dict[str, Union[Any, Callable[[], Any]]] = {}, ) -> AsyncToolboxTool: raise NotImplementedError("Synchronous methods not supported by async client.") @@ -209,6 +226,7 @@ def load_toolset( auth_headers: Optional[dict[str, Callable[[], str]]] = None, bound_params: dict[str, Union[Any, Callable[[], Any]]] = {}, strict: bool = False, + secure_params: dict[str, Union[Any, Callable[[], Any]]] = {}, ) -> list[AsyncToolboxTool]: raise NotImplementedError("Synchronous methods not supported by async client.") diff --git a/packages/toolbox-langchain/src/toolbox_langchain/async_tools.py b/packages/toolbox-langchain/src/toolbox_langchain/async_tools.py index 9aaf055e6..4fff83a96 100644 --- a/packages/toolbox-langchain/src/toolbox_langchain/async_tools.py +++ b/packages/toolbox-langchain/src/toolbox_langchain/async_tools.py @@ -167,6 +167,51 @@ def bind_param( """ return self.bind_params({param_name: param_value}) + def bind_secure_params( + self, + bound_secure_params: dict[str, Union[Any, Callable[[], Any]]], + ) -> "AsyncToolboxTool": + """ + Registers values or functions to retrieve the value for the + corresponding bound secure parameters. + + Args: + bound_secure_params: A dictionary of the bound secure parameter name to the + value or function of the bound secure value. + + Returns: + A new AsyncToolboxTool instance that is a deep copy of the current + instance, with added bound secure params. + + Raises: + ValueError: If any of the provided bound secure params is already bound. + """ + new_core_tool = self.__core_tool.bind_secure_params(bound_secure_params) + return AsyncToolboxTool(core_tool=new_core_tool) + + def bind_secure_param( + self, + param_name: str, + param_value: Union[Any, Callable[[], Any]], + ) -> "AsyncToolboxTool": + """ + Registers a value or a function to retrieve the value for a given bound + secure parameter. + + Args: + param_name: The name of the bound secure parameter. + param_value: The value of the bound secure parameter, or a callable that + returns the value. + + Returns: + A new AsyncToolboxTool instance that is a deep copy of the current + instance, with added bound secure param. + + Raises: + ValueError: If the provided bound secure param is already bound. + """ + return self.bind_secure_params({param_name: param_value}) + def add_telemetry_attributes( self, telemetry_attributes: TelemetryAttributes ) -> "AsyncToolboxTool": diff --git a/packages/toolbox-langchain/src/toolbox_langchain/client.py b/packages/toolbox-langchain/src/toolbox_langchain/client.py index d088ae683..a60d4b891 100644 --- a/packages/toolbox-langchain/src/toolbox_langchain/client.py +++ b/packages/toolbox-langchain/src/toolbox_langchain/client.py @@ -58,6 +58,7 @@ async def aload_tool( auth_headers: Optional[dict[str, Callable[[], str]]] = None, bound_params: dict[str, Union[Any, Callable[[], Any]]] = {}, telemetry_attributes: Optional[TelemetryAttributes] = None, + secure_params: dict[str, Union[Any, Callable[[], Any]]] = {}, ) -> ToolboxTool: """ Loads the tool with the given tool name from the Toolbox service. @@ -72,6 +73,8 @@ async def aload_tool( bound values. telemetry_attributes: Optional telemetry attributes (model, user id, agent id) sent to the server with every tool invocation. + secure_params: An optional mapping of secure parameter names to their + bound values. Returns: A tool loaded from the Toolbox. @@ -102,11 +105,16 @@ async def aload_tool( ) auth_token_getters = auth_headers + kwargs: dict[str, Any] = {} + if secure_params: + kwargs["secure_params"] = secure_params + core_tool = await to_thread( self.__core_client.load_tool, name=tool_name, auth_token_getters=auth_token_getters, bound_params=bound_params, + **kwargs, ) if telemetry_attributes is not None: core_tool = core_tool.add_telemetry_attributes(telemetry_attributes) @@ -121,6 +129,7 @@ async def aload_toolset( bound_params: dict[str, Union[Any, Callable[[], Any]]] = {}, strict: bool = False, telemetry_attributes: Optional[TelemetryAttributes] = None, + secure_params: dict[str, Union[Any, Callable[[], Any]]] = {}, ) -> list[ToolboxTool]: """ Loads tools from the Toolbox service, optionally filtered by toolset @@ -136,12 +145,14 @@ async def aload_toolset( bound_params: An optional mapping of parameter names to their bound values. strict: If True, raises an error if *any* loaded tool instance fails - to utilize all of the given parameters or auth tokens. (if any + to utilize all of the given parameters, auth tokens, or secure parameters (if any provided). If False (default), raises an error only if a - user-provided parameter or auth token cannot be applied to *any* + user-provided parameter, auth token, or secure parameter cannot be applied to *any* loaded tool across the set. telemetry_attributes: Optional telemetry attributes (model, user id, agent id) sent to the server with every tool invocation. + secure_params: An optional mapping of secure parameter names to their + bound values. Returns: A list of all tools loaded from the Toolbox. @@ -172,12 +183,17 @@ async def aload_toolset( ) auth_token_getters = auth_headers + kwargs: dict[str, Any] = {} + if secure_params: + kwargs["secure_params"] = secure_params + core_tools = await to_thread( self.__core_client.load_toolset, name=toolset_name, auth_token_getters=auth_token_getters, bound_params=bound_params, strict=strict, + **kwargs, ) tools = [] @@ -195,6 +211,7 @@ def load_tool( auth_headers: Optional[dict[str, Callable[[], str]]] = None, bound_params: dict[str, Union[Any, Callable[[], Any]]] = {}, telemetry_attributes: Optional[TelemetryAttributes] = None, + secure_params: dict[str, Union[Any, Callable[[], Any]]] = {}, ) -> ToolboxTool: """ Loads the tool with the given tool name from the Toolbox service. @@ -209,6 +226,8 @@ def load_tool( bound values. telemetry_attributes: Optional telemetry attributes (model, user id, agent id) sent to the server with every tool invocation. + secure_params: An optional mapping of secure parameter names to their + bound values. Returns: A tool loaded from the Toolbox. @@ -239,10 +258,15 @@ def load_tool( ) auth_token_getters = auth_headers + kwargs: dict[str, Any] = {} + if secure_params: + kwargs["secure_params"] = secure_params + core_sync_tool = self.__core_client.load_tool( name=tool_name, auth_token_getters=auth_token_getters, bound_params=bound_params, + **kwargs, ) if telemetry_attributes is not None: core_sync_tool = core_sync_tool.add_telemetry_attributes( @@ -259,6 +283,7 @@ def load_toolset( bound_params: dict[str, Union[Any, Callable[[], Any]]] = {}, strict: bool = False, telemetry_attributes: Optional[TelemetryAttributes] = None, + secure_params: dict[str, Union[Any, Callable[[], Any]]] = {}, ) -> list[ToolboxTool]: """ Loads tools from the Toolbox service, optionally filtered by toolset @@ -274,12 +299,14 @@ def load_toolset( bound_params: An optional mapping of parameter names to their bound values. strict: If True, raises an error if *any* loaded tool instance fails - to utilize all of the given parameters or auth tokens. (if any + to utilize all of the given parameters, auth tokens, or secure parameters (if any provided). If False (default), raises an error only if a - user-provided parameter or auth token cannot be applied to *any* + user-provided parameter, auth token, or secure parameter cannot be applied to *any* loaded tool across the set. telemetry_attributes: Optional telemetry attributes (model, user id, agent id) sent to the server with every tool invocation. + secure_params: An optional mapping of secure parameter names to their + bound values. Returns: A list of all tools loaded from the Toolbox. @@ -310,11 +337,16 @@ def load_toolset( ) auth_token_getters = auth_headers + kwargs: dict[str, Any] = {} + if secure_params: + kwargs["secure_params"] = secure_params + core_sync_tools = self.__core_client.load_toolset( name=toolset_name, auth_token_getters=auth_token_getters, bound_params=bound_params, strict=strict, + **kwargs, ) tools = [] diff --git a/packages/toolbox-langchain/src/toolbox_langchain/tools.py b/packages/toolbox-langchain/src/toolbox_langchain/tools.py index cb2f9d1ed..90e2ee7cb 100644 --- a/packages/toolbox-langchain/src/toolbox_langchain/tools.py +++ b/packages/toolbox-langchain/src/toolbox_langchain/tools.py @@ -153,6 +153,51 @@ def bind_param( """ return self.bind_params({param_name: param_value}) + def bind_secure_params( + self, + bound_secure_params: dict[str, Union[Any, Callable[[], Any]]], + ) -> "ToolboxTool": + """ + Registers values or functions to retrieve the value for the + corresponding bound secure parameters. + + Args: + bound_secure_params: A dictionary of the bound secure parameter name to the + value or function of the bound secure value. + + Returns: + A new ToolboxTool instance that is a deep copy of the current + instance, with added bound secure params. + + Raises: + ValueError: If any of the provided bound secure params is already bound. + """ + new_core_tool = self.__core_tool.bind_secure_params(bound_secure_params) + return ToolboxTool(core_tool=new_core_tool) + + def bind_secure_param( + self, + param_name: str, + param_value: Union[Any, Callable[[], Any]], + ) -> "ToolboxTool": + """ + Registers a value or a function to retrieve the value for a given bound + secure parameter. + + Args: + param_name: The name of the bound secure parameter. + param_value: The value of the bound secure parameter, or a callable that + returns the value. + + Returns: + A new ToolboxTool instance that is a deep copy of the current + instance, with added bound secure param. + + Raises: + ValueError: If the provided bound secure param is already bound. + """ + return self.bind_secure_params({param_name: param_value}) + def add_telemetry_attributes( self, telemetry_attributes: TelemetryAttributes ) -> "ToolboxTool": diff --git a/packages/toolbox-langchain/tests/test_async_client.py b/packages/toolbox-langchain/tests/test_async_client.py index f2632670d..a0b6d3e24 100644 --- a/packages/toolbox-langchain/tests/test_async_client.py +++ b/packages/toolbox-langchain/tests/test_async_client.py @@ -63,7 +63,9 @@ def mock_session(self): def mock_core_client_instance(self, mock_session): mock = AsyncMock(spec=ToolboxCoreClient) - async def mock_load_tool_impl(name, auth_token_getters, bound_params): + async def mock_load_tool_impl( + name, auth_token_getters, bound_params, secure_params={} + ): tool_schema_dict = MANIFEST_JSON["tools"].get(name) if not tool_schema_dict: raise ValueError(f"Tool '{name}' not in mock manifest_dict") @@ -83,7 +85,7 @@ async def mock_load_tool_impl(name, auth_token_getters, bound_params): mock.load_tool = AsyncMock(side_effect=mock_load_tool_impl) async def mock_load_toolset_impl( - name, auth_token_getters, bound_params, strict + name, auth_token_getters, bound_params, strict, secure_params={} ): core_tools_list = [] for tool_name_iter, tool_schema_dict in MANIFEST_JSON["tools"].items(): @@ -423,3 +425,35 @@ async def test_telemetry_enabled_forwarded( ) call_kwargs = mock_core_client_constructor.call_args[1] assert call_kwargs["telemetry_enabled"] == telemetry_enabled + + async def test_aload_tool_with_secure_params( + self, mock_client, mock_core_client_instance + ): + sec_params = {"api_key": "secret_key"} + tool = await mock_client.aload_tool("test_tool_1", secure_params=sec_params) + + assert isinstance(tool, AsyncToolboxTool) + mock_core_client_instance.load_tool.assert_called_once_with( + name="test_tool_1", + auth_token_getters={}, + bound_params={}, + secure_params=sec_params, + ) + + async def test_aload_toolset_with_secure_params( + self, mock_client, mock_core_client_instance + ): + sec_params = {"api_key": "secret_key"} + tools = await mock_client.aload_toolset( + "my_set", secure_params=sec_params, strict=True + ) + + assert len(tools) == 2 + assert isinstance(tools[0], AsyncToolboxTool) + mock_core_client_instance.load_toolset.assert_called_once_with( + name="my_set", + auth_token_getters={}, + bound_params={}, + strict=True, + secure_params=sec_params, + ) diff --git a/packages/toolbox-langchain/tests/test_async_tools.py b/packages/toolbox-langchain/tests/test_async_tools.py index d4d624174..6381c3db5 100644 --- a/packages/toolbox-langchain/tests/test_async_tools.py +++ b/packages/toolbox-langchain/tests/test_async_tools.py @@ -176,6 +176,57 @@ async def test_toolbox_tool_bind_params(self, toolbox_tool, params_to_bind): for bound_param_name in params_to_bind.keys(): assert bound_param_name not in new_core_tool_signature_params + async def test_toolbox_tool_bind_param(self, toolbox_tool): + new_langchain_tool = toolbox_tool.bind_param("param1", "bound-value") + assert isinstance( + new_langchain_tool._AsyncToolboxTool__core_tool, ToolboxCoreTool + ) + new_core_tool_signature_params = ( + new_langchain_tool._AsyncToolboxTool__core_tool.__signature__.parameters + ) + assert "param1" not in new_core_tool_signature_params + + async def test_toolbox_tool_bind_secure_params(self, toolbox_tool): + original_core_tool = toolbox_tool._AsyncToolboxTool__core_tool + with patch.object( + original_core_tool, + "bind_secure_params", + wraps=original_core_tool.bind_secure_params, + ) as mock_core_bind: + # Inject a secure param to test binding + from toolbox_core.protocol import ParameterSchema + + original_core_tool._ToolboxTool__secure_params = [ + ParameterSchema( + name="api_key", + type="string", + description="key", + required=True, + ) + ] + new_langchain_tool = toolbox_tool.bind_secure_params({"api_key": "sec123"}) + mock_core_bind.assert_called_once_with({"api_key": "sec123"}) + assert isinstance( + new_langchain_tool._AsyncToolboxTool__core_tool, ToolboxCoreTool + ) + + async def test_toolbox_tool_bind_secure_param(self, toolbox_tool): + original_core_tool = toolbox_tool._AsyncToolboxTool__core_tool + from toolbox_core.protocol import ParameterSchema + + original_core_tool._ToolboxTool__secure_params = [ + ParameterSchema( + name="api_key", type="string", description="key", required=True + ) + ] + new_langchain_tool = toolbox_tool.bind_secure_param("api_key", "sec123") + assert isinstance( + new_langchain_tool._AsyncToolboxTool__core_tool, ToolboxCoreTool + ) + assert new_langchain_tool._AsyncToolboxTool__core_tool._bound_secure_params == { + "api_key": "sec123" + } + async def test_toolbox_tool_bind_params_invalid(self, toolbox_tool): with pytest.raises( ValueError, match="unable to bind parameters: no parameter named param3" diff --git a/packages/toolbox-langchain/tests/test_client.py b/packages/toolbox-langchain/tests/test_client.py index 1b5275224..62c60f313 100644 --- a/packages/toolbox-langchain/tests/test_client.py +++ b/packages/toolbox-langchain/tests/test_client.py @@ -573,3 +573,39 @@ def test_telemetry_enabled_forwarded( ToolboxClient(URL, telemetry_enabled=telemetry_enabled) call_kwargs = mock_core_client_constructor.call_args[1] assert call_kwargs["telemetry_enabled"] == telemetry_enabled + + @patch("toolbox_langchain.client.ToolboxCoreSyncClient") + def test_load_tool_with_secure_params(self, mock_core_client_constructor): + mock_core_client = mock_core_client_constructor.return_value + mock_core_client.load_tool.return_value = create_mock_core_sync_tool() + + client = ToolboxClient(URL) + sec_params = {"api_key": "secret_key"} + tool = client.load_tool("my_tool", secure_params=sec_params) + + assert isinstance(tool, ToolboxTool) + mock_core_client.load_tool.assert_called_once_with( + name="my_tool", + auth_token_getters={}, + bound_params={}, + secure_params=sec_params, + ) + + @patch("toolbox_langchain.client.ToolboxCoreSyncClient") + def test_load_toolset_with_secure_params(self, mock_core_client_constructor): + mock_core_client = mock_core_client_constructor.return_value + mock_core_client.load_toolset.return_value = [create_mock_core_sync_tool()] + + client = ToolboxClient(URL) + sec_params = {"api_key": "secret_key"} + tools = client.load_toolset("my_set", secure_params=sec_params, strict=True) + + assert len(tools) == 1 + assert isinstance(tools[0], ToolboxTool) + mock_core_client.load_toolset.assert_called_once_with( + name="my_set", + auth_token_getters={}, + bound_params={}, + strict=True, + secure_params=sec_params, + ) diff --git a/packages/toolbox-langchain/tests/test_tools.py b/packages/toolbox-langchain/tests/test_tools.py index 354cf77c3..3890ce793 100644 --- a/packages/toolbox-langchain/tests/test_tools.py +++ b/packages/toolbox-langchain/tests/test_tools.py @@ -139,6 +139,8 @@ def mock_core_tool(self, tool_schema_dict): return_value=new_mock_instance_for_methods ) sync_mock.bind_params = Mock(return_value=new_mock_instance_for_methods) + sync_mock.bind_secure_param = Mock(return_value=new_mock_instance_for_methods) + sync_mock.bind_secure_params = Mock(return_value=new_mock_instance_for_methods) sync_mock.add_telemetry_attributes = Mock( return_value=new_mock_instance_for_methods ) @@ -178,6 +180,8 @@ def mock_core_sync_auth_tool(self, auth_tool_schema_dict): return_value=new_mock_instance_for_methods ) sync_mock.bind_params = Mock(return_value=new_mock_instance_for_methods) + sync_mock.bind_secure_param = Mock(return_value=new_mock_instance_for_methods) + sync_mock.bind_secure_params = Mock(return_value=new_mock_instance_for_methods) return sync_mock @pytest.fixture @@ -234,6 +238,25 @@ def test_toolbox_tool_bind_param(self, toolbox_tool, mock_core_tool): assert isinstance(new_langchain_tool, ToolboxTool) assert new_langchain_tool._ToolboxTool__core_tool == returned_core_tool_mock + def test_toolbox_tool_bind_secure_param(self, toolbox_tool, mock_core_tool): + returned_core_tool_mock = mock_core_tool.bind_secure_params.return_value + new_langchain_tool = toolbox_tool.bind_secure_param("api_key", "secret123") + + mock_core_tool.bind_secure_params.assert_called_once_with( + {"api_key": "secret123"} + ) + assert isinstance(new_langchain_tool, ToolboxTool) + assert new_langchain_tool._ToolboxTool__core_tool == returned_core_tool_mock + + def test_toolbox_tool_bind_secure_params(self, toolbox_tool, mock_core_tool): + returned_core_tool_mock = mock_core_tool.bind_secure_params.return_value + sec_dict = {"api_key": "secret123", "db_pass": "pass"} + new_langchain_tool = toolbox_tool.bind_secure_params(sec_dict) + + mock_core_tool.bind_secure_params.assert_called_once_with(sec_dict) + assert isinstance(new_langchain_tool, ToolboxTool) + assert new_langchain_tool._ToolboxTool__core_tool == returned_core_tool_mock + @pytest.mark.parametrize( "auth_token_getters", [