diff --git a/src/google/adk/flows/llm_flows/agent_transfer.py b/src/google/adk/flows/llm_flows/agent_transfer.py index 037b8c6d..ea144bf7 100644 --- a/src/google/adk/flows/llm_flows/agent_transfer.py +++ b/src/google/adk/flows/llm_flows/agent_transfer.py @@ -24,9 +24,8 @@ from typing_extensions import override from ...agents.invocation_context import InvocationContext from ...events.event import Event from ...models.llm_request import LlmRequest -from ...tools.function_tool import FunctionTool from ...tools.tool_context import ToolContext -from ...tools.transfer_to_agent_tool import transfer_to_agent +from ...tools.transfer_to_agent_tool import TransferToAgentTool from ._base_llm_processor import BaseLlmRequestProcessor if typing.TYPE_CHECKING: @@ -50,13 +49,18 @@ class _AgentTransferLlmRequestProcessor(BaseLlmRequestProcessor): if not transfer_targets: return + transfer_to_agent_tool = TransferToAgentTool( + agent_names=[agent.name for agent in transfer_targets] + ) + llm_request.append_instructions([ _build_target_agents_instructions( - invocation_context.agent, transfer_targets + transfer_to_agent_tool.name, + invocation_context.agent, + transfer_targets, ) ]) - transfer_to_agent_tool = FunctionTool(func=transfer_to_agent) tool_context = ToolContext(invocation_context) await transfer_to_agent_tool.process_llm_request( tool_context=tool_context, llm_request=llm_request @@ -80,10 +84,13 @@ line_break = '\n' def _build_target_agents_instructions( - agent: LlmAgent, target_agents: list[BaseAgent] + tool_name: str, + agent: LlmAgent, + target_agents: list[BaseAgent], ) -> str: # Build list of available agent names for the NOTE - # target_agents already includes parent agent if applicable, so no need to add it again + # target_agents already includes parent agent if applicable, + # so no need to add it again available_agent_names = [target_agent.name for target_agent in target_agents] # Sort for consistency @@ -101,15 +108,16 @@ You have a list of other agents to transfer to: _build_target_agents_info(target_agent) for target_agent in target_agents ])} -If you are the best to answer the question according to your description, you -can answer it. +If you are the best to answer the question according to your description, +you can answer it. If another agent is better for answering the question according to its -description, call `{_TRANSFER_TO_AGENT_FUNCTION_NAME}` function to transfer the -question to that agent. When transferring, do not generate any text other than -the function call. +description, call `{tool_name}` function to transfer the question to that +agent. When transferring, do not generate any text other than the function +call. -**NOTE**: the only available agents for `{_TRANSFER_TO_AGENT_FUNCTION_NAME}` function are {formatted_agent_names}. +**NOTE**: the only available agents for `{tool_name}` function are +{formatted_agent_names}. """ if agent.parent_agent and not agent.disallow_transfer_to_parent: @@ -119,9 +127,6 @@ If neither you nor the other agents are best for the question, transfer to your return si -_TRANSFER_TO_AGENT_FUNCTION_NAME = transfer_to_agent.__name__ - - def _get_transfer_targets(agent: LlmAgent) -> list[BaseAgent]: from ...agents.llm_agent import LlmAgent diff --git a/src/google/adk/tools/__init__.py b/src/google/adk/tools/__init__.py index 1777bd93..d359abb7 100644 --- a/src/google/adk/tools/__init__.py +++ b/src/google/adk/tools/__init__.py @@ -37,6 +37,7 @@ if TYPE_CHECKING: from .preload_memory_tool import preload_memory_tool as preload_memory from .tool_context import ToolContext from .transfer_to_agent_tool import transfer_to_agent + from .transfer_to_agent_tool import TransferToAgentTool from .url_context_tool import url_context from .vertex_ai_search_tool import VertexAiSearchTool @@ -75,6 +76,10 @@ _LAZY_MAPPING = { 'preload_memory': ('.preload_memory_tool', 'preload_memory_tool'), 'ToolContext': ('.tool_context', 'ToolContext'), 'transfer_to_agent': ('.transfer_to_agent_tool', 'transfer_to_agent'), + 'TransferToAgentTool': ( + '.transfer_to_agent_tool', + 'TransferToAgentTool', + ), 'url_context': ('.url_context_tool', 'url_context'), 'VertexAiSearchTool': ('.vertex_ai_search_tool', 'VertexAiSearchTool'), 'MCPToolset': ('.mcp_tool.mcp_toolset', 'MCPToolset'), diff --git a/src/google/adk/tools/transfer_to_agent_tool.py b/src/google/adk/tools/transfer_to_agent_tool.py index 99ee234b..2124e6aa 100644 --- a/src/google/adk/tools/transfer_to_agent_tool.py +++ b/src/google/adk/tools/transfer_to_agent_tool.py @@ -14,6 +14,12 @@ from __future__ import annotations +from typing import Optional + +from google.genai import types +from typing_extensions import override + +from .function_tool import FunctionTool from .tool_context import ToolContext @@ -23,7 +29,61 @@ def transfer_to_agent(agent_name: str, tool_context: ToolContext) -> None: This tool hands off control to another agent when it's more suitable to answer the user's question according to the agent's description. + Note: + For most use cases, you should use TransferToAgentTool instead of this + function directly. TransferToAgentTool provides additional enum constraints + that prevent LLMs from hallucinating invalid agent names. + Args: agent_name: the agent name to transfer to. """ tool_context.actions.transfer_to_agent = agent_name + + +class TransferToAgentTool(FunctionTool): + """A specialized FunctionTool for agent transfer with enum constraints. + + This tool enhances the base transfer_to_agent function by adding JSON Schema + enum constraints to the agent_name parameter. This prevents LLMs from + hallucinating invalid agent names by restricting choices to only valid agents. + + Attributes: + agent_names: List of valid agent names that can be transferred to. + """ + + def __init__( + self, + agent_names: list[str], + ): + """Initialize the TransferToAgentTool. + + Args: + agent_names: List of valid agent names that can be transferred to. + """ + super().__init__(func=transfer_to_agent) + self._agent_names = agent_names + + @override + def _get_declaration(self) -> Optional[types.FunctionDeclaration]: + """Add enum constraint to the agent_name parameter. + + Returns: + FunctionDeclaration with enum constraint on agent_name parameter. + """ + function_decl = super()._get_declaration() + if not function_decl: + return function_decl + + # Handle parameters (types.Schema object) + if function_decl.parameters: + agent_name_schema = function_decl.parameters.properties.get('agent_name') + if agent_name_schema: + agent_name_schema.enum = self._agent_names + + # Handle parameters_json_schema (dict) + if function_decl.parameters_json_schema: + properties = function_decl.parameters_json_schema.get('properties', {}) + if 'agent_name' in properties: + properties['agent_name']['enum'] = self._agent_names + + return function_decl diff --git a/tests/unittests/flows/llm_flows/test_agent_transfer_system_instructions.py b/tests/unittests/flows/llm_flows/test_agent_transfer_system_instructions.py index be97a627..b180a589 100644 --- a/tests/unittests/flows/llm_flows/test_agent_transfer_system_instructions.py +++ b/tests/unittests/flows/llm_flows/test_agent_transfer_system_instructions.py @@ -126,15 +126,16 @@ Agent name: peer_agent Agent description: Peer agent -If you are the best to answer the question according to your description, you -can answer it. +If you are the best to answer the question according to your description, +you can answer it. If another agent is better for answering the question according to its -description, call `transfer_to_agent` function to transfer the -question to that agent. When transferring, do not generate any text other than -the function call. +description, call `transfer_to_agent` function to transfer the question to that +agent. When transferring, do not generate any text other than the function +call. -**NOTE**: the only available agents for `transfer_to_agent` function are `a_agent`, `m_agent`, `parent_agent`, `peer_agent`, `z_agent`. +**NOTE**: the only available agents for `transfer_to_agent` function are +`a_agent`, `m_agent`, `parent_agent`, `peer_agent`, `z_agent`. If neither you nor the other agents are best for the question, transfer to your parent agent parent_agent.""" @@ -189,15 +190,16 @@ Agent name: agent2 Agent description: Second sub-agent -If you are the best to answer the question according to your description, you -can answer it. +If you are the best to answer the question according to your description, +you can answer it. If another agent is better for answering the question according to its -description, call `transfer_to_agent` function to transfer the -question to that agent. When transferring, do not generate any text other than -the function call. +description, call `transfer_to_agent` function to transfer the question to that +agent. When transferring, do not generate any text other than the function +call. -**NOTE**: the only available agents for `transfer_to_agent` function are `agent1`, `agent2`.""" +**NOTE**: the only available agents for `transfer_to_agent` function are +`agent1`, `agent2`.""" assert expected_content in instructions @@ -248,15 +250,16 @@ Agent name: parent_agent Agent description: Parent agent -If you are the best to answer the question according to your description, you -can answer it. +If you are the best to answer the question according to your description, +you can answer it. If another agent is better for answering the question according to its -description, call `transfer_to_agent` function to transfer the -question to that agent. When transferring, do not generate any text other than -the function call. +description, call `transfer_to_agent` function to transfer the question to that +agent. When transferring, do not generate any text other than the function +call. -**NOTE**: the only available agents for `transfer_to_agent` function are `parent_agent`, `sub_agent`. +**NOTE**: the only available agents for `transfer_to_agent` function are +`parent_agent`, `sub_agent`. If neither you nor the other agents are best for the question, transfer to your parent agent parent_agent.""" diff --git a/tests/unittests/tools/test_build_function_declaration.py b/tests/unittests/tools/test_build_function_declaration.py index f57c3d38..8a562c7c 100644 --- a/tests/unittests/tools/test_build_function_declaration.py +++ b/tests/unittests/tools/test_build_function_declaration.py @@ -411,3 +411,19 @@ def test_function_with_no_response_annotations(): # Changed: Now uses Any type instead of NULL for no return annotation assert function_decl.response is not None assert function_decl.response.type is None # Any type maps to None in schema + + +def test_transfer_to_agent_tool_with_enum_constraint(): + """Test TransferToAgentTool adds enum constraint to agent_name.""" + from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool + + agent_names = ['agent_a', 'agent_b', 'agent_c'] + tool = TransferToAgentTool(agent_names=agent_names) + + function_decl = tool._get_declaration() + + assert function_decl.name == 'transfer_to_agent' + assert function_decl.parameters.type == 'OBJECT' + assert function_decl.parameters.properties['agent_name'].type == 'STRING' + assert function_decl.parameters.properties['agent_name'].enum == agent_names + assert 'tool_context' not in function_decl.parameters.properties diff --git a/tests/unittests/tools/test_transfer_to_agent_tool.py b/tests/unittests/tools/test_transfer_to_agent_tool.py new file mode 100644 index 00000000..14b7b3ab --- /dev/null +++ b/tests/unittests/tools/test_transfer_to_agent_tool.py @@ -0,0 +1,164 @@ +# Copyright 2025 Google LLC +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Tests for TransferToAgentTool enum constraint functionality.""" + +from unittest.mock import patch + +from google.adk.tools.function_tool import FunctionTool +from google.adk.tools.transfer_to_agent_tool import TransferToAgentTool +from google.genai import types + + +def test_transfer_to_agent_tool_enum_constraint(): + """Test that TransferToAgentTool adds enum constraint to agent_name.""" + agent_names = ['agent_a', 'agent_b', 'agent_c'] + tool = TransferToAgentTool(agent_names=agent_names) + + decl = tool._get_declaration() + + assert decl is not None + assert decl.name == 'transfer_to_agent' + assert decl.parameters is not None + assert decl.parameters.type == types.Type.OBJECT + assert 'agent_name' in decl.parameters.properties + + agent_name_schema = decl.parameters.properties['agent_name'] + assert agent_name_schema.type == types.Type.STRING + assert agent_name_schema.enum == agent_names + + # Verify that agent_name is marked as required + assert decl.parameters.required == ['agent_name'] + + +def test_transfer_to_agent_tool_single_agent(): + """Test TransferToAgentTool with a single agent.""" + tool = TransferToAgentTool(agent_names=['single_agent']) + + decl = tool._get_declaration() + + assert decl is not None + agent_name_schema = decl.parameters.properties['agent_name'] + assert agent_name_schema.enum == ['single_agent'] + + +def test_transfer_to_agent_tool_multiple_agents(): + """Test TransferToAgentTool with multiple agents.""" + agent_names = ['agent_1', 'agent_2', 'agent_3', 'agent_4', 'agent_5'] + tool = TransferToAgentTool(agent_names=agent_names) + + decl = tool._get_declaration() + + assert decl is not None + agent_name_schema = decl.parameters.properties['agent_name'] + assert agent_name_schema.enum == agent_names + assert len(agent_name_schema.enum) == 5 + + +def test_transfer_to_agent_tool_empty_list(): + """Test TransferToAgentTool with an empty agent list.""" + tool = TransferToAgentTool(agent_names=[]) + + decl = tool._get_declaration() + + assert decl is not None + agent_name_schema = decl.parameters.properties['agent_name'] + assert agent_name_schema.enum == [] + + +def test_transfer_to_agent_tool_preserves_description(): + """Test that TransferToAgentTool preserves the original description.""" + tool = TransferToAgentTool(agent_names=['agent_a', 'agent_b']) + + decl = tool._get_declaration() + + assert decl is not None + assert decl.description is not None + assert 'Transfer the question to another agent' in decl.description + + +def test_transfer_to_agent_tool_preserves_parameter_type(): + """Test that TransferToAgentTool preserves the parameter type.""" + tool = TransferToAgentTool(agent_names=['agent_a']) + + decl = tool._get_declaration() + + assert decl is not None + agent_name_schema = decl.parameters.properties['agent_name'] + # Should still be a string type, just with enum constraint + assert agent_name_schema.type == types.Type.STRING + + +def test_transfer_to_agent_tool_no_extra_parameters(): + """Test that TransferToAgentTool doesn't add extra parameters.""" + tool = TransferToAgentTool(agent_names=['agent_a']) + + decl = tool._get_declaration() + + assert decl is not None + # Should only have agent_name parameter (tool_context is ignored) + assert len(decl.parameters.properties) == 1 + assert 'agent_name' in decl.parameters.properties + assert 'tool_context' not in decl.parameters.properties + + +def test_transfer_to_agent_tool_maintains_inheritance(): + """Test that TransferToAgentTool inherits from FunctionTool correctly.""" + tool = TransferToAgentTool(agent_names=['agent_a']) + + assert isinstance(tool, FunctionTool) + assert hasattr(tool, '_get_declaration') + assert hasattr(tool, 'process_llm_request') + + +def test_transfer_to_agent_tool_handles_parameters_json_schema(): + """Test that TransferToAgentTool handles parameters_json_schema format.""" + agent_names = ['agent_x', 'agent_y', 'agent_z'] + + # Create a mock FunctionDeclaration with parameters_json_schema + mock_decl = type('MockDecl', (), {})() + mock_decl.parameters = None # No Schema object + mock_decl.parameters_json_schema = { + 'type': 'object', + 'properties': { + 'agent_name': { + 'type': 'string', + 'description': 'Agent name to transfer to', + } + }, + 'required': ['agent_name'], + } + + # Temporarily patch FunctionTool._get_declaration + with patch.object( + FunctionTool, + '_get_declaration', + return_value=mock_decl, + ): + tool = TransferToAgentTool(agent_names=agent_names) + result = tool._get_declaration() + + # Verify enum was added to parameters_json_schema + assert result.parameters_json_schema is not None + assert 'agent_name' in result.parameters_json_schema['properties'] + assert ( + result.parameters_json_schema['properties']['agent_name']['enum'] + == agent_names + ) + assert ( + result.parameters_json_schema['properties']['agent_name']['type'] + == 'string' + ) + # Verify required field is preserved + assert result.parameters_json_schema['required'] == ['agent_name']