Skip to content

feat(toolbox-core): add basic implementation #103

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 5 commits into from
Mar 30, 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
9 changes: 9 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -3,3 +3,12 @@

# vscode
.vscode/

# python
env
venv
*.pyc
.python-version
**.egg-info/
__pycache__/**

8 changes: 5 additions & 3 deletions packages/toolbox-core/pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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",
Expand Down Expand Up @@ -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"]
Expand Down
2 changes: 2 additions & 0 deletions packages/toolbox-core/requirements.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
aiohttp==3.11.14
pydantic==2.10.6
4 changes: 2 additions & 2 deletions packages/toolbox-core/src/toolbox_core/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]
146 changes: 142 additions & 4 deletions packages/toolbox-core/src/toolbox_core/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,153 @@
# 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,
# 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 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


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
72 changes: 72 additions & 0 deletions packages/toolbox-core/src/toolbox_core/protocol.py
Original file line number Diff line number Diff line change
@@ -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]
13 changes: 13 additions & 0 deletions packages/toolbox-core/src/toolbox_core/py.typed
Original file line number Diff line number Diff line change
@@ -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.
96 changes: 96 additions & 0 deletions packages/toolbox-core/src/toolbox_core/tool.py
Original file line number Diff line number Diff line change
@@ -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)
Loading