You've already forked adk-python
mirror of
https://github.com/encounter/adk-python.git
synced 2026-03-30 10:57:20 -07:00
refactor(tools): Update ToolboxToolset to wrap toolbox-adk
### Description of Change
**Problem:**
The `ToolboxToolset` was implemented directly within `adk-python`, leading to code duplication and potential drift from the core `toolbox-adk` implementation.
**Solution:**
Refactored `ToolboxToolset` to act as a rigorous wrapper around the `toolbox-adk` package, which delegates all functionality to `toolbox_adk.ToolboxToolset`.
### Testing Plan
**Unit Tests:**
- [x] I have added or updated unit tests for my change.
- [x] All unit tests pass locally.
Summary:
- Verified initialization flows through to `toolbox-adk`.
- Verified `auth_token_getters` are correctly propagated.
- Verified type hints are static-analysis friendly.
**Manual End-to-End (E2E) Tests:**
Manually verified standard toolbox loading and execution with the new wrapper:
```python
from google.adk.tools import ToolboxToolset
from toolbox_adk import CredentialStrategy
# Loading with toolset_name
ts = ToolboxToolset(
server_url='http://localhost:8080',
toolset_name='calculator',
credentials=CredentialStrategy.toolbox_identity()
)
tools = await ts.get_tools()
print(f'Loaded {len(tools)} tools')
```
### Checklist
- [x] I have read the [CONTRIBUTING.md](https://github.com/google/adk-python/blob/main/CONTRIBUTING.md) document.
- [x] I have performed a self-review of my own code.
- [x] I have commented my code, particularly in hard-to-understand areas.
- [x] I have added tests that prove my fix is effective or that my feature works.
- [x] New and existing unit tests pass locally with my changes.
- [x] I have manually tested my changes end-to-end.
- [x] Any dependent changes have been merged and published in downstream modules.
PiperOrigin-RevId: 855798474
This commit is contained in:
committed by
Copybara-Service
parent
ab62b1bffd
commit
277084e313
+1
-1
@@ -157,7 +157,7 @@ extensions = [
|
||||
"llama-index-readers-file>=0.4.0", # For retrieval using LlamaIndex.
|
||||
"llama-index-embeddings-google-genai>=0.3.0", # For files retrieval using LlamaIndex.
|
||||
"lxml>=5.3.0", # For load_web_page tool.
|
||||
"toolbox-core>=0.1.0", # For tools.toolbox_toolset.ToolboxToolset
|
||||
"toolbox-adk>=0.1.0", # For tools.toolbox_toolset.ToolboxToolset
|
||||
]
|
||||
|
||||
otel-gcp = ["opentelemetry-instrumentation-google-genai>=0.3b0, <1.0.0"]
|
||||
|
||||
@@ -12,29 +12,41 @@
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
from typing import Callable
|
||||
from typing import List
|
||||
from typing import Mapping
|
||||
from typing import Optional
|
||||
from typing import TYPE_CHECKING
|
||||
from typing import Union
|
||||
|
||||
import toolbox_core as toolbox
|
||||
from typing_extensions import override
|
||||
|
||||
from ..agents.readonly_context import ReadonlyContext
|
||||
from .base_tool import BaseTool
|
||||
from .base_toolset import BaseToolset
|
||||
from .function_tool import FunctionTool
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from toolbox_adk import CredentialConfig
|
||||
|
||||
|
||||
class ToolboxToolset(BaseToolset):
|
||||
"""A class that provides access to toolbox toolsets.
|
||||
|
||||
This class acts as a bridge to the `toolbox-adk` package.
|
||||
You must install `toolbox-adk` to use this class.
|
||||
|
||||
Example:
|
||||
```python
|
||||
toolbox_toolset = ToolboxToolset("http://127.0.0.1:5000",
|
||||
toolset_name="my-toolset")
|
||||
from toolbox_adk import CredentialStrategy
|
||||
|
||||
toolbox_toolset = ToolboxToolset(
|
||||
server_url="http://127.0.0.1:5000",
|
||||
# toolset_name and tool_names are optional. If omitted, all tools are
|
||||
loaded.
|
||||
credentials=CredentialStrategy.toolbox_identity()
|
||||
)
|
||||
```
|
||||
"""
|
||||
@@ -44,64 +56,58 @@ class ToolboxToolset(BaseToolset):
|
||||
server_url: str,
|
||||
toolset_name: Optional[str] = None,
|
||||
tool_names: Optional[List[str]] = None,
|
||||
auth_token_getters: Optional[dict[str, Callable[[], str]]] = None,
|
||||
auth_token_getters: Optional[Mapping[str, Callable[[], str]]] = None,
|
||||
bound_params: Optional[
|
||||
Mapping[str, Union[Callable[[], Any], Any]]
|
||||
] = None,
|
||||
credentials: Optional[CredentialConfig] = None,
|
||||
additional_headers: Optional[Mapping[str, str]] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Args:
|
||||
|
||||
server_url: The URL of the toolbox server.
|
||||
toolset_name: The name of the toolbox toolset to load.
|
||||
tool_names: The names of the tools to load.
|
||||
auth_token_getters: A mapping of authentication service names to
|
||||
callables that return the corresponding authentication token. see:
|
||||
https://github.com/googleapis/mcp-toolbox-sdk-python/tree/main/packages/toolbox-core#authenticating-tools
|
||||
for details.
|
||||
bound_params: A mapping of parameter names to bind to specific values or
|
||||
callables that are called to produce values as needed. see:
|
||||
https://github.com/googleapis/mcp-toolbox-sdk-python/tree/main/packages/toolbox-core#binding-parameter-values
|
||||
for details.
|
||||
The resulting ToolboxToolset will contain both tools loaded by tool_names
|
||||
and toolset_name.
|
||||
server_url: The URL of the toolbox server.
|
||||
toolset_name: The name of the toolbox toolset to load.
|
||||
tool_names: The names of the tools to load.
|
||||
auth_token_getters: (Deprecated) Map of auth token getters.
|
||||
bound_params: Parameters to bind to the tools.
|
||||
credentials: (Optional) toolbox_adk.CredentialConfig object.
|
||||
additional_headers: (Optional) Static headers dictionary.
|
||||
**kwargs: Additional arguments passed to the underlying
|
||||
toolbox_adk.ToolboxToolset.
|
||||
"""
|
||||
if not tool_names and not toolset_name:
|
||||
raise ValueError("tool_names and toolset_name cannot both be None")
|
||||
if not toolset_name and not tool_names:
|
||||
raise ValueError(
|
||||
"Either 'toolset_name' or 'tool_names' must be provided."
|
||||
)
|
||||
|
||||
try:
|
||||
from toolbox_adk import ToolboxToolset as RealToolboxToolset # pylint: disable=import-outside-toplevel
|
||||
except ImportError as exc:
|
||||
raise ImportError(
|
||||
"ToolboxToolset requires the 'toolbox-adk' package. "
|
||||
"Please install it using `pip install toolbox-adk`."
|
||||
) from exc
|
||||
|
||||
super().__init__()
|
||||
self._server_url = server_url
|
||||
self._toolbox_client = toolbox.ToolboxClient(server_url)
|
||||
self._toolset_name = toolset_name
|
||||
self._tool_names = tool_names
|
||||
self._auth_token_getters = auth_token_getters or {}
|
||||
self._bound_params = bound_params or {}
|
||||
|
||||
self._delegate = RealToolboxToolset(
|
||||
server_url=server_url,
|
||||
toolset_name=toolset_name,
|
||||
tool_names=tool_names,
|
||||
credentials=credentials,
|
||||
additional_headers=additional_headers,
|
||||
bound_params=bound_params,
|
||||
auth_token_getters=auth_token_getters,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@override
|
||||
async def get_tools(
|
||||
self, readonly_context: Optional[ReadonlyContext] = None
|
||||
) -> list[BaseTool]:
|
||||
tools = []
|
||||
if self._toolset_name:
|
||||
tools.extend([
|
||||
FunctionTool(tool)
|
||||
for tool in await self._toolbox_client.load_toolset(
|
||||
self._toolset_name,
|
||||
auth_token_getters=self._auth_token_getters,
|
||||
bound_params=self._bound_params,
|
||||
)
|
||||
])
|
||||
if self._tool_names:
|
||||
tools.extend([
|
||||
FunctionTool(
|
||||
await self._toolbox_client.load_tool(
|
||||
tool_name,
|
||||
auth_token_getters=self._auth_token_getters,
|
||||
bound_params=self._bound_params,
|
||||
)
|
||||
)
|
||||
for tool_name in self._tool_names
|
||||
])
|
||||
return tools
|
||||
return await self._delegate.get_tools(readonly_context)
|
||||
|
||||
@override
|
||||
async def close(self):
|
||||
self._toolbox_client.close()
|
||||
await self._delegate.close()
|
||||
|
||||
Reference in New Issue
Block a user