feat(toolbox-llamaindex): Add client headers to Toolbox (#265)

* feat(toolbox-llamaindex): Add client headers to Toolbox

* lint

* rename package

Author: twishabansal

packages/toolbox-llamaindex/README.md

@@ -18,6 +18,12 @@ applications, enabling advanced orchestration and interaction with GenAI models.
 - [Use with LlamaIndex](#use-with-llamaindex)
     - [Maintain state](#maintain-state)
 - [Manual usage](#manual-usage)
+- [Client to Server Authentication](#client-to-server-authentication)
+    - [When is Client-to-Server Authentication Needed?](#when-is-client-to-server-authentication-needed)
+    - [How it works](#how-it-works)
+    - [Configuration](#configuration)
+    - [Authenticating with Google Cloud Servers](#authenticating-with-google-cloud-servers)
+    - [Step by Step Guide for Cloud Run](#step-by-step-guide-for-cloud-run)
 - [Authenticating Tools](#authenticating-tools)
     - [Supported Authentication Mechanisms](#supported-authentication-mechanisms)
     - [Configure Tools](#configure-tools)
@@ -33,14 +39,6 @@ applications, enabling advanced orchestration and interaction with GenAI models.
 
 <!-- /TOC -->
 
-## Installation
-
-```bash
-pip install toolbox-llamaindex
-```
-
-## Quickstart
-
 Here's a minimal example to get you started using
 [LlamaIndex](https://docs.llamaindex.ai/en/stable/#getting-started):
 
@@ -168,6 +166,86 @@ result = tools[0].call(name="Alice", age=30)
 This is useful for testing tools or when you need precise control over tool
 execution outside of an agent framework.
 
+## Client to Server Authentication
+
+This section describes how to authenticate the ToolboxClient itself when
+connecting to a Toolbox server instance that requires authentication. This is
+crucial for securing your Toolbox server endpoint, especially when deployed on
+platforms like Cloud Run, GKE,  or any environment where unauthenticated access is restricted.
+
+This client-to-server authentication ensures that the Toolbox server can verify the identity of the client making the request before any tool is loaded or called. It is different from [Authenticating Tools](#authenticating-tools), which deals with providing credentials for specific tools within an already connected Toolbox session.
+
+### When is Client-to-Server Authentication Needed?
+
+You'll need this type of authentication if your Toolbox server is configured to deny unauthenticated requests. For example:
+
+- Your Toolbox server is deployed on Cloud Run and configured to "Require authentication."
+- Your server is behind an Identity-Aware Proxy (IAP) or a similar authentication layer.
+- You have custom authentication middleware on your self-hosted Toolbox server.
+
+Without proper client authentication in these scenarios, attempts to connect or
+make calls (like `load_tool`) will likely fail with `Unauthorized` errors.
+
+### How it works
+
+The `ToolboxClient` (and `ToolboxSyncClient`) allows you to specify functions (or coroutines for the async client) that dynamically generate HTTP headers for every request sent to the Toolbox server. The most common use case is to add an Authorization header with a bearer token (e.g., a Google ID token).
+
+These header-generating functions are called just before each request, ensuring
+that fresh credentials or header values can be used.
+
+### Configuration
+
+You can configure these dynamic headers in two ways:
+
+1. **During Client Initialization**
+
+    ```python
+    from toolbox_llamaindex import ToolboxClient
+
+    client = ToolboxClient("toolbox-url", headers={"header1": header1_getter, "header2": header2_getter, ...})
+    ```
+
+1. **After Client Initialization**
+
+    ```python
+    from toolbox_llamaindex import ToolboxClient
+
+    client = ToolboxClient("toolbox-url")
+    client.add_headers({"header1": header1_getter, "header2": header2_getter, ...})
+    ```
+
+### Authenticating with Google Cloud Servers
+
+For Toolbox servers hosted on Google Cloud (e.g., Cloud Run) and requiring
+`Google ID token` authentication, the helper module
+[auth_methods](src/toolbox_core/auth_methods.py) provides utility functions.
+
+### Step by Step Guide for Cloud Run
+
+1. **Configure Permissions**: [Grant](https://cloud.google.com/run/docs/securing/managing-access#service-add-principals) the `roles/run.invoker` IAM role on the Cloud
+   Run service to the principal. This could be your `user account email` or a
+   `service account`.
+2. **Configure Credentials**
+    - Local Development: Set up
+   [ADC](https://cloud.google.com/docs/authentication/set-up-adc-local-dev-environment).
+    - Google Cloud Environments: When running within Google Cloud (e.g., Compute
+      Engine, GKE, another Cloud Run service, Cloud Functions), ADC is typically
+      configured automatically, using the environment's default service account.
+3. **Connect to the Toolbox Server**
+
+    ```python
+    from toolbox_core import auth_methods
+
+    auth_token_provider = auth_methods.aget_google_id_token # can also use sync method
+    client = ToolboxClient(
+        URL,
+        client_headers={"Authorization": auth_token_provider},
+    )
+    tools = await client.load_toolset()
+
+    # Now, you can use the client as usual.
+    ```
+
 ## Authenticating Tools
 
 > [!WARNING]

packages/toolbox-llamaindex/src/toolbox_llamaindex/async_client.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, Callable, Optional, Union
+from typing import Any, Awaitable, Callable, Mapping, Optional, Union
 from warnings import warn
 
 from aiohttp import ClientSession
@@ -30,6 +30,9 @@ def __init__(
         self,
         url: str,
         session: ClientSession,
+        client_headers: Optional[
+            Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]], str]]
+        ] = None,
     ):
         """
         Initializes the AsyncToolboxClient for the Toolbox service at the given URL.
@@ -38,7 +41,9 @@ def __init__(
             url: The base URL of the Toolbox service.
             session: An HTTP client session.
         """
-        self.__core_client = ToolboxCoreClient(url=url, session=session)
+        self.__core_client = ToolboxCoreClient(
+            url=url, session=session, client_headers=client_headers
+        )
 
     async def aload_tool(
         self,
@@ -185,3 +190,16 @@ 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)

packages/toolbox-llamaindex/src/toolbox_llamaindex/client.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from asyncio import to_thread
-from typing import Any, Callable, Optional, Union
+from typing import Any, Awaitable, Callable, Mapping, Optional, Union
 from warnings import warn
 
 from toolbox_core.sync_client import ToolboxSyncClient as ToolboxCoreSyncClient
@@ -27,14 +27,19 @@ class ToolboxClient:
     def __init__(
         self,
         url: str,
+        client_headers: Optional[
+            Mapping[str, Union[Callable[[], str], Callable[[], Awaitable[str]], str]]
+        ] = None,
     ) -> None:
         """
         Initializes the ToolboxClient for the Toolbox service at the given URL.
 
         Args:
             url: The base URL of the Toolbox service.
         """
-        self.__core_client = ToolboxCoreSyncClient(url=url)
+        self.__core_client = ToolboxCoreSyncClient(
+            url=url, client_headers=client_headers
+        )
 
     async def aload_tool(
         self,
@@ -287,3 +292,16 @@ 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)

packages/toolbox-llamaindex/tests/test_async_client.py

@@ -339,3 +339,22 @@ async def test_load_toolset_not_implemented(self, mock_client):
         assert "Synchronous methods not supported by async client." in str(
             excinfo.value
         )
+
+    @patch("toolbox_llamaindex.async_client.ToolboxCoreClient")
+    async def test_init_with_client_headers(
+        self, mock_core_client_constructor, mock_session
+    ):
+        """Tests that client_headers are passed to the core client during initialization."""
+        headers = {"X-Test-Header": "value"}
+        AsyncToolboxClient(URL, session=mock_session, client_headers=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
+        )

packages/toolbox-llamaindex/tests/test_client.py

@@ -422,3 +422,21 @@ async def test_aload_toolset_with_args(
             bound_params=bound_params,
             strict=True,
         )
+
+    @patch("toolbox_llamaindex.client.ToolboxCoreSyncClient")
+    def test_init_with_client_headers(self, mock_core_client_constructor):
+        """Tests that client_headers are passed to the core client during initialization."""
+        headers = {"X-Test-Header": "value"}
+        ToolboxClient(URL, client_headers=headers)
+        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)