feat: Add enum constraint to agent_name for transfer_to_agent

Merge https://github.com/google/adk-python/pull/2437

Current implementation of `transfer_to_agent` doesn't enforce strict constraints on agent names, we could use JSON Schema's enum definition to implement stricter constraints.

Co-authored-by: Xuan Yang <xygoogle@google.com>
COPYBARA_INTEGRATE_REVIEW=https://github.com/google/adk-python/pull/2437 from qieqieplus:main 052e8e73b9d61c0998573a2077f15864873d0dd7
PiperOrigin-RevId: 836410397

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
 

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'),

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

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."""
 

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

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']