From 058d9d1701f06cfed0fe804a5bb42136a3fbdb49 Mon Sep 17 00:00:00 2001 From: Douwe Maan Date: Fri, 9 Oct 2026 21:47:24 +0000 Subject: [PATCH] Keep a held Gemini filler open until the model carries on, and place inputs as they're sent when nothing is owed --- .../pydantic_ai/realtime/_google_lifecycle.py | 37 ++++++++++++------- tests/realtime/test_session_core.py | 32 ++++++++++++++++ 2 files changed, 55 insertions(+), 14 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/realtime/_google_lifecycle.py b/pydantic_ai_slim/pydantic_ai/realtime/_google_lifecycle.py index 1196fff8a6..7a17d88b1d 100644 --- a/pydantic_ai_slim/pydantic_ai/realtime/_google_lifecycle.py +++ b/pydantic_ai_slim/pydantic_ai/realtime/_google_lifecycle.py @@ -100,12 +100,13 @@ class GeminiLifecycle: self._audio_since_reply = True return if self._reply_owed(): + # It reaches the provider while the model works on a reply, which it follows. (A response that answers + # it places it ahead of itself, should that come first.) self._held.append(input_id) - elif not isinstance(content, (str, ToolResult)): + else: # Nothing is going on: it joins the conversation as it arrives. self._pending.append(InputAdded(input_id=input_id)) if isinstance(content, (str, ToolResult)): - # It asks for a reply: the response that answers it places it, ahead of that response. self._unanswered.append(input_id) def input_failed(self, input_id: InputId) -> None: @@ -137,16 +138,23 @@ class GeminiLifecycle: def _event(self, event: RealtimeCodecEvent) -> tuple[list[LifecycleEvent], list[LifecycleEvent]]: before: list[LifecycleEvent] = [] after: list[LifecycleEvent] = [] + output = isinstance(event, (AudioDelta, PartStartEvent, ToolCall)) or ( + isinstance(event, OutputTranscript) and bool(event.text) + ) if self._deferred is not None: - if isinstance(event, ToolCall): - # The tool call the model was stalling for: it belongs to the filler's response. + if isinstance(event, (ToolCall, ResponseDone)): + # The tool call the model was stalling for belongs to the filler's response, and so does the next + # boundary (handled below). self._deferred = None - else: - # The stall didn't end in the tool call it was for: the filler was a response of its own. - self._end_deferred(before) - if isinstance(event, (AudioDelta, PartStartEvent, ToolCall)) or ( - isinstance(event, OutputTranscript) and event.text - ): + elif output: + # The stall didn't end in the tool call it was for: the filler was a response of its own, which + # this one carries on. + self._end_deferred(before, continued=True) + elif isinstance(event, RealtimeResponseInterruptedEvent): + # The user cut in: the exchange is over. + self._end_deferred(before, continued=False) + # Anything else (the user's transcript, usage) leaves it held: the exchange isn't over. + if output: self._ensure_response(before) if isinstance(event, ToolCall) and event.response_usage_follows: self._calls_awaiting_usage = True @@ -277,13 +285,14 @@ class GeminiLifecycle: self._close_turn(events) self._place_held(events) - def _end_deferred(self, events: list[LifecycleEvent]) -> None: + def _end_deferred(self, events: list[LifecycleEvent], *, continued: bool) -> None: deferred, self._deferred = self._deferred, None assert deferred is not None - continued = self._open + held = self._open self._end(events, status='completed', finish_reason=deferred.finish_reason or 'stop') - # The exchange goes on past it, in whatever the model does next. - self._continues = continued + if continued: + # The exchange goes on past it, in the response starting now. + self._continues = held # --- user turns and inputs ---------------------------------------------------------------------- diff --git a/tests/realtime/test_session_core.py b/tests/realtime/test_session_core.py index 09dfa8e458..6a183a0611 100644 --- a/tests/realtime/test_session_core.py +++ b/tests/realtime/test_session_core.py @@ -49,6 +49,7 @@ from pydantic_ai.realtime._core import ( ToolReturned, TranscriptOverdue, ) +from pydantic_ai.realtime._google_lifecycle import GeminiLifecycle from pydantic_ai.realtime._lifecycle import ( InputAdded, InputLost, @@ -1127,3 +1128,34 @@ def test_a_turn_recorded_after_the_reply_to_it_is_evicted_before_that_reply() -> ) turn, response = session_core._turns['u1'], session_core._responses['r1'] # pyright: ignore[reportPrivateUsage] assert list(session_core._recorded_audio) == [turn, response] # pyright: ignore[reportPrivateUsage] + + +def test_a_held_filler_keeps_the_wait_through_a_transcript_until_the_model_carries_on() -> None: + """Gemini's extended-thinking model holds its filler open (`IN_PROGRESS`): the user's words then don't end the wait.""" + tracker = GeminiLifecycle(transcribes=True) + session_core = core() + + def message(*codec: Any) -> None: + for event, stale in tracker.message(list(codec)): + if not stale: + session_core.apply(event) + + session_core.apply(InputSent(input_id=0, request=text_request('Weather?'), solicits=True)) + tracker.input_sent(0, 'Weather?') + message(OutputTranscript('Let me check.')) + message(SessionUsage(RequestUsage(input_tokens=1)), ResponseDone(more_expected=True)) + wait = session_core.wait_tokens() + message(InputTranscript('Hm')) + assert session_core.still_owed(wait) != frozenset() + message(OutputTranscript('Still checking.')) + assert session_core.still_owed(wait) != frozenset() + message(SessionUsage(RequestUsage(input_tokens=2)), ResponseDone()) + assert session_core.still_owed(wait) == frozenset() + assert summary(session_core.all_messages()) == snapshot( + [ + '{prompt:Weather?}', + 'None [assistant:Let me check.] complete stop', + '{user:Hm}', + 'None [assistant:Still checking.] complete stop', + ] + )