Skip to content
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
45 changes: 43 additions & 2 deletions packages/toolbox-adk/src/toolbox_adk/tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
OAuth2Auth,
)
from google.adk.auth.auth_tool import AuthConfig
from google.genai.types import FunctionDeclaration, Type, Schema
from google.adk.tools.base_tool import BaseTool
from google.adk.tools.tool_context import ToolContext
from toolbox_core.tool import ToolboxTool as CoreToolboxTool
Expand Down Expand Up @@ -79,17 +80,57 @@ def __init__(
self._post_hook = post_hook
self._auth_config = auth_config


def _param_type_to_schema_type(self, param_type: str) -> Type:
type_map = {
"string": Type.STRING,
"integer": Type.INTEGER,
"number": Type.NUMBER,
"boolean": Type.BOOLEAN,
"array": Type.ARRAY,
"object": Type.OBJECT,
}
return type_map.get(param_type, Type.STRING)

@override
def _get_declaration(self) -> Optional[FunctionDeclaration]:
"""Gets the function declaration for the tool."""
properties = {}
required = []

# We do not use `google.genai.types.FunctionDeclaration.from_callable`
# here because it explicitly drops argument descriptions from the schema
# properties, lumping them all into the root description instead.
if hasattr(self._core_tool, '_params') and self._core_tool._params:
for param in self._core_tool._params:
properties[param.name] = Schema(
type=self._param_type_to_schema_type(param.type),
description=param.description or ""
)
if param.required:
required.append(param.name)

parameters = Schema(
type=Type.OBJECT,
properties=properties,
required=required
) if properties else None

return FunctionDeclaration(
name=self.name,
description=self.description,
parameters=parameters
)

@override
async def run_async(
self,
args: Dict[str, Any],
tool_context: ToolContext,
) -> Any:
# 1. Pre-hook
if self._pre_hook:
await self._pre_hook(tool_context, args)

# 2. ADK Auth Integration (3LO)
# Check if USER_IDENTITY is configured
reset_token = None

Expand Down
24 changes: 20 additions & 4 deletions packages/toolbox-adk/tests/integration/test_integration.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@
import os
from typing import Any, Optional
from inspect import signature, Parameter
from unittest.mock import MagicMock
from unittest.mock import MagicMock, AsyncMock

import pytest
from pydantic import ValidationError
Expand Down Expand Up @@ -60,6 +60,14 @@ async def test_load_toolset_and_run(self):
tool = next((t for t in tools if t.name == "get-row-by-id"), None)
assert tool is not None
assert isinstance(tool, ToolboxTool)

# Verify the function declaration schema builds correctly end-to-end
declaration = tool._get_declaration()
assert declaration is not None
assert declaration.name == "get-row-by-id"
assert declaration.parameters is not None
assert hasattr(declaration.parameters, 'properties')
assert "id" in declaration.parameters.properties

# Run it
ctx = MagicMock()
Expand Down Expand Up @@ -173,19 +181,27 @@ async def test_3lo_flow_simulation(self):
tool = tools[0]
assert isinstance(tool, ToolboxTool)
assert tool.name == "get-n-rows"

# Verify the function declaration schema builds correctly end-to-end
declaration = tool._get_declaration()
assert declaration is not None
assert declaration.name == "get-n-rows"
assert "num_rows" in declaration.parameters.properties

# Create a mock context that behaves like ADK's ReadonlyContext
mock_ctx_first = MagicMock()
# Simulate "No Auth Response Found"
mock_ctx_first.get_auth_response.return_value = None
mock_cred_service_first = AsyncMock()
mock_cred_service_first.load_credential.return_value = None
mock_ctx_first._invocation_context = MagicMock()
mock_ctx_first._invocation_context.credential_service = mock_cred_service_first

print("Running tool first time (expecting auth request)...")
result_first = await tool.run_async({"num_rows": "1"}, mock_ctx_first)

# The wrapper should catch the missing creds and request them.
assert (
result_first is None
), "Tool should return None to signal auth requirement"
assert result_first is None, "Tool should return None sig for auth requirement"
mock_ctx_first.request_credential.assert_called_once()

# Inspect the requested config
Expand Down
65 changes: 65 additions & 0 deletions packages/toolbox-adk/tests/unit/test_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

from toolbox_adk.credentials import CredentialConfig, CredentialType
from toolbox_adk.tool import ToolboxTool
from google.genai.types import Type


class TestToolboxTool:
Expand Down Expand Up @@ -267,6 +268,70 @@ async def test_3lo_exception_fallback(self):
# Should catch RuntimeError, call request_credential, and return None
assert result is None
ctx.request_credential.assert_called_once()

def test_param_type_to_schema_type(self):
core_tool = MagicMock()
core_tool.__name__ = "mock_tool"
core_tool.__doc__ = "mock doc"
tool = ToolboxTool(core_tool)

assert tool._param_type_to_schema_type("string") == Type.STRING
assert tool._param_type_to_schema_type("integer") == Type.INTEGER
assert tool._param_type_to_schema_type("boolean") == Type.BOOLEAN
assert tool._param_type_to_schema_type("number") == Type.NUMBER
assert tool._param_type_to_schema_type("array") == Type.ARRAY
assert tool._param_type_to_schema_type("object") == Type.OBJECT
assert tool._param_type_to_schema_type("unknown") == Type.STRING

def test_get_declaration(self):
# Create a mock for core tool parameters
class MockParam:
def __init__(self, name, param_type, description, required):
self.name = name
self.type = param_type
self.description = description
self.required = required

core_tool = MagicMock()
core_tool.__name__ = "mock_tool"
core_tool.__doc__ = "mock doc"
core_tool._params = [
MockParam("city", "string", "The city name", True),
MockParam("count", "integer", "Number of results", False)
]

tool = ToolboxTool(core_tool)
declaration = tool._get_declaration()

assert declaration.name == "mock_tool"
assert declaration.description == "mock doc"

parameters = declaration.parameters
assert parameters is not None
assert parameters.type == Type.OBJECT
assert "city" in parameters.properties
assert "count" in parameters.properties

assert parameters.properties["city"].type == Type.STRING
assert parameters.properties["city"].description == "The city name"

assert parameters.properties["count"].type == Type.INTEGER
assert parameters.properties["count"].description == "Number of results"

assert parameters.required == ["city"]

def test_get_declaration_no_params(self):
core_tool = MagicMock()
core_tool.__name__ = "mock_tool"
core_tool.__doc__ = "mock doc"
core_tool._params = []

tool = ToolboxTool(core_tool)
declaration = tool._get_declaration()

assert declaration.name == "mock_tool"
assert declaration.description == "mock doc"
assert getattr(declaration, "parameters", None) is None

def test_init_defaults(self):
# Test initialization with minimal tool metadata checks
Expand Down
Loading