Make additional_headers kwargs

This commit is contained in:
Xuan Yang
2025-10-21 12:44:50 -07:00
parent ace3e5a3a5
commit b8644ef32c
3 changed files with 8 additions and 2 deletions
@@ -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__(
@@ -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)
@@ -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]