fix: Add response schema for agent tool function declaration even when it's return None

PiperOrigin-RevId: 783610166

Author: seanzhougoogle

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
 

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