Skip to content

Commit 3741b74

Browse files
feat(tools): add opt-in flag to auto-attach load_artifacts to AgentTool sub-agents
Artifacts saved by a parent agent are forwarded to sub-agents via ForwardingArtifactService, but the bytes are only injected into an LLM agent's request when it explicitly calls the load_artifacts tool. This is easy to miss, especially when AgentTool wraps a composite agent like ParallelAgent with several sub-agents. Adds include_load_artifacts_tool (default False) to AgentTool that recursively attaches load_artifacts_tool to the wrapped agent and all of its sub-agents. Closes #3232
1 parent a238884 commit 3741b74

2 files changed

Lines changed: 127 additions & 0 deletions

File tree

src/google/adk/tools/agent_tool.py

Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,27 @@ def _get_input_schema(agent: BaseAgent) -> Optional[type[BaseModel]]:
7979
return None
8080

8181

82+
def _ensure_load_artifacts_tool(agent: BaseAgent) -> None:
83+
"""Recursively attaches `load_artifacts_tool` to `agent` and its sub-agents.
84+
85+
Only `LlmAgent` instances can hold tools, so non-LLM agents (e.g.
86+
`ParallelAgent`, `SequentialAgent`) are skipped but their sub-agents are
87+
still visited. This is a no-op if the tool is already present, so it is
88+
safe to call on every `AgentTool.run_async()`.
89+
"""
90+
from ..agents.llm_agent import LlmAgent
91+
from .load_artifacts_tool import load_artifacts_tool
92+
93+
if isinstance(agent, LlmAgent) and not any(
94+
getattr(tool, 'name', None) == load_artifacts_tool.name
95+
for tool in agent.tools
96+
):
97+
agent.tools.append(load_artifacts_tool)
98+
99+
for sub_agent in agent.sub_agents:
100+
_ensure_load_artifacts_tool(sub_agent)
101+
102+
82103
def _get_output_schema(agent: BaseAgent) -> Optional[SchemaType]:
83104
"""Extracts the output_schema from an agent.
84105
@@ -118,6 +139,14 @@ class AgentTool(BaseTool):
118139
to the agent's runner. When True (default), the agent will inherit all
119140
plugins from its parent. Set to False to run the agent with an isolated
120141
plugin environment.
142+
include_load_artifacts_tool: Whether to automatically attach the
143+
`load_artifacts` tool to the wrapped agent and all of its sub-agents
144+
(recursively). Artifacts saved by the parent agent are always
145+
forwarded to sub-agents via `ForwardingArtifactService`, but the
146+
bytes are only injected into an LLM agent's request when it calls
147+
`load_artifacts`. Defaults to False to avoid sending large payloads
148+
on every turn; set to True so sub-agents can see artifacts without
149+
having to add the tool to each agent manually.
121150
"""
122151

123152
def __init__(
@@ -127,11 +156,13 @@ def __init__(
127156
*,
128157
include_plugins: bool = True,
129158
propagate_grounding_metadata: bool = False,
159+
include_load_artifacts_tool: bool = False,
130160
):
131161
self.agent = agent
132162
self.skip_summarization: bool = skip_summarization
133163
self.include_plugins = include_plugins
134164
self.propagate_grounding_metadata = propagate_grounding_metadata
165+
self.include_load_artifacts_tool = include_load_artifacts_tool
135166

136167
super().__init__(name=agent.name, description=agent.description)
137168

@@ -214,6 +245,9 @@ async def run_async(
214245
if self.skip_summarization:
215246
tool_context.actions.skip_summarization = True
216247

248+
if self.include_load_artifacts_tool:
249+
_ensure_load_artifacts_tool(self.agent)
250+
217251
input_schema = _get_input_schema(self.agent)
218252
if input_schema:
219253
input_value = input_schema.model_validate(args)

tests/unittests/tools/test_agent_tool.py

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
from google.adk.agents.callback_context import CallbackContext
2222
from google.adk.agents.invocation_context import InvocationContext
2323
from google.adk.agents.llm_agent import Agent
24+
from google.adk.agents.parallel_agent import ParallelAgent
2425
from google.adk.agents.run_config import RunConfig
2526
from google.adk.agents.sequential_agent import SequentialAgent
2627
from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService
@@ -35,6 +36,7 @@
3536
from google.adk.runners import Runner
3637
from google.adk.sessions.in_memory_session_service import InMemorySessionService
3738
from google.adk.tools.agent_tool import AgentTool
39+
from google.adk.tools.load_artifacts_tool import load_artifacts_tool
3840
from google.adk.tools.tool_context import ToolContext
3941
from google.adk.utils.variant_utils import GoogleLLMVariant
4042
from google.genai import types
@@ -1148,6 +1150,97 @@ async def test_run_async_skips_thought_parts():
11481150
assert result == '42'
11491151

11501152

1153+
def test_include_load_artifacts_tool_default_false():
1154+
"""By default, load_artifacts is not added to the wrapped agent."""
1155+
mock_model = testing_utils.MockModel.create(
1156+
responses=[function_call_no_schema, 'response1', 'response2']
1157+
)
1158+
tool_agent = Agent(name='tool_agent', model=mock_model)
1159+
root_agent = Agent(
1160+
name='root_agent',
1161+
model=mock_model,
1162+
tools=[AgentTool(agent=tool_agent)],
1163+
)
1164+
1165+
runner = testing_utils.InMemoryRunner(root_agent)
1166+
runner.run('test1')
1167+
1168+
assert all(tool.name != 'load_artifacts' for tool in tool_agent.tools)
1169+
1170+
1171+
def test_include_load_artifacts_tool_true_adds_to_wrapped_agent():
1172+
"""When enabled, load_artifacts is attached to the wrapped LlmAgent."""
1173+
mock_model = testing_utils.MockModel.create(
1174+
responses=[function_call_no_schema, 'response1', 'response2']
1175+
)
1176+
tool_agent = Agent(name='tool_agent', model=mock_model)
1177+
root_agent = Agent(
1178+
name='root_agent',
1179+
model=mock_model,
1180+
tools=[AgentTool(agent=tool_agent, include_load_artifacts_tool=True)],
1181+
)
1182+
1183+
runner = testing_utils.InMemoryRunner(root_agent)
1184+
runner.run('test1')
1185+
1186+
assert any(tool.name == 'load_artifacts' for tool in tool_agent.tools)
1187+
1188+
1189+
def test_include_load_artifacts_tool_true_adds_to_sub_agents_recursively():
1190+
"""When enabled, load_artifacts is attached to sub-agents of a composite
1191+
wrapped agent (e.g. ParallelAgent), not just the top-level agent.
1192+
"""
1193+
sub_agent_1 = Agent(
1194+
name='sub_agent_1',
1195+
model=testing_utils.MockModel.create(responses=['sub_response_1']),
1196+
)
1197+
sub_agent_2 = Agent(
1198+
name='sub_agent_2',
1199+
model=testing_utils.MockModel.create(responses=['sub_response_2']),
1200+
)
1201+
parallel_agent = ParallelAgent(
1202+
name='parallel_tool_agent', sub_agents=[sub_agent_1, sub_agent_2]
1203+
)
1204+
1205+
function_call_for_parallel = Part.from_function_call(
1206+
name='parallel_tool_agent', args={'request': 'test1'}
1207+
)
1208+
mock_model_root = testing_utils.MockModel.create(
1209+
responses=[function_call_for_parallel, 'response2']
1210+
)
1211+
root_agent = Agent(
1212+
name='root_agent',
1213+
model=mock_model_root,
1214+
tools=[AgentTool(agent=parallel_agent, include_load_artifacts_tool=True)],
1215+
)
1216+
1217+
runner = testing_utils.InMemoryRunner(root_agent)
1218+
runner.run('test1')
1219+
1220+
assert any(tool.name == 'load_artifacts' for tool in sub_agent_1.tools)
1221+
assert any(tool.name == 'load_artifacts' for tool in sub_agent_2.tools)
1222+
1223+
1224+
def test_include_load_artifacts_tool_does_not_duplicate_existing_tool():
1225+
"""If the wrapped agent already has load_artifacts, it is not duplicated."""
1226+
mock_model = testing_utils.MockModel.create(
1227+
responses=[function_call_no_schema, 'response1', 'response2']
1228+
)
1229+
tool_agent = Agent(
1230+
name='tool_agent', model=mock_model, tools=[load_artifacts_tool]
1231+
)
1232+
root_agent = Agent(
1233+
name='root_agent',
1234+
model=mock_model,
1235+
tools=[AgentTool(agent=tool_agent, include_load_artifacts_tool=True)],
1236+
)
1237+
1238+
runner = testing_utils.InMemoryRunner(root_agent)
1239+
runner.run('test1')
1240+
1241+
assert sum(tool.name == 'load_artifacts' for tool in tool_agent.tools) == 1
1242+
1243+
11511244
class TestAgentToolWithCompositeAgents:
11521245
"""Tests for AgentTool wrapping composite agents (SequentialAgent, etc.)."""
11531246

0 commit comments

Comments
 (0)