Skip to content

Commit b3d0759

Browse files
wzhang2705GWeale
authored andcommitted
feat: Preserve transcription event order in conversation trajectory
Refactor live inference to emit synthetic text events for output transcriptions as they arrive (streaming), rather than accumulating all transcriptions and appending a single synthetic event after the turn completes. This preserves event ordering and enables downstream consumers to process transcriptions incrementally. Co-authored-by: Wenhua Zhang <wenhzhang@google.com> Change-Id: Idca79968729f2145ebf1f1c993d79ab7f6a3bd78
1 parent cbd14eb commit b3d0759

2 files changed

Lines changed: 74 additions & 22 deletions

File tree

src/google/adk/evaluation/evaluation_generator.py

Lines changed: 18 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -383,6 +383,7 @@ async def _generate_inferences_for_single_user_invocation_live(
383383
current_invocation_id: str,
384384
turn_complete_event: asyncio.Event,
385385
live_timeout_seconds: int,
386+
agent_name: str = _DEFAULT_AUTHOR,
386387
) -> AsyncGenerator[Event, None]:
387388
"""Generates inferences for a single user invocation in live mode."""
388389
yield Event(
@@ -408,6 +409,22 @@ async def _generate_inferences_for_single_user_invocation_live(
408409
event = await event_queue.get()
409410
if event.invocation_id == current_invocation_id:
410411
yield event
412+
# Emit a synthetic text event for each transcription, preserving
413+
# the order in which events are received.
414+
if (
415+
event.author != _USER_AUTHOR
416+
and event.output_transcription
417+
and event.output_transcription.text
418+
and event.partial
419+
):
420+
yield Event(
421+
content=Content(
422+
role="model",
423+
parts=[types.Part(text=event.output_transcription.text)],
424+
),
425+
author=agent_name,
426+
invocation_id=current_invocation_id,
427+
)
411428

412429
@staticmethod
413430
async def _generate_inferences_from_root_agent_live(
@@ -495,31 +512,10 @@ async def _generate_inferences_from_root_agent_live(
495512
current_invocation_id=live_session.current_invocation_id,
496513
turn_complete_event=live_session.turn_complete_event,
497514
live_timeout_seconds=live_timeout_seconds,
515+
agent_name=runner.agent.name,
498516
):
499517
events.append(event)
500518

501-
turn_transcription = ""
502-
for evt in events:
503-
if (
504-
evt.invocation_id == live_session.current_invocation_id
505-
and evt.author != _USER_AUTHOR
506-
and evt.output_transcription
507-
):
508-
if not evt.partial and evt.output_transcription.text:
509-
turn_transcription = evt.output_transcription.text
510-
else:
511-
turn_transcription += evt.output_transcription.text
512-
if turn_transcription:
513-
synthetic_event = Event(
514-
content=Content(
515-
role="model",
516-
parts=[types.Part(text=turn_transcription)],
517-
),
518-
author=runner.agent.name,
519-
invocation_id=live_session.current_invocation_id,
520-
)
521-
events.append(synthetic_event)
522-
523519
if live_session.live_finished.is_set():
524520
logger.info("Live session finished signal detected.")
525521
break

tests/unittests/evaluation/test_evaluation_generator.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -453,6 +453,62 @@ async def test_generate_inferences_live(self, mocker):
453453
with pytest.raises(StopAsyncIteration):
454454
await gen.__anext__()
455455

456+
@pytest.mark.asyncio
457+
async def test_generate_inferences_live_with_synthetic_events(self, mocker):
458+
"""Tests live inference generation with synthetic events."""
459+
mock_live_request_queue = mocker.MagicMock()
460+
event_queue = asyncio.Queue()
461+
turn_complete_event = asyncio.Event()
462+
463+
user_content = types.Content(parts=[types.Part(text="User query")])
464+
invocation_id = "inv1"
465+
466+
transcription = types.Transcription(text="Partial transcription")
467+
partial_event = Event(
468+
author="agent",
469+
content=types.Content(parts=[]),
470+
invocation_id=invocation_id,
471+
output_transcription=transcription,
472+
partial=True,
473+
)
474+
475+
gen = EvaluationGenerator._generate_inferences_for_single_user_invocation_live(
476+
live_request_queue=mock_live_request_queue,
477+
event_queue=event_queue,
478+
user_message=user_content,
479+
current_invocation_id=invocation_id,
480+
turn_complete_event=turn_complete_event,
481+
live_timeout_seconds=300,
482+
agent_name="custom_agent_name",
483+
)
484+
485+
# First yield should be the user message
486+
first_event = await gen.__anext__()
487+
assert first_event.author == "user"
488+
assert first_event.content == user_content
489+
assert first_event.invocation_id == invocation_id
490+
491+
# Mock turn_complete_event.wait to avoid blocking
492+
turn_complete_event.wait = mocker.AsyncMock()
493+
494+
# Put the partial event in the queue
495+
await event_queue.put(partial_event)
496+
497+
# Now advance
498+
second_event = await gen.__anext__()
499+
assert second_event == partial_event
500+
501+
# Next should be the synthetic event
502+
third_event = await gen.__anext__()
503+
assert third_event.author == "custom_agent_name"
504+
assert third_event.invocation_id == invocation_id
505+
assert third_event.content.role == "model"
506+
assert third_event.content.parts[0].text == "Partial transcription"
507+
508+
# The generator should be exhausted now
509+
with pytest.raises(StopAsyncIteration):
510+
await gen.__anext__()
511+
456512

457513
@pytest.fixture
458514
def mock_runner(mocker):

0 commit comments

Comments
 (0)