add headers

Author: twishabansal

packages/toolbox-langchain/src/toolbox_langchain/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, Callable, Optional, Union, Mapping, Awaitable
 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,7 @@ 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 +188,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)

packages/toolbox-langchain/src/toolbox_langchain/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, Callable, Optional, Union, Mapping, Awaitable
 from warnings import warn
 
 from toolbox_core.sync_client import ToolboxSyncClient as ToolboxCoreSyncClient
@@ -26,14 +26,17 @@ 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,
@@ -286,3 +289,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)