diff --git a/src/google/adk/tools/google_api_tool/google_api_tool.py b/src/google/adk/tools/google_api_tool/google_api_tool.py index b51f13f5..04d1ebb4 100644 --- a/src/google/adk/tools/google_api_tool/google_api_tool.py +++ b/src/google/adk/tools/google_api_tool/google_api_tool.py @@ -39,6 +39,7 @@ class GoogleApiTool(BaseTool): client_id: Optional[str] = None, client_secret: Optional[str] = None, service_account: Optional[ServiceAccount] = None, + *, additional_headers: Optional[Dict[str, str]] = None, ): super().__init__( diff --git a/src/google/adk/tools/google_api_tool/google_api_toolset.py b/src/google/adk/tools/google_api_tool/google_api_toolset.py index 5127f66d..714e6542 100644 --- a/src/google/adk/tools/google_api_tool/google_api_toolset.py +++ b/src/google/adk/tools/google_api_tool/google_api_toolset.py @@ -59,6 +59,7 @@ class GoogleApiToolset(BaseToolset): tool_filter: Optional[Union[ToolPredicate, List[str]]] = None, service_account: Optional[ServiceAccount] = None, tool_name_prefix: Optional[str] = None, + *, additional_headers: Optional[Dict[str, str]] = None, ): super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix) @@ -81,7 +82,7 @@ class GoogleApiToolset(BaseToolset): self._client_id, self._client_secret, self._service_account, - self._additional_headers, + additional_headers=self._additional_headers, ) for tool in await self._openapi_toolset.get_tools(readonly_context) if self._is_tool_selected(tool, readonly_context) diff --git a/tests/unittests/tools/google_api_tool/test_google_api_toolset.py b/tests/unittests/tools/google_api_tool/test_google_api_toolset.py index f8cf756c..5da0cb4b 100644 --- a/tests/unittests/tools/google_api_tool/test_google_api_toolset.py +++ b/tests/unittests/tools/google_api_tool/test_google_api_toolset.py @@ -214,7 +214,11 @@ class TestGoogleApiToolset: for i, rest_tool in enumerate(mock_rest_api_tools): mock_google_api_tool_class.assert_any_call( - rest_tool, client_id, client_secret, sa_mock, additional_headers + rest_tool, + client_id, + client_secret, + sa_mock, + additional_headers=additional_headers, ) assert tools[i] is mock_google_api_tool_instances[i]