Skip to content

feat: Add convenience methods for adding single auth token getter to tools #245

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 8 commits into from
May 20, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion packages/toolbox-core/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -381,7 +381,7 @@ loaded. This modifies the specific tool instance.
toolbox = ToolboxClient("http://127.0.0.1:5000")
tool = await toolbox.load_tool("my-tool")

auth_tool = tool.add_auth_token_getters({"my_auth": get_auth_token}) # Single token
auth_tool = tool.add_auth_token_getter("my_auth", get_auth_token) # Single token

# OR

Expand Down Expand Up @@ -459,6 +459,10 @@ specific tool instance.
toolbox = ToolboxClient("http://127.0.0.1:5000")
tool = await toolbox.load_tool("my-tool")

bound_tool = tool.bind_param("param", "value")

# OR

bound_tool = tool.bind_params({"param": "value"})
```

Expand Down
53 changes: 44 additions & 9 deletions packages/toolbox-core/src/toolbox_core/sync_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -139,35 +139,65 @@ def add_auth_token_getters(
auth_token_getters: Mapping[str, Callable[[], str]],
) -> "ToolboxSyncTool":
"""
Registers an auth token getter function that is used for AuthServices when tools
are invoked.
Registers auth token getter functions that are used for AuthServices
when tools are invoked.

Args:
auth_token_getters: A mapping of authentication service names to
callables that return the corresponding authentication token.

Returns:
A new ToolboxSyncTool instance with the specified authentication token
getters registered.
"""
A new ToolboxSyncTool instance with the specified authentication
token getters registered.

Raises:
ValueError: If an auth source has already been registered either to
the tool or to the corresponding client.

"""
new_async_tool = self.__async_tool.add_auth_token_getters(auth_token_getters)
return ToolboxSyncTool(new_async_tool, self.__loop, self.__thread)

def add_auth_token_getter(
self, auth_source: str, get_id_token: Callable[[], str]
) -> "ToolboxSyncTool":
"""
Registers an auth token getter function that is used for AuthService
when tools are invoked.

Args:
auth_source: The name of the authentication source.
get_id_token: A function that returns the ID token.

Returns:
A new ToolboxSyncTool instance with the specified authentication
token getter registered.

Raises:
ValueError: If the auth source has already been registered either to
the tool or to the corresponding client.

"""
return self.add_auth_token_getters({auth_source: get_id_token})

def bind_params(
self, bound_params: Mapping[str, Union[Callable[[], Any], Any]]
) -> "ToolboxSyncTool":
"""
Binds parameters to values or callables that produce values.

Args:
bound_params: A mapping of parameter names to values or callables that
produce values.
bound_params: A mapping of parameter names to values or callables
that produce values.

Returns:
A new ToolboxSyncTool instance with the specified parameters bound.
"""

Raises:
ValueError: If a parameter is already bound or is not defined by the
tool's definition.

"""
new_async_tool = self.__async_tool.bind_params(bound_params)
return ToolboxSyncTool(new_async_tool, self.__loop, self.__thread)

Expand All @@ -177,7 +207,7 @@ def bind_param(
param_value: Union[Callable[[], Any], Any],
) -> "ToolboxSyncTool":
"""
Binds a parameter to the value or callables that produce it.
Binds a parameter to the value or callable that produce the value.

Args:
param_name: The name of the bound parameter.
Expand All @@ -186,5 +216,10 @@ def bind_param(

Returns:
A new ToolboxSyncTool instance with the specified parameter bound.

Raises:
ValueError: If the parameter is already bound or is not defined by
the tool's definition.

"""
return self.bind_params({param_name: param_value})
46 changes: 39 additions & 7 deletions packages/toolbox-core/src/toolbox_core/tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -281,8 +281,8 @@ def add_auth_token_getters(
auth_token_getters: Mapping[str, Callable[[], str]],
) -> "ToolboxTool":
"""
Registers an auth token getter function that is used for AuthServices when tools
are invoked.
Registers auth token getter functions that are used for AuthServices
when tools are invoked.

Args:
auth_token_getters: A mapping of authentication service names to
Expand All @@ -292,9 +292,9 @@ def add_auth_token_getters(
A new ToolboxTool instance with the specified authentication token
getters registered.

Raises
ValueError: If the auth source has already been registered either
to the tool or to the corresponding client.
Raises:
ValueError: If an auth source has already been registered either to
the tool or to the corresponding client.
"""

# throw an error if the authentication source is already registered
Expand Down Expand Up @@ -346,6 +346,28 @@ def add_auth_token_getters(
required_authz_tokens=tuple(new_req_authz_tokens),
)

def add_auth_token_getter(
self, auth_source: str, get_id_token: Callable[[], str]
) -> "ToolboxTool":
"""
Registers an auth token getter function that is used for AuthService
when tools are invoked.

Args:
auth_source: The name of the authentication source.
get_id_token: A function that returns the ID token.

Returns:
A new ToolboxTool instance with the specified authentication token
getter registered.

Raises:
ValueError: If the auth source has already been registered either to
the tool or to the corresponding client.

"""
return self.add_auth_token_getters({auth_source: get_id_token})

def bind_params(
self, bound_params: Mapping[str, Union[Callable[[], Any], Any]]
) -> "ToolboxTool":
Expand All @@ -358,6 +380,11 @@ def bind_params(

Returns:
A new ToolboxTool instance with the specified parameters bound.

Raises:
ValueError: If a parameter is already bound or is not defined by the
tool's definition.

"""
param_names = set(p.name for p in self.__params)
for name in bound_params.keys():
Expand Down Expand Up @@ -389,14 +416,19 @@ def bind_param(
param_value: Union[Callable[[], Any], Any],
) -> "ToolboxTool":
"""
Binds a parameter to the value or callables that produce it.
Binds a parameter to the value or callable that produce the value.

Args:
param_name: The name of the bound parameter.
param_value: The value of the bound parameter, or a callable that
returns the value.

Returns:
A new ToolboxTool instance with the specified parameters bound.
A new ToolboxTool instance with the specified parameter bound.

Raises:
ValueError: If the parameter is already bound or is not defined by
the tool's definition.

"""
return self.bind_params({param_name: param_value})
33 changes: 33 additions & 0 deletions packages/toolbox-core/tests/test_sync_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -237,6 +237,39 @@ def test_toolbox_sync_tool_add_auth_token_getters(
)


def test_toolbox_sync_tool_add_auth_token_getter(
toolbox_sync_tool: ToolboxSyncTool,
mock_async_tool: MagicMock,
event_loop: asyncio.AbstractEventLoop,
mock_thread: MagicMock,
):
"""Tests the add_auth_token_getter method."""
auth_service = "service1"
auth_token_getter = lambda: "token1"

new_mock_async_tool = mock_async_tool.add_auth_token_getters.return_value
new_mock_async_tool.__name__ = "new_async_tool_with_auth"

new_sync_tool = toolbox_sync_tool.add_auth_token_getter(
auth_service, auth_token_getter
)

mock_async_tool.add_auth_token_getters.assert_called_once_with(
{auth_service: auth_token_getter}
)

assert isinstance(new_sync_tool, ToolboxSyncTool)
assert new_sync_tool is not toolbox_sync_tool
assert new_sync_tool._ToolboxSyncTool__async_tool is new_mock_async_tool
assert new_sync_tool._ToolboxSyncTool__loop is event_loop # Should be the same loop
assert (
new_sync_tool._ToolboxSyncTool__thread is mock_thread
) # Should be the same thread
assert (
new_sync_tool.__qualname__ == f"ToolboxSyncTool.{new_mock_async_tool.__name__}"
)


def test_toolbox_sync_tool_bind_params(
toolbox_sync_tool: ToolboxSyncTool,
mock_async_tool: MagicMock,
Expand Down
32 changes: 32 additions & 0 deletions packages/toolbox-core/tests/test_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -502,6 +502,38 @@ def test_add_auth_token_getters_unused_token(
tool_instance.add_auth_token_getters(unused_auth_getters)


def test_add_auth_token_getter_unused_token(
http_session: ClientSession,
sample_tool_params: list[ParameterSchema],
sample_tool_description: str,
unused_auth_getters: Mapping[str, Callable[[], str]],
):
"""
Tests ValueError when add_auth_token_getters is called with a getter for
an unused authentication service.
"""
tool_instance = ToolboxTool(
session=http_session,
base_url=HTTPS_BASE_URL,
name=TEST_TOOL_NAME,
description=sample_tool_description,
params=sample_tool_params,
required_authn_params={},
required_authz_tokens=[],
auth_service_token_getters={},
bound_params={},
client_headers={},
)

expected_error_message = "Authentication source\(s\) \`unused-auth-service\` unused by tool \`sample_tool\`."

with pytest.raises(ValueError, match=expected_error_message):
tool_instance.add_auth_token_getter(
next(iter(unused_auth_getters)),
unused_auth_getters[next(iter(unused_auth_getters))],
)


def test_toolbox_tool_underscore_name_property(toolbox_tool: ToolboxTool):
"""Tests the _name property."""
assert toolbox_tool._name == TEST_TOOL_NAME
Expand Down