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

fix(vertex_ai): drop a cleared turn's queued transcripts and carry its billed seconds

Browse source
This commit is contained in:
mateo-berri 2026-09-18 17:30:57 -07:00
parent dc11c34e3a
commit 8b1f78fa08
5 changed files with 104 additions and 17 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -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())

1
litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py
View file

@ -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 _:

1
litellm/types/llms/vertex_ai_speech_to_text.py
View file

@ -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 = (

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

@ -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"]]

9
tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py
View file

@ -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))