Skip to content

Commit 381d44c

Browse files
seanzhougooglecopybara-github
authored andcommitted
feat: Add get_auth_config method to toolset to expose auth requirement of the toolset
Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com> PiperOrigin-RevId: 862597425
1 parent 4341839 commit 381d44c

5 files changed

Lines changed: 65 additions & 0 deletions

File tree

src/google/adk/tools/apihub_tool/apihub_toolset.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,7 @@
2424
from ...agents.readonly_context import ReadonlyContext
2525
from ...auth.auth_credential import AuthCredential
2626
from ...auth.auth_schemes import AuthScheme
27+
from ...auth.auth_tool import AuthConfig
2728
from .._gemini_schema_util import _to_snake_case
2829
from ..base_toolset import BaseToolset
2930
from ..base_toolset import ToolPredicate
@@ -188,3 +189,13 @@ def _prepare_toolset(self) -> None:
188189
async def close(self):
189190
if self._openapi_toolset:
190191
await self._openapi_toolset.close()
192+
193+
@override
194+
def get_auth_config(self) -> AuthConfig | None:
195+
"""Returns the auth config for this toolset."""
196+
if self._auth_scheme is None:
197+
return None
198+
return AuthConfig(
199+
auth_scheme=self._auth_scheme,
200+
raw_auth_credential=self._auth_credential,
201+
)

src/google/adk/tools/application_integration_tool/application_integration_toolset.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from ...auth.auth_credential import ServiceAccount
2929
from ...auth.auth_credential import ServiceAccountCredential
3030
from ...auth.auth_schemes import AuthScheme
31+
from ...auth.auth_tool import AuthConfig
3132
from ..base_toolset import BaseToolset
3233
from ..base_toolset import ToolPredicate
3334
from ..openapi_tool.auth.auth_helpers import service_account_scheme_credential
@@ -278,3 +279,13 @@ async def get_tools(
278279
async def close(self) -> None:
279280
if self._openapi_toolset:
280281
await self._openapi_toolset.close()
282+
283+
@override
284+
def get_auth_config(self) -> AuthConfig | None:
285+
"""Returns the auth config for this toolset."""
286+
if self._auth_scheme is None:
287+
return None
288+
return AuthConfig(
289+
auth_scheme=self._auth_scheme,
290+
raw_auth_credential=self._auth_credential,
291+
)

src/google/adk/tools/base_toolset.py

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
from typing import Union
2929

3030
from ..agents.readonly_context import ReadonlyContext
31+
from ..auth.auth_tool import AuthConfig
3132
from .base_tool import BaseTool
3233

3334
if TYPE_CHECKING:
@@ -204,3 +205,21 @@ async def process_llm_request(
204205
llm_request: The outgoing LLM request, mutable this method.
205206
"""
206207
pass
208+
209+
def get_auth_config(self) -> Optional[AuthConfig]:
210+
"""Returns the auth config for this toolset. ADK will make sure the
211+
'exchanged_auth_credential' field in the config is populated with
212+
ready-to-use credential (e.g. oauth token for OAuth flow) before calling
213+
get_tools method or execute any tools returned by this toolset. Thus toolset
214+
can use this credential either for tool listing or tool calling. If tool
215+
calling needs a different credential from ADK client, call
216+
tool_context.request_credential in the tool.
217+
218+
Toolsets that support authentication should override this method to return
219+
an AuthConfig constructed from their auth_scheme, auth_credential, and
220+
optional credential_key parameters.
221+
222+
Returns:
223+
AuthConfig if the toolset has authentication configured, None otherwise.
224+
"""
225+
return None

src/google/adk/tools/mcp_tool/mcp_toolset.py

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
from ...agents.readonly_context import ReadonlyContext
4040
from ...auth.auth_credential import AuthCredential
4141
from ...auth.auth_schemes import AuthScheme
42+
from ...auth.auth_tool import AuthConfig
4243
from ..base_tool import BaseTool
4344
from ..base_toolset import BaseToolset
4445
from ..base_toolset import ToolPredicate
@@ -284,6 +285,16 @@ async def close(self) -> None:
284285
# Log the error but don't re-raise to avoid blocking shutdown
285286
print(f"Warning: Error during McpToolset cleanup: {e}", file=self._errlog)
286287

288+
@override
289+
def get_auth_config(self) -> AuthConfig | None:
290+
"""Returns the auth config for this toolset."""
291+
if self._auth_scheme is None:
292+
return None
293+
return AuthConfig(
294+
auth_scheme=self._auth_scheme,
295+
raw_auth_credential=self._auth_credential,
296+
)
297+
287298
@override
288299
@classmethod
289300
def from_config(

src/google/adk/tools/openapi_tool/openapi_spec_parser/openapi_toolset.py

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@
3232
from ....agents.readonly_context import ReadonlyContext
3333
from ....auth.auth_credential import AuthCredential
3434
from ....auth.auth_schemes import AuthScheme
35+
from ....auth.auth_tool import AuthConfig
3536
from ...base_toolset import BaseToolset
3637
from ...base_toolset import ToolPredicate
3738
from .openapi_spec_parser import OpenApiSpecParser
@@ -128,6 +129,8 @@ def __init__(
128129
"""
129130
super().__init__(tool_filter=tool_filter, tool_name_prefix=tool_name_prefix)
130131
self._header_provider = header_provider
132+
self._auth_scheme = auth_scheme
133+
self._auth_credential = auth_credential
131134
if not spec_dict:
132135
spec_dict = self._load_spec(spec_str, spec_str_type)
133136
self._ssl_verify = ssl_verify
@@ -211,3 +214,13 @@ def _parse(self, openapi_spec_dict: Dict[str, Any]) -> List[RestApiTool]:
211214
@override
212215
async def close(self):
213216
pass
217+
218+
@override
219+
def get_auth_config(self) -> AuthConfig | None:
220+
"""Returns the auth config for this toolset."""
221+
if self._auth_scheme is None:
222+
return None
223+
return AuthConfig(
224+
auth_scheme=self._auth_scheme,
225+
raw_auth_credential=self._auth_credential,
226+
)

0 commit comments

Comments
 (0)