diff --git a/src/google/adk/agents/llm_agent.py b/src/google/adk/agents/llm_agent.py index 09837e00..9a964786 100644 --- a/src/google/adk/agents/llm_agent.py +++ b/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 diff --git a/src/google/adk/flows/llm_flows/base_llm_flow.py b/src/google/adk/flows/llm_flows/base_llm_flow.py index 7ff1b4e5..5c5c7ec2 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/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, diff --git a/src/google/adk/tools/google_search_agent_tool.py b/src/google/adk/tools/google_search_agent_tool.py new file mode 100644 index 00000000..77cb6fed --- /dev/null +++ b/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 diff --git a/tests/unittests/agents/test_llm_agent_fields.py b/tests/unittests/agents/test_llm_agent_fields.py index e62cf4e8..bac0aace 100644 --- a/tests/unittests/agents/test_llm_agent_fields.py +++ b/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' diff --git a/tests/unittests/flows/llm_flows/test_base_llm_flow.py b/tests/unittests/flows/llm_flows/test_base_llm_flow.py index 8ae88536..1b310f87 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/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() diff --git a/tests/unittests/tools/test_google_search_agent_tool.py b/tests/unittests/tools/test_google_search_agent_tool.py new file mode 100644 index 00000000..dc9d9604 --- /dev/null +++ b/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