From 377b5a9b78fed898a9141a57210b0a8c2fb52488 Mon Sep 17 00:00:00 2001 From: "Xiang (Sean) Zhou" Date: Thu, 17 Jul 2025 09:50:56 -0700 Subject: [PATCH] fix: Add response schema for agent tool function declaration even when it's return None PiperOrigin-RevId: 784216811 --- src/google/adk/tools/agent_tool.py | 12 ++ tests/unittests/tools/test_agent_tool.py | 146 +++++++++++++++++++++++ 2 files changed, 158 insertions(+) diff --git a/src/google/adk/tools/agent_tool.py b/src/google/adk/tools/agent_tool.py index 4488e7fcf..2bbe8a9c7 100644 --- a/src/google/adk/tools/agent_tool.py +++ b/src/google/adk/tools/agent_tool.py @@ -61,6 +61,7 @@ def populate_name(cls, data: Any) -> Any: @override def _get_declaration(self) -> types.FunctionDeclaration: from ..agents.llm_agent import LlmAgent + from ..utils.variant_utils import GoogleLLMVariant if isinstance(self.agent, LlmAgent) and self.agent.input_schema: result = _automatic_function_calling_util.build_function_declaration( @@ -80,6 +81,17 @@ def _get_declaration(self) -> types.FunctionDeclaration: description=self.agent.description, name=self.name, ) + + # Set response schema for non-GEMINI_API variants + if self._api_variant != GoogleLLMVariant.GEMINI_API: + # Determine response type based on agent's output schema + if isinstance(self.agent, LlmAgent) and self.agent.output_schema: + # Agent has structured output schema - response is an object + result.response = types.Schema(type=types.Type.OBJECT) + else: + # Agent returns text - response is a string + result.response = types.Schema(type=types.Type.STRING) + result.name = self.name return result diff --git a/tests/unittests/tools/test_agent_tool.py b/tests/unittests/tools/test_agent_tool.py index eaef30d8d..8e2035eed 100644 --- a/tests/unittests/tools/test_agent_tool.py +++ b/tests/unittests/tools/test_agent_tool.py @@ -16,6 +16,8 @@ from google.adk.agents import SequentialAgent from google.adk.agents.callback_context import CallbackContext from google.adk.tools.agent_tool import AgentTool +from google.adk.utils.variant_utils import GoogleLLMVariant +from google.genai import types from google.genai.types import Part from pydantic import BaseModel from pytest import mark @@ -209,3 +211,147 @@ class CustomOutput(BaseModel): # The second request is the tool agent request. assert mock_model.requests[1].config.response_schema == CustomOutput assert mock_model.requests[1].config.response_mime_type == 'application/json' + + +@mark.parametrize( + 'env_variables', + [ + 'VERTEX', # Test VERTEX_AI variant + ], + indirect=True, +) +def test_agent_tool_response_schema_no_output_schema_vertex_ai(): + """Test AgentTool with no output schema has string response schema for VERTEX_AI.""" + tool_agent = Agent( + name='tool_agent', + model=testing_utils.MockModel.create(responses=['test response']), + ) + + agent_tool = AgentTool(agent=tool_agent) + declaration = agent_tool._get_declaration() + + assert declaration.name == 'tool_agent' + assert declaration.parameters.type == 'OBJECT' + assert declaration.parameters.properties['request'].type == 'STRING' + # Should have string response schema for VERTEX_AI + assert declaration.response is not None + assert declaration.response.type == types.Type.STRING + + +@mark.parametrize( + 'env_variables', + [ + 'VERTEX', # Test VERTEX_AI variant + ], + indirect=True, +) +def test_agent_tool_response_schema_with_output_schema_vertex_ai(): + """Test AgentTool with output schema has object response schema for VERTEX_AI.""" + + class CustomOutput(BaseModel): + custom_output: str + + tool_agent = Agent( + name='tool_agent', + model=testing_utils.MockModel.create(responses=['test response']), + output_schema=CustomOutput, + ) + + agent_tool = AgentTool(agent=tool_agent) + declaration = agent_tool._get_declaration() + + assert declaration.name == 'tool_agent' + # Should have object response schema for VERTEX_AI when output_schema exists + assert declaration.response is not None + assert declaration.response.type == types.Type.OBJECT + + +@mark.parametrize( + 'env_variables', + [ + 'GOOGLE_AI', # Test GEMINI_API variant + ], + indirect=True, +) +def test_agent_tool_response_schema_gemini_api(): + """Test AgentTool with GEMINI_API variant has no response schema.""" + + class CustomOutput(BaseModel): + custom_output: str + + tool_agent = Agent( + name='tool_agent', + model=testing_utils.MockModel.create(responses=['test response']), + output_schema=CustomOutput, + ) + + agent_tool = AgentTool(agent=tool_agent) + declaration = agent_tool._get_declaration() + + assert declaration.name == 'tool_agent' + # GEMINI_API should not have response schema + assert declaration.response is None + + +@mark.parametrize( + 'env_variables', + [ + 'VERTEX', # Test VERTEX_AI variant + ], + indirect=True, +) +def test_agent_tool_response_schema_with_input_schema_vertex_ai(): + """Test AgentTool with input and output schemas for VERTEX_AI.""" + + class CustomInput(BaseModel): + custom_input: str + + class CustomOutput(BaseModel): + custom_output: str + + tool_agent = Agent( + name='tool_agent', + model=testing_utils.MockModel.create(responses=['test response']), + input_schema=CustomInput, + output_schema=CustomOutput, + ) + + agent_tool = AgentTool(agent=tool_agent) + declaration = agent_tool._get_declaration() + + assert declaration.name == 'tool_agent' + assert declaration.parameters.type == 'OBJECT' + assert declaration.parameters.properties['custom_input'].type == 'STRING' + # Should have object response schema for VERTEX_AI when output_schema exists + assert declaration.response is not None + assert declaration.response.type == types.Type.OBJECT + + +@mark.parametrize( + 'env_variables', + [ + 'VERTEX', # Test VERTEX_AI variant + ], + indirect=True, +) +def test_agent_tool_response_schema_with_input_schema_no_output_vertex_ai(): + """Test AgentTool with input schema but no output schema for VERTEX_AI.""" + + class CustomInput(BaseModel): + custom_input: str + + tool_agent = Agent( + name='tool_agent', + model=testing_utils.MockModel.create(responses=['test response']), + input_schema=CustomInput, + ) + + agent_tool = AgentTool(agent=tool_agent) + declaration = agent_tool._get_declaration() + + assert declaration.name == 'tool_agent' + assert declaration.parameters.type == 'OBJECT' + assert declaration.parameters.properties['custom_input'].type == 'STRING' + # Should have string response schema for VERTEX_AI when no output_schema + assert declaration.response is not None + assert declaration.response.type == types.Type.STRING