scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 90

fix(vertex_ai): prune a discarded turn's id once its marker is delivered

Browse source
This commit is contained in:
mateo-berri 2026-09-18 17:44:28 -07:00
parent 8b1f78fa08
commit 3911d62bbe
2 changed files with 21 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
litellm/llms/vertex_ai/audio_transcription/realtime_backend.py
View file

@ -87,10 +87,16 @@ class _TurnResult:
@dataclass(frozen=True, slots=True)
class _TurnDiscarded:
pass
turn: int
_OutboxItem = str | _TurnResult | _StreamFailure | _Closed
@dataclass(frozen=True, slots=True)
class _TurnDiscardedEvent:
turn: int
event: str
_OutboxItem = str | _TurnResult | _TurnDiscardedEvent | _StreamFailure | _Closed
def open_speech_client(target: SpeechStreamingTarget, access_token: str) -> SpeechStreamingClient:
@ -304,6 +310,9 @@ class SpeechStreamingBackend:
raise _normal_closure()
case _TurnResult():
return None if item.turn in self._discarded_turns else item.event
case _TurnDiscardedEvent():
self._discarded_turns -= {item.turn}
return item.event
case str():
return item
case _:
@ -345,7 +354,10 @@ class SpeechStreamingBackend:
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()
_TurnDiscardedEvent(
turn=link.turn,
event=VertexSpeechStreamingTurnDiscarded(billed_seconds=self._billed_before).model_dump_json(),
)
)
case _:
assert_never(link)
@ -396,10 +408,11 @@ class SpeechStreamingBackend:
await self._link(_TURN_FINISHED_EVENT)
async def _discard_turn(self) -> None:
turn: Final = self._turn
streams: Final = self._turn
turn: Final = self._turn_index
self._turn = ()
self._discarded_turns |= {self._turn_index}
self._discarded_turns |= {turn}
self._turn_index += 1
for stream in turn:
for stream in streams:
stream.cancel()
await self._link(_TurnDiscarded())
await self._link(_TurnDiscarded(turn=turn))

1
tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py
View file

@ -260,6 +260,7 @@ async def test_discard_turn_drops_its_queued_results_and_keeps_google_billed_sec
await _until(lambda: len(client.streams[0]) == 3)
await backend.send(DISCARD_TURN)
assert await _recv(backend) == {"kind": "turn_discarded", "billed_seconds": 2.0}
assert backend._discarded_turns == frozenset()
await backend.send(b"\x03\x03")
fresh = await _recv(backend)
assert fresh["results"] == [{"transcript": "fresh", "is_final": True}]