From 9fc8f99b7a6edfe90ebffc87579da21020ed6e73 Mon Sep 17 00:00:00 2001 From: Jason Dai Date: Tue, 9 Jun 2026 22:00:44 -0700 Subject: [PATCH] fix: GenAI Client(evals) - Strip ADK-injected params from plain callable tools PiperOrigin-RevId: 929596761 --- agentplatform/_genai/types/evals.py | 44 ++++++++--- tests/unit/agentplatform/genai/test_evals.py | 83 ++++++++++++++++++++ vertexai/_genai/types/evals.py | 44 ++++++++--- 3 files changed, 153 insertions(+), 18 deletions(-) diff --git a/agentplatform/_genai/types/evals.py b/agentplatform/_genai/types/evals.py index 59c59b8948..4a919d8a5a 100644 --- a/agentplatform/_genai/types/evals.py +++ b/agentplatform/_genai/types/evals.py @@ -96,17 +96,43 @@ 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 + 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 index 1b05cb4b0b..3a5ff27c74 100644 --- 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 index a7abbc9266..bcb54308bc 100644 --- a/vertexai/_genai/types/evals.py +++ b/vertexai/_genai/types/evals.py @@ -96,17 +96,43 @@ 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 + 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.