Skip to content

Commit ab62b1b

Browse files
seanzhougooglecopybara-github
authored andcommitted
fix: Use the agent name as the author of the audio event
Co-authored-by: Xiang (Sean) Zhou <seanzhougoogle@google.com> PiperOrigin-RevId: 855789317
1 parent f668a5d commit ab62b1b

File tree

2 files changed

+58
-1
lines changed

2 files changed

+58
-1
lines changed

src/google/adk/flows/llm_flows/audio_cache_manager.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -181,10 +181,16 @@ async def _flush_cache_to_services(
181181
artifact_ref = f'artifact://{invocation_context.app_name}/{invocation_context.user_id}/{invocation_context.session.id}/_adk_live/{filename}#{revision_id}'
182182

183183
# Create event with file data reference to add to session
184+
# For model events, author should be the agent name, not the role
185+
author = (
186+
invocation_context.agent.name
187+
if audio_cache[0].role == 'model'
188+
else audio_cache[0].role
189+
)
184190
audio_event = Event(
185191
id=Event.new_id(),
186192
invocation_id=invocation_context.invocation_id,
187-
author=audio_cache[0].role,
193+
author=author,
188194
content=types.Content(
189195
role=audio_cache[0].role,
190196
parts=[

tests/unittests/flows/llm_flows/test_audio_cache_manager.py

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -387,3 +387,54 @@ async def test_filename_uses_first_chunk_timestamp(self):
387387
assert filename.startswith(
388388
f'adk_live_audio_storage_input_audio_{expected_timestamp_ms}'
389389
)
390+
391+
@pytest.mark.asyncio
392+
async def test_flush_event_author_for_user_audio(self):
393+
"""Test that flushed user audio events have 'user' as author."""
394+
invocation_context = await testing_utils.create_invocation_context(
395+
testing_utils.create_test_agent()
396+
)
397+
398+
# Set up mock artifact service
399+
mock_artifact_service = AsyncMock()
400+
mock_artifact_service.save_artifact.return_value = 123
401+
invocation_context.artifact_service = mock_artifact_service
402+
403+
# Cache user input audio
404+
input_blob = types.Blob(data=b'user_audio_data', mime_type='audio/pcm')
405+
self.manager.cache_audio(invocation_context, input_blob, 'input')
406+
407+
# Flush cache and get events
408+
events = await self.manager.flush_caches(
409+
invocation_context, flush_user_audio=True, flush_model_audio=False
410+
)
411+
412+
# Verify event author is 'user' for user audio
413+
assert len(events) == 1
414+
assert events[0].author == 'user'
415+
assert events[0].content.role == 'user'
416+
417+
@pytest.mark.asyncio
418+
async def test_flush_event_author_for_model_audio(self):
419+
"""Test that flushed model audio events have agent name as author, not 'model'."""
420+
agent = testing_utils.create_test_agent(name='my_test_agent')
421+
invocation_context = await testing_utils.create_invocation_context(agent)
422+
423+
# Set up mock artifact service
424+
mock_artifact_service = AsyncMock()
425+
mock_artifact_service.save_artifact.return_value = 123
426+
invocation_context.artifact_service = mock_artifact_service
427+
428+
# Cache model output audio
429+
output_blob = types.Blob(data=b'model_audio_data', mime_type='audio/wav')
430+
self.manager.cache_audio(invocation_context, output_blob, 'output')
431+
432+
# Flush cache and get events
433+
events = await self.manager.flush_caches(
434+
invocation_context, flush_user_audio=False, flush_model_audio=True
435+
)
436+
437+
# Verify event author is agent name (not 'model') for model audio
438+
assert len(events) == 1
439+
assert events[0].author == 'my_test_agent' # Agent name, not 'model'
440+
assert events[0].content.role == 'model' # Role is still 'model'

0 commit comments

Comments
 (0)