From 8b1f78fa08fb6322398ff78b27646ce15e6684d9 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:30:57 -0700 Subject: [PATCH] fix(vertex_ai): drop a cleared turn's queued transcripts and carry its billed seconds --- .../audio_transcription/realtime_backend.py | 59 +++++++++++++++---- .../realtime_transformation.py | 1 + .../types/llms/vertex_ai_speech_to_text.py | 1 + .../test_vertex_ai_realtime_backend.py | 51 ++++++++++++++-- .../test_vertex_ai_realtime_transformation.py | 9 ++- 5 files changed, 104 insertions(+), 17 deletions(-) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py index 0c16fea9e9f..4c8338c027e 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py @@ -45,7 +45,6 @@ _LINK_QUEUE_SIZE: Final = 64 _CLOSE_REASON_MAX_CHARS: Final = 120 _CONFIGURED_EVENT: Final = VertexSpeechStreamingConfigured().model_dump_json() _TURN_FINISHED_EVENT: Final = VertexSpeechStreamingTurnFinished().model_dump_json() -_TURN_DISCARDED_EVENT: Final = VertexSpeechStreamingTurnDiscarded().model_dump_json() _COMMAND_ADAPTER: Final = TypeAdapter[VertexSpeechStreamingCommandUnion](VertexSpeechStreamingCommand) _TIMEDELTA_ADAPTER: Final = TypeAdapter(timedelta) _SPEECH_EVENTS: Final[MappingProxyType[str, Literal["begin", "end"]]] = MappingProxyType( @@ -80,6 +79,20 @@ class _Closed: pass +@dataclass(frozen=True, slots=True) +class _TurnResult: + turn: int + event: str + + +@dataclass(frozen=True, slots=True) +class _TurnDiscarded: + pass + + +_OutboxItem = str | _TurnResult | _StreamFailure | _Closed + + def open_speech_client(target: SpeechStreamingTarget, access_token: str) -> SpeechStreamingClient: try: from google.api_core.client_options import ClientOptions @@ -146,10 +159,12 @@ class _RecognizeStream: request_type: "type[StreamingRecognizeRequest]", first_request: "StreamingRecognizeRequest", opened_at: float, + turn: int, ) -> None: self._client: Final = client self._request_type: Final = request_type self.opened_at: Final = opened_at + self.turn: Final = turn self._requests: Final[asyncio.Queue[StreamingRecognizeRequest | None]] = asyncio.Queue( maxsize=REQUEST_QUEUE_SIZE ) @@ -177,7 +192,7 @@ class _RecognizeStream: self._closed = True await self._client.transport.close() - async def relay(self, outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", billed_before: float) -> float: + async def relay(self, outbox: asyncio.Queue[_OutboxItem], billed_before: float) -> float: if self._cancelled: await self.close() return 0.0 @@ -193,12 +208,14 @@ class _RecognizeStream: await self.close() return self.billed_seconds - async def _forward(self, outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", billed_before: float) -> None: + async def _forward(self, outbox: asyncio.Queue[_OutboxItem], billed_before: float) -> None: try: responses: Final = await self._client.streaming_recognize(self._drain()) async for response in responses: self._note(response) - await outbox.put(_response_event(response, billed_before + self.billed_seconds)) + await outbox.put( + _TurnResult(turn=self.turn, event=_response_event(response, billed_before + self.billed_seconds)) + ) except Exception as e: # noqa: BLE001 # task boundary: a swallowed failure would hang the client session verbose_logger.warning("Google Speech-to-Text streaming failed: %s", e) await outbox.put(_StreamFailure(reason=f"Google Speech-to-Text streaming failed: {e}")) @@ -214,6 +231,9 @@ class _RecognizeStream: yield request +_Link = _RecognizeStream | str | _TurnDiscarded + + class SpeechStreamingBackend: def __init__( self, @@ -229,11 +249,13 @@ class SpeechStreamingBackend: self._clock: Final = clock self._rotation_seconds: Final = rotation_seconds self._rotation_deadline_seconds: Final = rotation_deadline_seconds - self._outbox: Final[asyncio.Queue[str | _StreamFailure | _Closed]] = asyncio.Queue(maxsize=OUTBOX_SIZE) - self._links: Final[asyncio.Queue[_RecognizeStream | str]] = asyncio.Queue(maxsize=_LINK_QUEUE_SIZE) + self._outbox: Final[asyncio.Queue[_OutboxItem]] = asyncio.Queue(maxsize=OUTBOX_SIZE) + self._links: Final[asyncio.Queue[_Link]] = asyncio.Queue(maxsize=_LINK_QUEUE_SIZE) self._pump: asyncio.Task[None] | None = None self._config: StreamingRecognitionConfig | None = None self._turn: tuple[_RecognizeStream, ...] = () + self._turn_index: int = 0 + self._discarded_turns: frozenset[int] = frozenset() self._billed_before: float = 0.0 self._closed: bool = False @@ -267,9 +289,12 @@ class SpeechStreamingBackend: assert_never(command) async def recv(self, decode: bool | None = None) -> str | bytes: - if self._closed and self._outbox.empty(): - raise _normal_closure() - item: Final = await self._outbox.get() + while not (self._closed and self._outbox.empty()): + if (event := self._deliverable(await self._outbox.get())) is not None: + return event + raise _normal_closure() + + def _deliverable(self, item: _OutboxItem) -> str | None: match item: case _StreamFailure(): raise ConnectionClosedError( @@ -277,6 +302,8 @@ class SpeechStreamingBackend: ) case _Closed(): raise _normal_closure() + case _TurnResult(): + return None if item.turn in self._discarded_turns else item.event case str(): return item case _: @@ -301,7 +328,7 @@ class SpeechStreamingBackend: if isinstance(link, _RecognizeStream): await link.close() - async def _link(self, item: _RecognizeStream | str) -> None: + async def _link(self, item: _Link) -> None: if self._pump is None: self._pump = asyncio.create_task(self._pump_links()) await self._links.put(item) @@ -310,12 +337,16 @@ class SpeechStreamingBackend: while True: await self._relay(await self._links.get()) - async def _relay(self, link: _RecognizeStream | str) -> None: + async def _relay(self, link: _Link) -> None: match link: case str(): await self._outbox.put(link) case _RecognizeStream(): self._billed_before += await link.relay(self._outbox, self._billed_before) + case _TurnDiscarded(): + await self._outbox.put( + VertexSpeechStreamingTurnDiscarded(billed_seconds=self._billed_before).model_dump_json() + ) case _: assert_never(link) @@ -351,6 +382,7 @@ class SpeechStreamingBackend: request_type=StreamingRecognizeRequest, first_request=StreamingRecognizeRequest(recognizer=self._target.recognizer, streaming_config=config), opened_at=self._clock(), + turn=self._turn_index, ) await self._link(stream) return stream @@ -358,6 +390,7 @@ class SpeechStreamingBackend: async def _finish_turn(self) -> None: turn: Final = self._turn self._turn = () + self._turn_index += 1 if turn: await turn[-1].half_close() await self._link(_TURN_FINISHED_EVENT) @@ -365,6 +398,8 @@ class SpeechStreamingBackend: async def _discard_turn(self) -> None: turn: Final = self._turn self._turn = () + self._discarded_turns |= {self._turn_index} + self._turn_index += 1 for stream in turn: stream.cancel() - await self._link(_TURN_DISCARDED_EVENT) + await self._link(_TurnDiscarded()) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py index dab2e980fd0..ac23901accb 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py @@ -236,6 +236,7 @@ class ChirpEventTransformer: case VertexSpeechStreamingTurnFinished(): return self._finish_turn() case VertexSpeechStreamingTurnDiscarded(): + self._billed_seconds = max(self._billed_seconds, frame.billed_seconds) self._turn = None return () case _: diff --git a/litellm/types/llms/vertex_ai_speech_to_text.py b/litellm/types/llms/vertex_ai_speech_to_text.py index e954960d41f..d07a5bbc192 100644 --- a/litellm/types/llms/vertex_ai_speech_to_text.py +++ b/litellm/types/llms/vertex_ai_speech_to_text.py @@ -93,6 +93,7 @@ class VertexSpeechStreamingTurnFinished(BaseModel): class VertexSpeechStreamingTurnDiscarded(BaseModel): model_config = ConfigDict(frozen=True) kind: Literal["turn_discarded"] = "turn_discarded" + billed_seconds: float VertexSpeechStreamingEventUnion = ( diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py index d6f65c90806..15601c5ca6c 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py @@ -1,6 +1,6 @@ import asyncio import json -from collections.abc import AsyncIterator, Sequence +from collections.abc import AsyncIterator, Callable, Sequence from dataclasses import replace from datetime import timedelta from typing import Final @@ -128,6 +128,14 @@ async def _configure(backend: SpeechStreamingBackend) -> None: assert await _recv(backend) == {"kind": "configured"} +async def _until(condition: Callable[[], bool]) -> None: + async def poll() -> None: + while not condition(): + await asyncio.sleep(0) + + await asyncio.wait_for(poll(), timeout=2) + + def _audio(stream: list[StreamingRecognizeRequest]) -> list[bytes]: return [bytes(request.audio) for request in stream[1:]] @@ -221,7 +229,7 @@ async def test_turn_commands_without_audio_answer_immediately(): await backend.send(FINISH_TURN) assert await _recv(backend) == {"kind": "turn_finished"} await backend.send(DISCARD_TURN) - assert await _recv(backend) == {"kind": "turn_discarded"} + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 0.0} @pytest.mark.asyncio @@ -232,12 +240,47 @@ async def test_discard_turn_cancels_the_open_stream_and_the_next_turn_starts_fre await backend.send(b"\x01\x01") assert await _transcript(backend) == "draft" await backend.send(DISCARD_TURN) - assert await _recv(backend) == {"kind": "turn_discarded"} + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 0.0} await backend.send(b"\x02\x02") assert await _transcript(backend) == "again" assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01"], [b"\x02\x02"]] +@pytest.mark.asyncio +async def test_discard_turn_drops_its_queued_results_and_keeps_google_billed_seconds(): + client = _FakeSpeechClient( + [_response("draft"), _response("leftover", is_final=True, billed=2.0)], + [_response("fresh", is_final=True, billed=1.0)], + ) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "draft" + await backend.send(b"\x02\x02") + await _until(lambda: len(client.streams[0]) == 3) + await backend.send(DISCARD_TURN) + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 2.0} + await backend.send(b"\x03\x03") + fresh = await _recv(backend) + assert fresh["results"] == [{"transcript": "fresh", "is_final": True}] + assert fresh["billed_seconds"] == 3.0 + + +@pytest.mark.asyncio +async def test_discard_turn_keeps_the_queued_results_of_the_turn_finished_before_it(): + client = _FakeSpeechClient([_response("one", is_final=True, billed=2.0)], [_response("two")]) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + await backend.send(FINISH_TURN) + await backend.send(b"\x02\x02") + await _until(lambda: len(client.streams) == 2 and len(client.streams[1]) == 2) + await backend.send(DISCARD_TURN) + assert await _transcript(backend) == "one" + assert await _recv(backend) == {"kind": "turn_finished"} + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 2.0} + + @pytest.mark.asyncio async def test_billed_seconds_accumulate_across_turns(): client = _FakeSpeechClient( @@ -425,7 +468,7 @@ async def test_discard_turn_cancels_every_stream_of_the_turn(): now[0] = 240.0 await backend.send(b"\x02\x02") await backend.send(DISCARD_TURN) - assert await _recv(backend) == {"kind": "turn_discarded"} + assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 0.0} await backend.send(b"\x03\x03") assert await _transcript(backend) == "fresh" assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01"], [b"\x03\x03"]] diff --git a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py index 719dd621c82..84c3a4e244a 100644 --- a/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py @@ -317,13 +317,20 @@ def test_manual_turns_complete_on_commit_without_speech_events(): def test_clear_discards_the_open_turn(): config = _configured(turn_detection=None) draft = _backend_events(config, _response(("draft", False))) - assert _backend_events(config, VertexSpeechStreamingTurnDiscarded()) == [] + assert _backend_events(config, VertexSpeechStreamingTurnDiscarded(billed_seconds=0.0)) == [] assert _backend_events(config, VertexSpeechStreamingTurnFinished()) == [] fresh = _backend_events(config, _response(("again", False))) assert fresh[0]["delta"] == "again" assert fresh[0]["item_id"] != draft[0]["item_id"] +def test_cleared_audio_keeps_google_billed_seconds_for_the_close_flush(): + config = _configured(turn_detection=None) + assert _backend_events(config, _response(("draft", False), billed_seconds=1.0)) != [] + assert _backend_events(config, VertexSpeechStreamingTurnDiscarded(billed_seconds=2.5)) == [] + assert config.unbilled_usage_on_session_close(MODEL) == {"type": "duration", "seconds": 2.5} + + def test_usage_is_billed_once_across_turns_and_flushed_on_close(): config = _configured() first = _backend_events(config, _response(("one", True), billed_seconds=2.0))