Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 20 additions & 2 deletions src/google/adk/tools/_automatic_function_calling_util.py
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@
from typing import Any
from typing import Callable
from typing import Dict
from typing import Literal
from typing import Optional
from typing import Union

Expand Down Expand Up @@ -329,7 +328,26 @@ def from_function_with_options(
return declaration

return_annotation = inspect.signature(func).return_annotation
if return_annotation is inspect._empty:

# Handle functions with no return annotation or that return None
if (
return_annotation is inspect._empty
or return_annotation is None
or return_annotation is type(None)
):
# Create a response schema for None/null return
return_value = inspect.Parameter(
'return_value',
inspect.Parameter.POSITIONAL_OR_KEYWORD,
annotation=None,
)
declaration.response = (
_function_parameter_parse_util._parse_schema_from_parameter(
variant,
return_value,
func.__name__,
)
)
return declaration

return_value = inspect.Parameter(
Expand Down
12 changes: 12 additions & 0 deletions src/google/adk/tools/agent_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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

Expand Down
146 changes: 146 additions & 0 deletions tests/unittests/tools/test_agent_tool.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
117 changes: 117 additions & 0 deletions tests/unittests/tools/test_build_function_declaration.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@

from google.adk.tools import _automatic_function_calling_util
from google.adk.tools.agent_tool import ToolContext
from google.adk.utils.variant_utils import GoogleLLMVariant
from google.genai import types
# TODO: crewai requires python 3.10 as minimum
# from crewai_tools import FileReadTool
from pydantic import BaseModel
Expand Down Expand Up @@ -262,3 +264,118 @@ def simple_function(input_str: List[CustomInput]) -> str:
# assert function_decl.name == 'directory_read_tool'
# assert function_decl.parameters.type == 'OBJECT'
# assert function_decl.parameters.properties['file_path'].type == 'STRING'


def test_function_no_return_annotation_gemini_api():
"""Test function with no return annotation using GEMINI_API variant."""

def function_no_return(param: str):
"""A function with no return annotation."""
return None

function_decl = _automatic_function_calling_util.build_function_declaration(
func=function_no_return, variant=GoogleLLMVariant.GEMINI_API
)

assert function_decl.name == 'function_no_return'
assert function_decl.parameters.type == 'OBJECT'
assert function_decl.parameters.properties['param'].type == 'STRING'
# GEMINI_API should not have response schema
assert function_decl.response is None


def test_function_no_return_annotation_vertex_ai():
"""Test function with no return annotation using VERTEX_AI variant."""

def function_no_return(param: str):
"""A function with no return annotation."""
return None

function_decl = _automatic_function_calling_util.build_function_declaration(
func=function_no_return, variant=GoogleLLMVariant.VERTEX_AI
)

assert function_decl.name == 'function_no_return'
assert function_decl.parameters.type == 'OBJECT'
assert function_decl.parameters.properties['param'].type == 'STRING'
# VERTEX_AI should have response schema for None return
assert function_decl.response is not None
assert function_decl.response.type == types.Type.NULL


def test_function_explicit_none_return_vertex_ai():
"""Test function with explicit None return annotation using VERTEX_AI variant."""

def function_none_return(param: str) -> None:
"""A function that explicitly returns None."""
pass

function_decl = _automatic_function_calling_util.build_function_declaration(
func=function_none_return, variant=GoogleLLMVariant.VERTEX_AI
)

assert function_decl.name == 'function_none_return'
assert function_decl.parameters.type == 'OBJECT'
assert function_decl.parameters.properties['param'].type == 'STRING'
# VERTEX_AI should have response schema for explicit None return
assert function_decl.response is not None
assert function_decl.response.type == types.Type.NULL


def test_function_explicit_none_return_gemini_api():
"""Test function with explicit None return annotation using GEMINI_API variant."""

def function_none_return(param: str) -> None:
"""A function that explicitly returns None."""
pass

function_decl = _automatic_function_calling_util.build_function_declaration(
func=function_none_return, variant=GoogleLLMVariant.GEMINI_API
)

assert function_decl.name == 'function_none_return'
assert function_decl.parameters.type == 'OBJECT'
assert function_decl.parameters.properties['param'].type == 'STRING'
# GEMINI_API should not have response schema
assert function_decl.response is None


def test_function_regular_return_type_vertex_ai():
"""Test function with regular return type using VERTEX_AI variant."""

def function_string_return(param: str) -> str:
"""A function that returns a string."""
return param

function_decl = _automatic_function_calling_util.build_function_declaration(
func=function_string_return, variant=GoogleLLMVariant.VERTEX_AI
)

assert function_decl.name == 'function_string_return'
assert function_decl.parameters.type == 'OBJECT'
assert function_decl.parameters.properties['param'].type == 'STRING'
# VERTEX_AI should have response schema for string return
assert function_decl.response is not None
assert function_decl.response.type == types.Type.STRING


def test_transfer_to_agent_like_function():
"""Test a function similar to transfer_to_agent that caused the original issue."""

def transfer_to_agent(agent_name: str, tool_context: ToolContext):
"""Transfer the question to another agent."""
tool_context.actions.transfer_to_agent = agent_name

function_decl = _automatic_function_calling_util.build_function_declaration(
func=transfer_to_agent,
ignore_params=['tool_context'],
variant=GoogleLLMVariant.VERTEX_AI,
)

assert function_decl.name == 'transfer_to_agent'
assert function_decl.parameters.type == 'OBJECT'
assert function_decl.parameters.properties['agent_name'].type == 'STRING'
assert 'tool_context' not in function_decl.parameters.properties
# This should now have a response schema for VERTEX_AI variant
assert function_decl.response is not None
assert function_decl.response.type == types.Type.NULL
Loading