diff --git a/.gitignore b/.gitignore index 04fd2c0b..a5c5ce0f 100644 --- a/.gitignore +++ b/.gitignore @@ -3,3 +3,12 @@ # vscode .vscode/ + +# python +env +venv +*.pyc +.python-version +**.egg-info/ +__pycache__/** + diff --git a/packages/toolbox-core/pyproject.toml b/packages/toolbox-core/pyproject.toml index 980360ae..edc45a8a 100644 --- a/packages/toolbox-core/pyproject.toml +++ b/packages/toolbox-core/pyproject.toml @@ -9,9 +9,10 @@ authors = [ {name = "Google LLC", email = "googleapis-packages@google.com"} ] -# TODO: Add deps -#dependencies = [ -#] +dependencies = [ + "pydantic>=2.7.0,<3.0.0", + "aiohttp>=3.8.6,<4.0.0", +] classifiers = [ "Intended Audience :: Developers", @@ -43,6 +44,7 @@ test = [ "isort==6.0.1", "mypy==1.15.0", "pytest==8.3.5", + "pytest-aioresponses==0.3.0" ] [build-system] requires = ["setuptools"] diff --git a/packages/toolbox-core/requirements.txt b/packages/toolbox-core/requirements.txt index e69de29b..439ae59e 100644 --- a/packages/toolbox-core/requirements.txt +++ b/packages/toolbox-core/requirements.txt @@ -0,0 +1,2 @@ +aiohttp==3.11.14 +pydantic==2.10.6 diff --git a/packages/toolbox-core/src/toolbox_core/__init__.py b/packages/toolbox-core/src/toolbox_core/__init__.py index 433463de..882b2046 100644 --- a/packages/toolbox-core/src/toolbox_core/__init__.py +++ b/packages/toolbox-core/src/toolbox_core/__init__.py @@ -12,6 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .client import DummyClass +from .client import ToolboxClient -__all__ = ["DummyClass"] +__all__ = ["ToolboxClient"] diff --git a/packages/toolbox-core/src/toolbox_core/client.py b/packages/toolbox-core/src/toolbox_core/client.py index 3caf79e1..00960690 100644 --- a/packages/toolbox-core/src/toolbox_core/client.py +++ b/packages/toolbox-core/src/toolbox_core/client.py @@ -4,7 +4,7 @@ # 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 +# 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, @@ -12,7 +12,146 @@ # See the License for the specific language governing permissions and # limitations under the License. +from typing import Optional -class DummyClass: - def __init__(self): - self.val = "dummy value" +from aiohttp import ClientSession + +from .protocol import ManifestSchema, ToolSchema +from .tool import ToolboxTool + +# Add random comment + +class ToolboxClient: + """ + An asynchronous client for interacting with a Toolbox service. + + Provides methods to discover and load tools defined by a remote Toolbox + service endpoint. It manages an underlying `aiohttp.ClientSession`. + """ + + __base_url: str + __session: ClientSession + + def __init__( + self, + url: str, + session: Optional[ClientSession] = None, + ): + """ + Initializes the ToolboxClient. + + Args: + url: The base URL for the Toolbox service API (e.g., "http://localhost:8000"). + session: An optional existing `aiohttp.ClientSession` to use. + If None (default), a new session is created internally. Note that + if a session is provided, its lifecycle (including closing) + should typically be managed externally. + """ + self.__base_url = url + + # If no aiohttp.ClientSession is provided, make our own + if session is None: + session = ClientSession() + self.__session = session + + def __parse_tool(self, name: str, schema: ToolSchema) -> ToolboxTool: + """Internal helper to create a callable tool from its schema.""" + tool = ToolboxTool( + session=self.__session, + base_url=self.__base_url, + name=name, + desc=schema.description, + params=[p.to_param() for p in schema.parameters], + ) + return tool + + async def __aenter__(self): + """ + Enter the runtime context related to this client instance. + + Allows the client to be used as an asynchronous context manager + (e.g., `async with ToolboxClient(...) as client:`). + + Returns: + self: The client instance itself. + """ + return self + + async def __aexit__(self, exc_type, exc_val, exc_tb): + """ + Exit the runtime context and close the internally managed session. + + Allows the client to be used as an asynchronous context manager + (e.g., `async with ToolboxClient(...) as client:`). + """ + await self.close() + + async def close(self): + """ + Asynchronously 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. + """ + await self.__session.close() + + async def load_tool( + self, + name: str, + ) -> ToolboxTool: + """ + Asynchronously loads a tool from the server. + + Retrieves the schema for the specified tool from the Toolbox server and + returns a callable object (`ToolboxTool`) that can be used to invoke the + tool remotely. + + Args: + name: The unique name or identifier of the tool to load. + + Returns: + ToolboxTool: A callable object representing the loaded tool, ready + for execution. The specific arguments and behavior of the callable + depend on the tool itself. + + """ + + # request the definition of the tool from the server + url = f"{self.__base_url}/api/tool/{name}" + async with self.__session.get(url) as response: + json = await response.json() + manifest: ManifestSchema = ManifestSchema(**json) + + # parse the provided definition to a tool + if name not in manifest.tools: + # TODO: Better exception + raise Exception(f"Tool '{name}' not found!") + tool = self.__parse_tool(name, manifest.tools[name]) + + return tool + + async def load_toolset( + self, + name: str, + ) -> list[ToolboxTool]: + """ + Asynchronously fetches a toolset and loads all tools defined within it. + + Args: + name: Name of the toolset to load tools. + + Returns: + list[ToolboxTool]: A list of callables, one for each tool defined + in the toolset. + """ + # Request the definition of the tool from the server + url = f"{self.__base_url}/api/toolset/{name}" + async with self.__session.get(url) 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) for n, s in manifest.tools.items()] + return tools diff --git a/packages/toolbox-core/src/toolbox_core/protocol.py b/packages/toolbox-core/src/toolbox_core/protocol.py new file mode 100644 index 00000000..5f1cc86d --- /dev/null +++ b/packages/toolbox-core/src/toolbox_core/protocol.py @@ -0,0 +1,72 @@ +# 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. + +from inspect import Parameter +from typing import Optional, Type + +from pydantic import BaseModel + + +class ParameterSchema(BaseModel): + """ + Schema for a tool parameter. + """ + + name: str + type: str + description: str + authSources: Optional[list[str]] = None + items: Optional["ParameterSchema"] = None + + def __get_type(self) -> Type: + if self.type == "string": + return str + elif self.type == "integer": + return int + elif self.type == "float": + return float + elif self.type == "boolean": + return bool + elif self.type == "array": + if self.items is None: + raise Exception("Unexpected value: type is 'list' but items is None") + return list[self._items.to_type()] # type: ignore + + raise ValueError(f"Unsupported schema type: {self.type}") + + def to_param(self) -> Parameter: + return Parameter( + self.name, + Parameter.POSITIONAL_OR_KEYWORD, + annotation=self.__get_type(), + ) + + +class ToolSchema(BaseModel): + """ + Schema for a tool. + """ + + description: str + parameters: list[ParameterSchema] + authRequired: list[str] = [] + + +class ManifestSchema(BaseModel): + """ + Schema for the Toolbox manifest. + """ + + serverVersion: str + tools: dict[str, ToolSchema] diff --git a/packages/toolbox-core/src/toolbox_core/py.typed b/packages/toolbox-core/src/toolbox_core/py.typed index e69de29b..0a2669d7 100644 --- a/packages/toolbox-core/src/toolbox_core/py.typed +++ b/packages/toolbox-core/src/toolbox_core/py.typed @@ -0,0 +1,13 @@ +# 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. diff --git a/packages/toolbox-core/src/toolbox_core/tool.py b/packages/toolbox-core/src/toolbox_core/tool.py new file mode 100644 index 00000000..48e4626c --- /dev/null +++ b/packages/toolbox-core/src/toolbox_core/tool.py @@ -0,0 +1,96 @@ +# 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. + + +from inspect import Parameter, Signature +from typing import Any + +from aiohttp import ClientSession + + +class ToolboxTool: + """ + A callable proxy object representing a specific tool on a remote Toolbox server. + + Instances of this class behave like asynchronous 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. + """ + + __url: str + __session: ClientSession + __signature__: Signature + + def __init__( + self, + session: ClientSession, + base_url: str, + name: str, + desc: str, + params: list[Parameter], + ): + """ + Initializes a callable that will trigger the tool invocation through the Toolbox server. + + Args: + session: The `aiohttp.ClientSession` used for making API requests. + base_url: The base URL of the Toolbox server API. + name: The name of the remote tool. + desc: The description of the remote tool (used as its docstring). + params: A list of `inspect.Parameter` objects defining the tool's + arguments and their types/defaults. + """ + + # used to invoke the toolbox API + self.__session = session + self.__url = f"{base_url}/api/tool/{name}/invoke" + + # the following properties are set to help anyone that might inspect it determine + self.__name__ = name + self.__doc__ = desc + self.__signature__ = Signature(parameters=params, return_annotation=str) + self.__annotations__ = {p.name: p.annotation for p in params} + # TODO: self.__qualname__ ?? + + async def __call__(self, *args: Any, **kwargs: Any) -> str: + """ + Asynchronously 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. + """ + all_args = self.__signature__.bind(*args, **kwargs) + all_args.apply_defaults() # Include default values if not provided + payload = all_args.arguments + + async with self.__session.post( + self.__url, + json=payload, + ) as resp: + ret = await resp.json() + if "error" in ret: + # TODO: better error + raise Exception(ret["error"]) + return ret.get("result", ret) diff --git a/packages/toolbox-core/tests/test_client.py b/packages/toolbox-core/tests/test_client.py new file mode 100644 index 00000000..b19c575b --- /dev/null +++ b/packages/toolbox-core/tests/test_client.py @@ -0,0 +1,108 @@ +# 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 inspect + +import pytest + +from toolbox_core import ToolboxClient +from toolbox_core.protocol import ManifestSchema, ParameterSchema, ToolSchema + +TEST_BASE_URL = "http://toolbox.example.com" + + +@pytest.fixture() +def test_tool_str(): + return ToolSchema( + description="Test Tool 1 Description", + parameters=[ + ParameterSchema( + name="param1", type="string", description="Description of Param1" + ) + ], + ) + + +@pytest.fixture() +def test_tool_int_bool(): + """Fixture for the second test tool schema.""" + return ToolSchema( + description="Test Tool 2 Description", + parameters=[ + ParameterSchema(name="argA", type="integer", description="Argument A"), + ParameterSchema(name="argB", type="boolean", description="Argument B"), + ], + ) + + +@pytest.mark.asyncio +async def test_load_tool_success(aioresponses, test_tool_str): + """ + Tests successfully loading a tool when the API returns a valid manifest. + """ + # Mock out responses from server + TOOL_NAME = "test_tool_1" + manifest = ManifestSchema(serverVersion="0.0.0", tools={TOOL_NAME: test_tool_str}) + aioresponses.get( + f"{TEST_BASE_URL}/api/tool/{TOOL_NAME}", + payload=manifest.model_dump(), + status=200, + ) + aioresponses.post( + f"{TEST_BASE_URL}/api/tool/{TOOL_NAME}/invoke", + payload={"result": "ok"}, + status=200, + ) + + async with ToolboxClient(TEST_BASE_URL) as client: + # Load a Tool + loaded_tool = await client.load_tool(TOOL_NAME) + + # Assertions + assert callable(loaded_tool) + # Assert introspection attributes are set correctly + assert loaded_tool.__name__ == TOOL_NAME + assert loaded_tool.__doc__ == test_tool_str.description + + # Assert signature inspection + sig = inspect.signature(loaded_tool) + assert list(sig.parameters.keys()) == [p.name for p in test_tool_str.parameters] + + assert await loaded_tool("some value") == "ok" + + +@pytest.mark.asyncio +async def test_load_toolset_success(aioresponses, test_tool_str, test_tool_int_bool): + """Tests successfully loading a toolset with multiple tools.""" + TOOLSET_NAME = "my_toolset" + TOOL1 = "tool1" + TOOL2 = "tool2" + manifest = ManifestSchema( + serverVersion="0.0.0", tools={TOOL1: test_tool_str, TOOL2: test_tool_int_bool} + ) + aioresponses.get( + f"{TEST_BASE_URL}/api/toolset/{TOOLSET_NAME}", + payload=manifest.model_dump(), + status=200, + ) + + async with ToolboxClient(TEST_BASE_URL) as client: + tools = await client.load_toolset(TOOLSET_NAME) + + assert isinstance(tools, list) + assert len(tools) == len(manifest.tools) + + # Check if tools were created correctly + assert {t.__name__ for t in tools} == manifest.tools.keys() diff --git a/packages/toolbox-core/tests/test_core.py b/packages/toolbox-core/tests/test_core.py deleted file mode 100644 index 76ce659f..00000000 --- a/packages/toolbox-core/tests/test_core.py +++ /dev/null @@ -1,20 +0,0 @@ -# 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. -from toolbox_core import DummyClass - - -class TestToolboxCore: - def test_dummy(self): - dummy_variable = DummyClass() - assert dummy_variable.val == "dummy value"