ADK changes

PiperOrigin-RevId: 814319961

Author: xuanyang15

src/google/adk/agents/llm_agent.py

@@ -112,8 +112,22 @@
 
 
 async def _convert_tool_union_to_tools(
-    tool_union: ToolUnion, ctx: ReadonlyContext
+    tool_union: ToolUnion,
+    ctx: ReadonlyContext,
+    model: Union[str, BaseLlm],
+    multiple_tools: bool = False,
 ) -> list[BaseTool]:
+  from ..tools.google_search_tool import google_search
+
+  # Wrap google_search tool with AgentTool if there are multiple tools because
+  # the built-in tools cannot be used together with other tools.
+  # TODO(b/448114567): Remove once the workaround is no longer needed.
+  if multiple_tools and tool_union is google_search:
+    from ..tools.google_search_agent_tool import create_google_search_agent
+    from ..tools.google_search_agent_tool import GoogleSearchAgentTool
+
+    return [GoogleSearchAgentTool(create_google_search_agent(model))]
+
   if isinstance(tool_union, BaseTool):
     return [tool_union]
   if callable(tool_union):
@@ -462,8 +476,16 @@ async def canonical_tools(
     This method is only for use by Agent Development Kit.
     """
     resolved_tools = []
+    # We may need to wrap some built-in tools if there are other tools
+    # because the built-in tools cannot be used together with other tools.
+    # TODO(b/448114567): Remove once the workaround is no longer needed.
+    multiple_tools = len(self.tools) > 1
     for tool_union in self.tools:
-      resolved_tools.extend(await _convert_tool_union_to_tools(tool_union, ctx))
+      resolved_tools.extend(
+          await _convert_tool_union_to_tools(
+              tool_union, ctx, self.model, multiple_tools
+          )
+      )
     return resolved_tools
 
   @property

src/google/adk/flows/llm_flows/base_llm_flow.py

@@ -45,6 +45,7 @@
 from ...telemetry.tracing import trace_send_data
 from ...telemetry.tracing import tracer
 from ...tools.base_toolset import BaseToolset
+from ...tools.google_search_tool import google_search
 from ...tools.tool_context import ToolContext
 from ...utils.context_utils import Aclosing
 from .audio_cache_manager import AudioCacheManager
@@ -442,6 +443,11 @@ async def _preprocess_async(
           yield event
 
     # Run processors for tools.
+
+    # We may need to wrap some built-in tools if there are other tools
+    # because the built-in tools cannot be used together with other tools.
+    # TODO(b/448114567): Remove once the workaround is no longer needed.
+    multiple_tools = len(agent.tools) > 1
     for tool_union in agent.tools:
       tool_context = ToolContext(invocation_context)
 
@@ -455,7 +461,10 @@ async def _preprocess_async(
 
       # Then process all tools from this tool union
       tools = await _convert_tool_union_to_tools(
-          tool_union, ReadonlyContext(invocation_context)
+          tool_union,
+          ReadonlyContext(invocation_context),
+          llm_request.model,
+          multiple_tools,
       )
       for tool in tools:
         await tool.process_llm_request(
@@ -818,6 +827,26 @@ async def _handle_after_model_callback(
 
     agent = invocation_context.agent
 
+    # Add grounding metadata to the response if needed.
+    # TODO(b/448114567): Remove this function once the workaround is no longer needed.
+    async def _maybe_add_grounding_metadata(
+        response: Optional[LlmResponse] = None,
+    ) -> Optional[LlmResponse]:
+      readonly_context = ReadonlyContext(invocation_context)
+      tools = await agent.canonical_tools(readonly_context)
+      if not any(tool.name == 'google_search_agent' for tool in tools):
+        return response
+      ground_metadata = invocation_context.session.state.get(
+          'temp:_adk_grounding_metadata', None
+      )
+      if not ground_metadata:
+        return response
+
+      if not response:
+        response = llm_response
+      response.grounding_metadata = ground_metadata
+      return response
+
     callback_context = CallbackContext(
         invocation_context, event_actions=model_response_event.actions
     )
@@ -830,20 +859,21 @@ async def _handle_after_model_callback(
         )
     )
     if callback_response:
-      return callback_response
+      return await _maybe_add_grounding_metadata(callback_response)
 
     # If no overrides are provided from the plugins, further run the canonical
     # callbacks.
     if not agent.canonical_after_model_callbacks:
-      return
+      return await _maybe_add_grounding_metadata()
     for callback in agent.canonical_after_model_callbacks:
       callback_response = callback(
           callback_context=callback_context, llm_response=llm_response
       )
       if inspect.isawaitable(callback_response):
         callback_response = await callback_response
       if callback_response:
-        return callback_response
+        return await _maybe_add_grounding_metadata(callback_response)
+    return await _maybe_add_grounding_metadata()
 
   def _finalize_model_response_event(
       self,

src/google/adk/tools/google_search_agent_tool.py

@@ -0,0 +1,140 @@
+# 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.
+
+from __future__ import annotations
+
+from typing import Any
+from typing import Union
+
+from google.genai import types
+from typing_extensions import override
+
+from ..agents.llm_agent import LlmAgent
+from ..memory.in_memory_memory_service import InMemoryMemoryService
+from ..models.base_llm import BaseLlm
+from ..utils.context_utils import Aclosing
+from ._forwarding_artifact_service import ForwardingArtifactService
+from .agent_tool import AgentTool
+from .google_search_tool import google_search
+from .tool_context import ToolContext
+
+
+def create_google_search_agent(model: Union[str, BaseLlm]) -> LlmAgent:
+  """Create a sub-agent that only uses google_search tool."""
+  return LlmAgent(
+      name='google_search_agent',
+      model=model,
+      description=(
+          'An agent for performing Google search using the `google_search` tool'
+      ),
+      instruction="""
+        You are a specialized Google search agent.
+
+        When given a search query, use the `google_search` tool to find the related information.
+      """,
+      tools=[google_search],
+  )
+
+
+class GoogleSearchAgentTool(AgentTool):
+  """A tool that wraps a sub-agent that only uses google_search tool.
+
+  This is a workaround to support using google_search tool with other tools.
+  TODO(b/448114567): Remove once the workaround is no longer needed.
+
+  Attributes:
+    model: The model to use for the sub-agent.
+  """
+
+  def __init__(self, agent: LlmAgent):
+    self.agent = agent
+    super().__init__(agent=self.agent)
+
+  @override
+  async def run_async(
+      self,
+      *,
+      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 isinstance(self.agent, LlmAgent) and self.agent.input_schema:
+      input_value = self.agent.input_schema.model_validate(args)
+      content = types.Content(
+          role='user',
+          parts=[
+              types.Part.from_text(
+                  text=input_value.model_dump_json(exclude_none=True)
+              )
+          ],
+      )
+    else:
+      content = types.Content(
+          role='user',
+          parts=[types.Part.from_text(text=args['request'])],
+      )
+    runner = Runner(
+        app_name=self.agent.name,
+        agent=self.agent,
+        artifact_service=ForwardingArtifactService(tool_context),
+        session_service=InMemorySessionService(),
+        memory_service=InMemoryMemoryService(),
+        credential_service=tool_context._invocation_context.credential_service,
+        plugins=list(tool_context._invocation_context.plugin_manager.plugins),
+    )
+
+    state_dict = {
+        k: v
+        for k, v in tool_context.state.to_dict().items()
+        if not k.startswith('_adk')  # Filter out adk internal states
+    }
+    session = await runner.session_service.create_session(
+        app_name=self.agent.name,
+        user_id=tool_context._invocation_context.user_id,
+        state=state_dict,
+    )
+
+    last_content = None
+    last_grounding_metadata = None
+    async with Aclosing(
+        runner.run_async(
+            user_id=session.user_id, session_id=session.id, new_message=content
+        )
+    ) as agen:
+      async for event in agen:
+        # Forward state delta to parent session.
+        if event.actions.state_delta:
+          tool_context.state.update(event.actions.state_delta)
+        if event.content:
+          last_content = event.content
+          last_grounding_metadata = event.grounding_metadata
+
+    if not last_content:
+      return ''
+    merged_text = '\n'.join(p.text for p in last_content.parts if p.text)
+    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)
+    else:
+      tool_result = merged_text
+
+    if last_grounding_metadata:
+      tool_context.state['temp:_adk_grounding_metadata'] = (
+          last_grounding_metadata
+      )
+    return tool_result

tests/unittests/agents/test_llm_agent_fields.py

@@ -26,6 +26,7 @@
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.registry import LLMRegistry
 from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.google_search_tool import google_search
 from google.genai import types
 from pydantic import BaseModel
 import pytest
@@ -279,3 +280,63 @@ def test_allow_transfer_by_default():
 
   assert not agent.disallow_transfer_to_parent
   assert not agent.disallow_transfer_to_peers
+
+
+# TODO(b/448114567): Remove TestCanonicalTools once the workaround
+# is no longer needed.
+class TestCanonicalTools:
+  """Unit tests for canonical_tools in LlmAgent."""
+
+  @staticmethod
+  def _my_tool(sides: int) -> int:
+    return sides
+
+  async def test_handle_google_search_with_other_tools(self):
+    """Test that google_search is wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+            google_search,
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 2
+    assert tools[0].name == '_my_tool'
+    assert tools[1].name == 'google_search_agent'
+    assert tools[1].__class__.__name__ == 'GoogleSearchAgentTool'
+
+  async def test_handle_google_search_only(self):
+    """Test that google_search is not wrapped into an agent."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            google_search,
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 1
+    assert tools[0].name == 'google_search'
+    assert tools[0].__class__.__name__ == 'GoogleSearchTool'
+
+  async def test_no_google_search(self):
+    """Test other tools are not affected."""
+    agent = LlmAgent(
+        name='test_agent',
+        model='gemini-pro',
+        tools=[
+            self._my_tool,
+        ],
+    )
+    ctx = await _create_readonly_context(agent)
+    tools = await agent.canonical_tools(ctx)
+
+    assert len(tools) == 1
+    assert tools[0].name == '_my_tool'
+    assert tools[0].__class__.__name__ == 'FunctionTool'