Skip to content

Commit a43b905

Browse files
committed
feat(agents): annotate events with the LoopAgent iteration
LoopAgent re-runs its sub-agents each iteration, but events from different iterations are indistinguishable in the stream (same author, branch and invocation_id), so consumers cannot tell iterations apart — e.g. a ParallelAgent nested in a LoopAgent, whose branch ids repeat every iteration. Record the current 0-based iteration under Event.custom_metadata['loop_iteration'] as each sub-agent event bubbles up. Additive and behavior-preserving; setdefault means the nearest enclosing loop wins when loops are nested. Fixes #6266.
1 parent b2dda6e commit a43b905

2 files changed

Lines changed: 77 additions & 0 deletions

File tree

src/google/adk/agents/loop_agent.py

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -79,6 +79,27 @@ class LoopAgent(BaseAgent):
7979
escalates.
8080
"""
8181

82+
LOOP_ITERATION_KEY: ClassVar[str] = 'loop_iteration'
83+
"""Key under ``Event.custom_metadata`` recording the 0-based iteration of the
84+
nearest enclosing ``LoopAgent`` during which the event was produced.
85+
86+
Without it, events produced in different iterations are indistinguishable
87+
(same ``author``, ``branch`` and ``invocation_id``), so consumers cannot tell
88+
loop iterations apart — e.g. a ``ParallelAgent`` nested in a ``LoopAgent``,
89+
whose branch ids repeat every iteration. See
90+
https://github.com/google/adk-python/issues/6266.
91+
"""
92+
93+
def _annotate_loop_iteration(self, event: Event, iteration: int) -> None:
94+
"""Record the current loop iteration on ``event.custom_metadata``.
95+
96+
Uses ``setdefault`` so that when ``LoopAgent``s are nested the nearest
97+
enclosing loop (which annotates first, as the event bubbles up) wins.
98+
"""
99+
metadata = dict(event.custom_metadata) if event.custom_metadata else {}
100+
metadata.setdefault(self.LOOP_ITERATION_KEY, iteration)
101+
event.custom_metadata = metadata
102+
82103
@override
83104
async def _run_async_impl(
84105
self, ctx: InvocationContext
@@ -112,6 +133,7 @@ async def _run_async_impl(
112133

113134
async with Aclosing(sub_agent.run_async(ctx)) as agen:
114135
async for event in agen:
136+
self._annotate_loop_iteration(event, times_looped)
115137
yield event
116138
if event.actions.escalate:
117139
should_exit = True

tests/unittests/agents/test_loop_agent.py

Lines changed: 55 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -139,6 +139,61 @@ async def test_run_async(request: pytest.FixtureRequest, resumable: bool):
139139
assert simplified_events == expected_events
140140

141141

142+
@pytest.mark.asyncio
143+
async def test_run_async_annotates_loop_iteration(
144+
request: pytest.FixtureRequest,
145+
):
146+
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
147+
loop_agent = LoopAgent(
148+
name=f'{request.function.__name__}_test_loop_agent',
149+
max_iterations=3,
150+
sub_agents=[agent],
151+
)
152+
parent_ctx = await _create_parent_invocation_context(
153+
request.function.__name__, loop_agent, resumable=False
154+
)
155+
156+
events = [e async for e in loop_agent.run_async(parent_ctx)]
157+
158+
sub_agent_events = [e for e in events if e.author == agent.name]
159+
assert len(sub_agent_events) == 3
160+
assert [
161+
e.custom_metadata[LoopAgent.LOOP_ITERATION_KEY]
162+
for e in sub_agent_events
163+
] == [0, 1, 2]
164+
165+
166+
@pytest.mark.asyncio
167+
async def test_run_async_loop_iteration_survives_nested_loops(
168+
request: pytest.FixtureRequest,
169+
):
170+
# An inner LoopAgent annotates first (as events bubble up), so its iteration
171+
# wins over the outer loop's for events it produced.
172+
agent = _TestingAgent(name=f'{request.function.__name__}_test_agent')
173+
inner = LoopAgent(
174+
name=f'{request.function.__name__}_inner_loop',
175+
max_iterations=2,
176+
sub_agents=[agent],
177+
)
178+
outer = LoopAgent(
179+
name=f'{request.function.__name__}_outer_loop',
180+
max_iterations=2,
181+
sub_agents=[inner],
182+
)
183+
parent_ctx = await _create_parent_invocation_context(
184+
request.function.__name__, outer, resumable=False
185+
)
186+
187+
events = [e async for e in outer.run_async(parent_ctx)]
188+
189+
sub_agent_events = [e for e in events if e.author == agent.name]
190+
# outer x2 * inner x2 = 4 events; each carries the inner loop's iteration.
191+
assert [
192+
e.custom_metadata[LoopAgent.LOOP_ITERATION_KEY]
193+
for e in sub_agent_events
194+
] == [0, 1, 0, 1]
195+
196+
142197
@pytest.mark.asyncio
143198
async def test_resume_async(request: pytest.FixtureRequest):
144199
agent_1 = _TestingAgent(name=f'{request.function.__name__}_test_agent_1')

0 commit comments

Comments
 (0)