diff --git a/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py b/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py index 4f2ebff63..c232644f3 100644 --- a/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py +++ b/packages/toolbox-llamaindex/src/toolbox_llamaindex/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 at least one provided parameter or auth token (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-llamaindex/src/toolbox_llamaindex/async_tools.py b/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_tools.py index 94a8f07b6..398a0f5bc 100644 --- a/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_tools.py +++ b/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_tools.py @@ -175,7 +175,7 @@ def bind_param( returns the value. Returns: - A new ToolboxTool instance that is a deep copy of the current + A new AsyncToolboxTool instance that is a deep copy of the current instance, with added bound param. Raises: @@ -183,6 +183,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-llamaindex/src/toolbox_llamaindex/client.py b/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py index 6e06ced11..65b33999e 100644 --- a/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py +++ b/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py @@ -59,6 +59,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. @@ -73,6 +74,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. @@ -103,11 +106,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) @@ -122,6 +130,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 @@ -137,12 +146,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 at least one provided parameter or auth token (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. @@ -173,12 +184,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 = [] @@ -196,6 +212,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. @@ -210,6 +227,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. @@ -240,10 +259,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( @@ -260,6 +284,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 @@ -275,12 +300,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 at least one provided parameter or auth token (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. @@ -311,11 +338,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-llamaindex/src/toolbox_llamaindex/tools.py b/packages/toolbox-llamaindex/src/toolbox_llamaindex/tools.py index d34fa61be..6157286ae 100644 --- a/packages/toolbox-llamaindex/src/toolbox_llamaindex/tools.py +++ b/packages/toolbox-llamaindex/src/toolbox_llamaindex/tools.py @@ -175,6 +175,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-llamaindex/tests/test_async_client.py b/packages/toolbox-llamaindex/tests/test_async_client.py index f34eefb80..28df75297 100644 --- a/packages/toolbox-llamaindex/tests/test_async_client.py +++ b/packages/toolbox-llamaindex/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-llamaindex/tests/test_async_tools.py b/packages/toolbox-llamaindex/tests/test_async_tools.py index 0be4ffd5b..b26675d2c 100644 --- a/packages/toolbox-llamaindex/tests/test_async_tools.py +++ b/packages/toolbox-llamaindex/tests/test_async_tools.py @@ -177,6 +177,58 @@ 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_llamaindex_tool = toolbox_tool.bind_param("param1", "bound-value") + assert isinstance( + new_llamaindex_tool._AsyncToolboxTool__core_tool, ToolboxCoreTool + ) + new_core_tool_signature_params = ( + new_llamaindex_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: + from toolbox_core.protocol import ParameterSchema + + original_core_tool._ToolboxTool__secure_params = [ + ParameterSchema( + name="api_key", + type="string", + description="key", + required=True, + ) + ] + new_llamaindex_tool = toolbox_tool.bind_secure_params({"api_key": "sec123"}) + mock_core_bind.assert_called_once_with({"api_key": "sec123"}) + assert isinstance( + new_llamaindex_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_llamaindex_tool = toolbox_tool.bind_secure_param("api_key", "sec123") + assert isinstance( + new_llamaindex_tool._AsyncToolboxTool__core_tool, ToolboxCoreTool + ) + assert ( + new_llamaindex_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-llamaindex/tests/test_client.py b/packages/toolbox-llamaindex/tests/test_client.py index c903b944c..2f9b56afb 100644 --- a/packages/toolbox-llamaindex/tests/test_client.py +++ b/packages/toolbox-llamaindex/tests/test_client.py @@ -635,3 +635,39 @@ async def test_aload_toolset_with_deprecated_args( bound_params=bound_params, strict=False, ) + + @patch("toolbox_llamaindex.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_llamaindex.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-llamaindex/tests/test_tools.py b/packages/toolbox-llamaindex/tests/test_tools.py index 5e8085093..dee1484ef 100644 --- a/packages/toolbox-llamaindex/tests/test_tools.py +++ b/packages/toolbox-llamaindex/tests/test_tools.py @@ -141,6 +141,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 ) @@ -180,6 +182,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 @@ -236,6 +240,25 @@ def test_toolbox_tool_bind_param(self, toolbox_tool, mock_core_tool): assert isinstance(new_llamaindex_tool, ToolboxTool) assert new_llamaindex_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_llamaindex_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_llamaindex_tool, ToolboxTool) + assert new_llamaindex_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_llamaindex_tool = toolbox_tool.bind_secure_params(sec_dict) + + mock_core_tool.bind_secure_params.assert_called_once_with(sec_dict) + assert isinstance(new_llamaindex_tool, ToolboxTool) + assert new_llamaindex_tool._ToolboxTool__core_tool == returned_core_tool_mock + @pytest.mark.parametrize( "auth_token_getters", [