fix(toolbox-sdk)!: deprecate 'add_auth_headers' in favor of 'add_auth_tokens' (#170)

Client code should not be concerned with how authentication tokens are
handled internally. This commit refactors method and parameter names to
use the more abstract term "token" instead of "header," improving the
developer experience and hiding implementation details.

Author: anubhav756

README.md

@@ -204,7 +204,7 @@ configuring tools for authenticated parameters.
 
 ### Configure SDK for Authentication
 
-Provide the `auth_headers` parameter to the `load_tool` or `load_toolset` calls
+Provide the `auth_tokens` parameter to the `load_tool` or `load_toolset` calls
 with a dictionary. The keys of this dictionary should match the names of the
 authentication sources configured in your tools file (e.g., `my_auth_service`),
 and the values should be callable functions (e.g., lambdas or regular functions)
@@ -213,30 +213,30 @@ that return the ID token of the logged-in user.
 Here's an example:
 
 ```py
-def get_auth_header():
+def get_auth_token():
     # ... Logic to retrieve ID token (e.g., from local storage, OAuth flow)
     # This example just returns a placeholder. Replace with your actual token retrieval.
     return "YOUR_ID_TOKEN"
 
 toolbox = ToolboxClient("http://localhost:5000")
 
-tools = toolbox.load_toolset(auth_headers={ "my_auth_service": get_auth_header })
+tools = toolbox.load_toolset(auth_tokens={ "my_auth_service": get_auth_token })
 
 # OR
 
-tool = toolbox.load_tool("my_tool", auth_headers={ "my_auth_service": get_auth_header })
+tool = toolbox.load_tool("my_tool", auth_tokens={ "my_auth_service": get_auth_token })
 ```
 
-Alternatively, you can call the `add_auth_header` method to configure
+Alternatively, you can call the `add_auth_token` method to configure
 authentication separately.
 
 ```py
-toolbox.add_auth_header("my_auth_service", get_auth_header)
+toolbox.add_auth_token("my_auth_service", get_auth_token)
 ```
 
 > [!NOTE]
-> Authentication headers added via `load_tool`, `load_toolset`, or
-> `add_auth_header` apply to all subsequent tool invocations, regardless of when
+> Authentication tokens added via `load_tool`, `load_toolset`, or
+> `add_auth_token` apply to all subsequent tool invocations, regardless of when
 > the tool was loaded. This ensures a consistent authentication context.
 
 ### Complete Example
@@ -245,7 +245,7 @@ toolbox.add_auth_header("my_auth_service", get_auth_header)
 import asyncio
 from toolbox_langchain_sdk import ToolboxClient
 
-async def get_auth_header():
+async def get_auth_token():
     # Replace with your actual ID token retrieval logic.
     # For example, using a library like google-auth
     # from google.oauth2 import id_token
@@ -257,7 +257,7 @@ async def get_auth_header():
 
 async def main():
     toolbox = ToolboxClient("http://localhost:5000")
-    toolbox.add_auth_header("my_auth_service", get_auth_header)
+    toolbox.add_auth_token("my_auth_service", get_auth_token)
     tools = await toolbox.load_toolset()
     result = await tools[0].arun({"input": "some input"})
     print(result)

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "PyYAML>=6.0.1,<7.0.0",
     "pydantic>=2.7.0,<3.0.0",
     "aiohttp>=3.8.6,<4.0.0",
+    "deprecated>=1.1.0,<2.0.0",
 ]
 
 classifiers = [

requirements.txt

@@ -2,3 +2,4 @@ langchain-core==0.3.21
 PyYAML==6.0.2
 pydantic==2.10.2
 aiohttp==3.11.7
+deprecated==1.2.15

src/toolbox_langchain_sdk/client.py

@@ -1,8 +1,9 @@
 import asyncio
-import warnings
 from typing import Any, Callable, Optional, Type
+from warnings import warn
 
 from aiohttp import ClientSession
+from deprecated import deprecated
 from langchain_core.tools import StructuredTool
 from pydantic import BaseModel
 
@@ -183,13 +184,17 @@ def _process_auth_params(self, manifest: ManifestSchema) -> None:
             # If none of the permitted auth sources of a parameter are
             # registered, raise a warning message to the user.
             if not self._validate_auth(tool_name):
-                warnings.warn(
+                warn(
                     f"Some parameters of tool {tool_name} require authentication, but no valid auth sources are registered. Please register the required sources before use."
                 )
 
+    @deprecated("Please use `add_auth_token` instead.")
     def add_auth_header(
         self, auth_source: str, get_id_token: Callable[[], str]
     ) -> None:
+        self.add_auth_token(auth_source, get_id_token)
+
+    def add_auth_token(self, auth_source: str, get_id_token: Callable[[], str]) -> None:
         """
         Registers a function to retrieve an ID token for a given authentication
         source.
@@ -201,23 +206,39 @@ def add_auth_header(
         self._id_token_getters[auth_source] = get_id_token
 
     async def load_tool(
-        self, tool_name: str, auth_headers: dict[str, Callable[[], str]] = {}
+        self,
+        tool_name: str,
+        auth_tokens: dict[str, Callable[[], str]] = {},
+        auth_headers: Optional[dict[str, Callable[[], str]]] = None,
     ) -> StructuredTool:
         """
         Loads the tool, with the given tool name, from the Toolbox service.
 
         Args:
             tool_name: The name of the tool to load.
-            auth_headers: A mapping of authentication source names to
+            auth_tokens: A mapping of authentication source names to
                 functions that retrieve ID tokens. If provided, these will
                 override or be added to the existing ID token getters.
                 Default: Empty.
 
         Returns:
             A tool loaded from the Toolbox
         """
-        for auth_source, get_id_token in auth_headers.items():
-            self.add_auth_header(auth_source, get_id_token)
+        if auth_headers:
+            if auth_tokens:
+                warn(
+                    "Both `auth_tokens` and `auth_headers` are provided. `auth_headers` is deprecated, and `auth_tokens` will be used.",
+                    DeprecationWarning,
+                )
+            else:
+                warn(
+                    "Argument `auth_headers` is deprecated. Use `auth_tokens` instead.",
+                    DeprecationWarning,
+                )
+                auth_tokens = auth_headers
+
+        for auth_source, get_id_token in auth_tokens.items():
+            self.add_auth_token(auth_source, get_id_token)
 
         manifest: ManifestSchema = await self._load_tool_manifest(tool_name)
 
@@ -228,7 +249,8 @@ async def load_tool(
     async def load_toolset(
         self,
         toolset_name: Optional[str] = None,
-        auth_headers: dict[str, Callable[[], str]] = {},
+        auth_tokens: dict[str, Callable[[], str]] = {},
+        auth_headers: Optional[dict[str, Callable[[], str]]] = None,
     ) -> list[StructuredTool]:
         """
         Loads tools from the Toolbox service, optionally filtered by toolset
@@ -237,16 +259,29 @@ async def load_toolset(
         Args:
             toolset_name: The name of the toolset to load.
                 Default: None. If not provided, then all the tools are loaded.
-            auth_headers: A mapping of authentication source names to
+            auth_tokens: A mapping of authentication source names to
                 functions that retrieve ID tokens. If provided, these will
                 override or be added to the existing ID token getters.
                 Default: Empty.
 
         Returns:
             A list of all tools loaded from the Toolbox.
         """
-        for auth_source, get_id_token in auth_headers.items():
-            self.add_auth_header(auth_source, get_id_token)
+        if auth_headers:
+            if auth_tokens:
+                warn(
+                    "Both `auth_tokens` and `auth_headers` are provided. `auth_headers` is deprecated, and `auth_tokens` will be used.",
+                    DeprecationWarning,
+                )
+            else:
+                warn(
+                    "Argument `auth_headers` is deprecated. Use `auth_tokens` instead.",
+                    DeprecationWarning,
+                )
+                auth_tokens = auth_headers
+
+        for auth_source, get_id_token in auth_tokens.items():
+            self.add_auth_token(auth_source, get_id_token)
 
         tools: list[StructuredTool] = []
         manifest: ManifestSchema = await self._load_toolset_manifest(toolset_name)

src/toolbox_langchain_sdk/utils.py

@@ -1,8 +1,10 @@
 import json
 import warnings
 from typing import Any, Callable, Optional, Type, cast
+from warnings import warn
 
 from aiohttp import ClientSession
+from deprecated import deprecated
 from pydantic import BaseModel, Field, create_model
 
 
@@ -100,22 +102,27 @@ def _parse_type(type_: str) -> Any:
         raise ValueError(f"Unsupported schema type: {type_}")
 
 
+@deprecated("Please use `_get_auth_tokens` instead.")
 def _get_auth_headers(id_token_getters: dict[str, Callable[[], str]]) -> dict[str, str]:
+    return _get_auth_tokens(id_token_getters)
+
+
+def _get_auth_tokens(id_token_getters: dict[str, Callable[[], str]]) -> dict[str, str]:
     """
     Gets id tokens for the given auth sources in the getters map and returns
-    headers to be included in tool invocation.
+    tokens to be included in tool invocation.
 
     Args:
         id_token_getters: A dict that maps auth source names to the functions
         that return its ID token.
 
     Returns:
-        A dictionary of headers to be included in the tool invocation.
+        A dictionary of tokens to be included in the tool invocation.
     """
-    auth_headers = {}
+    auth_tokens = {}
     for auth_source, get_id_token in id_token_getters.items():
-        auth_headers[f"{auth_source}_token"] = get_id_token()
-    return auth_headers
+        auth_tokens[f"{auth_source}_token"] = get_id_token()
+    return auth_tokens
 
 
 async def _invoke_tool(
@@ -141,20 +148,20 @@ async def _invoke_tool(
         invocation.
     """
     url = f"{url}/api/tool/{tool_name}/invoke"
-    auth_headers = _get_auth_headers(id_token_getters)
+    auth_tokens = _get_auth_tokens(id_token_getters)
 
     # ID tokens contain sensitive user information (claims). Transmitting these
     # over HTTP exposes the data to interception and unauthorized access. Always
     # use HTTPS to ensure secure communication and protect user privacy.
-    if auth_headers and not url.startswith("https://"):
-        warnings.warn(
+    if auth_tokens and not url.startswith("https://"):
+        warn(
             "Sending ID token over HTTP. User data may be exposed. Use HTTPS for secure communication."
         )
 
     async with session.post(
         url,
         json=_convert_none_to_empty_string(data),
-        headers=auth_headers,
+        headers=auth_tokens,
     ) as response:
         response.raise_for_status()
         return await response.json()