diff --git a/packages/toolbox-core/src/toolbox_core/__init__.py b/packages/toolbox-core/src/toolbox_core/__init__.py index 882b2046..7ea4a145 100644 --- a/packages/toolbox-core/src/toolbox_core/__init__.py +++ b/packages/toolbox-core/src/toolbox_core/__init__.py @@ -13,5 +13,6 @@ # limitations under the License. from .client import ToolboxClient +from .sync_client import ToolboxSyncClient -__all__ = ["ToolboxClient"] +__all__ = ["ToolboxClient", "ToolboxSyncClient"] diff --git a/packages/toolbox-core/src/toolbox_core/sync_client.py b/packages/toolbox-core/src/toolbox_core/sync_client.py new file mode 100644 index 00000000..36877223 --- /dev/null +++ b/packages/toolbox-core/src/toolbox_core/sync_client.py @@ -0,0 +1,149 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +from threading import Thread +from typing import Any, Awaitable, Callable, Mapping, Optional, TypeVar, Union + +from aiohttp import ClientSession + +from .client import ToolboxClient +from .sync_tool import ToolboxSyncTool + +T = TypeVar("T") + + +class ToolboxSyncClient: + """ + An synchronous client for interacting with a Toolbox service. + + Provides methods to discover and load tools defined by a remote Toolbox + service endpoint. + """ + + __loop: Optional[asyncio.AbstractEventLoop] = None + __thread: Optional[Thread] = None + + def __init__( + self, + url: str, + ): + """ + Initializes the ToolboxSyncClient. + + Args: + url: The base URL for the Toolbox service API (e.g., "http://localhost:5000"). + """ + # Running a loop in a background thread allows us to support async + # methods from non-async environments. + if self.__class__.__loop is None: + loop = asyncio.new_event_loop() + thread = Thread(target=loop.run_forever, daemon=True) + thread.start() + self.__class__.__thread = thread + self.__class__.__loop = loop + + async def create_client(): + return ToolboxClient(url) + + # Ignoring type since we're already checking the existence of a loop above. + self.__async_client = asyncio.run_coroutine_threadsafe( + create_client(), self.__class__.__loop # type: ignore + ).result() + + def close(self): + """ + Synchronously closes the underlying client session. Doing so will cause + any tools created by this Client to cease to function. + + If the session was provided externally during initialization, the caller + is responsible for its lifecycle, but calling close here will still + attempt to close it. + """ + coro = self.__async_client.close() + asyncio.run_coroutine_threadsafe(coro, self.__loop).result() + + def load_tool( + self, + name: str, + auth_token_getters: dict[str, Callable[[], str]] = {}, + bound_params: Mapping[str, Union[Callable[[], Any], Any]] = {}, + ) -> ToolboxSyncTool: + """ + Synchronously loads a tool from the server. + + Retrieves the schema for the specified tool from the Toolbox server and + returns a callable object (`ToolboxSyncTool`) that can be used to invoke the + tool remotely. + + Args: + name: The unique name or identifier of the tool to load. + 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. + + Returns: + ToolboxSyncTool: A callable object representing the loaded tool, ready + for execution. The specific arguments and behavior of the callable + depend on the tool itself. + """ + coro = self.__async_client.load_tool(name, auth_token_getters, bound_params) + + # We have already created a new loop in the init method in case it does not already exist + async_tool = asyncio.run_coroutine_threadsafe(coro, self.__loop).result() # type: ignore + + if not self.__loop or not self.__thread: + raise ValueError("Background loop or thread cannot be None.") + return ToolboxSyncTool(async_tool, self.__loop, self.__thread) + + def load_toolset( + self, + name: str, + auth_token_getters: dict[str, Callable[[], str]] = {}, + bound_params: Mapping[str, Union[Callable[[], Any], Any]] = {}, + ) -> list[ToolboxSyncTool]: + """ + Synchronously fetches a toolset and loads all tools defined within it. + + Args: + name: Name of the toolset to load tools. + 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. + + Returns: + list[ToolboxSyncTool]: A list of callables, one for each tool defined + in the toolset. + """ + coro = self.__async_client.load_toolset(name, auth_token_getters, bound_params) + + # We have already created a new loop in the init method in case it does not already exist + async_tools = asyncio.run_coroutine_threadsafe(coro, self.__loop).result() # type: ignore + + if not self.__loop or not self.__thread: + raise ValueError("Background loop or thread cannot be None.") + return [ + ToolboxSyncTool(async_tool, self.__loop, self.__thread) + for async_tool in async_tools + ] + + def __enter__(self): + """Enter the runtime context related to this client instance.""" + return self + + def __exit__(self, exc_type, exc_val, exc_tb): + """Exit the runtime context and close the client session.""" + self.close() diff --git a/packages/toolbox-core/src/toolbox_core/sync_tool.py b/packages/toolbox-core/src/toolbox_core/sync_tool.py new file mode 100644 index 00000000..ad24e8e3 --- /dev/null +++ b/packages/toolbox-core/src/toolbox_core/sync_tool.py @@ -0,0 +1,142 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +from asyncio import AbstractEventLoop +from inspect import Signature +from threading import Thread +from typing import Any, Callable, Mapping, TypeVar, Union + +from .tool import ToolboxTool + +T = TypeVar("T") + + +class ToolboxSyncTool: + """ + A callable proxy object representing a specific tool on a remote Toolbox server. + + Instances of this class behave like synchronous functions. When called, they + send a request to the corresponding tool's endpoint on the Toolbox server with + the provided arguments. + + It utilizes Python's introspection features (`__name__`, `__doc__`, + `__signature__`, `__annotations__`) so that standard tools like `help()` + and `inspect` work as expected. + """ + + def __init__( + self, async_tool: ToolboxTool, loop: AbstractEventLoop, thread: Thread + ): + """ + Initializes a callable that will trigger the tool invocation through the + Toolbox server. + + Args: + async_tool: An instance of the asynchronous ToolboxTool. + loop: The event loop used to run asynchronous tasks. + thread: The thread to run blocking operations in. + """ + + if not isinstance(async_tool, ToolboxTool): + raise TypeError("async_tool must be an instance of ToolboxTool") + + self.__async_tool = async_tool + self.__loop = loop + self.__thread = thread + + # NOTE: We cannot define __qualname__ as a @property here. + # Properties are designed to compute values dynamically when accessed on an *instance* (using 'self'). + # However, Python needs the class's __qualname__ attribute to be a plain string + # *before* any instances exist, specifically when the 'class ToolboxSyncTool:' statement + # itself is being processed during module import or class definition. + # Defining __qualname__ as a property leads to a TypeError because the class object needs + # a string value immediately, not a descriptor that evaluates later. + self.__qualname__ = f"{self.__class__.__qualname__}.{self.__name__}" + + @property + def __name__(self) -> str: + return self.__async_tool.__name__ + + @property + def __doc__(self) -> Union[str, None]: # type: ignore[override] + # Standard Python object attributes like __doc__ are technically "writable". + # But not defining a setter function makes this a read-only property. + # Mypy flags this issue in the type checks. + return self.__async_tool.__doc__ + + @property + def __signature__(self) -> Signature: + return self.__async_tool.__signature__ + + @property + def __annotations__(self) -> dict[str, Any]: # type: ignore[override] + # Standard Python object attributes like __doc__ are technically "writable". + # But not defining a setter function makes this a read-only property. + # Mypy flags this issue in the type checks. + return self.__async_tool.__annotations__ + + def __call__(self, *args: Any, **kwargs: Any) -> str: + """ + Synchronously calls the remote tool with the provided arguments. + + Validates arguments against the tool's signature, then sends them + as a JSON payload in a POST request to the tool's invoke URL. + + Args: + *args: Positional arguments for the tool. + **kwargs: Keyword arguments for the tool. + + Returns: + The string result returned by the remote tool execution. + """ + coro = self.__async_tool(*args, **kwargs) + return asyncio.run_coroutine_threadsafe(coro, self.__loop).result() + + def add_auth_token_getters( + self, + auth_token_getters: Mapping[str, Callable[[], str]], + ) -> "ToolboxSyncTool": + """ + Registers an auth token getter function that is 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. + """ + + new_async_tool = self.__async_tool.add_auth_token_getters(auth_token_getters) + return ToolboxSyncTool(new_async_tool, self.__loop, self.__thread) + + def bind_parameters( + 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. + + Returns: + A new ToolboxSyncTool instance with the specified parameters bound. + """ + + new_async_tool = self.__async_tool.bind_parameters(bound_params) + return ToolboxSyncTool(new_async_tool, self.__loop, self.__thread) diff --git a/packages/toolbox-core/tests/test_sync_e2e.py b/packages/toolbox-core/tests/test_sync_e2e.py new file mode 100644 index 00000000..055ca977 --- /dev/null +++ b/packages/toolbox-core/tests/test_sync_e2e.py @@ -0,0 +1,183 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest + +from toolbox_core.sync_client import ToolboxSyncClient +from toolbox_core.sync_tool import ToolboxSyncTool + + +# --- Shared Fixtures Defined at Module Level --- +@pytest.fixture(scope="module") +def toolbox(): + """Creates a ToolboxSyncClient instance shared by all tests in this module.""" + toolbox = ToolboxSyncClient("http://localhost:5000") + try: + yield toolbox + finally: + toolbox.close() + + +@pytest.fixture(scope="function") +def get_n_rows_tool(toolbox: ToolboxSyncClient) -> ToolboxSyncTool: + """Load the 'get-n-rows' tool using the shared toolbox client.""" + tool = toolbox.load_tool("get-n-rows") + assert tool.__name__ == "get-n-rows" + return tool + + +@pytest.mark.usefixtures("toolbox_server") +class TestBasicE2E: + @pytest.mark.parametrize( + "toolset_name, expected_length, expected_tools", + [ + ("my-toolset", 1, ["get-row-by-id"]), + ("my-toolset-2", 2, ["get-n-rows", "get-row-by-id"]), + ], + ) + def test_load_toolset_specific( + self, + toolbox: ToolboxSyncClient, + toolset_name: str, + expected_length: int, + expected_tools: list[str], + ): + """Load a specific toolset""" + toolset = toolbox.load_toolset(toolset_name) + assert len(toolset) == expected_length + tool_names = {tool.__name__ for tool in toolset} + assert tool_names == set(expected_tools) + + def test_run_tool(self, get_n_rows_tool: ToolboxSyncTool): + """Invoke a tool.""" + response = get_n_rows_tool(num_rows="2") + + assert isinstance(response, str) + assert "row1" in response + assert "row2" in response + assert "row3" not in response + + def test_run_tool_missing_params(self, get_n_rows_tool: ToolboxSyncTool): + """Invoke a tool with missing params.""" + with pytest.raises(TypeError, match="missing a required argument: 'num_rows'"): + get_n_rows_tool() + + def test_run_tool_wrong_param_type(self, get_n_rows_tool: ToolboxSyncTool): + """Invoke a tool with wrong param type.""" + with pytest.raises( + Exception, + match=r"num_rows\s+Input should be a valid string\s+\[type=string_type,\s+input_value=2,\s+input_type=int\]", + ): + get_n_rows_tool(num_rows=2) + + +@pytest.mark.usefixtures("toolbox_server") +class TestBindParams: + def test_bind_params( + self, toolbox: ToolboxSyncClient, get_n_rows_tool: ToolboxSyncTool + ): + """Bind a param to an existing tool.""" + new_tool = get_n_rows_tool.bind_parameters({"num_rows": "3"}) + response = new_tool() + assert isinstance(response, str) + assert "row1" in response + assert "row2" in response + assert "row3" in response + assert "row4" not in response + + def test_bind_params_callable( + self, toolbox: ToolboxSyncClient, get_n_rows_tool: ToolboxSyncTool + ): + """Bind a callable param to an existing tool.""" + new_tool = get_n_rows_tool.bind_parameters({"num_rows": lambda: "3"}) + response = new_tool() + assert isinstance(response, str) + assert "row1" in response + assert "row2" in response + assert "row3" in response + assert "row4" not in response + + +@pytest.mark.usefixtures("toolbox_server") +class TestAuth: + def test_run_tool_unauth_with_auth( + self, toolbox: ToolboxSyncClient, auth_token2: str + ): + """Tests running a tool that doesn't require auth, with auth provided.""" + tool = toolbox.load_tool( + "get-row-by-id", auth_token_getters={"my-test-auth": lambda: auth_token2} + ) + response = tool(id="2") + assert "row2" in response + + def test_run_tool_no_auth(self, toolbox: ToolboxSyncClient): + """Tests running a tool requiring auth without providing auth.""" + tool = toolbox.load_tool("get-row-by-id-auth") + with pytest.raises( + Exception, + match="tool invocation not authorized. Please make sure your specify correct auth headers", + ): + tool(id="2") + + def test_run_tool_wrong_auth(self, toolbox: ToolboxSyncClient, auth_token2: str): + """Tests running a tool with incorrect auth. The tool + requires a different authentication than the one provided.""" + tool = toolbox.load_tool("get-row-by-id-auth") + auth_tool = tool.add_auth_token_getters({"my-test-auth": lambda: auth_token2}) + with pytest.raises( + Exception, + match="tool invocation not authorized", + ): + auth_tool(id="2") + + def test_run_tool_auth(self, toolbox: ToolboxSyncClient, auth_token1: str): + """Tests running a tool with correct auth.""" + tool = toolbox.load_tool("get-row-by-id-auth") + auth_tool = tool.add_auth_token_getters({"my-test-auth": lambda: auth_token1}) + response = auth_tool(id="2") + assert "row2" in response + + def test_run_tool_param_auth_no_auth(self, toolbox: ToolboxSyncClient): + """Tests running a tool with a param requiring auth, without auth.""" + tool = toolbox.load_tool("get-row-by-email-auth") + with pytest.raises( + Exception, + match="One or more of the following authn services are required to invoke this tool: my-test-auth", + ): + tool() + + def test_run_tool_param_auth(self, toolbox: ToolboxSyncClient, auth_token1: str): + """Tests running a tool with a param requiring auth, with correct auth.""" + tool = toolbox.load_tool( + "get-row-by-email-auth", + auth_token_getters={"my-test-auth": lambda: auth_token1}, + ) + response = tool() + assert "row4" in response + assert "row5" in response + assert "row6" in response + + def test_run_tool_param_auth_no_field( + self, toolbox: ToolboxSyncClient, auth_token1: str + ): + """Tests running a tool with a param requiring auth, with insufficient auth.""" + tool = toolbox.load_tool( + "get-row-by-content-auth", + auth_token_getters={"my-test-auth": lambda: auth_token1}, + ) + with pytest.raises( + Exception, + match="no field named row_data in claims", + ): + tool()