From a047039cf9ba56a271b104dd7f6e9fe6c58b1303 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Tue, 10 Jun 2025 20:33:15 +0530 Subject: [PATCH 1/3] chore!: remove addHeaders feature --- packages/toolbox-core/README.md | 23 ++----- .../toolbox-core/src/toolbox_core/client.py | 24 ------- .../src/toolbox_core/sync_client.py | 17 ----- packages/toolbox-core/tests/test_client.py | 44 ------------ .../toolbox-core/tests/test_sync_client.py | 67 ------------------- packages/toolbox-langchain/README.md | 27 +++----- .../src/toolbox_langchain/async_client.py | 15 ----- .../src/toolbox_langchain/client.py | 17 +---- .../toolbox-langchain/tests/test_client.py | 9 --- packages/toolbox-llamaindex/README.md | 27 +++----- .../src/toolbox_llamaindex/async_client.py | 13 ---- .../src/toolbox_llamaindex/client.py | 15 +---- .../tests/test_async_client.py | 8 --- .../toolbox-llamaindex/tests/test_client.py | 9 --- 14 files changed, 25 insertions(+), 290 deletions(-) diff --git a/packages/toolbox-core/README.md b/packages/toolbox-core/README.md index 7da6f149..c4461df3 100644 --- a/packages/toolbox-core/README.md +++ b/packages/toolbox-core/README.md @@ -300,24 +300,15 @@ that fresh credentials or header values can be used. ### Configuration -You can configure these dynamic headers in two ways: +You can configure these dynamic headers as seen below: -1. **During Client Initialization** - - ```python - from toolbox_core import ToolboxClient - - async with ToolboxClient("toolbox-url", client_headers={"header1": header1_getter, "header2": header2_getter, ...}) as client: - ``` - -1. **After Client Initialization** - - ```python - from toolbox_core import ToolboxClient +```python +from toolbox_core import ToolboxClient - async with ToolboxClient("toolbox-url") as client: - client.add_headers({"header1": header1_getter, "header2": header2_getter, ...}) - ``` +async with ToolboxClient("toolbox-url", client_headers={"header1": header1_getter, "header2": header2_getter, ...}) as client: + # Use client + pass +``` ### Authenticating with Google Cloud Servers diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 1fcea46c..24fec26a 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -330,27 +330,3 @@ async def load_toolset( ) return tools - - def add_headers( - self, - headers: Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]]]], - ) -> None: - """ - Add headers to be included in each request sent through this client. - - Args: - headers: Headers to include in each request sent through this client. - - Raises: - ValueError: If any of the headers are already registered in the client. - """ - existing_headers = self.__client_headers.keys() - incoming_headers = headers.keys() - duplicates = existing_headers & incoming_headers - if duplicates: - raise ValueError( - f"Client header(s) `{', '.join(duplicates)}` already registered in the client." - ) - - merged_headers = {**self.__client_headers, **headers} - self.__client_headers = merged_headers diff --git a/packages/toolbox-core/src/toolbox_core/sync_client.py b/packages/toolbox-core/src/toolbox_core/sync_client.py index 7754468f..dcc1008b 100644 --- a/packages/toolbox-core/src/toolbox_core/sync_client.py +++ b/packages/toolbox-core/src/toolbox_core/sync_client.py @@ -153,23 +153,6 @@ def load_toolset( for async_tool in async_tools ] - def add_headers( - self, - headers: Mapping[ - str, Union[Callable[[], str], Callable[[], Awaitable[str]], str] - ], - ) -> None: - """ - Add headers to be included in each request sent through this client. - - Args: - headers: Headers to include in each request sent through this client. - - Raises: - ValueError: If any of the headers are already registered in the client. - """ - self.__async_client.add_headers(headers) - def __enter__(self): """Enter the runtime context related to this client instance.""" return self diff --git a/packages/toolbox-core/tests/test_client.py b/packages/toolbox-core/tests/test_client.py index 1b36fe0d..b92e78a1 100644 --- a/packages/toolbox-core/tests/test_client.py +++ b/packages/toolbox-core/tests/test_client.py @@ -1417,50 +1417,6 @@ async def test_load_toolset_with_headers( assert len(tools) == 1 assert tools[0].__name__ == tool_name - @pytest.mark.asyncio - async def test_add_headers_success( - self, aioresponses, test_tool_str, static_header - ): - """Tests adding headers after client initialization.""" - tool_name = "tool_after_add_headers" - manifest = ManifestSchema( - serverVersion="0.0.0", tools={tool_name: test_tool_str} - ) - expected_payload = {"result": "added_ok"} - - get_callback = self.create_callback_factory( - expected_header=static_header, - callback_payload=manifest.model_dump(), - ) - aioresponses.get(f"{TEST_BASE_URL}/api/tool/{tool_name}", callback=get_callback) - - post_callback = self.create_callback_factory( - expected_header=static_header, - callback_payload=expected_payload, - ) - aioresponses.post( - f"{TEST_BASE_URL}/api/tool/{tool_name}/invoke", callback=post_callback - ) - - async with ToolboxClient(TEST_BASE_URL) as client: - client.add_headers(static_header) - assert client._ToolboxClient__client_headers == static_header - - tool = await client.load_tool(tool_name) - result = await tool(param1="test") - assert result == expected_payload["result"] - - @pytest.mark.asyncio - async def test_add_headers_duplicate_fail(self, static_header): - """Tests that adding a duplicate header via add_headers raises - ValueError.""" - async with ToolboxClient(TEST_BASE_URL, client_headers=static_header) as client: - with pytest.raises( - ValueError, - match=f"Client header\\(s\\) `X-Static-Header` already registered", - ): - await client.add_headers(static_header) - @pytest.mark.asyncio async def test_client_header_auth_token_conflict_fail( self, aioresponses, test_tool_auth diff --git a/packages/toolbox-core/tests/test_sync_client.py b/packages/toolbox-core/tests/test_sync_client.py index 32634d53..ad887211 100644 --- a/packages/toolbox-core/tests/test_sync_client.py +++ b/packages/toolbox-core/tests/test_sync_client.py @@ -349,73 +349,6 @@ def test_load_tool_raises_if_loop_or_thread_none(self): client.close() # Clean up manually created client # sync_client_environment will handle the final cleanup of original_class_loop/thread. - -class TestSyncClientHeaders: - """Additive tests for client header functionality specific to ToolboxSyncClient if any, - or counterparts to async client header tests.""" - - def test_sync_add_headers_success( - self, aioresponses, test_tool_str_schema, sync_client - ): - tool_name = "tool_after_add_headers_sync" - manifest = ManifestSchema( - serverVersion="0.0.0", tools={tool_name: test_tool_str_schema} - ) - expected_payload = {"result": "added_sync_ok"} - headers_to_add = {"X-Custom-SyncHeader": "sync_value"} - - def get_callback(url, **kwargs): - # The sync_client might have default headers. Check ours are present. - assert kwargs.get("headers") is not None - for key, value in headers_to_add.items(): - assert kwargs["headers"].get(key) == value - return CallbackResult(status=200, payload=manifest.model_dump()) - - aioresponses.get(f"{TEST_BASE_URL}/api/tool/{tool_name}", callback=get_callback) - - def post_callback(url, **kwargs): - assert kwargs.get("headers") is not None - for key, value in headers_to_add.items(): - assert kwargs["headers"].get(key) == value - return CallbackResult(status=200, payload=expected_payload) - - aioresponses.post( - f"{TEST_BASE_URL}/api/tool/{tool_name}/invoke", callback=post_callback - ) - - sync_client.add_headers(headers_to_add) - tool = sync_client.load_tool(tool_name) - result = tool(param1="test") - assert result == expected_payload["result"] - - def test_sync_add_headers_duplicate_fail(self): - """Tests that adding a duplicate header via add_headers raises ValueError (from async client).""" - initial_headers = {"X-Initial-Header": "initial_value"} - mock_async_client = AsyncMock(spec=ToolboxClient) - - # Configure add_headers to simulate the ValueError from ToolboxClient - def mock_add_headers(headers): - # Simulate ToolboxClient's check - if "X-Initial-Header" in headers: - raise ValueError( - "Client header(s) `X-Initial-Header` already registered" - ) - - mock_async_client.add_headers = Mock(side_effect=mock_add_headers) - - with patch( - "toolbox_core.sync_client.ToolboxClient", return_value=mock_async_client - ): - with ToolboxSyncClient( - TEST_BASE_URL, client_headers=initial_headers - ) as client: - with pytest.raises( - ValueError, - match="Client header\\(s\\) `X-Initial-Header` already registered", - ): - client.add_headers({"X-Initial-Header": "another_value"}) - - class TestSyncAuth: @pytest.fixture def expected_header_token(self): diff --git a/packages/toolbox-langchain/README.md b/packages/toolbox-langchain/README.md index 478bd087..42146d7c 100644 --- a/packages/toolbox-langchain/README.md +++ b/packages/toolbox-langchain/README.md @@ -231,27 +231,16 @@ that fresh credentials or header values can be used. ### Configuration -You can configure these dynamic headers in two ways: +You can configure these dynamic headers as follows: -1. **During Client Initialization** - - ```python - from toolbox_langchain import ToolboxClient - - client = ToolboxClient( - "toolbox-url", - client_headers={"header1": header1_getter, "header2": header2_getter, ...} - ) - ``` - -1. **After Client Initialization** - - ```python - from toolbox_langchain import ToolboxClient +```python +from toolbox_langchain import ToolboxClient - client = ToolboxClient("toolbox-url") - client.add_headers({"header1": header1_getter, "header2": header2_getter, ...}) - ``` +client = ToolboxClient( + "toolbox-url", + client_headers={"header1": header1_getter, "header2": header2_getter, ...} +) +``` ### Authenticating with Google Cloud Servers diff --git a/packages/toolbox-langchain/src/toolbox_langchain/async_client.py b/packages/toolbox-langchain/src/toolbox_langchain/async_client.py index a29bfb85..4757188d 100644 --- a/packages/toolbox-langchain/src/toolbox_langchain/async_client.py +++ b/packages/toolbox-langchain/src/toolbox_langchain/async_client.py @@ -190,18 +190,3 @@ def load_toolset( strict: bool = False, ) -> list[AsyncToolboxTool]: raise NotImplementedError("Synchronous methods not supported by async client.") - - def add_headers( - self, - headers: Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]]]], - ) -> None: - """ - Add headers to be included in each request sent through this client. - - Args: - headers: Headers to include in each request sent through this client. - - Raises: - ValueError: If any of the headers are already registered in the client. - """ - self.__core_client.add_headers(headers) diff --git a/packages/toolbox-langchain/src/toolbox_langchain/client.py b/packages/toolbox-langchain/src/toolbox_langchain/client.py index 19cb60d0..f9e76799 100644 --- a/packages/toolbox-langchain/src/toolbox_langchain/client.py +++ b/packages/toolbox-langchain/src/toolbox_langchain/client.py @@ -290,19 +290,4 @@ def load_toolset( tools = [] for core_sync_tool in core_sync_tools: tools.append(ToolboxTool(core_tool=core_sync_tool)) - return tools - - def add_headers( - self, - headers: Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]]]], - ) -> None: - """ - Add headers to be included in each request sent through this client. - - Args: - headers: Headers to include in each request sent through this client. - - Raises: - ValueError: If any of the headers are already registered in the client. - """ - self.__core_client.add_headers(headers) + return tools \ No newline at end of file diff --git a/packages/toolbox-langchain/tests/test_client.py b/packages/toolbox-langchain/tests/test_client.py index bf42c870..73ecdb5d 100644 --- a/packages/toolbox-langchain/tests/test_client.py +++ b/packages/toolbox-langchain/tests/test_client.py @@ -429,12 +429,3 @@ def test_init_with_client_headers(self, mock_core_client_constructor): mock_core_client_constructor.assert_called_once_with( url=URL, client_headers=headers ) - - @patch("toolbox_langchain.client.ToolboxCoreSyncClient") - def test_add_headers(self, mock_core_client_constructor): - """Tests that add_headers calls the core client's add_headers.""" - mock_core_instance = mock_core_client_constructor.return_value - client = ToolboxClient(URL) - headers = {"X-Another-Header": "dynamic_value"} - client.add_headers(headers) - mock_core_instance.add_headers.assert_called_once_with(headers) diff --git a/packages/toolbox-llamaindex/README.md b/packages/toolbox-llamaindex/README.md index c7539231..db0b1658 100644 --- a/packages/toolbox-llamaindex/README.md +++ b/packages/toolbox-llamaindex/README.md @@ -204,27 +204,16 @@ that fresh credentials or header values can be used. ### Configuration -You can configure these dynamic headers in two ways: +You can configure these dynamic headers as follows: -1. **During Client Initialization** - - ```python - from toolbox_llamaindex import ToolboxClient - - client = ToolboxClient( - "toolbox-url", - client_headers={"header1": header1_getter, "header2": header2_getter, ...} - ) - ``` - -1. **After Client Initialization** - - ```python - from toolbox_llamaindex import ToolboxClient +```python +from toolbox_llamaindex import ToolboxClient - client = ToolboxClient("toolbox-url") - client.add_headers({"header1": header1_getter, "header2": header2_getter, ...}) - ``` +client = ToolboxClient( + "toolbox-url", + client_headers={"header1": header1_getter, "header2": header2_getter, ...} +) +``` ### Authenticating with Google Cloud Servers diff --git a/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py b/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py index fa1f3807..def070f5 100644 --- a/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py +++ b/packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py @@ -190,16 +190,3 @@ def load_toolset( strict: bool = False, ) -> list[AsyncToolboxTool]: raise NotImplementedError("Synchronous methods not supported by async client.") - - def add_headers( - self, - headers: Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]]]], - ) -> None: - """ - Add headers to be included in each request sent through this client. - Args: - headers: Headers to include in each request sent through this client. - Raises: - ValueError: If any of the headers are already registered in the client. - """ - self.__core_client.add_headers(headers) diff --git a/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py b/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py index ef277909..b4fe22f1 100644 --- a/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py +++ b/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py @@ -291,17 +291,4 @@ def load_toolset( tools = [] for core_sync_tool in core_sync_tools: tools.append(ToolboxTool(core_tool=core_sync_tool)) - return tools - - def add_headers( - self, - headers: Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]]]], - ) -> None: - """ - Add headers to be included in each request sent through this client. - Args: - headers: Headers to include in each request sent through this client. - Raises: - ValueError: If any of the headers are already registered in the client. - """ - self.__core_client.add_headers(headers) + return tools \ No newline at end of file diff --git a/packages/toolbox-llamaindex/tests/test_async_client.py b/packages/toolbox-llamaindex/tests/test_async_client.py index cc798715..6cefad47 100644 --- a/packages/toolbox-llamaindex/tests/test_async_client.py +++ b/packages/toolbox-llamaindex/tests/test_async_client.py @@ -350,11 +350,3 @@ async def test_init_with_client_headers( mock_core_client_constructor.assert_called_once_with( url=URL, session=mock_session, client_headers=headers ) - - async def test_add_headers(self, mock_client): - """Tests that add_headers calls the core client's add_headers.""" - headers = {"X-Another-Header": lambda: "dynamic_value"} - mock_client.add_headers(headers) - mock_client._AsyncToolboxClient__core_client.add_headers.assert_called_once_with( - headers - ) diff --git a/packages/toolbox-llamaindex/tests/test_client.py b/packages/toolbox-llamaindex/tests/test_client.py index 0c525b8d..b059e108 100644 --- a/packages/toolbox-llamaindex/tests/test_client.py +++ b/packages/toolbox-llamaindex/tests/test_client.py @@ -431,12 +431,3 @@ def test_init_with_client_headers(self, mock_core_client_constructor): mock_core_client_constructor.assert_called_once_with( url=URL, client_headers=headers ) - - @patch("toolbox_llamaindex.client.ToolboxCoreSyncClient") - def test_add_headers(self, mock_core_client_constructor): - """Tests that add_headers calls the core client's add_headers.""" - mock_core_instance = mock_core_client_constructor.return_value - client = ToolboxClient(URL) - headers = {"X-Another-Header": "dynamic_value"} - client.add_headers(headers) - mock_core_instance.add_headers.assert_called_once_with(headers) From a9c33c99da571fdcba2ec5a4012de78f8a09218e Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Tue, 10 Jun 2025 20:37:17 +0530 Subject: [PATCH 2/3] lint --- packages/toolbox-core/tests/test_sync_client.py | 1 + packages/toolbox-langchain/src/toolbox_langchain/client.py | 2 +- packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py | 2 +- 3 files changed, 3 insertions(+), 2 deletions(-) diff --git a/packages/toolbox-core/tests/test_sync_client.py b/packages/toolbox-core/tests/test_sync_client.py index ad887211..ebc7f479 100644 --- a/packages/toolbox-core/tests/test_sync_client.py +++ b/packages/toolbox-core/tests/test_sync_client.py @@ -349,6 +349,7 @@ def test_load_tool_raises_if_loop_or_thread_none(self): client.close() # Clean up manually created client # sync_client_environment will handle the final cleanup of original_class_loop/thread. + class TestSyncAuth: @pytest.fixture def expected_header_token(self): diff --git a/packages/toolbox-langchain/src/toolbox_langchain/client.py b/packages/toolbox-langchain/src/toolbox_langchain/client.py index f9e76799..00dc9c4d 100644 --- a/packages/toolbox-langchain/src/toolbox_langchain/client.py +++ b/packages/toolbox-langchain/src/toolbox_langchain/client.py @@ -290,4 +290,4 @@ def load_toolset( tools = [] for core_sync_tool in core_sync_tools: tools.append(ToolboxTool(core_tool=core_sync_tool)) - return tools \ No newline at end of file + return tools diff --git a/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py b/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py index b4fe22f1..3cc8ea68 100644 --- a/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py +++ b/packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py @@ -291,4 +291,4 @@ def load_toolset( tools = [] for core_sync_tool in core_sync_tools: tools.append(ToolboxTool(core_tool=core_sync_tool)) - return tools \ No newline at end of file + return tools From d0ad8af6ad5e3c5c215577105c3896fd9deaf9c0 Mon Sep 17 00:00:00 2001 From: Twisha Bansal Date: Tue, 10 Jun 2025 20:50:51 +0530 Subject: [PATCH 3/3] remove files --- packages/toolbox-core/README.md | 23 +++++-- .../toolbox-core/src/toolbox_core/client.py | 24 +++++++ .../src/toolbox_core/sync_client.py | 17 +++++ packages/toolbox-core/tests/test_client.py | 44 +++++++++++++ .../toolbox-core/tests/test_sync_client.py | 66 +++++++++++++++++++ packages/toolbox-langchain/README.md | 27 +++++--- .../src/toolbox_langchain/async_client.py | 15 +++++ .../src/toolbox_langchain/client.py | 15 +++++ .../toolbox-langchain/tests/test_client.py | 9 +++ 9 files changed, 225 insertions(+), 15 deletions(-) diff --git a/packages/toolbox-core/README.md b/packages/toolbox-core/README.md index c4461df3..7da6f149 100644 --- a/packages/toolbox-core/README.md +++ b/packages/toolbox-core/README.md @@ -300,15 +300,24 @@ that fresh credentials or header values can be used. ### Configuration -You can configure these dynamic headers as seen below: +You can configure these dynamic headers in two ways: -```python -from toolbox_core import ToolboxClient +1. **During Client Initialization** -async with ToolboxClient("toolbox-url", client_headers={"header1": header1_getter, "header2": header2_getter, ...}) as client: - # Use client - pass -``` + ```python + from toolbox_core import ToolboxClient + + async with ToolboxClient("toolbox-url", client_headers={"header1": header1_getter, "header2": header2_getter, ...}) as client: + ``` + +1. **After Client Initialization** + + ```python + from toolbox_core import ToolboxClient + + async with ToolboxClient("toolbox-url") as client: + client.add_headers({"header1": header1_getter, "header2": header2_getter, ...}) + ``` ### Authenticating with Google Cloud Servers diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 24fec26a..1fcea46c 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -330,3 +330,27 @@ async def load_toolset( ) return tools + + def add_headers( + self, + headers: Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]]]], + ) -> None: + """ + Add headers to be included in each request sent through this client. + + Args: + headers: Headers to include in each request sent through this client. + + Raises: + ValueError: If any of the headers are already registered in the client. + """ + existing_headers = self.__client_headers.keys() + incoming_headers = headers.keys() + duplicates = existing_headers & incoming_headers + if duplicates: + raise ValueError( + f"Client header(s) `{', '.join(duplicates)}` already registered in the client." + ) + + merged_headers = {**self.__client_headers, **headers} + self.__client_headers = merged_headers diff --git a/packages/toolbox-core/src/toolbox_core/sync_client.py b/packages/toolbox-core/src/toolbox_core/sync_client.py index dcc1008b..7754468f 100644 --- a/packages/toolbox-core/src/toolbox_core/sync_client.py +++ b/packages/toolbox-core/src/toolbox_core/sync_client.py @@ -153,6 +153,23 @@ def load_toolset( for async_tool in async_tools ] + def add_headers( + self, + headers: Mapping[ + str, Union[Callable[[], str], Callable[[], Awaitable[str]], str] + ], + ) -> None: + """ + Add headers to be included in each request sent through this client. + + Args: + headers: Headers to include in each request sent through this client. + + Raises: + ValueError: If any of the headers are already registered in the client. + """ + self.__async_client.add_headers(headers) + def __enter__(self): """Enter the runtime context related to this client instance.""" return self diff --git a/packages/toolbox-core/tests/test_client.py b/packages/toolbox-core/tests/test_client.py index b92e78a1..1b36fe0d 100644 --- a/packages/toolbox-core/tests/test_client.py +++ b/packages/toolbox-core/tests/test_client.py @@ -1417,6 +1417,50 @@ async def test_load_toolset_with_headers( assert len(tools) == 1 assert tools[0].__name__ == tool_name + @pytest.mark.asyncio + async def test_add_headers_success( + self, aioresponses, test_tool_str, static_header + ): + """Tests adding headers after client initialization.""" + tool_name = "tool_after_add_headers" + manifest = ManifestSchema( + serverVersion="0.0.0", tools={tool_name: test_tool_str} + ) + expected_payload = {"result": "added_ok"} + + get_callback = self.create_callback_factory( + expected_header=static_header, + callback_payload=manifest.model_dump(), + ) + aioresponses.get(f"{TEST_BASE_URL}/api/tool/{tool_name}", callback=get_callback) + + post_callback = self.create_callback_factory( + expected_header=static_header, + callback_payload=expected_payload, + ) + aioresponses.post( + f"{TEST_BASE_URL}/api/tool/{tool_name}/invoke", callback=post_callback + ) + + async with ToolboxClient(TEST_BASE_URL) as client: + client.add_headers(static_header) + assert client._ToolboxClient__client_headers == static_header + + tool = await client.load_tool(tool_name) + result = await tool(param1="test") + assert result == expected_payload["result"] + + @pytest.mark.asyncio + async def test_add_headers_duplicate_fail(self, static_header): + """Tests that adding a duplicate header via add_headers raises + ValueError.""" + async with ToolboxClient(TEST_BASE_URL, client_headers=static_header) as client: + with pytest.raises( + ValueError, + match=f"Client header\\(s\\) `X-Static-Header` already registered", + ): + await client.add_headers(static_header) + @pytest.mark.asyncio async def test_client_header_auth_token_conflict_fail( self, aioresponses, test_tool_auth diff --git a/packages/toolbox-core/tests/test_sync_client.py b/packages/toolbox-core/tests/test_sync_client.py index ebc7f479..32634d53 100644 --- a/packages/toolbox-core/tests/test_sync_client.py +++ b/packages/toolbox-core/tests/test_sync_client.py @@ -350,6 +350,72 @@ def test_load_tool_raises_if_loop_or_thread_none(self): # sync_client_environment will handle the final cleanup of original_class_loop/thread. +class TestSyncClientHeaders: + """Additive tests for client header functionality specific to ToolboxSyncClient if any, + or counterparts to async client header tests.""" + + def test_sync_add_headers_success( + self, aioresponses, test_tool_str_schema, sync_client + ): + tool_name = "tool_after_add_headers_sync" + manifest = ManifestSchema( + serverVersion="0.0.0", tools={tool_name: test_tool_str_schema} + ) + expected_payload = {"result": "added_sync_ok"} + headers_to_add = {"X-Custom-SyncHeader": "sync_value"} + + def get_callback(url, **kwargs): + # The sync_client might have default headers. Check ours are present. + assert kwargs.get("headers") is not None + for key, value in headers_to_add.items(): + assert kwargs["headers"].get(key) == value + return CallbackResult(status=200, payload=manifest.model_dump()) + + aioresponses.get(f"{TEST_BASE_URL}/api/tool/{tool_name}", callback=get_callback) + + def post_callback(url, **kwargs): + assert kwargs.get("headers") is not None + for key, value in headers_to_add.items(): + assert kwargs["headers"].get(key) == value + return CallbackResult(status=200, payload=expected_payload) + + aioresponses.post( + f"{TEST_BASE_URL}/api/tool/{tool_name}/invoke", callback=post_callback + ) + + sync_client.add_headers(headers_to_add) + tool = sync_client.load_tool(tool_name) + result = tool(param1="test") + assert result == expected_payload["result"] + + def test_sync_add_headers_duplicate_fail(self): + """Tests that adding a duplicate header via add_headers raises ValueError (from async client).""" + initial_headers = {"X-Initial-Header": "initial_value"} + mock_async_client = AsyncMock(spec=ToolboxClient) + + # Configure add_headers to simulate the ValueError from ToolboxClient + def mock_add_headers(headers): + # Simulate ToolboxClient's check + if "X-Initial-Header" in headers: + raise ValueError( + "Client header(s) `X-Initial-Header` already registered" + ) + + mock_async_client.add_headers = Mock(side_effect=mock_add_headers) + + with patch( + "toolbox_core.sync_client.ToolboxClient", return_value=mock_async_client + ): + with ToolboxSyncClient( + TEST_BASE_URL, client_headers=initial_headers + ) as client: + with pytest.raises( + ValueError, + match="Client header\\(s\\) `X-Initial-Header` already registered", + ): + client.add_headers({"X-Initial-Header": "another_value"}) + + class TestSyncAuth: @pytest.fixture def expected_header_token(self): diff --git a/packages/toolbox-langchain/README.md b/packages/toolbox-langchain/README.md index 42146d7c..478bd087 100644 --- a/packages/toolbox-langchain/README.md +++ b/packages/toolbox-langchain/README.md @@ -231,16 +231,27 @@ that fresh credentials or header values can be used. ### Configuration -You can configure these dynamic headers as follows: +You can configure these dynamic headers in two ways: -```python -from toolbox_langchain import ToolboxClient +1. **During Client Initialization** -client = ToolboxClient( - "toolbox-url", - client_headers={"header1": header1_getter, "header2": header2_getter, ...} -) -``` + ```python + from toolbox_langchain import ToolboxClient + + client = ToolboxClient( + "toolbox-url", + client_headers={"header1": header1_getter, "header2": header2_getter, ...} + ) + ``` + +1. **After Client Initialization** + + ```python + from toolbox_langchain import ToolboxClient + + client = ToolboxClient("toolbox-url") + client.add_headers({"header1": header1_getter, "header2": header2_getter, ...}) + ``` ### Authenticating with Google Cloud Servers diff --git a/packages/toolbox-langchain/src/toolbox_langchain/async_client.py b/packages/toolbox-langchain/src/toolbox_langchain/async_client.py index 4757188d..a29bfb85 100644 --- a/packages/toolbox-langchain/src/toolbox_langchain/async_client.py +++ b/packages/toolbox-langchain/src/toolbox_langchain/async_client.py @@ -190,3 +190,18 @@ def load_toolset( strict: bool = False, ) -> list[AsyncToolboxTool]: raise NotImplementedError("Synchronous methods not supported by async client.") + + def add_headers( + self, + headers: Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]]]], + ) -> None: + """ + Add headers to be included in each request sent through this client. + + Args: + headers: Headers to include in each request sent through this client. + + Raises: + ValueError: If any of the headers are already registered in the client. + """ + self.__core_client.add_headers(headers) diff --git a/packages/toolbox-langchain/src/toolbox_langchain/client.py b/packages/toolbox-langchain/src/toolbox_langchain/client.py index 00dc9c4d..19cb60d0 100644 --- a/packages/toolbox-langchain/src/toolbox_langchain/client.py +++ b/packages/toolbox-langchain/src/toolbox_langchain/client.py @@ -291,3 +291,18 @@ def load_toolset( for core_sync_tool in core_sync_tools: tools.append(ToolboxTool(core_tool=core_sync_tool)) return tools + + def add_headers( + self, + headers: Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]]]], + ) -> None: + """ + Add headers to be included in each request sent through this client. + + Args: + headers: Headers to include in each request sent through this client. + + Raises: + ValueError: If any of the headers are already registered in the client. + """ + self.__core_client.add_headers(headers) diff --git a/packages/toolbox-langchain/tests/test_client.py b/packages/toolbox-langchain/tests/test_client.py index 73ecdb5d..bf42c870 100644 --- a/packages/toolbox-langchain/tests/test_client.py +++ b/packages/toolbox-langchain/tests/test_client.py @@ -429,3 +429,12 @@ def test_init_with_client_headers(self, mock_core_client_constructor): mock_core_client_constructor.assert_called_once_with( url=URL, client_headers=headers ) + + @patch("toolbox_langchain.client.ToolboxCoreSyncClient") + def test_add_headers(self, mock_core_client_constructor): + """Tests that add_headers calls the core client's add_headers.""" + mock_core_instance = mock_core_client_constructor.return_value + client = ToolboxClient(URL) + headers = {"X-Another-Header": "dynamic_value"} + client.add_headers(headers) + mock_core_instance.add_headers.assert_called_once_with(headers)