fix: GenAI Client(evals) - Strip ADK-injected params from plain callable tools

jsondai · copybara-github · commit 551c5bd20380 · 2026-06-15T17:40:30.000-07:00
PiperOrigin-RevId: 932761778
diff --git a/agentplatform/_genai/types/evals.py b/agentplatform/_genai/types/evals.py
@@ -96,17 +96,48 @@ def _get_tool_declarations_from_agent(agent: Any) -> genai_types.ToolListUnion:
                     tool_declarations.append({"function_declarations": [declaration]})
                 continue
 
-            tool_declarations.append(
-                {
-                    "function_declarations": [
-                        genai_types.FunctionDeclaration.from_callable_with_api_option(
-                            callable=tool
-                        )
-                    ]
-                }
-            )
+            declaration = AgentConfig._get_declaration_from_callable(tool)
+            if declaration is not None:
+                tool_declarations.append({"function_declarations": [declaration]})
         return tool_declarations
 
+    @staticmethod
+    def _get_declaration_from_callable(
+        tool: Any,
+    ) -> Optional[genai_types.FunctionDeclaration]:
+        """Builds a function declaration for a plain callable tool.
+
+        ADK agents store plain Python functions in `agent.tools` and only wrap
+        them in `FunctionTool` lazily at runtime. Such functions often take
+        ADK-injected parameters (e.g. `tool_context: ToolContext`) that the
+        generic `google-genai` schema generator rejects. When google-adk is
+        available, wrap the callable in ADK's `FunctionTool` so its declaration
+        logic strips those injected parameters. Otherwise, fall back to the
+        generic generator.
+
+        Args:
+          tool: A plain callable tool from an agent's `tools` list.
+
+        Returns:
+          The function declaration for the tool, or None if the tool has no
+          declaration.
+        """
+        # pylint: disable=g-import-not-at-top,protected-access
+        # The returned FunctionDeclaration may populate either `parameters` or
+        # `parameters_json_schema` depending on the installed google-adk version
+        # and the JSON_SCHEMA_FOR_FUNC_DECL feature flag (default-on in adk>=2.2).
+        # A future adk major version will drop `parameters`, so downstream
+        # consumers must handle both fields.
+        try:
+            from google.adk.tools.function_tool import FunctionTool
+
+            return FunctionTool(func=tool)._get_declaration()  # type: ignore[no-any-return]
+        except ImportError:
+            pass
+        return genai_types.FunctionDeclaration.from_callable_with_api_option(
+            callable=tool
+        )
+
     @classmethod
     def from_agent(cls, agent: Any) -> "AgentConfig":
         """Creates an AgentConfig from an ADK agent.
diff --git a/tests/unit/agentplatform/genai/test_evals.py b/tests/unit/agentplatform/genai/test_evals.py
@@ -14,6 +14,7 @@
 #
 # pylint: disable=protected-access,bad-continuation,
 import base64
+import builtins
 import importlib
 import json
 import os
@@ -6012,6 +6013,88 @@ def my_search_tool(query: str) -> str:
             mock_function_declaration
         ]
 
+    def test_load_from_agent_plain_callable_wraps_in_adk_function_tool(self):
+        """Plain callables with ToolContext params are declared via ADK FunctionTool."""
+
+        def memorize(key: str, value: str, tool_context: "ToolContext"):  # noqa: F821
+            tool_context.state[key] = value
+            return {"status": "ok"}
+
+        mock_declaration = mock.Mock(spec=genai_types.FunctionDeclaration)
+        mock_function_tool_cls = mock.MagicMock()
+        mock_function_tool_cls.return_value._get_declaration.return_value = (
+            mock_declaration
+        )
+        mock_modules = {
+            "google.adk": mock.MagicMock(),
+            "google.adk.tools": mock.MagicMock(),
+            "google.adk.tools.function_tool": mock.MagicMock(
+                FunctionTool=mock_function_tool_cls
+            ),
+        }
+
+        mock_agent = mock.Mock()
+        mock_agent.name = "mock_agent"
+        mock_agent.instruction = "mock instruction"
+        mock_agent.description = "mock description"
+        mock_agent.tools = [memorize]
+        mock_agent.sub_agents = []
+
+        with (
+            mock.patch.object(
+                genai_types.FunctionDeclaration, "from_callable_with_api_option"
+            ) as mock_from_callable,
+            mock.patch.dict(sys.modules, mock_modules),
+        ):
+            agent_info = agentplatform_genai_types.evals.AgentInfo.load_from_agent(
+                agent=mock_agent,
+            )
+
+        assert agent_info.agents["mock_agent"].tools[0].function_declarations == [
+            mock_declaration
+        ]
+        mock_function_tool_cls.assert_called_once_with(func=memorize)
+        mock_from_callable.assert_not_called()
+
+    def test_load_from_agent_plain_callable_falls_back_without_adk(self):
+        """When google-adk is unavailable, plain callables use from_callable."""
+
+        def my_plain_tool(query: str) -> str:
+            return query
+
+        mock_declaration = mock.Mock(spec=genai_types.FunctionDeclaration)
+
+        mock_agent = mock.Mock()
+        mock_agent.name = "mock_agent"
+        mock_agent.instruction = "mock instruction"
+        mock_agent.description = "mock description"
+        mock_agent.tools = [my_plain_tool]
+        mock_agent.sub_agents = []
+
+        real_import = builtins.__import__
+
+        def _no_adk_import(name, *args, **kwargs):
+            if name == "google.adk.tools.function_tool":
+                raise ImportError("google-adk not installed")
+            return real_import(name, *args, **kwargs)
+
+        with (
+            mock.patch.object(
+                genai_types.FunctionDeclaration,
+                "from_callable_with_api_option",
+                return_value=mock_declaration,
+            ) as mock_from_callable,
+            mock.patch.object(builtins, "__import__", side_effect=_no_adk_import),
+        ):
+            agent_info = agentplatform_genai_types.evals.AgentInfo.load_from_agent(
+                agent=mock_agent,
+            )
+
+        assert agent_info.agents["mock_agent"].tools[0].function_declarations == [
+            mock_declaration
+        ]
+        mock_from_callable.assert_called_once_with(callable=my_plain_tool)
+
 
 class TestValidateDatasetAgentData:
     """Unit tests for the _validate_dataset_agent_data function."""
diff --git a/vertexai/_genai/types/evals.py b/vertexai/_genai/types/evals.py
@@ -96,17 +96,48 @@ def _get_tool_declarations_from_agent(agent: Any) -> genai_types.ToolListUnion:
                     tool_declarations.append({"function_declarations": [declaration]})
                 continue
 
-            tool_declarations.append(
-                {
-                    "function_declarations": [
-                        genai_types.FunctionDeclaration.from_callable_with_api_option(
-                            callable=tool
-                        )
-                    ]
-                }
-            )
+            declaration = AgentConfig._get_declaration_from_callable(tool)
+            if declaration is not None:
+                tool_declarations.append({"function_declarations": [declaration]})
         return tool_declarations
 
+    @staticmethod
+    def _get_declaration_from_callable(
+        tool: Any,
+    ) -> Optional[genai_types.FunctionDeclaration]:
+        """Builds a function declaration for a plain callable tool.
+
+        ADK agents store plain Python functions in `agent.tools` and only wrap
+        them in `FunctionTool` lazily at runtime. Such functions often take
+        ADK-injected parameters (e.g. `tool_context: ToolContext`) that the
+        generic `google-genai` schema generator rejects. When google-adk is
+        available, wrap the callable in ADK's `FunctionTool` so its declaration
+        logic strips those injected parameters. Otherwise, fall back to the
+        generic generator.
+
+        Args:
+          tool: A plain callable tool from an agent's `tools` list.
+
+        Returns:
+          The function declaration for the tool, or None if the tool has no
+          declaration.
+        """
+        # pylint: disable=g-import-not-at-top,protected-access
+        # The returned FunctionDeclaration may populate either `parameters` or
+        # `parameters_json_schema` depending on the installed google-adk version
+        # and the JSON_SCHEMA_FOR_FUNC_DECL feature flag (default-on in adk>=2.2).
+        # A future adk major version will drop `parameters`, so downstream
+        # consumers must handle both fields.
+        try:
+            from google.adk.tools.function_tool import FunctionTool
+
+            return FunctionTool(func=tool)._get_declaration()  # type: ignore[no-any-return]
+        except ImportError:
+            pass
+        return genai_types.FunctionDeclaration.from_callable_with_api_option(
+            callable=tool
+        )
+
     @classmethod
     def from_agent(cls, agent: Any) -> "AgentConfig":
         """Creates an AgentConfig from an ADK agent.
