Skip to content

Commit b4ce8ef

Browse files
committed
feat: Add convenience methods for adding single auth token getter to tools
1 parent fcdfdae commit b4ce8ef

File tree

2 files changed

+44
-0
lines changed

2 files changed

+44
-0
lines changed

packages/toolbox-core/src/toolbox_core/sync_tool.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,28 @@ def add_auth_token_getters(
154154
new_async_tool = self.__async_tool.add_auth_token_getters(auth_token_getters)
155155
return ToolboxSyncTool(new_async_tool, self.__loop, self.__thread)
156156

157+
def add_auth_token_getter(
158+
self, auth_source: str, get_id_token: Callable[[], str]
159+
) -> "ToolboxSyncTool":
160+
"""
161+
Registers auth token getter functions that are used for AuthServices
162+
when tools are invoked.
163+
164+
Args:
165+
auth_source: The name of the authentication source.
166+
get_id_token: A function that returns the ID token.
167+
168+
Returns:
169+
A new ToolboxSyncTool instance with the specified authentication
170+
token getters registered.
171+
172+
Raises:
173+
ValueError: If the auth source has already been registered either to
174+
the tool or to the corresponding client.
175+
176+
"""
177+
return self.add_auth_token_getters({auth_source: get_id_token})
178+
157179
def bind_params(
158180
self, bound_params: Mapping[str, Union[Callable[[], Any], Any]]
159181
) -> "ToolboxSyncTool":

packages/toolbox-core/src/toolbox_core/tool.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -346,6 +346,28 @@ def add_auth_token_getters(
346346
required_authz_tokens=tuple(new_req_authz_tokens),
347347
)
348348

349+
def add_auth_token_getter(
350+
self, auth_source: str, get_id_token: Callable[[], str]
351+
) -> "ToolboxTool":
352+
"""
353+
Registers auth token getter functions that are used for AuthServices
354+
when tools are invoked.
355+
356+
Args:
357+
auth_source: The name of the authentication source.
358+
get_id_token: A function that returns the ID token.
359+
360+
Returns:
361+
A new ToolboxTool instance with the specified authentication token
362+
getters registered.
363+
364+
Raises
365+
ValueError: If the auth source has already been registered either
366+
to the tool or to the corresponding client.
367+
368+
"""
369+
return self.add_auth_token_getters({auth_source: get_id_token})
370+
349371
def bind_params(
350372
self, bound_params: Mapping[str, Union[Callable[[], Any], Any]]
351373
) -> "ToolboxTool":

0 commit comments

Comments
 (0)