ADK changes

PiperOrigin-RevId: 814319961

src/google/adk/agents/llm_agent.py

@@ -112,8 +112,22 @@ ToolUnion: TypeAlias = Union[Callable, BaseTool, BaseToolset]
 
 
 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 @@ class LlmAgent(BaseAgent):
     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_call_llm
 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 @@ class BaseLlmFlow(ABC):
           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 @@ class BaseLlmFlow(ABC):
 
       # 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 @@ class BaseLlmFlow(ABC):
 
     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,12 +859,12 @@ class BaseLlmFlow(ABC):
         )
     )
     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
@@ -843,7 +872,8 @@ class BaseLlmFlow(ABC):
       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.agents.readonly_context import ReadonlyContext
 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'

tests/unittests/flows/llm_flows/test_base_llm_flow.py

@@ -14,13 +14,18 @@
 
 """Unit tests for BaseLlmFlow toolset integration."""
 
+from typing import Optional
 from unittest.mock import AsyncMock
 
+from google.adk.agents.callback_context import CallbackContext
 from google.adk.agents.llm_agent import Agent
+from google.adk.events.event import Event
 from google.adk.flows.llm_flows.base_llm_flow import BaseLlmFlow
 from google.adk.models.llm_request import LlmRequest
 from google.adk.models.llm_response import LlmResponse
+from google.adk.plugins.base_plugin import BasePlugin
 from google.adk.tools.base_toolset import BaseToolset
+from google.adk.tools.google_search_tool import google_search
 from google.genai import types
 import pytest
 
@@ -148,3 +153,222 @@ async def test_preprocess_handles_mixed_tools_and_toolsets():
   # Verify that process_llm_request was called on both tools and toolsets
   assert mock_tool.process_llm_request_called
   assert mock_toolset.process_llm_request_called
+
+
+# TODO(b/448114567): Remove the following test_preprocess_with_google_search
+# tests once the workaround is no longer needed.
+@pytest.mark.asyncio
+async def test_preprocess_with_google_search_only():
+  """Test _preprocess_async with only the google_search tool."""
+  agent = Agent(name='test_agent', model='gemini-pro', tools=[google_search])
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+  flow = BaseLlmFlowForTesting()
+  llm_request = LlmRequest(model='gemini-pro')
+  async for _ in flow._preprocess_async(invocation_context, llm_request):
+    pass
+
+  assert len(llm_request.config.tools) == 1
+  assert llm_request.config.tools[0].google_search is not None
+
+
+@pytest.mark.asyncio
+async def test_preprocess_with_google_search_workaround():
+  """Test _preprocess_async with google_search and another tool."""
+
+  def _my_tool(sides: int) -> int:
+    """A simple tool."""
+    return sides
+
+  agent = Agent(
+      name='test_agent', model='gemini-pro', tools=[_my_tool, google_search]
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, user_content='test message'
+  )
+  flow = BaseLlmFlowForTesting()
+  llm_request = LlmRequest(model='gemini-pro')
+  async for _ in flow._preprocess_async(invocation_context, llm_request):
+    pass
+
+  assert len(llm_request.config.tools) == 1
+  declarations = llm_request.config.tools[0].function_declarations
+  assert len(declarations) == 2
+  assert {d.name for d in declarations} == {'_my_tool', 'google_search_agent'}
+
+
+# TODO(b/448114567): Remove the following
+# test_handle_after_model_callback_grounding tests once the workaround
+# is no longer needed.
+def dummy_tool():
+  pass
+
+
+@pytest.mark.parametrize(
+    'tools, state_metadata, expect_metadata',
+    [
+        ([], None, False),
+        ([google_search, dummy_tool], {'foo': 'bar'}, True),
+        ([dummy_tool], {'foo': 'bar'}, False),
+        ([google_search, dummy_tool], None, False),
+    ],
+    ids=[
+        'no_search_no_grounding',
+        'with_search_with_grounding',
+        'no_search_with_grounding',
+        'with_search_no_grounding',
+    ],
+)
+@pytest.mark.asyncio
+async def test_handle_after_model_callback_grounding_with_no_callbacks(
+    tools, state_metadata, expect_metadata
+):
+  """Test handling grounding metadata when there are no callbacks."""
+  agent = Agent(name='test_agent', tools=tools)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  if state_metadata:
+    invocation_context.session.state['temp:_adk_grounding_metadata'] = (
+        state_metadata
+    )
+
+  llm_response = LlmResponse(
+      content=types.Content(parts=[types.Part.from_text(text='response')])
+  )
+  event = Event(
+      id=Event.new_id(),
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+  )
+  flow = BaseLlmFlowForTesting()
+
+  result = await flow._handle_after_model_callback(
+      invocation_context, llm_response, event
+  )
+
+  if expect_metadata:
+    llm_response.grounding_metadata = state_metadata
+    assert result == llm_response
+  else:
+    assert result is None
+
+
+@pytest.mark.parametrize(
+    'tools, state_metadata, expect_metadata',
+    [
+        ([], None, False),
+        ([google_search, dummy_tool], {'foo': 'bar'}, True),
+        ([dummy_tool], {'foo': 'bar'}, False),
+        ([google_search, dummy_tool], None, False),
+    ],
+    ids=[
+        'no_search_no_grounding',
+        'with_search_with_grounding',
+        'no_search_with_grounding',
+        'with_search_no_grounding',
+    ],
+)
+@pytest.mark.asyncio
+async def test_handle_after_model_callback_grounding_with_callback_override(
+    tools, state_metadata, expect_metadata
+):
+  """Test handling grounding metadata when there is a callback override."""
+  agent_response = LlmResponse(
+      content=types.Content(parts=[types.Part.from_text(text='agent')])
+  )
+  agent_callback = AsyncMock(return_value=agent_response)
+
+  agent = Agent(
+      name='test_agent', tools=tools, after_model_callback=[agent_callback]
+  )
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent
+  )
+  if state_metadata:
+    invocation_context.session.state['temp:_adk_grounding_metadata'] = (
+        state_metadata
+    )
+
+  llm_response = LlmResponse(
+      content=types.Content(parts=[types.Part.from_text(text='response')])
+  )
+  event = Event(
+      id=Event.new_id(),
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+  )
+  flow = BaseLlmFlowForTesting()
+
+  result = await flow._handle_after_model_callback(
+      invocation_context, llm_response, event
+  )
+
+  if expect_metadata:
+    agent_response.grounding_metadata = state_metadata
+
+  assert result == agent_response
+  agent_callback.assert_called_once()
+
+
+@pytest.mark.parametrize(
+    'tools, state_metadata, expect_metadata',
+    [
+        ([], None, False),
+        ([google_search, dummy_tool], {'foo': 'bar'}, True),
+        ([dummy_tool], {'foo': 'bar'}, False),
+        ([google_search, dummy_tool], None, False),
+    ],
+    ids=[
+        'no_search_no_grounding',
+        'with_search_with_grounding',
+        'no_search_with_grounding',
+        'with_search_no_grounding',
+    ],
+)
+@pytest.mark.asyncio
+async def test_handle_after_model_callback_grounding_with_plugin_override(
+    tools, state_metadata, expect_metadata
+):
+  """Test handling grounding metadata when there is a plugin override."""
+  plugin_response = LlmResponse(
+      content=types.Content(parts=[types.Part.from_text(text='plugin')])
+  )
+
+  class _MockPlugin(BasePlugin):
+
+    def __init__(self):
+      super().__init__(name='mock_plugin')
+
+    after_model_callback = AsyncMock(return_value=plugin_response)
+
+  plugin = _MockPlugin()
+  agent = Agent(name='test_agent', tools=tools)
+  invocation_context = await testing_utils.create_invocation_context(
+      agent=agent, plugins=[plugin]
+  )
+  if state_metadata:
+    invocation_context.session.state['temp:_adk_grounding_metadata'] = (
+        state_metadata
+    )
+
+  llm_response = LlmResponse(
+      content=types.Content(parts=[types.Part.from_text(text='response')])
+  )
+  event = Event(
+      id=Event.new_id(),
+      invocation_id=invocation_context.invocation_id,
+      author=agent.name,
+  )
+  flow = BaseLlmFlowForTesting()
+
+  result = await flow._handle_after_model_callback(
+      invocation_context, llm_response, event
+  )
+
+  if expect_metadata:
+    plugin_response.grounding_metadata = state_metadata
+
+  assert result == plugin_response
+  plugin.after_model_callback.assert_called_once()

tests/unittests/tools/test_google_search_agent_tool.py

@@ -0,0 +1,139 @@
+# 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 google.adk.agents.invocation_context import InvocationContext
+from google.adk.agents.llm_agent import Agent
+from google.adk.models.llm_response import LlmResponse
+from google.adk.sessions.in_memory_session_service import InMemorySessionService
+from google.adk.tools.google_search_agent_tool import GoogleSearchAgentTool
+from google.adk.tools.tool_context import ToolContext
+from google.genai import types
+from google.genai.types import Part
+from pytest import mark
+
+from .. import testing_utils
+
+function_call_no_schema = Part.from_function_call(
+    name='tool_agent', args={'request': 'test1'}
+)
+
+grounding_metadata = types.GroundingMetadata(web_search_queries=['test query'])
+
+
+# TODO(b/448114567): Remove test_grounding_metadata_ tests once the workaround
+# is no longer needed.
+
+
+@mark.asyncio
+async def test_grounding_metadata_is_stored_in_state_during_invocation():
+  """Verify grounding_metadata is stored in the state during invocation."""
+
+  # Mock model for the tool_agent that returns grounding_metadata
+  tool_agent_model = testing_utils.MockModel.create(
+      responses=[
+          LlmResponse(
+              content=types.Content(
+                  parts=[Part.from_text(text='response from tool')]
+              ),
+              grounding_metadata=grounding_metadata,
+          )
+      ]
+  )
+
+  tool_agent = Agent(
+      name='tool_agent',
+      model=tool_agent_model,
+  )
+
+  agent_tool = GoogleSearchAgentTool(agent=tool_agent)
+
+  session_service = InMemorySessionService()
+  session = await session_service.create_session(
+      app_name='test_app', user_id='test_user'
+  )
+
+  invocation_context = InvocationContext(
+      invocation_id='invocation_id',
+      agent=tool_agent,
+      session=session,
+      session_service=session_service,
+  )
+  tool_context = ToolContext(invocation_context=invocation_context)
+  tool_result = await agent_tool.run_async(
+      args=function_call_no_schema.function_call.args, tool_context=tool_context
+  )
+
+  # Verify the tool result
+  assert tool_result == 'response from tool'
+
+  # Verify grounding_metadata is stored in the state
+  assert tool_context.state['temp:_adk_grounding_metadata'] == (
+      grounding_metadata
+  )
+
+
+@mark.asyncio
+async def test_grounding_metadata_is_not_stored_in_state_after_invocation():
+  """Verify grounding_metadata is not stored in the state after invocation."""
+
+  # Mock model for the tool_agent that returns grounding_metadata
+  tool_agent_model = testing_utils.MockModel.create(
+      responses=[
+          LlmResponse(
+              content=types.Content(
+                  parts=[Part.from_text(text='response from tool')]
+              ),
+              grounding_metadata=grounding_metadata,
+          )
+      ]
+  )
+
+  tool_agent = Agent(
+      name='tool_agent',
+      model=tool_agent_model,
+  )
+
+  # Mock model for the root_agent
+  root_agent_model = testing_utils.MockModel.create(
+      responses=[
+          function_call_no_schema,  # Call the tool_agent
+          'Final response from root',
+      ]
+  )
+
+  root_agent = Agent(
+      name='root_agent',
+      model=root_agent_model,
+      tools=[GoogleSearchAgentTool(agent=tool_agent)],
+  )
+
+  runner = testing_utils.InMemoryRunner(root_agent)
+  events = runner.run('test input')
+
+  # Find the function response event
+  function_response_event = None
+  for event in events:
+    if event.get_function_responses():
+      function_response_event = event
+      break
+
+  # Verify the function response
+  assert function_response_event is not None
+  function_responses = function_response_event.get_function_responses()
+  assert len(function_responses) == 1
+  tool_output = function_responses[0].response
+  assert tool_output == {'result': 'response from tool'}
+
+  # Verify grounding_metadata is not stored in the root_agent's state
+  assert 'temp:_adk_grounding_metadata' not in runner.session.state