fix(adk): implement _get_declaration to natively support ADK schema builder (#534)

* fix(adk): Implement generating `FunctionDeclaration` from tool parameters.

* test(adk): fix test_3lo_flow_simulation by properly mocking credential_service

* chore: improve comments

* chore: delint

* chore: move `google.genai.types` import to top of file.

Author: anubhav756

packages/toolbox-adk/src/toolbox_adk/tool.py

@@ -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
@@ -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
 

packages/toolbox-adk/tests/integration/test_integration.py

@@ -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
@@ -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()
@@ -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

packages/toolbox-adk/tests/unit/test_tool.py

@@ -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:
@@ -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