|
19 | 19 | from google.adk.agents.base_agent import BaseAgent |
20 | 20 | from google.adk.agents.context import Context |
21 | 21 | from google.adk.agents.invocation_context import InvocationContext as BaseInvocationContext |
| 22 | +from google.adk.agents.llm_agent import LlmAgent |
| 23 | +from google.adk.agents.parallel_agent import ParallelAgent |
| 24 | +from google.adk.agents.sequential_agent import SequentialAgent |
22 | 25 | from google.adk.events.event import Event |
23 | 26 | from google.adk.sessions.in_memory_session_service import InMemorySessionService |
24 | 27 | from google.adk.sessions.session import Session |
|
27 | 30 | from google.genai import types |
28 | 31 | import pytest |
29 | 32 |
|
| 33 | +from tests.unittests import testing_utils |
| 34 | + |
30 | 35 | from .workflow_testing_utils import create_runner |
31 | 36 | from .workflow_testing_utils import InputCapturingNode |
32 | 37 | from .workflow_testing_utils import simplify_events_with_node |
@@ -165,3 +170,116 @@ async def test_run_async_with_agent_node_piping_data( |
165 | 170 | # AgentNode does not record content as node output, so the next node |
166 | 171 | # receives None as input. |
167 | 172 | assert node_b.received_inputs == [None] |
| 173 | + |
| 174 | + |
| 175 | +async def _run_workflow(wf: Workflow) -> list[Event]: |
| 176 | + ss = InMemorySessionService() |
| 177 | + runner = create_runner(app_name='test', node=wf, session_service=ss) |
| 178 | + session = await ss.create_session(app_name='test', user_id='u') |
| 179 | + msg = types.Content(parts=[types.Part(text='start')], role='user') |
| 180 | + return [ |
| 181 | + event |
| 182 | + async for event in runner.run_async( |
| 183 | + user_id='u', session_id=session.id, new_message=msg |
| 184 | + ) |
| 185 | + ] |
| 186 | + |
| 187 | + |
| 188 | +@pytest.mark.asyncio |
| 189 | +async def test_node_after_parallel_agent_node_sees_every_branch(): |
| 190 | + """A node after a ParallelAgent node runs on the workflow's own branch. |
| 191 | +
|
| 192 | + The ParallelAgent relays each sub-agent's events on that sub-agent's |
| 193 | + branch. Its successor must not inherit whichever of those branches spoke |
| 194 | + last: an agent there reads history filtered by branch, so it would see only |
| 195 | + that one sub-agent's answer, and its own replies would land on that |
| 196 | + sub-agent's branch. |
| 197 | + """ |
| 198 | + researcher_a = LlmAgent( |
| 199 | + name='researcher_a', |
| 200 | + model=testing_utils.MockModel.create(responses=['finding from a']), |
| 201 | + ) |
| 202 | + researcher_b = LlmAgent( |
| 203 | + name='researcher_b', |
| 204 | + model=testing_utils.MockModel.create(responses=['finding from b']), |
| 205 | + ) |
| 206 | + research = ParallelAgent( |
| 207 | + name='research', sub_agents=[researcher_a, researcher_b] |
| 208 | + ) |
| 209 | + summarizer_model = testing_utils.MockModel.create(responses=['summary']) |
| 210 | + report = SequentialAgent( |
| 211 | + name='report', |
| 212 | + sub_agents=[LlmAgent(name='summarizer', model=summarizer_model)], |
| 213 | + ) |
| 214 | + |
| 215 | + events = await _run_workflow( |
| 216 | + Workflow(name='wf', edges=[(START, research, report)]) |
| 217 | + ) |
| 218 | + |
| 219 | + summarizer_request = ' '.join( |
| 220 | + part.text |
| 221 | + for content in summarizer_model.requests[0].contents |
| 222 | + for part in content.parts or [] |
| 223 | + if part.text |
| 224 | + ) |
| 225 | + assert 'finding from a' in summarizer_request |
| 226 | + assert 'finding from b' in summarizer_request |
| 227 | + assert {event.author: event.branch for event in events if event.content} == { |
| 228 | + 'researcher_a': 'research.researcher_a', |
| 229 | + 'researcher_b': 'research.researcher_b', |
| 230 | + 'summarizer': None, |
| 231 | + } |
| 232 | + |
| 233 | + |
| 234 | +class _Dispatcher(BaseAgent): |
| 235 | + """Runs its sub-agent on a sub-branch and hands back the final answer.""" |
| 236 | + |
| 237 | + async def _run_async_impl( |
| 238 | + self, ctx: BaseInvocationContext |
| 239 | + ) -> AsyncGenerator[Event, None]: |
| 240 | + target = self.sub_agents[0] |
| 241 | + parent_branch = ctx.branch |
| 242 | + ctx.branch = f'{self.name}.{target.name}' |
| 243 | + async for event in target.run_async(ctx): |
| 244 | + if event.is_final_response(): |
| 245 | + event.branch = parent_branch |
| 246 | + yield event |
| 247 | + |
| 248 | + |
| 249 | +@pytest.mark.asyncio |
| 250 | +async def test_answer_relayed_onto_parent_branch_stays_there(): |
| 251 | + """A relayed answer re-stamped onto the parent branch is not pulled back. |
| 252 | +
|
| 253 | + The dispatcher keeps its sub-agent's tool traffic on a sub-branch and hands |
| 254 | + the final answer back on its own branch, so later agents see the answer |
| 255 | + but not the tool calls. The tool events must not move the dispatcher node |
| 256 | + onto the sub-branch, or the branch-less answer gets stamped with it. |
| 257 | + """ |
| 258 | + |
| 259 | + def lookup_arrival() -> dict[str, str]: |
| 260 | + """Looks up the arrival time.""" |
| 261 | + return {'arrival': 'Tuesday 14:05'} |
| 262 | + |
| 263 | + bookings = LlmAgent( |
| 264 | + name='bookings', |
| 265 | + model=testing_utils.MockModel.create( |
| 266 | + responses=[ |
| 267 | + types.Part.from_function_call(name='lookup_arrival', args={}), |
| 268 | + 'You land on Tuesday at 14:05.', |
| 269 | + ] |
| 270 | + ), |
| 271 | + tools=[lookup_arrival], |
| 272 | + ) |
| 273 | + dispatcher = _Dispatcher(name='dispatcher', sub_agents=[bookings]) |
| 274 | + |
| 275 | + events = await _run_workflow(Workflow(name='wf', edges=[(START, dispatcher)])) |
| 276 | + |
| 277 | + bookings_events = [e for e in events if e.author == 'bookings' and e.content] |
| 278 | + assert [e.branch for e in bookings_events] == [ |
| 279 | + 'dispatcher.bookings', |
| 280 | + 'dispatcher.bookings', |
| 281 | + None, |
| 282 | + ] |
| 283 | + assert bookings_events[-1].content.parts[0].text == ( |
| 284 | + 'You land on Tuesday at 14:05.' |
| 285 | + ) |
0 commit comments