diff --git a/contributing/samples/agent_tool_with_grounding_metadata/__init__.py b/contributing/samples/agent_tool_with_grounding_metadata/__init__.py new file mode 100644 index 00000000..4015e47d --- /dev/null +++ b/contributing/samples/agent_tool_with_grounding_metadata/__init__.py @@ -0,0 +1,15 @@ +# Copyright 2026 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 . import agent diff --git a/contributing/samples/agent_tool_with_grounding_metadata/agent.py b/contributing/samples/agent_tool_with_grounding_metadata/agent.py new file mode 100644 index 00000000..521038d8 --- /dev/null +++ b/contributing/samples/agent_tool_with_grounding_metadata/agent.py @@ -0,0 +1,74 @@ +# Copyright 2026 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. + +import os +import random + +from dotenv import load_dotenv +from google.adk import Agent +from google.adk.tools.agent_tool import AgentTool +from google.adk.tools.tool_context import ToolContext +from google.adk.tools.vertex_ai_search_tool import VertexAiSearchTool + +load_dotenv(override=True) + +VERTEXAI_DATASTORE_ID = os.getenv("VERTEXAI_DATASTORE_ID") +if not VERTEXAI_DATASTORE_ID: + raise ValueError("VERTEXAI_DATASTORE_ID environment variable not set") + + +def roll_die(sides: int, tool_context: ToolContext) -> int: + """Roll a die and return the rolled result. + + Args: + sides: The integer number of sides the die has. + + Returns: + An integer of the result of rolling the die. + """ + result = random.randint(1, sides) + if "rolls" not in tool_context.state: + tool_context.state["rolls"] = [] + + tool_context.state["rolls"] = tool_context.state["rolls"] + [result] + return result + + +vertex_ai_search_agent = Agent( + model="gemini-3-flash-preview", + name="vertex_ai_search_agent", + description="An agent for performing Vertex AI search.", + tools=[ + VertexAiSearchTool( + data_store_id=VERTEXAI_DATASTORE_ID, + ) + ], +) + +root_agent = Agent( + model="gemini-3.1-pro-preview", + name="hello_world_agent", + description="A hello world agent with multiple tools.", + instruction=""" + You are a helpful assistant which can help user to roll dice and search for information. + - Use `roll_die` tool to roll dice. + - Use `vertex_ai_search_agent` to search for Google Agent Development Kit (ADK) information in the datastore. + """, + tools=[ + roll_die, + AgentTool( + agent=vertex_ai_search_agent, propagate_grounding_metadata=True + ), + ], +) 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 bd0037bd..0f2d3600 100644 --- a/src/google/adk/flows/llm_flows/base_llm_flow.py +++ b/src/google/adk/flows/llm_flows/base_llm_flow.py @@ -259,7 +259,9 @@ async def _handle_after_model_callback( tools = await agent.canonical_tools(readonly_context) invocation_context.canonical_tools_cache = tools - if not any(tool.name == 'google_search_agent' for tool in tools): + if not any( + getattr(tool, 'propagate_grounding_metadata', False) for tool in tools + ): return response ground_metadata = invocation_context.session.state.get( 'temp:_adk_grounding_metadata', None diff --git a/src/google/adk/tools/agent_tool.py b/src/google/adk/tools/agent_tool.py index f53c18df..8a43eb63 100644 --- a/src/google/adk/tools/agent_tool.py +++ b/src/google/adk/tools/agent_tool.py @@ -113,10 +113,12 @@ class AgentTool(BaseTool): skip_summarization: bool = False, *, include_plugins: bool = True, + propagate_grounding_metadata: bool = False, ): self.agent = agent self.skip_summarization: bool = skip_summarization self.include_plugins = include_plugins + self.propagate_grounding_metadata = propagate_grounding_metadata super().__init__(name=agent.name, description=agent.description) @@ -247,6 +249,7 @@ class AgentTool(BaseTool): ) 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 @@ -258,6 +261,7 @@ class AgentTool(BaseTool): tool_context.state.update(event.actions.state_delta) if event.content: last_content = event.content + last_grounding_metadata = event.grounding_metadata # Clean up runner resources (especially MCP sessions) # to avoid "Attempted to exit cancel scope in a different task" errors @@ -273,6 +277,12 @@ class AgentTool(BaseTool): tool_result = validate_schema(output_schema, merged_text) else: tool_result = merged_text + + if self.propagate_grounding_metadata and last_grounding_metadata: + tool_context.state['temp:_adk_grounding_metadata'] = ( + last_grounding_metadata + ) + return tool_result @override diff --git a/src/google/adk/tools/google_search_agent_tool.py b/src/google/adk/tools/google_search_agent_tool.py index 7ed09c79..f21e915c 100644 --- a/src/google/adk/tools/google_search_agent_tool.py +++ b/src/google/adk/tools/google_search_agent_tool.py @@ -14,21 +14,12 @@ 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._schema_utils import validate_schema -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: @@ -60,80 +51,4 @@ class GoogleSearchAgentTool(AgentTool): 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 last_content is None or last_content.parts is None: - 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 = validate_schema(self.agent.output_schema, merged_text) - else: - tool_result = merged_text - - if last_grounding_metadata: - tool_context.state['temp:_adk_grounding_metadata'] = ( - last_grounding_metadata - ) - return tool_result + super().__init__(agent=self.agent, propagate_grounding_metadata=True) 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 3dfadbca..9578a3f5 100644 --- a/tests/unittests/flows/llm_flows/test_base_llm_flow.py +++ b/tests/unittests/flows/llm_flows/test_base_llm_flow.py @@ -284,7 +284,6 @@ async def test_handle_after_model_callback_grounding_with_no_callbacks( invocation_id=invocation_context.invocation_id, author=agent.name, ) - flow = BaseLlmFlowForTesting() result = await _handle_after_model_callback( invocation_context, llm_response, event @@ -341,7 +340,6 @@ async def test_handle_after_model_callback_grounding_with_callback_override( invocation_id=invocation_context.invocation_id, author=agent.name, ) - flow = BaseLlmFlowForTesting() result = await _handle_after_model_callback( invocation_context, llm_response, event @@ -403,7 +401,6 @@ async def test_handle_after_model_callback_grounding_with_plugin_override( invocation_id=invocation_context.invocation_id, author=agent.name, ) - flow = BaseLlmFlowForTesting() result = await _handle_after_model_callback( invocation_context, llm_response, event @@ -430,6 +427,7 @@ async def test_handle_after_model_callback_caches_canonical_tools(): def __init__(self): super().__init__(name='google_search_agent', description='Mock search') + self.propagate_grounding_metadata = True async def call(self, **kwargs): return 'mock result' @@ -459,7 +457,6 @@ async def test_handle_after_model_callback_caches_canonical_tools(): invocation_id=invocation_context.invocation_id, author=agent.name, ) - flow = BaseLlmFlowForTesting() # Call _handle_after_model_callback multiple times with the same context result1 = await _handle_after_model_callback(