Skip to content

feat: Implement strict validation for client while loading tools #205

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 9 commits into from
May 12, 2025
98 changes: 88 additions & 10 deletions packages/toolbox-core/src/toolbox_core/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -157,6 +157,9 @@ async def load_tool(
for execution. The specific arguments and behavior of the callable
depend on the tool itself.

Raises:
ValueError: If the loaded tool instance fails to utilize at least
one provided parameter or auth token (if any provided).
"""
# Resolve client headers
resolved_headers = {
Expand All @@ -174,55 +177,130 @@ async def load_tool(
if name not in manifest.tools:
# TODO: Better exception
raise Exception(f"Tool '{name}' not found!")
tool, _, _ = self.__parse_tool(
tool, used_auth_keys, used_bound_keys = self.__parse_tool(
name,
manifest.tools[name],
auth_token_getters,
bound_params,
self.__client_headers,
)

provided_auth_keys = set(auth_token_getters.keys())
provided_bound_keys = set(bound_params.keys())

unused_auth = provided_auth_keys - used_auth_keys
unused_bound = provided_bound_keys - used_bound_keys

if unused_auth or unused_bound:
error_messages = []
if unused_auth:
error_messages.append(f"unused auth tokens: {', '.join(unused_auth)}")
if unused_bound:
error_messages.append(
f"unused bound parameters: {', '.join(unused_bound)}"
)
raise ValueError(
f"Validation failed for tool '{name}': { '; '.join(error_messages) }."
)

return tool

async def load_toolset(
self,
name: Optional[str] = None,
auth_token_getters: dict[str, Callable[[], str]] = {},
bound_params: Mapping[str, Union[Callable[[], Any], Any]] = {},
strict: bool = False,
) -> list[ToolboxTool]:
"""
Asynchronously fetches a toolset and loads all tools defined within it.

Args:
name: Name of the toolset to load tools.
name: Name of the toolset to load. If None, loads the default toolset.
auth_token_getters: A mapping of authentication service names to
callables that return the corresponding authentication token.
bound_params: A mapping of parameter names to bind to specific values or
callables that are called to produce values as needed.
strict: If True, raises an error if *any* loaded tool instance fails
to utilize at least one provided parameter or auth token (if any
provided). If False (default), raises an error only if a
user-provided parameter or auth token cannot be applied to *any*
loaded tool across the set.

Returns:
list[ToolboxTool]: A list of callables, one for each tool defined
in the toolset.

Raises:
ValueError: If validation fails based on the `strict` flag.
"""

# Resolve client headers
original_headers = self.__client_headers
resolved_headers = {
header_name: await resolve_value(original_headers[header_name])
for header_name in original_headers
}
# Request the definition of the tool from the server
# Request the definition of the toolset from the server
url = f"{self.__base_url}/api/toolset/{name or ''}"
async with self.__session.get(url, headers=resolved_headers) as response:
json = await response.json()
manifest: ManifestSchema = ManifestSchema(**json)

# parse each tools name and schema into a list of ToolboxTools
tools = [
self.__parse_tool(
n, s, auth_token_getters, bound_params, self.__client_headers
)[0]
for n, s in manifest.tools.items()
]
tools: list[ToolboxTool] = []
overall_used_auth_keys: set[str] = set()
overall_used_bound_params: set[str] = set()
provided_auth_keys = set(auth_token_getters.keys())
provided_bound_keys = set(bound_params.keys())

# parse each tool's name and schema into a list of ToolboxTools
for tool_name, schema in manifest.tools.items():
tool, used_auth_keys, used_bound_keys = self.__parse_tool(
tool_name,
schema,
auth_token_getters,
bound_params,
self.__client_headers,
)
tools.append(tool)

if strict:
unused_auth = provided_auth_keys - used_auth_keys
unused_bound = provided_bound_keys - used_bound_keys
if unused_auth or unused_bound:
error_messages = []
if unused_auth:
error_messages.append(
f"unused auth tokens: {', '.join(unused_auth)}"
)
if unused_bound:
error_messages.append(
f"unused bound parameters: {', '.join(unused_bound)}"
)
raise ValueError(
f"Validation failed for tool '{tool_name}': { '; '.join(error_messages) }."
)
else:
overall_used_auth_keys.update(used_auth_keys)
overall_used_bound_params.update(used_bound_keys)

unused_auth = provided_auth_keys - overall_used_auth_keys
unused_bound = provided_bound_keys - overall_used_bound_params

if unused_auth or unused_bound:
error_messages = []
if unused_auth:
error_messages.append(
f"unused auth tokens could not be applied to any tool: {', '.join(unused_auth)}"
)
if unused_bound:
error_messages.append(
f"unused bound parameters could not be applied to any tool: {', '.join(unused_bound)}"
)
raise ValueError(
f"Validation failed for toolset '{name or 'default'}': { '; '.join(error_messages) }."
)

return tools

async def add_headers(
Expand Down
Loading