|
21 | 21 | from google.adk.agents.callback_context import CallbackContext |
22 | 22 | from google.adk.agents.invocation_context import InvocationContext |
23 | 23 | from google.adk.agents.llm_agent import Agent |
| 24 | +from google.adk.agents.parallel_agent import ParallelAgent |
24 | 25 | from google.adk.agents.run_config import RunConfig |
25 | 26 | from google.adk.agents.sequential_agent import SequentialAgent |
26 | 27 | from google.adk.artifacts.in_memory_artifact_service import InMemoryArtifactService |
|
35 | 36 | from google.adk.runners import Runner |
36 | 37 | from google.adk.sessions.in_memory_session_service import InMemorySessionService |
37 | 38 | from google.adk.tools.agent_tool import AgentTool |
| 39 | +from google.adk.tools.load_artifacts_tool import load_artifacts_tool |
38 | 40 | from google.adk.tools.tool_context import ToolContext |
39 | 41 | from google.adk.utils.variant_utils import GoogleLLMVariant |
40 | 42 | from google.genai import types |
@@ -1148,6 +1150,97 @@ async def test_run_async_skips_thought_parts(): |
1148 | 1150 | assert result == '42' |
1149 | 1151 |
|
1150 | 1152 |
|
| 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 | + |
1151 | 1244 | class TestAgentToolWithCompositeAgents: |
1152 | 1245 | """Tests for AgentTool wrapping composite agents (SequentialAgent, etc.).""" |
1153 | 1246 |
|
|
0 commit comments