fix: recursively extract input/output schema for AgentTool

Fixes: https://github.com/google/adk-python/issues/4154

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 859440231

src/google/adk/tools/agent_tool.py

@@ -15,9 +15,11 @@
 from __future__ import annotations
 
 from typing import Any
+from typing import Optional
 from typing import TYPE_CHECKING
 
 from google.genai import types
+from pydantic import BaseModel
 from pydantic import model_validator
 from typing_extensions import override
 
@@ -37,6 +39,56 @@ if TYPE_CHECKING:
   from ..agents.base_agent import BaseAgent
 
 
+def _get_input_schema(agent: BaseAgent) -> Optional[type[BaseModel]]:
+  """Extracts the input_schema from an agent.
+
+  For LlmAgent, returns its input_schema directly.
+  For agents with sub_agents, recursively searches the first sub-agent for an
+  input_schema.
+
+  Args:
+    agent: The agent to extract input_schema from.
+
+  Returns:
+    The input_schema if found, None otherwise.
+  """
+  from ..agents.llm_agent import LlmAgent
+
+  if isinstance(agent, LlmAgent):
+    return agent.input_schema
+
+  # For composite agents, check the first sub-agent
+  if agent.sub_agents:
+    return _get_input_schema(agent.sub_agents[0])
+
+  return None
+
+
+def _get_output_schema(agent: BaseAgent) -> Optional[type[BaseModel]]:
+  """Extracts the output_schema from an agent.
+
+  For LlmAgent, returns its output_schema directly.
+  For agents with sub_agents, recursively searches the last sub-agent for an
+  output_schema.
+
+  Args:
+    agent: The agent to extract output_schema from.
+
+  Returns:
+    The output_schema if found, None otherwise.
+  """
+  from ..agents.llm_agent import LlmAgent
+
+  if isinstance(agent, LlmAgent):
+    return agent.output_schema
+
+  # For composite agents, check the last sub-agent
+  if agent.sub_agents:
+    return _get_output_schema(agent.sub_agents[-1])
+
+  return None
+
+
 class AgentTool(BaseTool):
   """A tool that wraps an agent.
 
@@ -74,12 +126,14 @@ class AgentTool(BaseTool):
 
   @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:
+    input_schema = _get_input_schema(self.agent)
+    output_schema = _get_output_schema(self.agent)
+
+    if input_schema:
       result = _automatic_function_calling_util.build_function_declaration(
-          func=self.agent.input_schema, variant=self._api_variant
+          func=input_schema, variant=self._api_variant
       )
       # Override the description with the agent's description
       result.description = self.agent.description
@@ -114,7 +168,7 @@ class AgentTool(BaseTool):
     # 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:
+      if output_schema:
         # Agent has structured output schema - response is an object
         if is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL):
           result.response_json_schema = {'type': 'object'}
@@ -137,15 +191,15 @@ class AgentTool(BaseTool):
       args: dict[str, Any],
       tool_context: ToolContext,
   ) -> Any:
-    from ..agents.llm_agent import LlmAgent
     from ..runners import Runner
     from ..sessions.in_memory_session_service import InMemorySessionService
 
     if self.skip_summarization:
       tool_context.actions.skip_summarization = True
 
-    if isinstance(self.agent, LlmAgent) and self.agent.input_schema:
-      input_value = self.agent.input_schema.model_validate(args)
+    input_schema = _get_input_schema(self.agent)
+    if input_schema:
+      input_value = input_schema.model_validate(args)
       content = types.Content(
           role='user',
           parts=[
@@ -212,10 +266,11 @@ class AgentTool(BaseTool):
     merged_text = '\n'.join(
         p.text for p in last_content.parts if p.text and not p.thought
     )
-    if isinstance(self.agent, LlmAgent) and self.agent.output_schema:
-      tool_result = self.agent.output_schema.model_validate_json(
-          merged_text
-      ).model_dump(exclude_none=True)
+    output_schema = _get_output_schema(self.agent)
+    if output_schema:
+      tool_result = output_schema.model_validate_json(merged_text).model_dump(
+          exclude_none=True
+      )
     else:
       tool_result = merged_text
     return tool_result