diff --git a/litellm/llms/base_llm/realtime/transcription_protocol.py b/litellm/llms/base_llm/realtime/transcription_protocol.py index 05151bd4a78..c3264911d46 100644 --- a/litellm/llms/base_llm/realtime/transcription_protocol.py +++ b/litellm/llms/base_llm/realtime/transcription_protocol.py @@ -59,37 +59,46 @@ class TranscriptionSessionUpdate: return None if self.turn_detection is None else self.turn_detection.get("type") -def json_object(payload: str) -> Mapping[str, JsonValue]: +ProtocolErrorType = type[RealtimeTranscriptionProtocolError] + + +def json_object(payload: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError) -> Mapping[str, JsonValue]: try: value: Final = _JSON_ADAPTER.validate_json(payload) except ValidationError: - raise RealtimeTranscriptionProtocolError("invalid JSON object") from None + raise error("invalid JSON object") from None if not isinstance(value, dict): - raise RealtimeTranscriptionProtocolError("message must be a JSON object") + raise error("message must be a JSON object") return value -def json_mapping(value: JsonValue | None, name: str) -> Mapping[str, JsonValue]: +def json_mapping( + value: JsonValue | None, name: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError +) -> Mapping[str, JsonValue]: if value is None: return EMPTY_JSON_OBJECT if not isinstance(value, dict): - raise RealtimeTranscriptionProtocolError(f"{name} must be an object") + raise error(f"{name} must be an object") return value -def json_string(value: JsonValue | None, name: str) -> str | None: +def json_string( + value: JsonValue | None, name: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError +) -> str | None: if value is None: return None if not isinstance(value, str): - raise RealtimeTranscriptionProtocolError(f"{name} must be a string") + raise error(f"{name} must be a string") return value -def json_integer(value: JsonValue | None, name: str) -> int | None: +def json_integer( + value: JsonValue | None, name: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError +) -> int | None: if value is None: return None if isinstance(value, bool) or not isinstance(value, int): - raise RealtimeTranscriptionProtocolError(f"{name} must be an integer") + raise error(f"{name} must be an integer") return value @@ -97,49 +106,52 @@ def new_event_id() -> str: return f"event_{uuid.uuid4().hex}" -def parse_transcription_session_update(payload: str) -> TranscriptionSessionUpdate: - message: Final = json_object(payload) +def parse_transcription_session_update( + payload: str, error: ProtocolErrorType = RealtimeTranscriptionProtocolError +) -> TranscriptionSessionUpdate: + message: Final = json_object(payload, error) if message.get("type") not in SESSION_UPDATE_EVENT_TYPES: - raise RealtimeTranscriptionProtocolError("expected session.update") - session: Final = json_mapping(message.get("session"), "session") + raise error("expected session.update") + session: Final = json_mapping(message.get("session"), "session", error) if not session: - raise RealtimeTranscriptionProtocolError("session.update requires a session object") - audio: Final = json_mapping(session.get("audio"), "session.audio") - audio_input: Final = json_mapping(audio.get("input"), "session.audio.input") + raise error("session.update requires a session object") + audio: Final = json_mapping(session.get("audio"), "session.audio", error) + audio_input: Final = json_mapping(audio.get("input"), "session.audio.input", error) beta_transcription: Final = session.get("input_audio_transcription") ga_transcription: Final = audio_input.get("transcription") if beta_transcription is not None and ga_transcription is not None: - raise RealtimeTranscriptionProtocolError("input transcription must use either beta or GA layout") + raise error("input transcription must use either beta or GA layout") transcription: Final = json_mapping( beta_transcription if beta_transcription is not None else ga_transcription, "input audio transcription", + error, ) turn_detection_present: Final = "turn_detection" in session or "turn_detection" in audio_input turn_detection: Final = session.get("turn_detection", audio_input.get("turn_detection")) return TranscriptionSessionUpdate( - session_type=json_string(session.get("type"), "session.type"), - audio_format=_parse_audio_format(session, audio_input), - model=json_string(transcription.get("model"), "transcription model"), - language=json_string(transcription.get("language"), "language"), + session_type=json_string(session.get("type"), "session.type", error), + audio_format=_parse_audio_format(session, audio_input, error), + model=json_string(transcription.get("model"), "transcription model", error), + language=json_string(transcription.get("language"), "language", error), unsupported_transcription_keys=tuple( sorted(key for key in transcription if key not in _SUPPORTED_TRANSCRIPTION_KEYS) ), - turn_detection=None if turn_detection is None else json_mapping(turn_detection, "turn_detection"), + turn_detection=None if turn_detection is None else json_mapping(turn_detection, "turn_detection", error), turn_detection_disabled=turn_detection_present and turn_detection is None, ) def _parse_audio_format( - session: Mapping[str, JsonValue], audio_input: Mapping[str, JsonValue] + session: Mapping[str, JsonValue], audio_input: Mapping[str, JsonValue], error: ProtocolErrorType ) -> TranscriptionAudioFormat | None: beta_format: Final = session.get("input_audio_format") ga_format: Final = audio_input.get("format") if beta_format is not None and ga_format is not None: - raise RealtimeTranscriptionProtocolError("input audio format must use either beta or GA layout") + raise error("input audio format must use either beta or GA layout") if beta_format is not None: return TranscriptionAudioFormat( layout="beta", - encoding=json_string(beta_format, "session.input_audio_format"), + encoding=json_string(beta_format, "session.input_audio_format", error), rate=None, channels=None, ) @@ -147,26 +159,30 @@ def _parse_audio_format( return None if isinstance(ga_format, str): return TranscriptionAudioFormat(layout="ga", encoding=ga_format, rate=None, channels=None) - format_mapping: Final = json_mapping(ga_format, "session.audio.input.format") + format_mapping: Final = json_mapping(ga_format, "session.audio.input.format", error) return TranscriptionAudioFormat( layout="ga", - encoding=json_string(format_mapping.get("type"), "session.audio.input.format.type"), - rate=json_integer(format_mapping.get("rate"), "session.audio.input.format.rate"), - channels=json_integer(format_mapping.get("channels"), "session.audio.input.format.channels"), + encoding=json_string(format_mapping.get("type"), "session.audio.input.format.type", error), + rate=json_integer(format_mapping.get("rate"), "session.audio.input.format.rate", error), + channels=json_integer(format_mapping.get("channels"), "session.audio.input.format.channels", error), ) -def decode_pcm16_append(audio: JsonValue | None, max_encoded_bytes: int | None = None) -> bytes: +def decode_pcm16_append( + audio: JsonValue | None, + max_encoded_bytes: int | None = None, + error: ProtocolErrorType = RealtimeTranscriptionProtocolError, +) -> bytes: if not isinstance(audio, str): - raise RealtimeTranscriptionProtocolError("Audio must be a base64 string") + raise error("Audio must be a base64 string") if max_encoded_bytes is not None and len(audio) > max_encoded_bytes: - raise RealtimeTranscriptionProtocolError("Audio append exceeds the four-second backlog limit") + raise error("Audio append exceeds the four-second backlog limit") try: decoded: Final = base64.b64decode(audio, validate=True) except (binascii.Error, ValueError): - raise RealtimeTranscriptionProtocolError("Audio must be valid base64") from None + raise error("Audio must be valid base64") from None if len(decoded) % 2: - raise RealtimeTranscriptionProtocolError("PCM16 audio must contain complete samples") + raise error("PCM16 audio must contain complete samples") return decoded diff --git a/litellm/llms/base_llm/realtime/transformation.py b/litellm/llms/base_llm/realtime/transformation.py index e1b16a5985d..1f4ad29fa74 100644 --- a/litellm/llms/base_llm/realtime/transformation.py +++ b/litellm/llms/base_llm/realtime/transformation.py @@ -1,9 +1,10 @@ from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence from types import TracebackType -from typing import TYPE_CHECKING, Any, Protocol, Self +from typing import TYPE_CHECKING, Any, Protocol import httpx +from typing_extensions import Self from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents from litellm.types.realtime import ( diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py index f0ca448bea0..442c79255af 100644 --- a/litellm/llms/meta/realtime/transformation.py +++ b/litellm/llms/meta/realtime/transformation.py @@ -239,7 +239,7 @@ def _parse_mode(update: TranscriptionSessionUpdate) -> MuseMode: def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig: - update: Final = parse_transcription_session_update(payload) + update: Final = parse_transcription_session_update(payload, MuseProtocolError) if update.session_type not in (None, "transcription", "realtime"): raise MuseProtocolError("Muse Voice supports transcription sessions only") if update.unsupported_transcription_keys: @@ -486,7 +486,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig): model: str, session_configuration_request: str | None = None, ) -> tuple[str | bytes, ...]: - request: Final = json_object(message) + request: Final = json_object(message, MuseProtocolError) event_type: Final = request.get("type") if event_type in ("session.update", "transcription_session.update"): return self._configure(message, model) @@ -538,7 +538,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig): return result def _backend_events(self, payload: str) -> tuple[OpenAIRealtimeEvents, ...]: - frame: Final = json_object(payload) + frame: Final = json_object(payload, MuseProtocolError) session_id: Final = frame.get("sessionId") if session_id is None: return self._transformer.transform(frame) @@ -560,7 +560,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig): def _append_audio(self, request: Mapping[str, JsonValue]) -> tuple[bytes, ...]: config: Final = self._require_config() - audio: Final = decode_pcm16_append(request.get("audio"), config.max_encoded_append_bytes) + audio: Final = decode_pcm16_append(request.get("audio"), config.max_encoded_append_bytes, MuseProtocolError) buffered: Final = self._pending_audio + audio packet_end: Final = len(buffered) - len(buffered) % config.packet_bytes self._pending_audio = buffered[packet_end:] diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py index 0ecbe7209f6..87c72c193bc 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py @@ -4,9 +4,10 @@ from collections.abc import AsyncIterable, AsyncIterator, Awaitable, Callable from dataclasses import dataclass from datetime import timedelta from types import MappingProxyType, TracebackType -from typing import TYPE_CHECKING, Final, Literal, Protocol, Self +from typing import TYPE_CHECKING, Final, Literal, Protocol from pydantic import TypeAdapter +from typing_extensions import Self, assert_never from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK from websockets.frames import Close @@ -37,6 +38,10 @@ SPEECH_SDK_INSTALL_HINT: Final = ( ) STREAM_FAILURE_CLOSE_CODE: Final = 1011 STREAM_ROTATION_SECONDS: Final = 240.0 +STREAM_ROTATION_DEADLINE_SECONDS: Final = 280.0 +REQUEST_QUEUE_SIZE: Final = 64 +OUTBOX_SIZE: Final = 256 +_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() @@ -129,6 +134,10 @@ def _billed_seconds(response: "StreamingRecognizeResponse") -> float: return _TIMEDELTA_ADAPTER.validate_python(response.metadata.total_billed_duration).total_seconds() +def _normal_closure() -> ConnectionClosedOK: + return ConnectionClosedOK(rcvd=Close(1000, ""), sent=None) + + class _RecognizeStream: def __init__( self, @@ -136,63 +145,59 @@ class _RecognizeStream: client: SpeechStreamingClient, request_type: "type[StreamingRecognizeRequest]", first_request: "StreamingRecognizeRequest", - outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", - previous: "_RecognizeStream | None", opened_at: float, ) -> None: self._client: Final = client self._request_type: Final = request_type - self._outbox: Final = outbox - self._previous: _RecognizeStream | None = previous self.opened_at: Final = opened_at - self._requests: Final[asyncio.Queue[StreamingRecognizeRequest | None]] = asyncio.Queue() + self._requests: Final[asyncio.Queue[StreamingRecognizeRequest | None]] = asyncio.Queue( + maxsize=REQUEST_QUEUE_SIZE + ) self._requests.put_nowait(first_request) - self._billed_seconds: float = 0.0 - self._base_billed_seconds: float = 0.0 - self._task: Final = asyncio.create_task(self._run()) + self.speech_active: bool = False + self.billed_seconds: float = 0.0 + self._cancelled: bool = False + self._task: asyncio.Task[None] | None = None - @property - def billed_seconds(self) -> float: - return self._base_billed_seconds + self._billed_seconds + async def send_audio(self, audio: bytes) -> None: + await self._requests.put(self._request_type(audio=audio)) - def send_audio(self, audio: bytes) -> None: - self._requests.put_nowait(self._request_type(audio=audio)) + async def half_close(self) -> None: + await self._requests.put(None) - def half_close(self) -> None: - self._requests.put_nowait(None) + def cancel(self) -> None: + self._cancelled = True + if self._task is not None: + self._task.cancel() - async def wait(self) -> None: - await asyncio.gather(self._task, return_exceptions=True) + async def relay(self, outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", billed_before: float) -> float: + if self._cancelled: + return 0.0 + task: Final = asyncio.create_task(self._forward(outbox, billed_before)) + self._task = task + try: + await asyncio.wait((task,)) + except asyncio.CancelledError: + task.cancel() + await asyncio.wait((task,)) + raise + return self.billed_seconds - async def cancel(self) -> None: - self._task.cancel() - await self.wait() - - async def cancel_chain(self) -> None: - previous: Final = self._previous - self._task.cancel() - if previous is not None: - await previous.cancel_chain() - await self.wait() - - async def _run(self) -> None: - previous: Final = self._previous - if previous is not None: - await previous.wait() - self._base_billed_seconds = previous.billed_seconds - self._previous = None + async def _forward(self, outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", billed_before: float) -> None: try: responses: Final = await self._client.streaming_recognize(self._drain()) async for response in responses: - self._billed_seconds = max(self._billed_seconds, _billed_seconds(response)) - await self._outbox.put(_response_event(response, self.billed_seconds)) - await self._outbox.put(_TURN_FINISHED_EVENT) - except asyncio.CancelledError: - self._outbox.put_nowait(_TURN_DISCARDED_EVENT) - raise + self._note(response) + await outbox.put(_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 self._outbox.put(_StreamFailure(reason=f"Google Speech-to-Text streaming failed: {e}")) + await outbox.put(_StreamFailure(reason=f"Google Speech-to-Text streaming failed: {e}")) + + def _note(self, response: "StreamingRecognizeResponse") -> None: + activity: Final = _SPEECH_EVENTS.get(response.speech_event_type.name) + if activity is not None: + self.speech_active = activity == "begin" + self.billed_seconds = max(self.billed_seconds, _billed_seconds(response)) async def _drain(self) -> "AsyncIterator[StreamingRecognizeRequest]": while (request := await self._requests.get()) is not None: @@ -207,16 +212,21 @@ class SpeechStreamingBackend: client_factory: Callable[[SpeechStreamingTarget], SpeechStreamingClient] = open_speech_client, clock: Callable[[], float] = time.monotonic, rotation_seconds: float = STREAM_ROTATION_SECONDS, + rotation_deadline_seconds: float = STREAM_ROTATION_DEADLINE_SECONDS, ) -> None: self._target: Final = target self._client_factory: Final = client_factory self._clock: Final = clock self._rotation_seconds: Final = rotation_seconds - self._outbox: Final[asyncio.Queue[str | _StreamFailure | _Closed]] = asyncio.Queue() + 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._pump: asyncio.Task[None] | None = None self._client: SpeechStreamingClient | None = None self._config: StreamingRecognitionConfig | None = None - self._stream: _RecognizeStream | None = None - self._last_stream: _RecognizeStream | None = None + self._turn: tuple[_RecognizeStream, ...] = () + self._billed_before: float = 0.0 + self._closed: bool = False async def __aenter__(self) -> Self: return self @@ -230,20 +240,26 @@ class SpeechStreamingBackend: await self.close() async def send(self, message: str | bytes) -> None: + if self._closed: + raise _normal_closure() if isinstance(message, bytes): - self._send_audio(message) + await self._send_audio(message) return command: Final = _COMMAND_ADAPTER.validate_json(message) match command: case VertexSpeechStreamingConfigure(): self._config = _streaming_config(command) - await self._outbox.put(_CONFIGURED_EVENT) + await self._link(_CONFIGURED_EVENT) case VertexSpeechStreamingFinishTurn(): - self._finish_turn() + await self._finish_turn() case VertexSpeechStreamingDiscardTurn(): await self._discard_turn() + case _: + 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() match item: case _StreamFailure(): @@ -251,35 +267,67 @@ class SpeechStreamingBackend: rcvd=Close(STREAM_FAILURE_CLOSE_CODE, item.reason[:_CLOSE_REASON_MAX_CHARS]), sent=None ) case _Closed(): - raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=None) + raise _normal_closure() case str(): return item + case _: + assert_never(item) async def close(self) -> None: - self._stream = None - last_stream: Final = self._last_stream - self._last_stream = None - if last_stream is not None: - await last_stream.cancel_chain() + if self._closed: + return + self._closed = True + self._turn = () + pump: Final = self._pump + if pump is not None: + pump.cancel() + await asyncio.wait((pump,)) client: Final = self._client self._client = None if client is not None: await client.transport.close() - self._outbox.put_nowait(_Closed()) + if not self._outbox.full(): + self._outbox.put_nowait(_Closed()) - def _send_audio(self, audio: bytes) -> None: - self._rotate_expiring_stream() - stream: Final = self._stream if self._stream is not None else self._open_stream() - stream.send_audio(audio) + async def _link(self, item: _RecognizeStream | str) -> None: + if self._pump is None: + self._pump = asyncio.create_task(self._pump_links()) + await self._links.put(item) - def _rotate_expiring_stream(self) -> None: - stream: Final = self._stream - if stream is None or self._clock() - stream.opened_at < self._rotation_seconds: - return - self._stream = None - stream.half_close() + async def _pump_links(self) -> None: + while True: + await self._relay(await self._links.get()) - def _open_stream(self) -> _RecognizeStream: + async def _relay(self, link: _RecognizeStream | str) -> None: + match link: + case str(): + await self._outbox.put(link) + case _RecognizeStream(): + self._billed_before += await link.relay(self._outbox, self._billed_before) + case _: + assert_never(link) + + async def _send_audio(self, audio: bytes) -> None: + stream: Final = await self._turn_stream() + await stream.send_audio(audio) + + async def _turn_stream(self) -> _RecognizeStream: + current: Final = self._turn[-1] if self._turn else None + if current is not None and not self._expired(current): + return current + if current is not None: + await current.half_close() + stream: Final = await self._open_stream() + self._turn = (*self._turn, stream) + return stream + + def _expired(self, stream: _RecognizeStream) -> bool: + elapsed: Final = self._clock() - stream.opened_at + if elapsed >= self._rotation_deadline_seconds: + return True + return elapsed >= self._rotation_seconds and not stream.speech_active + + async def _open_stream(self) -> _RecognizeStream: from google.cloud.speech_v2.types import StreamingRecognizeRequest config: Final = self._config @@ -291,26 +339,21 @@ class SpeechStreamingBackend: client=self._client, request_type=StreamingRecognizeRequest, first_request=StreamingRecognizeRequest(recognizer=self._target.recognizer, streaming_config=config), - outbox=self._outbox, - previous=self._last_stream, opened_at=self._clock(), ) - self._stream = stream - self._last_stream = stream + await self._link(stream) return stream - def _finish_turn(self) -> None: - stream: Final = self._stream - self._stream = None - if stream is None: - self._outbox.put_nowait(_TURN_FINISHED_EVENT) - return - stream.half_close() + async def _finish_turn(self) -> None: + turn: Final = self._turn + self._turn = () + if turn: + await turn[-1].half_close() + await self._link(_TURN_FINISHED_EVENT) async def _discard_turn(self) -> None: - stream: Final = self._stream - self._stream = None - if stream is None: - await self._outbox.put(_TURN_DISCARDED_EVENT) - return - await stream.cancel() + turn: Final = self._turn + self._turn = () + for stream in turn: + stream.cancel() + await self._link(_TURN_DISCARDED_EVENT) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py index 5e2857f24c3..6ec7a21a134 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py @@ -3,7 +3,9 @@ from dataclasses import dataclass, replace from typing import Final from pydantic import JsonValue, TypeAdapter +from typing_extensions import assert_never +import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.audio_utils.utils import normalize_transcription_language_to_bcp47 @@ -56,7 +58,7 @@ DEFAULT_SAMPLE_RATE_HERTZ: Final = 24_000 MIN_SAMPLE_RATE_HERTZ: Final = 8_000 MAX_SAMPLE_RATE_HERTZ: Final = 48_000 MAX_AUDIO_MESSAGE_BYTES: Final = 25_000 -SPEECH_TO_TEXT_MODEL_PREFIX: Final = "chirp" +_SPEECH_TO_TEXT_ENDPOINTS: Final = frozenset({"/v1/audio/transcriptions", "/v1/realtime"}) _VERTEX_MODEL_PREFIX: Final = "vertex_ai/" _STREAMING_EVENT_ADAPTER: Final = TypeAdapter[VertexSpeechStreamingEventUnion](VertexSpeechStreamingEvent) _FINISH_TURN_COMMAND: Final = VertexSpeechStreamingFinishTurn().model_dump_json() @@ -99,7 +101,15 @@ class ChirpSessionConfig: def is_vertex_speech_to_text_model(model: str) -> bool: - return normalize_speech_to_text_model(model).startswith(SPEECH_TO_TEXT_MODEL_PREFIX) + try: + info: Final = litellm.get_model_info( + model=normalize_speech_to_text_model(model), custom_llm_provider="vertex_ai" + ) + except Exception: # noqa: BLE001 # get_model_info raises for unmapped models, which are not Speech-to-Text models + return False + if info.get("mode") != "audio_transcription": + return False + return _SPEECH_TO_TEXT_ENDPOINTS <= frozenset(info.get("supported_endpoints") or ()) def normalize_speech_to_text_model(model: str) -> str: @@ -116,7 +126,7 @@ def default_session_config(model: str) -> ChirpSessionConfig: def parse_chirp_session_update(payload: str, expected_model: str) -> ChirpSessionConfig: - update: Final = parse_transcription_session_update(payload) + update: Final = parse_transcription_session_update(payload, ChirpProtocolError) if update.session_type not in (None, "transcription", "realtime"): raise ChirpProtocolError("Speech-to-Text streaming supports transcription sessions only") if update.unsupported_transcription_keys: @@ -228,6 +238,8 @@ class ChirpEventTransformer: case VertexSpeechStreamingTurnDiscarded(): self._turn = None return () + case _: + assert_never(frame) def _response(self, frame: VertexSpeechStreamingResponse) -> tuple[OpenAIRealtimeEvents, ...]: self._billed_seconds = max(self._billed_seconds, frame.billed_seconds) @@ -367,7 +379,7 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): model: str, session_configuration_request: str | None = None, ) -> tuple[str | bytes, ...]: - request: Final = json_object(message) + request: Final = json_object(message, ChirpProtocolError) event_type: Final = request.get("type") if event_type in ("session.update", "transcription_session.update"): return self._configure(message, model) @@ -417,7 +429,7 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): def _append_audio(self, request: Mapping[str, JsonValue]) -> tuple[bytes, ...]: self._require_config() - audio: Final = decode_pcm16_append(request.get("audio")) + audio: Final = decode_pcm16_append(request.get("audio"), error=ChirpProtocolError) return tuple( audio[start : start + MAX_AUDIO_MESSAGE_BYTES] for start in range(0, len(audio), MAX_AUDIO_MESSAGE_BYTES) ) diff --git a/pyproject.toml b/pyproject.toml index c58fd1d2158..361a5309ef3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -277,6 +277,7 @@ ci = [ "langgraph>=1.2.4,<1.3.0", "langgraph-prebuilt>=1.1.0,<1.3.0", "claude-agent-sdk==0.1.44", + "google-cloud-speech==2.40.0", ] healthcheck = [ "httpx==0.28.1", diff --git a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py index 1262d124d89..a5d7e47fb65 100644 --- a/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py +++ b/tests/test_litellm/llms/meta/realtime/test_meta_realtime_transformation.py @@ -6,7 +6,6 @@ from unittest.mock import MagicMock import pytest -from litellm.llms.base_llm.realtime.transcription_protocol import RealtimeTranscriptionProtocolError from litellm.llms.meta.realtime.transformation import ( DEFAULT_MUSE_REALTIME_URL, MUSE_MODEL, @@ -161,7 +160,7 @@ def test_language_normalization_uses_official_muse_names(source: str, expected: ], ) def test_session_rejects_unsupported_audio_model_and_hints(session: dict[str, object], message: str): - with pytest.raises(RealtimeTranscriptionProtocolError, match=message): + with pytest.raises(MuseProtocolError, match=message): parse_session_update(_event("session.update", session={"type": "transcription", **session}), MUSE_MODEL) @@ -585,7 +584,7 @@ def test_pcm_is_packetized_into_raw_binary_frames(rate: int, packet_bytes: int): def test_invalid_audio_appends_are_rejected(audio: object, message: str): config = _configured() - with pytest.raises(RealtimeTranscriptionProtocolError, match=message): + with pytest.raises(MuseProtocolError, match=message): config.transform_realtime_request(_event("input_audio_buffer.append", audio=audio), MUSE_MODEL) 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 87ab3da37a5..49758e62415 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,3 +1,4 @@ +import asyncio import json from collections.abc import AsyncIterator, Sequence from datetime import timedelta @@ -13,7 +14,7 @@ from google.cloud.speech_v2.types import ( ) from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK -from litellm.llms.vertex_ai.audio_transcription.realtime_backend import SpeechStreamingBackend +from litellm.llms.vertex_ai.audio_transcription.realtime_backend import REQUEST_QUEUE_SIZE, SpeechStreamingBackend from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import SpeechStreamingTarget TARGET: Final = SpeechStreamingTarget( @@ -26,7 +27,7 @@ CONFIGURE: Final = json.dumps( ) FINISH_TURN: Final = json.dumps({"kind": "finish_turn"}) DISCARD_TURN: Final = json.dumps({"kind": "discard_turn"}) -ScriptItem = StreamingRecognizeResponse | Exception +ScriptItem = StreamingRecognizeResponse | Exception | asyncio.Event def _response( @@ -84,13 +85,16 @@ class _FakeSpeechClient: async for request in requests: received.append(request) if request.audio and script: - yield self._next(script) + yield await self._next(script) while script: - yield self._next(script) + yield await self._next(script) @staticmethod - def _next(script: list[ScriptItem]) -> StreamingRecognizeResponse: + async def _next(script: list[ScriptItem]) -> StreamingRecognizeResponse: item: Final = script.pop(0) + if isinstance(item, asyncio.Event): + await item.wait() + return await _FakeSpeechClient._next(script) if isinstance(item, Exception): raise item return item @@ -101,11 +105,18 @@ def _backend(client: _FakeSpeechClient, **kwargs: object) -> SpeechStreamingBack async def _recv(backend: SpeechStreamingBackend) -> dict[str, object]: - message: Final = await backend.recv() + message: Final = await asyncio.wait_for(backend.recv(), timeout=2) assert isinstance(message, str) return json.loads(message) +async def _transcript(backend: SpeechStreamingBackend) -> str: + event: Final = await _recv(backend) + assert event["kind"] == "response", event + (result,) = event["results"] + return result["transcript"] + + async def _configure(backend: SpeechStreamingBackend) -> None: await backend.send(CONFIGURE) assert await _recv(backend) == {"kind": "configured"} @@ -149,7 +160,9 @@ async def test_audio_streams_through_one_recognize_call_with_the_config_first(): @pytest.mark.asyncio async def test_voice_activity_events_are_relayed(): - client = _FakeSpeechClient([_response(None, event="SPEECH_ACTIVITY_BEGIN"), _response(None, event="SPEECH_ACTIVITY_END")]) + client = _FakeSpeechClient( + [_response(None, event="SPEECH_ACTIVITY_BEGIN"), _response(None, event="SPEECH_ACTIVITY_END")] + ) async with _backend(client) as backend: await _configure(backend) await backend.send(b"\x00\x00") @@ -181,16 +194,17 @@ async def test_stream_failure_closes_the_session_with_1011_and_the_reason(): @pytest.mark.asyncio -async def test_close_discards_the_open_turn_then_reports_a_normal_closure(): +async def test_close_reports_a_normal_closure_to_both_directions(): client = _FakeSpeechClient([_response("hi")]) backend = _backend(client) await _configure(backend) await backend.send(b"\x00\x00") - assert (await _recv(backend))["results"][0]["transcript"] == "hi" + assert await _transcript(backend) == "hi" await backend.close() - assert await _recv(backend) == {"kind": "turn_discarded"} with pytest.raises(ConnectionClosedOK): await backend.recv() + with pytest.raises(ConnectionClosedOK): + await backend.send(b"\x00\x00") assert client.transport.closed @@ -210,17 +224,19 @@ async def test_discard_turn_cancels_the_open_stream_and_the_next_turn_starts_fre async with _backend(client) as backend: await _configure(backend) await backend.send(b"\x01\x01") - assert (await _recv(backend))["results"][0]["transcript"] == "draft" + assert await _transcript(backend) == "draft" await backend.send(DISCARD_TURN) assert await _recv(backend) == {"kind": "turn_discarded"} await backend.send(b"\x02\x02") - assert (await _recv(backend))["results"][0]["transcript"] == "again" + 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_billed_seconds_accumulate_across_turns(): - client = _FakeSpeechClient([_response("one", is_final=True, billed=2.0)], [_response("two", is_final=True, billed=3.0)]) + client = _FakeSpeechClient( + [_response("one", is_final=True, billed=2.0)], [_response("two", is_final=True, billed=3.0)] + ) async with _backend(client) as backend: await _configure(backend) await backend.send(b"\x00\x00") @@ -236,26 +252,132 @@ async def test_billed_seconds_accumulate_across_turns(): @pytest.mark.asyncio -async def test_streams_rotate_before_the_five_minute_limit_without_losing_audio(): +async def test_streams_rotate_before_the_five_minute_limit_without_ending_the_turn(): now = [0.0] client = _FakeSpeechClient( [_response("first"), _response("first half", is_final=True, billed=239.0)], - [_response("second")], + [_response("second", billed=1.0)], ) async with _backend(client, clock=lambda: now[0], rotation_seconds=240.0) as backend: await _configure(backend) await backend.send(b"\x01\x01") - assert (await _recv(backend))["results"][0]["transcript"] == "first" + assert await _transcript(backend) == "first" now[0] = 239.0 await backend.send(b"\x02\x02") - assert (await _recv(backend))["results"][0]["transcript"] == "first half" + assert await _transcript(backend) == "first half" now[0] = 240.0 await backend.send(b"\x03\x03") - assert await _recv(backend) == {"kind": "turn_finished"} second = await _recv(backend) - assert second["results"][0]["transcript"] == "second" - assert second["billed_seconds"] == 239.0 + assert second["results"] == [{"transcript": "second", "is_final": False}] + assert second["billed_seconds"] == 240.0 await backend.send(FINISH_TURN) assert await _recv(backend) == {"kind": "turn_finished"} assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01", b"\x02\x02"], [b"\x03\x03"]] assert client.streams[1][0].streaming_config.config.model == "chirp_3" + + +@pytest.mark.asyncio +async def test_turn_finished_follows_results_that_arrive_after_a_rotation(): + now = [0.0] + client = _FakeSpeechClient( + [_response("one"), _response("one two", is_final=True)], + [_response("three")], + ) + async with _backend(client, clock=lambda: now[0], rotation_seconds=240.0) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "one" + now[0] = 240.0 + await backend.send(b"\x02\x02") + await backend.send(FINISH_TURN) + assert await _transcript(backend) == "one two" + assert await _transcript(backend) == "three" + assert await _recv(backend) == {"kind": "turn_finished"} + + +@pytest.mark.asyncio +async def test_rotation_waits_for_a_pause_in_speech(): + now = [0.0] + client = _FakeSpeechClient( + [ + _response(None, event="SPEECH_ACTIVITY_BEGIN"), + _response("still talking"), + _response("still talking", is_final=True, event="SPEECH_ACTIVITY_END"), + ], + [_response("next")], + ) + async with _backend( + client, clock=lambda: now[0], rotation_seconds=240.0, rotation_deadline_seconds=280.0 + ) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert (await _recv(backend))["speech_event"] == "begin" + now[0] = 250.0 + await backend.send(b"\x02\x02") + assert await _transcript(backend) == "still talking" + now[0] = 260.0 + await backend.send(b"\x03\x03") + assert (await _recv(backend))["speech_event"] == "end" + now[0] = 261.0 + await backend.send(b"\x04\x04") + assert await _transcript(backend) == "next" + assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01", b"\x02\x02", b"\x03\x03"], [b"\x04\x04"]] + + +@pytest.mark.asyncio +async def test_rotation_is_forced_at_the_deadline_during_continuous_speech(): + now = [0.0] + client = _FakeSpeechClient( + [_response(None, event="SPEECH_ACTIVITY_BEGIN"), _response("still talking")], + [_response("cut off")], + ) + async with _backend( + client, clock=lambda: now[0], rotation_seconds=240.0, rotation_deadline_seconds=280.0 + ) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert (await _recv(backend))["speech_event"] == "begin" + now[0] = 279.0 + await backend.send(b"\x02\x02") + assert await _transcript(backend) == "still talking" + now[0] = 280.0 + await backend.send(b"\x03\x03") + assert await _transcript(backend) == "cut off" + assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01", b"\x02\x02"], [b"\x03\x03"]] + + +@pytest.mark.asyncio +async def test_discard_turn_cancels_every_stream_of_the_turn(): + now = [0.0] + hold = asyncio.Event() + client = _FakeSpeechClient( + [_response("draft"), hold, _response("never delivered")], + [_response("fresh", is_final=True)], + ) + async with _backend(client, clock=lambda: now[0], rotation_seconds=240.0) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "draft" + now[0] = 240.0 + await backend.send(b"\x02\x02") + await backend.send(DISCARD_TURN) + assert await _recv(backend) == {"kind": "turn_discarded"} + 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"]] + + +@pytest.mark.asyncio +async def test_audio_sends_block_once_the_request_queue_is_full(): + hold = asyncio.Event() + client = _FakeSpeechClient([hold, _response("late", is_final=True)]) + async with _backend(client) as backend: + await _configure(backend) + for _ in range(REQUEST_QUEUE_SIZE + 1): + await backend.send(b"\x00\x00") + blocked = asyncio.create_task(backend.send(b"\x00\x00")) + await asyncio.sleep(0) + assert not blocked.done() + hold.set() + await asyncio.wait_for(blocked, timeout=2) + assert await _transcript(backend) == "late" 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 e294c7cad36..fcbfad8bf12 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 @@ -100,7 +100,10 @@ def _types(events: list[dict[str, object]]) -> list[object]: def _commands(config: VertexChirpRealtimeConfig, payload: str) -> list[object]: - return [json.loads(command) if isinstance(command, str) else command for command in config.transform_realtime_request(payload, MODEL)] + return [ + json.loads(command) if isinstance(command, str) else command + for command in config.transform_realtime_request(payload, MODEL) + ] @pytest.mark.parametrize( @@ -108,9 +111,11 @@ def _commands(config: VertexChirpRealtimeConfig, payload: str) -> list[object]: [ ("vertex_ai/chirp_3", True), ("chirp_3", True), - ("chirp_2", True), + ("chirp_2", False), ("gemini-live-2.5-flash", False), ("vertex_ai/gemini-2.0-flash-live-preview-04-09", False), + ("vertex_ai/gemini-3.5-transcribe-live-preview", False), + ("gemini-3.5-transcribe-preview", False), ], ) def test_is_vertex_speech_to_text_model(model: str, expected: bool): @@ -132,7 +137,11 @@ def test_beta_session_update_defaults_the_rate_and_auto_detects_the_language(): config = parse_chirp_session_update( _event( "transcription_session.update", - session={"input_audio_format": "pcm16", "input_audio_transcription": {"model": MODEL}, "turn_detection": None}, + session={ + "input_audio_format": "pcm16", + "input_audio_transcription": {"model": MODEL}, + "turn_detection": None, + }, ), MODEL, ) @@ -146,7 +155,10 @@ def test_beta_session_update_defaults_the_rate_and_auto_detects_the_language(): (_event("session.update", session={"type": "realtime_voice"}), "transcription sessions only"), (_ga_session_update(model="gemini-live-2.5-flash"), "cannot be changed"), (_event("session.update", session={"audio": {"input": {"format": {"type": "audio/pcmu"}}}}), "pcm16"), - (_event("session.update", session={"audio": {"input": {"format": {"type": "audio/pcm", "channels": 2}}}}), "mono"), + ( + _event("session.update", session={"audio": {"input": {"format": {"type": "audio/pcm", "channels": 2}}}}), + "mono", + ), (_ga_session_update(rate=4_000), "sample rates"), (_ga_session_update(rate=96_000), "sample rates"), (_ga_session_update(turn_detection="semantic_vad"), "server_vad"), @@ -169,7 +181,9 @@ def test_session_update_configures_once_and_later_updates_are_ignored(): def test_audio_and_commits_before_session_update_are_rejected(): config = _config() with pytest.raises(ChirpProtocolError, match=r"session\.update must configure"): - config.transform_realtime_request(_event("input_audio_buffer.append", audio=base64.b64encode(b"\x00\x00").decode()), MODEL) + config.transform_realtime_request( + _event("input_audio_buffer.append", audio=base64.b64encode(b"\x00\x00").decode()), MODEL + ) with pytest.raises(ChirpProtocolError, match=r"session\.update must configure"): config.transform_realtime_request(_event("input_audio_buffer.commit"), MODEL) @@ -177,8 +191,14 @@ def test_audio_and_commits_before_session_update_are_rejected(): def test_append_is_split_into_google_sized_chunks(): config = _configured() audio = bytes(range(256)) * 250 - chunks = config.transform_realtime_request(_event("input_audio_buffer.append", audio=base64.b64encode(audio).decode()), MODEL) - assert [len(chunk) for chunk in chunks] == [MAX_AUDIO_MESSAGE_BYTES, MAX_AUDIO_MESSAGE_BYTES, 64_000 - 2 * MAX_AUDIO_MESSAGE_BYTES] + chunks = config.transform_realtime_request( + _event("input_audio_buffer.append", audio=base64.b64encode(audio).decode()), MODEL + ) + assert [len(chunk) for chunk in chunks] == [ + MAX_AUDIO_MESSAGE_BYTES, + MAX_AUDIO_MESSAGE_BYTES, + 64_000 - 2 * MAX_AUDIO_MESSAGE_BYTES, + ] assert b"".join(chunk for chunk in chunks if isinstance(chunk, bytes)) == audio diff --git a/uv.lock b/uv.lock index de7ccb302e9..9d26d22c980 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-15T01:04:49.417319Z" +exclude-newer = "2026-09-15T21:32:43.124695Z" exclude-newer-span = "P3D" [manifest] @@ -4628,6 +4628,7 @@ ci = [ { name = "blockbuster" }, { name = "claude-agent-sdk" }, { name = "detect-secrets" }, + { name = "google-cloud-speech" }, { name = "google-generativeai" }, { name = "jsonlines" }, { name = "langchain" }, @@ -4822,6 +4823,7 @@ ci = [ { name = "blockbuster", specifier = "==1.5.26" }, { name = "claude-agent-sdk", specifier = "==0.1.44" }, { name = "detect-secrets", specifier = "==1.5.0" }, + { name = "google-cloud-speech", specifier = "==2.40.0" }, { name = "google-generativeai", specifier = "==0.8.6" }, { name = "jsonlines", specifier = "==4.0.0" }, { name = "langchain", specifier = "==1.3.9" },