From e82d15a3aa9ea7f24b420a0c3a9ef10c9f9cc7a2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 17 Sep 2026 17:57:40 -0700 Subject: [PATCH 1/7] feat(vertex_ai): stream Chirp speech-to-text over /v1/realtime Bridge OpenAI Realtime transcription sessions on vertex_ai/chirp_* models to Google Speech-to-Text v2 StreamingRecognize over gRPC, so partial and final transcripts stream back while audio is still being sent. Interim results become delta events, finals become completed events carrying billed seconds, the gRPC stream rotates at 240 s under Google's five-minute cap with billed time chained across rotations, and audio is split into 25 KB requests. The OpenAI transcription protocol helpers move into a shared module that Meta Muse now uses too, google-cloud-speech ships behind a new stt-vertex-chirp extra bundled into the proxy runtime, and the cost map lists /v1/realtime for chirp_3. --- .../litellm_core_utils/realtime_streaming.py | 4 +- .../realtime/transcription_protocol.py | 259 +++++++++++ .../llms/base_llm/realtime/transformation.py | 23 +- litellm/llms/custom_httpx/llm_http_handler.py | 7 +- litellm/llms/meta/realtime/transformation.py | 248 +++------- .../audio_transcription/realtime_backend.py | 316 +++++++++++++ .../realtime_transformation.py | 433 ++++++++++++++++++ .../audio_transcription/transformation.py | 7 +- ...odel_prices_and_context_window_backup.json | 3 +- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- litellm/realtime_api/main.py | 28 +- .../types/llms/vertex_ai_speech_to_text.py | 66 ++- model_prices_and_context_window.json | 3 +- pyproject.toml | 7 + .../llms/base_llm/realtime/__init__.py | 0 .../realtime/test_transcription_protocol.py | 128 ++++++ .../custom_httpx/test_llm_http_handler.py | 142 ++++++ .../test_meta_realtime_transformation.py | 5 +- .../test_vertex_ai_realtime_backend.py | 261 +++++++++++ .../test_vertex_ai_realtime_transformation.py | 355 ++++++++++++++ tests/test_litellm/realtime_api/test_main.py | 66 +++ uv.lock | 27 +- 22 files changed, 2180 insertions(+), 210 deletions(-) create mode 100644 litellm/llms/base_llm/realtime/transcription_protocol.py create mode 100644 litellm/llms/vertex_ai/audio_transcription/realtime_backend.py create mode 100644 litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py create mode 100644 tests/test_litellm/llms/base_llm/realtime/__init__.py create mode 100644 tests/test_litellm/llms/base_llm/realtime/test_transcription_protocol.py create mode 100644 tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py create mode 100644 tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index fa567bdf4c9..d2fbb26bb02 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -12,7 +12,7 @@ import litellm from litellm._logging import redact_internal_details_from_client_message, verbose_logger from litellm.constants import REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig +from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend from litellm.types.llms.openai import ( OpenAIRealtimeEvents, OpenAIRealtimeOutputItemDone, @@ -127,7 +127,7 @@ class RealTimeStreaming: def __init__( self, websocket: Any, - backend_ws: CLIENT_CONNECTION_CLASS, + backend_ws: CLIENT_CONNECTION_CLASS | RealtimeBackend, logging_obj: LiteLLMLogging, provider_config: BaseRealtimeConfig | None = None, model: str = "", diff --git a/litellm/llms/base_llm/realtime/transcription_protocol.py b/litellm/llms/base_llm/realtime/transcription_protocol.py new file mode 100644 index 00000000000..05151bd4a78 --- /dev/null +++ b/litellm/llms/base_llm/realtime/transcription_protocol.py @@ -0,0 +1,259 @@ +import base64 +import binascii +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final, Literal + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm._uuid import uuid +from litellm.types.llms.openai import ( + OpenAIRealtimeErrorEvent, + OpenAIRealtimeInputAudioBufferSpeechEvent, + OpenAIRealtimeInputAudioTranscriptionCompleted, + OpenAIRealtimeInputAudioTranscriptionDelta, + OpenAIRealtimeServerVadTurnDetection, + OpenAIRealtimeTranscriptionSession, + OpenAIRealtimeTranscriptionSessionCreated, + OpenAIRealtimeTranscriptionSettings, +) +from litellm.types.realtime import RealtimeInputAudioTranscriptionDurationUsage, RealtimeInputAudioTranscriptionUsage + +SESSION_UPDATE_EVENT_TYPES: Final = frozenset(("session.update", "transcription_session.update")) +PCM16_ENCODINGS: Final = frozenset(("pcm16", "audio/pcm")) +SERVER_VAD_TURN_DETECTION: Final[OpenAIRealtimeServerVadTurnDetection] = {"type": "server_vad"} +EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) +_SUPPORTED_TRANSCRIPTION_KEYS: Final = frozenset(("model", "language")) +_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) + + +class RealtimeTranscriptionProtocolError(ValueError): + pass + + +@dataclass(frozen=True, slots=True) +class TranscriptionAudioFormat: + layout: Literal["beta", "ga"] + encoding: str | None + rate: int | None + channels: int | None + + @property + def is_pcm16(self) -> bool: + return self.encoding in PCM16_ENCODINGS + + +@dataclass(frozen=True, slots=True) +class TranscriptionSessionUpdate: + session_type: str | None + audio_format: TranscriptionAudioFormat | None + model: str | None + language: str | None + unsupported_transcription_keys: tuple[str, ...] + turn_detection: Mapping[str, JsonValue] | None + turn_detection_disabled: bool + + @property + def turn_detection_type(self) -> JsonValue | None: + return None if self.turn_detection is None else self.turn_detection.get("type") + + +def json_object(payload: str) -> Mapping[str, JsonValue]: + try: + value: Final = _JSON_ADAPTER.validate_json(payload) + except ValidationError: + raise RealtimeTranscriptionProtocolError("invalid JSON object") from None + if not isinstance(value, dict): + raise RealtimeTranscriptionProtocolError("message must be a JSON object") + return value + + +def json_mapping(value: JsonValue | None, name: str) -> Mapping[str, JsonValue]: + if value is None: + return EMPTY_JSON_OBJECT + if not isinstance(value, dict): + raise RealtimeTranscriptionProtocolError(f"{name} must be an object") + return value + + +def json_string(value: JsonValue | None, name: str) -> str | None: + if value is None: + return None + if not isinstance(value, str): + raise RealtimeTranscriptionProtocolError(f"{name} must be a string") + return value + + +def json_integer(value: JsonValue | None, name: str) -> 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") + return value + + +def new_event_id() -> str: + return f"event_{uuid.uuid4().hex}" + + +def parse_transcription_session_update(payload: str) -> TranscriptionSessionUpdate: + message: Final = json_object(payload) + if message.get("type") not in SESSION_UPDATE_EVENT_TYPES: + raise RealtimeTranscriptionProtocolError("expected session.update") + session: Final = json_mapping(message.get("session"), "session") + 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") + 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") + transcription: Final = json_mapping( + beta_transcription if beta_transcription is not None else ga_transcription, + "input audio transcription", + ) + 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"), + 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_disabled=turn_detection_present and turn_detection is None, + ) + + +def _parse_audio_format( + session: Mapping[str, JsonValue], audio_input: Mapping[str, JsonValue] +) -> 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") + if beta_format is not None: + return TranscriptionAudioFormat( + layout="beta", + encoding=json_string(beta_format, "session.input_audio_format"), + rate=None, + channels=None, + ) + if ga_format is None: + 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") + 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"), + ) + + +def decode_pcm16_append(audio: JsonValue | None, max_encoded_bytes: int | None = None) -> bytes: + if not isinstance(audio, str): + raise RealtimeTranscriptionProtocolError("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") + try: + decoded: Final = base64.b64decode(audio, validate=True) + except (binascii.Error, ValueError): + raise RealtimeTranscriptionProtocolError("Audio must be valid base64") from None + if len(decoded) % 2: + raise RealtimeTranscriptionProtocolError("PCM16 audio must contain complete samples") + return decoded + + +def _transcription_settings(model: str, language: str | None) -> OpenAIRealtimeTranscriptionSettings: + if language is None: + model_only: Final[OpenAIRealtimeTranscriptionSettings] = {"model": model} + return model_only + with_language: Final[OpenAIRealtimeTranscriptionSettings] = {"model": model, "language": language} + return with_language + + +def transcription_session( + *, session_id: str, model: str, sample_rate: int, language: str | None, server_vad: bool +) -> OpenAIRealtimeTranscriptionSession: + settings: Final = _transcription_settings(model, language) + session: Final[OpenAIRealtimeTranscriptionSession] = { + "id": session_id, + "object": "realtime.transcription_session", + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": sample_rate}, + "transcription": settings, + "turn_detection": SERVER_VAD_TURN_DETECTION if server_vad else None, + } + }, + } + return session + + +def transcription_session_created_event( + session: OpenAIRealtimeTranscriptionSession, +) -> OpenAIRealtimeTranscriptionSessionCreated: + event: Final[OpenAIRealtimeTranscriptionSessionCreated] = { + "type": "session.created", + "event_id": new_event_id(), + "session": session, + } + return event + + +def error_event(message: str) -> OpenAIRealtimeErrorEvent: + event: Final[OpenAIRealtimeErrorEvent] = { + "type": "error", + "error": {"type": "server_error", "message": message}, + } + return event + + +def speech_event( + event_type: Literal["input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped"], item_id: str +) -> OpenAIRealtimeInputAudioBufferSpeechEvent: + event: Final[OpenAIRealtimeInputAudioBufferSpeechEvent] = { + "type": event_type, + "event_id": new_event_id(), + "item_id": item_id, + } + return event + + +def delta_event(item_id: str, delta: str) -> OpenAIRealtimeInputAudioTranscriptionDelta: + event: Final[OpenAIRealtimeInputAudioTranscriptionDelta] = { + "type": "conversation.item.input_audio_transcription.delta", + "event_id": new_event_id(), + "item_id": item_id, + "content_index": 0, + "delta": delta, + } + return event + + +def completed_event( + item_id: str, transcript: str, usage: RealtimeInputAudioTranscriptionUsage | None +) -> OpenAIRealtimeInputAudioTranscriptionCompleted: + event: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = { + "type": "conversation.item.input_audio_transcription.completed", + "event_id": new_event_id(), + "item_id": item_id, + "content_index": 0, + "transcript": transcript, + } + if usage is None: + return event + billed: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = {**event, "usage": usage} + return billed + + +def duration_usage(seconds: float) -> RealtimeInputAudioTranscriptionUsage: + usage: Final[RealtimeInputAudioTranscriptionDurationUsage] = {"type": "duration", "seconds": seconds} + return usage diff --git a/litellm/llms/base_llm/realtime/transformation.py b/litellm/llms/base_llm/realtime/transformation.py index e44cccc1a62..e1b16a5985d 100644 --- a/litellm/llms/base_llm/realtime/transformation.py +++ b/litellm/llms/base_llm/realtime/transformation.py @@ -1,6 +1,7 @@ from abc import ABC, abstractmethod from collections.abc import Mapping, Sequence -from typing import TYPE_CHECKING, Any +from types import TracebackType +from typing import TYPE_CHECKING, Any, Protocol, Self import httpx @@ -21,6 +22,23 @@ else: LiteLLMLoggingObj = Any +class RealtimeBackend(Protocol): + async def __aenter__(self) -> Self: ... + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: ... + + async def send(self, message: str | bytes) -> None: ... + + async def recv(self, decode: bool | None = None) -> str | bytes: ... + + async def close(self) -> None: ... + + class BaseRealtimeConfig(ABC): @abstractmethod def validate_environment( @@ -78,6 +96,9 @@ class BaseRealtimeConfig(ABC): def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None: return None + async def open_backend(self, url: str, headers: Mapping[str, str]) -> RealtimeBackend | None: + return None + def transform_session_created_event( self, model: str, diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 857adf5b9f1..75ecfed1044 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -6180,7 +6180,12 @@ class BaseLLMHTTPHandler: ssl_context = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT) ssl_context.check_hostname = False ssl_context.verify_mode = ssl.CERT_NONE - backend_ws: Final = await self._open_realtime_backend_ws(websockets, url, headers, ssl_context) + provider_backend: Final = await provider_config.open_backend(url, headers) + backend_ws: Final = ( + provider_backend + if provider_backend is not None + else await self._open_realtime_backend_ws(websockets, url, headers, ssl_context) + ) async with backend_ws: _request_data: Final[dict[str, object]] = {} if litellm_metadata: diff --git a/litellm/llms/meta/realtime/transformation.py b/litellm/llms/meta/realtime/transformation.py index 1b8943f0cee..f0ca448bea0 100644 --- a/litellm/llms/meta/realtime/transformation.py +++ b/litellm/llms/meta/realtime/transformation.py @@ -1,36 +1,41 @@ import asyncio -import base64 -import binascii import json import math import time from collections.abc import Awaitable, Callable, Iterator, Mapping from dataclasses import dataclass from types import MappingProxyType -from typing import Final, Literal +from typing import Final from urllib.parse import urlparse, urlunparse -from pydantic import JsonValue, TypeAdapter, ValidationError +from pydantic import JsonValue from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.realtime.transcription_protocol import ( + RealtimeTranscriptionProtocolError, + TranscriptionSessionUpdate, + completed_event, + decode_pcm16_append, + delta_event, + duration_usage, + error_event, + json_object, + parse_transcription_session_update, + speech_event, + transcription_session, + transcription_session_created_event, +) from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig from litellm.secret_managers.main import get_secret_str from litellm.types.llms.meta import MuseAudioEncoding, MuseHandshake, MuseMode, MuseSampleRate from litellm.types.llms.openai import ( - OpenAIRealtimeErrorEvent, OpenAIRealtimeEvents, - OpenAIRealtimeInputAudioBufferSpeechEvent, - OpenAIRealtimeInputAudioTranscriptionCompleted, - OpenAIRealtimeInputAudioTranscriptionDelta, - OpenAIRealtimeServerVadTurnDetection, OpenAIRealtimeTranscriptionSession, OpenAIRealtimeTranscriptionSessionCreated, - OpenAIRealtimeTranscriptionSettings, ) from litellm.types.realtime import ( - RealtimeInputAudioTranscriptionDurationUsage, RealtimeInputAudioTranscriptionUsage, RealtimeResponseTransformInput, RealtimeResponseTypedDict, @@ -98,17 +103,13 @@ _LANGUAGE_CODES: Final = MappingProxyType( "zh": "Mandarin Chinese", } ) -_SUPPORTED_TRANSCRIPTION_KEYS: Final = frozenset(("model", "language")) _MAX_AUDIO_BACKLOG_SECONDS: Final = 4 _PACKET_MS: Final = 80 _END_STREAM: Final = '{"type":"endStream"}' _PROVIDER_ERROR_MESSAGE: Final = "Meta Muse realtime transcription failed" -_JSON_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) -_EMPTY_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) -_SERVER_VAD: Final[OpenAIRealtimeServerVadTurnDetection] = {"type": "server_vad"} -class MuseProtocolError(ValueError): +class MuseProtocolError(RealtimeTranscriptionProtocolError): pass @@ -150,26 +151,13 @@ class MuseSessionConfig: return biased def openai_session(self, session_id: str) -> OpenAIRealtimeTranscriptionSession: - session: Final[OpenAIRealtimeTranscriptionSession] = { - "id": session_id, - "object": "realtime.transcription_session", - "type": "transcription", - "audio": { - "input": { - "format": {"type": "audio/pcm", "rate": self.sample_rate}, - "transcription": self._transcription_settings(), - "turn_detection": None if self.mode == "PUSH_TO_TALK" else _SERVER_VAD, - } - }, - } - return session - - def _transcription_settings(self) -> OpenAIRealtimeTranscriptionSettings: - base: Final[OpenAIRealtimeTranscriptionSettings] = {"model": self.model} - if not self.language_bias: - return base - localized: Final[OpenAIRealtimeTranscriptionSettings] = {**base, "language": self.language_bias[0]} - return localized + return transcription_session( + session_id=session_id, + model=self.model, + sample_rate=self.sample_rate, + language=self.language_bias[0] if self.language_bias else None, + server_vad=self.mode != "PUSH_TO_TALK", + ) _DEFAULT_SESSION_CONFIG: Final = MuseSessionConfig( @@ -177,40 +165,10 @@ _DEFAULT_SESSION_CONFIG: Final = MuseSessionConfig( ) -def _json_object(payload: str) -> Mapping[str, JsonValue]: - try: - value: Final = _JSON_ADAPTER.validate_json(payload) - except ValidationError: - raise MuseProtocolError("invalid JSON object") from None - if not isinstance(value, dict): - raise MuseProtocolError("message must be a JSON object") - return value - - -def _mapping(value: JsonValue | None, name: str) -> Mapping[str, JsonValue]: - if value is None: - return _EMPTY_OBJECT - if not isinstance(value, dict): - raise MuseProtocolError(f"{name} must be an object") - return value - - -def _string(value: JsonValue | None, name: str) -> str | None: - if value is None: - return None - if not isinstance(value, str): - raise MuseProtocolError(f"{name} must be a string") - return value - - def _normalize_model(model: str) -> str: return model.removeprefix("meta/").strip() -def _event_id() -> str: - return f"event_{uuid.uuid4().hex}" - - def normalize_language(language: str) -> str: value: Final = language.strip() if not value: @@ -254,138 +212,55 @@ def build_muse_realtime_url(api_base: str | None) -> str: return urlunparse((scheme, netloc, "/v1/asr/realtime", "", "", "")) -def _parse_sample_rate(session: Mapping[str, JsonValue]) -> MuseSampleRate: - beta_format: Final = session.get("input_audio_format") - audio: Final = _mapping(session.get("audio"), "session.audio") - audio_input: Final = _mapping(audio.get("input"), "session.audio.input") - ga_format: Final = audio_input.get("format") - if beta_format is not None and ga_format is not None: - raise MuseProtocolError("input audio format must use either beta or GA layout") - if beta_format is not None: - if beta_format != "pcm16": +def _parse_sample_rate(update: TranscriptionSessionUpdate) -> MuseSampleRate: + audio_format: Final = update.audio_format + if audio_format is None: + return 24_000 + if audio_format.layout == "beta": + if audio_format.encoding != "pcm16": raise MuseProtocolError("Muse Voice requires pcm16 input audio") return 24_000 - if ga_format is None: - return 24_000 - if isinstance(ga_format, str): - if ga_format != "pcm16": - raise MuseProtocolError("Muse Voice requires audio/pcm input audio") - return 24_000 - format_mapping: Final = _mapping(ga_format, "session.audio.input.format") - if format_mapping.get("type") != "audio/pcm": + if not audio_format.is_pcm16: raise MuseProtocolError("Muse Voice requires audio/pcm input audio") - channels: Final = format_mapping.get("channels", 1) - if isinstance(channels, bool) or channels != 1: + if audio_format.channels not in (None, 1): raise MuseProtocolError("Muse Voice requires mono input audio") - rate: Final = format_mapping.get("rate", 24_000) - if isinstance(rate, bool) or not isinstance(rate, int) or rate not in SUPPORTED_SAMPLE_RATES: + rate: Final = 24_000 if audio_format.rate is None else audio_format.rate + if rate not in SUPPORTED_SAMPLE_RATES: raise MuseProtocolError("Muse Voice supports PCM16 at 16000 Hz or 24000 Hz") return 16_000 if rate == 16_000 else 24_000 -def _parse_mode(session: Mapping[str, JsonValue], audio_input: Mapping[str, JsonValue]) -> MuseMode: - 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")) - if turn_detection_present and turn_detection is None: +def _parse_mode(update: TranscriptionSessionUpdate) -> MuseMode: + if update.turn_detection_disabled: return "PUSH_TO_TALK" - if turn_detection is None: - return "ENDPOINTING" - turn_detection_mapping: Final = _mapping(turn_detection, "turn_detection") - if turn_detection_mapping.get("type") not in (None, "server_vad"): + if update.turn_detection_type not in (None, "server_vad"): raise MuseProtocolError("Muse Voice supports server_vad turn detection or null") return "ENDPOINTING" def parse_session_update(payload: str, expected_model: str) -> MuseSessionConfig: - message: Final = _json_object(payload) - if message.get("type") not in ("session.update", "transcription_session.update"): - raise MuseProtocolError("expected session.update") - session: Final = _mapping(message.get("session"), "session") - if not session: - raise MuseProtocolError("session.update requires a session object") - if session.get("type") not in (None, "transcription", "realtime"): + update: Final = parse_transcription_session_update(payload) + if update.session_type not in (None, "transcription", "realtime"): raise MuseProtocolError("Muse Voice supports transcription sessions only") - audio: Final = _mapping(session.get("audio"), "session.audio") - audio_input: Final = _mapping(audio.get("input"), "session.audio.input") - 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 MuseProtocolError("input transcription must use either beta or GA layout") - transcription: Final = _mapping( - beta_transcription if beta_transcription is not None else ga_transcription, - "input audio transcription", - ) - unsupported: Final = tuple(sorted(key for key in transcription if key not in _SUPPORTED_TRANSCRIPTION_KEYS)) - if unsupported: - verbose_logger.warning("Meta realtime: dropping unsupported transcription settings %s", unsupported) - requested_model: Final = _string(transcription.get("model"), "transcription model") + if update.unsupported_transcription_keys: + verbose_logger.warning( + "Meta realtime: dropping unsupported transcription settings %s", update.unsupported_transcription_keys + ) normalized_model: Final = _normalize_model(expected_model) if normalized_model != MUSE_MODEL: raise MuseProtocolError("unsupported Meta realtime model") - if requested_model is not None and _normalize_model(requested_model) != normalized_model: + if update.model is not None and _normalize_model(update.model) != normalized_model: raise MuseProtocolError("realtime session model cannot be changed") - language: Final = _string(transcription.get("language"), "language") return MuseSessionConfig( model=normalized_model, - mode=_parse_mode(session, audio_input), - sample_rate=_parse_sample_rate(session), - language_bias=() if language is None else (normalize_language(language),), + mode=_parse_mode(update), + sample_rate=_parse_sample_rate(update), + language_bias=() if update.language is None else (normalize_language(update.language),), ) def session_created_event(config: MuseSessionConfig, session_id: str) -> OpenAIRealtimeTranscriptionSessionCreated: - event: Final[OpenAIRealtimeTranscriptionSessionCreated] = { - "type": "session.created", - "event_id": _event_id(), - "session": config.openai_session(session_id), - } - return event - - -def error_event(message: str) -> OpenAIRealtimeErrorEvent: - event: Final[OpenAIRealtimeErrorEvent] = { - "type": "error", - "error": {"type": "server_error", "message": message}, - } - return event - - -def _speech_event( - event_type: Literal["input_audio_buffer.speech_started", "input_audio_buffer.speech_stopped"], item_id: str -) -> OpenAIRealtimeInputAudioBufferSpeechEvent: - event: Final[OpenAIRealtimeInputAudioBufferSpeechEvent] = { - "type": event_type, - "event_id": _event_id(), - "item_id": item_id, - } - return event - - -def _delta_event(item_id: str, delta: str) -> OpenAIRealtimeInputAudioTranscriptionDelta: - event: Final[OpenAIRealtimeInputAudioTranscriptionDelta] = { - "type": "conversation.item.input_audio_transcription.delta", - "event_id": _event_id(), - "item_id": item_id, - "content_index": 0, - "delta": delta, - } - return event - - -def _completed_event( - item_id: str, transcript: str, usage: RealtimeInputAudioTranscriptionUsage | None -) -> OpenAIRealtimeInputAudioTranscriptionCompleted: - event: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = { - "type": "conversation.item.input_audio_transcription.completed", - "event_id": _event_id(), - "item_id": item_id, - "content_index": 0, - "transcript": transcript, - } - if usage is None: - return event - billed: Final[OpenAIRealtimeInputAudioTranscriptionCompleted] = {**event, "usage": usage} - return billed + return transcription_session_created_event(config.openai_session(session_id)) def _required_turn_id(message: Mapping[str, JsonValue], event: str) -> str: @@ -424,18 +299,18 @@ class _TurnState: has_content: Final = self.latest_partial is not None or self.final_text is not None if (self.started or has_content) and not self.start_emitted: self.start_emitted = True - yield _speech_event("input_audio_buffer.speech_started", self.item_id) + yield speech_event("input_audio_buffer.speech_started", self.item_id) if self.latest_partial is not None and self.final_text is None: delta: Final = _new_suffix(self.emitted_partial, self.latest_partial) if delta: self.emitted_partial = self.latest_partial - yield _delta_event(self.item_id, delta) + yield delta_event(self.item_id, delta) if self.stopped and not self.stopped_emitted: self.stopped_emitted = True - yield _speech_event("input_audio_buffer.speech_stopped", self.item_id) + yield speech_event("input_audio_buffer.speech_stopped", self.item_id) if self.final_text is not None and self.stopped_emitted and not self.completed_emitted: self.completed_emitted = True - yield _completed_event(self.item_id, self.final_text, take_usage()) + yield completed_event(self.item_id, self.final_text, take_usage()) class MuseEventTransformer: @@ -467,8 +342,7 @@ class MuseEventTransformer: if seconds <= 0: return None self._unbilled_seconds = 0.0 - usage: Final[RealtimeInputAudioTranscriptionDurationUsage] = {"type": "duration", "seconds": seconds} - return usage + return duration_usage(seconds) def _apply_turn_event(self, event_type: JsonValue | None, message: Mapping[str, JsonValue]) -> _TurnState | None: match event_type: @@ -612,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) event_type: Final = request.get("type") if event_type in ("session.update", "transcription_session.update"): return self._configure(message, model) @@ -664,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) session_id: Final = frame.get("sessionId") if session_id is None: return self._transformer.transform(frame) @@ -686,17 +560,7 @@ class MetaRealtimeConfig(BaseRealtimeConfig): def _append_audio(self, request: Mapping[str, JsonValue]) -> tuple[bytes, ...]: config: Final = self._require_config() - encoded: Final = request.get("audio") - if not isinstance(encoded, str): - raise MuseProtocolError("Audio must be a base64 string") - if len(encoded) > config.max_encoded_append_bytes: - raise MuseProtocolError("Audio append exceeds the four-second backlog limit") - try: - audio: Final = base64.b64decode(encoded, validate=True) - except (binascii.Error, ValueError): - raise MuseProtocolError("Audio must be valid base64") from None - if len(audio) % 2: - raise MuseProtocolError("PCM16 audio must contain complete samples") + audio: Final = decode_pcm16_append(request.get("audio"), config.max_encoded_append_bytes) 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 new file mode 100644 index 00000000000..0ecbe7209f6 --- /dev/null +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py @@ -0,0 +1,316 @@ +import asyncio +import time +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 pydantic import TypeAdapter +from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK +from websockets.frames import Close + +from litellm import verbose_logger +from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import SpeechStreamingTarget +from litellm.types.llms.vertex_ai_speech_to_text import ( + VertexSpeechStreamingCommand, + VertexSpeechStreamingCommandUnion, + VertexSpeechStreamingConfigure, + VertexSpeechStreamingConfigured, + VertexSpeechStreamingDiscardTurn, + VertexSpeechStreamingFinishTurn, + VertexSpeechStreamingResponse, + VertexSpeechStreamingResult, + VertexSpeechStreamingTurnDiscarded, + VertexSpeechStreamingTurnFinished, +) + +if TYPE_CHECKING: + from google.cloud.speech_v2.types import ( + StreamingRecognitionConfig, + StreamingRecognizeRequest, + StreamingRecognizeResponse, + ) + +SPEECH_SDK_INSTALL_HINT: Final = ( + "google-cloud-speech is not installed. Install with `pip install 'litellm[stt-vertex-chirp]'`." +) +STREAM_FAILURE_CLOSE_CODE: Final = 1011 +STREAM_ROTATION_SECONDS: Final = 240.0 +_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( + { + "SPEECH_ACTIVITY_BEGIN": "begin", + "SPEECH_ACTIVITY_END": "end", + "END_OF_SINGLE_UTTERANCE": "end", + } +) + + +class ClosableTransport(Protocol): + def close(self) -> Awaitable[None]: ... + + +class SpeechStreamingClient(Protocol): + def streaming_recognize( + self, requests: "AsyncIterator[StreamingRecognizeRequest] | None" = None + ) -> "Awaitable[AsyncIterable[StreamingRecognizeResponse]]": ... + + @property + def transport(self) -> ClosableTransport: ... + + +@dataclass(frozen=True, slots=True) +class _StreamFailure: + reason: str + + +@dataclass(frozen=True, slots=True) +class _Closed: + pass + + +def open_speech_client(target: SpeechStreamingTarget) -> SpeechStreamingClient: + try: + from google.api_core.client_options import ClientOptions + from google.cloud.speech_v2 import SpeechAsyncClient + from google.oauth2.credentials import Credentials + except ImportError as e: + raise ImportError(SPEECH_SDK_INSTALL_HINT) from e + return SpeechAsyncClient( + credentials=Credentials(token=target.access_token), + transport="grpc_asyncio", + client_options=ClientOptions(api_endpoint=target.api_endpoint), + ) + + +def _streaming_config(command: VertexSpeechStreamingConfigure) -> "StreamingRecognitionConfig": + from google.cloud.speech_v2.types import ( + ExplicitDecodingConfig, + RecognitionConfig, + StreamingRecognitionConfig, + StreamingRecognitionFeatures, + ) + + return StreamingRecognitionConfig( + config=RecognitionConfig( + explicit_decoding_config=ExplicitDecodingConfig( + encoding=ExplicitDecodingConfig.AudioEncoding.LINEAR16, + sample_rate_hertz=command.sample_rate_hertz, + audio_channel_count=1, + ), + model=command.model, + language_codes=command.language_codes, + ), + streaming_features=StreamingRecognitionFeatures(interim_results=True, enable_voice_activity_events=True), + ) + + +def _response_event(response: "StreamingRecognizeResponse", billed_seconds: float) -> str: + return VertexSpeechStreamingResponse( + speech_event=_SPEECH_EVENTS.get(response.speech_event_type.name, "none"), + results=tuple( + VertexSpeechStreamingResult( + transcript=result.alternatives[0].transcript if result.alternatives else "", + is_final=result.is_final, + ) + for result in response.results + ), + billed_seconds=billed_seconds, + ).model_dump_json() + + +def _billed_seconds(response: "StreamingRecognizeResponse") -> float: + return _TIMEDELTA_ADAPTER.validate_python(response.metadata.total_billed_duration).total_seconds() + + +class _RecognizeStream: + def __init__( + self, + *, + 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.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()) + + @property + def billed_seconds(self) -> float: + return self._base_billed_seconds + self._billed_seconds + + def send_audio(self, audio: bytes) -> None: + self._requests.put_nowait(self._request_type(audio=audio)) + + def half_close(self) -> None: + self._requests.put_nowait(None) + + async def wait(self) -> None: + await asyncio.gather(self._task, return_exceptions=True) + + 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 + 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 + 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}")) + + async def _drain(self) -> "AsyncIterator[StreamingRecognizeRequest]": + while (request := await self._requests.get()) is not None: + yield request + + +class SpeechStreamingBackend: + def __init__( + self, + target: SpeechStreamingTarget, + *, + client_factory: Callable[[SpeechStreamingTarget], SpeechStreamingClient] = open_speech_client, + clock: Callable[[], float] = time.monotonic, + rotation_seconds: float = STREAM_ROTATION_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._client: SpeechStreamingClient | None = None + self._config: StreamingRecognitionConfig | None = None + self._stream: _RecognizeStream | None = None + self._last_stream: _RecognizeStream | None = None + + async def __aenter__(self) -> Self: + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + await self.close() + + async def send(self, message: str | bytes) -> None: + if isinstance(message, bytes): + 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) + case VertexSpeechStreamingFinishTurn(): + self._finish_turn() + case VertexSpeechStreamingDiscardTurn(): + await self._discard_turn() + + async def recv(self, decode: bool | None = None) -> str | bytes: + item: Final = await self._outbox.get() + match item: + case _StreamFailure(): + raise ConnectionClosedError( + rcvd=Close(STREAM_FAILURE_CLOSE_CODE, item.reason[:_CLOSE_REASON_MAX_CHARS]), sent=None + ) + case _Closed(): + raise ConnectionClosedOK(rcvd=Close(1000, ""), sent=None) + case str(): + return 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() + client: Final = self._client + self._client = None + if client is not None: + await client.transport.close() + 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) + + 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() + + def _open_stream(self) -> _RecognizeStream: + from google.cloud.speech_v2.types import StreamingRecognizeRequest + + config: Final = self._config + if config is None: + raise RuntimeError("audio was sent before the Speech-to-Text stream was configured") + if self._client is None: + self._client = self._client_factory(self._target) + stream: Final = _RecognizeStream( + 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 + 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 _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() diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py new file mode 100644 index 00000000000..5e2857f24c3 --- /dev/null +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py @@ -0,0 +1,433 @@ +from collections.abc import Callable, Mapping +from dataclasses import dataclass, replace +from typing import Final + +from pydantic import JsonValue, TypeAdapter + +from litellm import verbose_logger +from litellm._uuid import uuid +from litellm.litellm_core_utils.audio_utils.utils import normalize_transcription_language_to_bcp47 +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.llms.base_llm.realtime.transcription_protocol import ( + RealtimeTranscriptionProtocolError, + TranscriptionAudioFormat, + TranscriptionSessionUpdate, + completed_event, + decode_pcm16_append, + delta_event, + duration_usage, + json_object, + parse_transcription_session_update, + speech_event, + transcription_session, + transcription_session_created_event, +) +from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig, RealtimeBackend +from litellm.llms.vertex_ai.audio_transcription.transformation import ( + AUTO_LANGUAGE_CODE, + DEFAULT_SPEECH_TO_TEXT_LOCATION, + speech_to_text_host, + validate_vertex_transcription_location, + validate_vertex_transcription_project_id, +) +from litellm.types.llms.openai import ( + OpenAIRealtimeEvents, + OpenAIRealtimeTranscriptionSession, + OpenAIRealtimeTranscriptionSessionCreated, +) +from litellm.types.llms.vertex_ai_speech_to_text import ( + VertexSpeechStreamingConfigure, + VertexSpeechStreamingConfigured, + VertexSpeechStreamingDiscardTurn, + VertexSpeechStreamingEvent, + VertexSpeechStreamingEventUnion, + VertexSpeechStreamingFinishTurn, + VertexSpeechStreamingResponse, + VertexSpeechStreamingTurnDiscarded, + VertexSpeechStreamingTurnFinished, +) +from litellm.types.realtime import ( + RealtimeInputAudioTranscriptionUsage, + RealtimeResponseTransformInput, + RealtimeResponseTypedDict, +) + +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" +_VERTEX_MODEL_PREFIX: Final = "vertex_ai/" +_STREAMING_EVENT_ADAPTER: Final = TypeAdapter[VertexSpeechStreamingEventUnion](VertexSpeechStreamingEvent) +_FINISH_TURN_COMMAND: Final = VertexSpeechStreamingFinishTurn().model_dump_json() +_DISCARD_TURN_COMMAND: Final = VertexSpeechStreamingDiscardTurn().model_dump_json() + + +class ChirpProtocolError(RealtimeTranscriptionProtocolError): + pass + + +@dataclass(frozen=True, slots=True) +class SpeechStreamingTarget: + api_endpoint: str + recognizer: str + access_token: str + + +@dataclass(frozen=True, slots=True) +class ChirpSessionConfig: + model: str + language: str | None + sample_rate: int + server_vad: bool + + def openai_session(self, session_id: str) -> OpenAIRealtimeTranscriptionSession: + return transcription_session( + session_id=session_id, + model=self.model, + sample_rate=self.sample_rate, + language=self.language, + server_vad=self.server_vad, + ) + + def configure_command(self) -> str: + return VertexSpeechStreamingConfigure( + model=self.model, + language_codes=(AUTO_LANGUAGE_CODE,) if self.language is None else (self.language,), + sample_rate_hertz=self.sample_rate, + ).model_dump_json() + + +def is_vertex_speech_to_text_model(model: str) -> bool: + return normalize_speech_to_text_model(model).startswith(SPEECH_TO_TEXT_MODEL_PREFIX) + + +def normalize_speech_to_text_model(model: str) -> str: + return model.removeprefix(_VERTEX_MODEL_PREFIX) + + +def default_session_config(model: str) -> ChirpSessionConfig: + return ChirpSessionConfig( + model=normalize_speech_to_text_model(model), + language=None, + sample_rate=DEFAULT_SAMPLE_RATE_HERTZ, + server_vad=True, + ) + + +def parse_chirp_session_update(payload: str, expected_model: str) -> ChirpSessionConfig: + update: Final = parse_transcription_session_update(payload) + if update.session_type not in (None, "transcription", "realtime"): + raise ChirpProtocolError("Speech-to-Text streaming supports transcription sessions only") + if update.unsupported_transcription_keys: + verbose_logger.debug( + "Speech-to-Text streaming: ignoring unsupported transcription settings %s", + update.unsupported_transcription_keys, + ) + model: Final = normalize_speech_to_text_model(expected_model) + if update.model is not None and normalize_speech_to_text_model(update.model) != model: + raise ChirpProtocolError("realtime session model cannot be changed") + return ChirpSessionConfig( + model=model, + language=None if update.language is None else normalize_transcription_language_to_bcp47(update.language), + sample_rate=_parse_sample_rate(update.audio_format), + server_vad=_parse_server_vad(update), + ) + + +def _parse_sample_rate(audio_format: TranscriptionAudioFormat | None) -> int: + if audio_format is None: + return DEFAULT_SAMPLE_RATE_HERTZ + if not audio_format.is_pcm16: + raise ChirpProtocolError("Speech-to-Text streaming requires pcm16 input audio") + if audio_format.channels not in (None, 1): + raise ChirpProtocolError("Speech-to-Text streaming requires mono input audio") + rate: Final = DEFAULT_SAMPLE_RATE_HERTZ if audio_format.rate is None else audio_format.rate + if not MIN_SAMPLE_RATE_HERTZ <= rate <= MAX_SAMPLE_RATE_HERTZ: + raise ChirpProtocolError( + f"Speech-to-Text streaming supports sample rates from {MIN_SAMPLE_RATE_HERTZ} Hz" + f" to {MAX_SAMPLE_RATE_HERTZ} Hz" + ) + return rate + + +def _parse_server_vad(update: TranscriptionSessionUpdate) -> bool: + if update.turn_detection_disabled: + return False + if update.turn_detection_type not in (None, "server_vad"): + raise ChirpProtocolError("Speech-to-Text streaming supports server_vad turn detection or null") + return True + + +def session_created_event(config: ChirpSessionConfig, session_id: str) -> OpenAIRealtimeTranscriptionSessionCreated: + return transcription_session_created_event(config.openai_session(session_id)) + + +def _normalize_word(word: str) -> str: + return "".join(char for char in word if char.isalnum()).casefold() + + +def new_words(previous: str, current: str) -> str: + previous_words: Final = previous.split() + current_words: Final = current.split() + common: Final = next( + ( + index + for index, (old, new) in enumerate(zip(previous_words, current_words, strict=False)) + if _normalize_word(old) != _normalize_word(new) + ), + min(len(previous_words), len(current_words)), + ) + appended: Final = " ".join(current_words[common:]) + if not appended: + return "" + return f" {appended}" if common else appended + + +def _join_transcript(committed: str, tail: str) -> str: + return " ".join(part for part in (committed, tail) if part) + + +@dataclass(frozen=True, slots=True) +class _Turn: + item_id: str + committed: str = "" + preview: str = "" + started_emitted: bool = False + stopped_emitted: bool = False + + +class ChirpEventTransformer: + def __init__(self, *, new_item_id: Callable[[], str] = lambda: f"item_{uuid.uuid4().hex}") -> None: + self._new_item_id: Final = new_item_id + self._config: ChirpSessionConfig | None = None + self._session_id: str | None = None + self._turn: _Turn | None = None + self._billed_seconds: float = 0.0 + self._reported_seconds: float = 0.0 + + def configure(self, config: ChirpSessionConfig, session_id: str) -> None: + self._config = config + self._session_id = session_id + + def take_unbilled_usage(self) -> RealtimeInputAudioTranscriptionUsage | None: + unreported: Final = self._billed_seconds - self._reported_seconds + if unreported <= 0: + return None + self._reported_seconds = self._billed_seconds + return duration_usage(unreported) + + def transform(self, frame: VertexSpeechStreamingEventUnion) -> tuple[OpenAIRealtimeEvents, ...]: + match frame: + case VertexSpeechStreamingConfigured(): + return (session_created_event(self._require_config(), self._require_session_id()),) + case VertexSpeechStreamingResponse(): + return self._response(frame) + case VertexSpeechStreamingTurnFinished(): + return self._finish_turn() + case VertexSpeechStreamingTurnDiscarded(): + self._turn = None + return () + + def _response(self, frame: VertexSpeechStreamingResponse) -> tuple[OpenAIRealtimeEvents, ...]: + self._billed_seconds = max(self._billed_seconds, frame.billed_seconds) + interim: Final = " ".join( + result.transcript.strip() for result in frame.results if not result.is_final and result.transcript.strip() + ) + finals: Final = tuple( + result.transcript.strip() for result in frame.results if result.is_final and result.transcript.strip() + ) + begin_events: Final = self._begin() if frame.speech_event == "begin" or interim or finals else () + interim_events: Final = self._hypothesis(interim) if interim else () + final_events: Final = tuple(event for final in finals for event in self._final(final)) + end_events: Final = self._stop() if frame.speech_event == "end" else () + return (*begin_events, *interim_events, *final_events, *end_events) + + def _begin(self) -> tuple[OpenAIRealtimeEvents, ...]: + if self._turn is None: + self._turn = _Turn(item_id=self._new_item_id()) + turn: Final = self._turn + if turn.started_emitted or not self._require_config().server_vad: + return () + self._turn = replace(turn, started_emitted=True) + return (speech_event("input_audio_buffer.speech_started", turn.item_id),) + + def _stop(self) -> tuple[OpenAIRealtimeEvents, ...]: + turn: Final = self._turn + if turn is None or turn.stopped_emitted or not self._require_config().server_vad: + return () + self._turn = replace(turn, stopped_emitted=True) + return (speech_event("input_audio_buffer.speech_stopped", turn.item_id),) + + def _hypothesis(self, text: str) -> tuple[OpenAIRealtimeEvents, ...]: + turn: Final = self._require_turn() + hypothesis: Final = _join_transcript(turn.committed, text) + delta: Final = new_words(turn.preview, hypothesis) + self._turn = replace(turn, preview=hypothesis) + return (delta_event(turn.item_id, delta),) if delta else () + + def _final(self, text: str) -> tuple[OpenAIRealtimeEvents, ...]: + turn: Final = self._require_turn() + committed: Final = _join_transcript(turn.committed, text) + delta: Final = new_words(turn.preview, committed) + self._turn = replace(turn, committed=committed, preview=committed) + delta_events: Final[tuple[OpenAIRealtimeEvents, ...]] = (delta_event(turn.item_id, delta),) if delta else () + if not self._require_config().server_vad: + return delta_events + return (*delta_events, *self._complete()) + + def _finish_turn(self) -> tuple[OpenAIRealtimeEvents, ...]: + if self._turn is None: + return () + return self._complete() + + def _complete(self) -> tuple[OpenAIRealtimeEvents, ...]: + turn: Final = self._require_turn() + stop_events: Final = self._stop() + transcript: Final = turn.committed or turn.preview + self._turn = None + return (*stop_events, completed_event(turn.item_id, transcript, self.take_unbilled_usage())) + + def _require_turn(self) -> _Turn: + if self._turn is None: + self._turn = _Turn(item_id=self._new_item_id()) + return self._turn + + def _require_config(self) -> ChirpSessionConfig: + if self._config is None: + raise ChirpProtocolError("session.update must configure the session before the backend responds") + return self._config + + def _require_session_id(self) -> str: + if self._session_id is None: + raise ChirpProtocolError("session.update must configure the session before the backend responds") + return self._session_id + + +def _default_backend_factory(target: SpeechStreamingTarget) -> RealtimeBackend: + from litellm.llms.vertex_ai.audio_transcription.realtime_backend import SpeechStreamingBackend + + return SpeechStreamingBackend(target) + + +class VertexChirpRealtimeConfig(BaseRealtimeConfig): + def __init__( + self, + *, + access_token: str, + project: str, + location: str | None, + backend_factory: Callable[[SpeechStreamingTarget], RealtimeBackend] = _default_backend_factory, + ) -> None: + self._access_token: Final = access_token + self._project: Final = validate_vertex_transcription_project_id(project) + self._location: Final = validate_vertex_transcription_location(location, DEFAULT_SPEECH_TO_TEXT_LOCATION) + self._backend_factory: Final = backend_factory + self._transformer: Final = ChirpEventTransformer() + self._config: ChirpSessionConfig | None = None + self._session_id: str | None = None + + def validate_environment( + self, + headers: dict[str, str], # mutable-ok: BaseRealtimeConfig contract + model: str, + api_key: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseRealtimeConfig contract + return headers + + def get_complete_url(self, api_base: str | None, model: str, api_key: str | None = None) -> str: + if not is_vertex_speech_to_text_model(model): + raise ValueError(f"Unsupported Speech-to-Text streaming model: {model}") + return _api_endpoint(api_base) if api_base else speech_to_text_host(self._location) + + async def open_backend(self, url: str, headers: Mapping[str, str]) -> RealtimeBackend | None: + return self._backend_factory( + SpeechStreamingTarget( + api_endpoint=url, + recognizer=f"projects/{self._project}/locations/{self._location}/recognizers/_", + access_token=self._access_token, + ) + ) + + def is_setup_message(self, msg_obj: Mapping[str, object]) -> bool: + return msg_obj.get("kind") == "configure" + + def transform_session_created_event( + self, + model: str, + logging_session_id: str, + session_configuration_request: str | None = None, + ) -> OpenAIRealtimeTranscriptionSessionCreated: + self._session_id = logging_session_id + return session_created_event(default_session_config(model), logging_session_id) + + def transform_realtime_request( + self, + message: str, + model: str, + session_configuration_request: str | None = None, + ) -> tuple[str | bytes, ...]: + request: Final = json_object(message) + event_type: Final = request.get("type") + if event_type in ("session.update", "transcription_session.update"): + return self._configure(message, model) + if event_type == "input_audio_buffer.append": + return self._append_audio(request) + if event_type in ("input_audio_buffer.commit", "input_audio_buffer.end"): + self._require_config() + return (_FINISH_TURN_COMMAND,) + if event_type == "input_audio_buffer.clear": + self._require_config() + return (_DISCARD_TURN_COMMAND,) + verbose_logger.debug("Speech-to-Text streaming: dropping unsupported client event %s", event_type) + return () + + def unbilled_usage_on_session_close(self, model: str) -> RealtimeInputAudioTranscriptionUsage | None: + return self._transformer.take_unbilled_usage() + + def transform_realtime_response( + self, + message: str | bytes, + model: str, + logging_obj: LiteLLMLoggingObj, + realtime_response_transform_input: RealtimeResponseTransformInput, + ) -> RealtimeResponseTypedDict: + frame: Final = _STREAMING_EVENT_ADAPTER.validate_json(message) + events: Final = list(self._transformer.transform(frame)) # mutable-ok: response field is a list + result: Final[RealtimeResponseTypedDict] = { + "response": events, + "current_output_item_id": realtime_response_transform_input.get("current_output_item_id"), + "current_response_id": realtime_response_transform_input.get("current_response_id"), + "current_delta_chunks": realtime_response_transform_input.get("current_delta_chunks"), + "current_conversation_id": realtime_response_transform_input.get("current_conversation_id"), + "current_item_chunks": realtime_response_transform_input.get("current_item_chunks"), + "current_delta_type": realtime_response_transform_input.get("current_delta_type"), + "session_configuration_request": realtime_response_transform_input.get("session_configuration_request"), + } + return result + + def _configure(self, message: str, model: str) -> tuple[str, ...]: + if self._config is not None: + verbose_logger.debug("Speech-to-Text streaming: ignoring session.update after the stream was configured") + return () + config: Final = parse_chirp_session_update(message, model) + self._config = config + self._transformer.configure(config, self._session_id or f"sess_{uuid.uuid4().hex}") + return (config.configure_command(),) + + def _append_audio(self, request: Mapping[str, JsonValue]) -> tuple[bytes, ...]: + self._require_config() + audio: Final = decode_pcm16_append(request.get("audio")) + return tuple( + audio[start : start + MAX_AUDIO_MESSAGE_BYTES] for start in range(0, len(audio), MAX_AUDIO_MESSAGE_BYTES) + ) + + def _require_config(self) -> ChirpSessionConfig: + if self._config is None: + raise ChirpProtocolError("session.update must configure the session before audio is sent") + return self._config + + +def _api_endpoint(api_base: str) -> str: + without_scheme: Final = api_base.split("://", 1)[-1] + return without_scheme.split("/", 1)[0] diff --git a/litellm/llms/vertex_ai/audio_transcription/transformation.py b/litellm/llms/vertex_ai/audio_transcription/transformation.py index db3504c9a6a..b1284e15def 100644 --- a/litellm/llms/vertex_ai/audio_transcription/transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/transformation.py @@ -42,6 +42,10 @@ def validate_vertex_transcription_location(location: str | None, default_locatio raise VertexAIError(status_code=400, message=str(e)) from e +def speech_to_text_host(location: str) -> str: + return "speech.googleapis.com" if location == "global" else f"{location}-speech.googleapis.com" + + def validate_vertex_transcription_project_id(project_id: str) -> str: if not project_id or ".." in project_id or any(c in project_id for c in _URL_UNSAFE_PROJECT_CHARS): raise VertexAIError(status_code=400, message=f"Invalid vertex_project format: {project_id!r}") @@ -122,8 +126,7 @@ class VertexAIAudioTranscriptionConfig(BaseAudioTranscriptionConfig, VertexBase) project_id: Final = validate_vertex_transcription_project_id( self.safe_get_vertex_ai_project(litellm_params) or self._resolve_project_id_from_credentials(litellm_params) ) - host: Final = "speech.googleapis.com" if location == "global" else f"{location}-speech.googleapis.com" - base_url: Final = (api_base or f"https://{host}").rstrip("/") + base_url: Final = (api_base or f"https://{speech_to_text_host(location)}").rstrip("/") return f"{base_url}/v2/projects/{project_id}/locations/{location}/recognizers/_:recognize" def _resolve_project_id_from_credentials(self, litellm_params: dict) -> str: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 55fa2ec3bf8..88c078df4f4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -48005,7 +48005,8 @@ "mode": "audio_transcription", "source": "https://cloud.google.com/speech-to-text/pricing", "supported_endpoints": [ - "/v1/audio/transcriptions" + "/v1/audio/transcriptions", + "/v1/realtime" ] }, "vertex_ai/claude-3-5-haiku": { diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 8a880b5832f..4f5d70d6871 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19346,7 +19346,7 @@ } } }, - "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " + "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" }, "500": { "content": { diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index d67e4555a29..ba95514f556 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -38,6 +38,10 @@ from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_pr from ..llms.bedrock.realtime.handler import BedrockRealtime from ..llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..llms.openai.realtime.handler import OpenAIRealtime +from ..llms.vertex_ai.audio_transcription.realtime_transformation import ( + VertexChirpRealtimeConfig, + is_vertex_speech_to_text_model, +) from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig from ..llms.vertex_ai.vertex_llm_base import VertexBase from ..llms.xai.realtime.handler import XAIRealtime @@ -539,8 +543,6 @@ async def _arealtime( or get_secret_str("VERTEXAI_LOCATION") ) - resolved_location: Final = vertex_llm_base.get_vertex_region(vertex_region=vertex_location, model=model) - ( access_token, resolved_project, @@ -551,10 +553,11 @@ async def _arealtime( timeout_seconds=REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS, ) - vertex_realtime_config: Final = VertexAIRealtimeConfig( + vertex_realtime_config: Final = _vertex_realtime_config( + model=model, access_token=access_token, project=resolved_project, - location=resolved_location, + location=vertex_location, ) await base_llm_http_handler.async_realtime( @@ -575,6 +578,18 @@ async def _arealtime( raise ValueError(f"Unsupported model: {model}") +def _vertex_realtime_config( + model: str, access_token: str, project: str, location: str | None +) -> VertexAIRealtimeConfig | VertexChirpRealtimeConfig: + if is_vertex_speech_to_text_model(model): + return VertexChirpRealtimeConfig(access_token=access_token, project=project, location=location) + return VertexAIRealtimeConfig( + access_token=access_token, + project=project, + location=vertex_llm_base.get_vertex_region(vertex_region=location, model=model), + ) + + def _is_transcription_only_realtime_model(model: str, custom_llm_provider: str) -> bool: try: model_info: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) @@ -682,6 +697,11 @@ async def _realtime_health_check( api_base=resolved_api_base or "https://api.x.ai/v1", query_params={"model": model} ) elif custom_llm_provider == "vertex_ai": + if is_vertex_speech_to_text_model(model): + raise ValueError( + f"Realtime health checks are not supported for Speech-to-Text streaming model {model};" + " health check it with mode audio_transcription" + ) vertex_model_params: Final = dict(resolved_params) resolved_location: Final = vertex_llm_base.get_vertex_region( vertex_region=VertexBase.safe_get_vertex_ai_location(vertex_model_params), diff --git a/litellm/types/llms/vertex_ai_speech_to_text.py b/litellm/types/llms/vertex_ai_speech_to_text.py index 8995d98385b..e954960d41f 100644 --- a/litellm/types/llms/vertex_ai_speech_to_text.py +++ b/litellm/types/llms/vertex_ai_speech_to_text.py @@ -1,4 +1,6 @@ -from pydantic import BaseModel +from typing import Annotated, Literal + +from pydantic import BaseModel, ConfigDict, Field from typing_extensions import TypedDict @@ -38,3 +40,65 @@ class VertexSpeechToTextResponseMetadata(BaseModel): class VertexSpeechToTextRecognizeResponse(BaseModel): results: list[VertexSpeechToTextResult] = [] metadata: VertexSpeechToTextResponseMetadata | None = None + + +class VertexSpeechStreamingConfigure(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["configure"] = "configure" + model: str + language_codes: tuple[str, ...] + sample_rate_hertz: int + + +class VertexSpeechStreamingFinishTurn(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["finish_turn"] = "finish_turn" + + +class VertexSpeechStreamingDiscardTurn(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["discard_turn"] = "discard_turn" + + +VertexSpeechStreamingCommandUnion = ( + VertexSpeechStreamingConfigure | VertexSpeechStreamingFinishTurn | VertexSpeechStreamingDiscardTurn +) +VertexSpeechStreamingCommand = Annotated[VertexSpeechStreamingCommandUnion, Field(discriminator="kind")] + + +class VertexSpeechStreamingResult(BaseModel): + model_config = ConfigDict(frozen=True) + transcript: str + is_final: bool + + +class VertexSpeechStreamingResponse(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["response"] = "response" + speech_event: Literal["none", "begin", "end"] + results: tuple[VertexSpeechStreamingResult, ...] + billed_seconds: float + + +class VertexSpeechStreamingConfigured(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["configured"] = "configured" + + +class VertexSpeechStreamingTurnFinished(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["turn_finished"] = "turn_finished" + + +class VertexSpeechStreamingTurnDiscarded(BaseModel): + model_config = ConfigDict(frozen=True) + kind: Literal["turn_discarded"] = "turn_discarded" + + +VertexSpeechStreamingEventUnion = ( + VertexSpeechStreamingResponse + | VertexSpeechStreamingConfigured + | VertexSpeechStreamingTurnFinished + | VertexSpeechStreamingTurnDiscarded +) +VertexSpeechStreamingEvent = Annotated[VertexSpeechStreamingEventUnion, Field(discriminator="kind")] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 55fa2ec3bf8..88c078df4f4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -48005,7 +48005,8 @@ "mode": "audio_transcription", "source": "https://cloud.google.com/speech-to-text/pricing", "supported_endpoints": [ - "/v1/audio/transcriptions" + "/v1/audio/transcriptions", + "/v1/realtime" ] }, "vertex_ai/claude-3-5-haiku": { diff --git a/pyproject.toml b/pyproject.toml index dfe84a28d52..c58fd1d2158 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -129,6 +129,12 @@ grpc = [ # Newest non-yanked release older than the 30-day cutoff. "grpcio==1.78.0", ] +stt-vertex-chirp = [ + # Google Cloud Speech-to-Text v2 streaming (gRPC) for Chirp models on + # /v1/realtime. Imported lazily inside the backend so litellm core stays + # usable without it. + "google-cloud-speech>=2.40.0,<3.0", +] stt-nvidia-riva = [ # NVIDIA Riva STT provider (gRPC). These are imported lazily inside the # provider handler so litellm core remains usable without them. @@ -152,6 +158,7 @@ proxy-runtime = [ # Keep these in a dedicated extra so uv-based images preserve the same # feature surface without forcing the base SDK install to grow. "google-cloud-aiplatform>=1.133.0,<2.0", + "google-cloud-speech>=2.40.0,<3.0", "google-genai>=1.37.0,<2.0", "anthropic[vertex]>=0.84.0,<1.0", "grpcio==1.78.0", diff --git a/tests/test_litellm/llms/base_llm/realtime/__init__.py b/tests/test_litellm/llms/base_llm/realtime/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/tests/test_litellm/llms/base_llm/realtime/test_transcription_protocol.py b/tests/test_litellm/llms/base_llm/realtime/test_transcription_protocol.py new file mode 100644 index 00000000000..d38cc480e85 --- /dev/null +++ b/tests/test_litellm/llms/base_llm/realtime/test_transcription_protocol.py @@ -0,0 +1,128 @@ +import base64 +import json + +import pytest + +from litellm.llms.base_llm.realtime.transcription_protocol import ( + RealtimeTranscriptionProtocolError, + completed_event, + decode_pcm16_append, + parse_transcription_session_update, + transcription_session, +) + + +def _session_update(session: dict[str, object]) -> str: + return json.dumps({"type": "session.update", "session": session}) + + +def test_ga_layout_parses_format_language_and_turn_detection(): + update = parse_transcription_session_update( + _session_update( + { + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": 16_000, "channels": 1}, + "transcription": {"model": "chirp_3", "language": "pt-BR", "prompt": "names"}, + "turn_detection": {"type": "server_vad", "threshold": 0.5}, + } + }, + } + ) + ) + assert update.session_type == "transcription" + assert update.audio_format is not None + assert (update.audio_format.layout, update.audio_format.rate, update.audio_format.channels) == ("ga", 16_000, 1) + assert update.audio_format.is_pcm16 + assert (update.model, update.language) == ("chirp_3", "pt-BR") + assert update.unsupported_transcription_keys == ("prompt",) + assert update.turn_detection_type == "server_vad" + assert not update.turn_detection_disabled + + +def test_beta_layout_parses_flat_fields(): + update = parse_transcription_session_update( + json.dumps( + { + "type": "transcription_session.update", + "session": { + "input_audio_format": "pcm16", + "input_audio_transcription": {"model": "whisper-1"}, + "turn_detection": None, + }, + } + ) + ) + assert update.audio_format is not None + assert (update.audio_format.layout, update.audio_format.encoding) == ("beta", "pcm16") + assert update.audio_format.is_pcm16 + assert update.model == "whisper-1" + assert update.turn_detection_disabled + + +def test_absent_turn_detection_is_not_disabled(): + update = parse_transcription_session_update(_session_update({"audio": {"input": {"transcription": {}}}})) + assert update.turn_detection is None + assert not update.turn_detection_disabled + + +@pytest.mark.parametrize( + ("payload", "message"), + [ + ("not json", "invalid JSON object"), + ("[]", "must be a JSON object"), + (json.dumps({"type": "response.create"}), "expected session.update"), + (_session_update({}), "requires a session object"), + (_session_update({"input_audio_format": "pcm16", "audio": {"input": {"format": "pcm16"}}}), "either beta or GA"), + (_session_update({"input_audio_transcription": {}, "audio": {"input": {"transcription": {}}}}), "either beta or GA"), + (_session_update({"audio": {"input": {"format": {"rate": "fast"}}}}), "must be an integer"), + (_session_update({"audio": {"input": {"format": {"rate": True}}}}), "must be an integer"), + (_session_update({"audio": {"input": {"transcription": {"language": 7}}}}), "must be a string"), + (_session_update({"audio": {"input": {"transcription": []}}}), "must be an object"), + ], +) +def test_malformed_session_updates_are_rejected(payload: str, message: str): + with pytest.raises(RealtimeTranscriptionProtocolError, match=message): + parse_transcription_session_update(payload) + + +def test_decode_pcm16_append_returns_the_raw_samples(): + assert decode_pcm16_append(base64.b64encode(b"\x01\x02\x03\x04").decode()) == b"\x01\x02\x03\x04" + + +@pytest.mark.parametrize( + ("audio", "message"), + [ + (None, "must be a base64 string"), + ("@@@", "must be valid base64"), + (base64.b64encode(b"\x01\x02\x03").decode(), "complete samples"), + ], +) +def test_decode_pcm16_append_rejects_bad_audio(audio: object, message: str): + with pytest.raises(RealtimeTranscriptionProtocolError, match=message): + decode_pcm16_append(audio) + + +def test_decode_pcm16_append_enforces_the_backlog_limit(): + with pytest.raises(RealtimeTranscriptionProtocolError, match="backlog limit"): + decode_pcm16_append(base64.b64encode(b"\x00" * 8).decode(), max_encoded_bytes=4) + + +def test_transcription_session_reflects_negotiated_settings(): + manual = transcription_session(session_id="sess_1", model="chirp_3", sample_rate=16_000, language=None, server_vad=False) + assert manual["id"] == "sess_1" + assert manual["audio"]["input"] == { + "format": {"type": "audio/pcm", "rate": 16_000}, + "transcription": {"model": "chirp_3"}, + "turn_detection": None, + } + vad = transcription_session(session_id="sess_1", model="chirp_3", sample_rate=24_000, language="en-US", server_vad=True) + assert vad["audio"]["input"]["transcription"] == {"model": "chirp_3", "language": "en-US"} + assert vad["audio"]["input"]["turn_detection"] == {"type": "server_vad"} + + +def test_completed_event_carries_usage_only_when_billed(): + assert "usage" not in completed_event("item_1", "hello", None) + billed = completed_event("item_1", "hello", {"type": "duration", "seconds": 2.5}) + assert (billed["item_id"], billed["transcript"], billed["usage"]) == ("item_1", "hello", {"type": "duration", "seconds": 2.5}) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 95dceccb2f5..1f71ffd43f6 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -1,4 +1,5 @@ import asyncio +import base64 import json import logging import threading @@ -2064,6 +2065,7 @@ async def _run_async_realtime_with_backend_failure(client_ws): provider_config = Mock() provider_config.get_complete_url.return_value = "wss://backend.example/live" provider_config.validate_environment.return_value = {} + provider_config.open_backend = AsyncMock(return_value=None) with patch.object( handler, @@ -3707,3 +3709,143 @@ def test_image_edit_handler_keeps_the_sync_transform(): assert config.transform_calls == ["sync"] assert captured["body"] == {"transformed_by": "sync"} assert response.data[0].b64_json == "sync" + + +class _ScriptedClientWebSocket(_FakeClientWebSocket): + def __init__(self, messages: list[str], last_event_type: str) -> None: + super().__init__() + self._messages: Final = list(messages) + self._last_event_type: Final = last_event_type + self._backend_done: Final = asyncio.Event() + + async def receive_text(self) -> str: + if self._messages: + return self._messages.pop(0) + await asyncio.wait_for(self._backend_done.wait(), timeout=5) + raise RuntimeError("client went away") + + async def send_text(self, payload: str) -> None: + await super().send_text(payload) + if json.loads(payload).get("type") == self._last_event_type: + self._backend_done.set() + + def sent_events(self) -> list[dict[str, object]]: + return [json.loads(payload) for name, payload in self.events if name == "send_text"] + + +@pytest.mark.asyncio +async def test_async_realtime_bridges_a_transcription_session_through_the_provider_backend(): + import websockets.exceptions # noqa: F401 # binds the submodule so async_realtime's except clause resolves, as in the proxy process + + from datetime import timedelta + + from google.cloud.speech_v2.types import ( + RecognitionResponseMetadata, + SpeechRecognitionAlternative, + StreamingRecognitionResult, + StreamingRecognizeResponse, + ) + + from litellm.llms.vertex_ai.audio_transcription.realtime_backend import SpeechStreamingBackend + from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig + + def google_response(transcript: str, is_final: bool, billed: float) -> StreamingRecognizeResponse: + return StreamingRecognizeResponse( + results=[ + StreamingRecognitionResult( + alternatives=[SpeechRecognitionAlternative(transcript=transcript)], is_final=is_final + ) + ], + metadata=RecognitionResponseMetadata(total_billed_duration=timedelta(seconds=billed)), + ) + + class FakeTransport: + async def close(self) -> None: + return None + + class FakeSpeechClient: + transport = FakeTransport() + + def __init__(self) -> None: + self.requests: Final[list[object]] = [] + + async def streaming_recognize(self, requests=None): + return self._respond(requests) + + async def _respond(self, requests): + script = [google_response("four score", False, 0.0), google_response("Four score and seven", True, 2.0)] + async for request in requests: + self.requests.append(request) + if request.audio and script: + yield script.pop(0) + + speech_client = FakeSpeechClient() + provider_config = VertexChirpRealtimeConfig( + access_token="token", + project="proj-1", + location="us", + backend_factory=lambda target: SpeechStreamingBackend(target, client_factory=lambda target: speech_client), + ) + audio = base64.b64encode(b"\x00\x01" * 800).decode() + client_ws = _ScriptedClientWebSocket( + [ + json.dumps( + { + "type": "session.update", + "session": { + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": 16000}, + "transcription": {"model": "chirp_3", "language": "en"}, + "turn_detection": {"type": "server_vad"}, + } + }, + }, + } + ), + json.dumps({"type": "input_audio_buffer.append", "audio": audio}), + json.dumps({"type": "input_audio_buffer.append", "audio": audio}), + json.dumps({"type": "input_audio_buffer.commit"}), + ], + last_event_type="conversation.item.input_audio_transcription.completed", + ) + logging_obj = Mock() + logging_obj.litellm_trace_id = "trace_1" + logging_obj.model_call_details = {} + logging_obj.dispatch_success_handlers = AsyncMock() + logging_obj.dispatch_failure_handlers = AsyncMock() + handler = BaseLLMHTTPHandler() + + with patch.object(handler, "_open_realtime_backend_ws", AsyncMock(side_effect=AssertionError("dialed a websocket"))) as dial: + await handler.async_realtime( + model="chirp_3", + websocket=client_ws, + logging_obj=logging_obj, + provider_config=provider_config, + headers={}, + query_params={"model": "chirp_3", "intent": "transcription"}, + ) + + dial.assert_not_awaited() + events = client_ws.sent_events() + assert [event["type"] for event in events] == [ + "session.created", + "session.updated", + "input_audio_buffer.speech_started", + "conversation.item.input_audio_transcription.delta", + "conversation.item.input_audio_transcription.delta", + "input_audio_buffer.speech_stopped", + "conversation.item.input_audio_transcription.completed", + ] + assert events[0]["session"]["audio"]["input"]["transcription"] == {"model": "chirp_3"} + assert events[1]["session"]["audio"]["input"] == { + "format": {"type": "audio/pcm", "rate": 16000}, + "transcription": {"model": "chirp_3", "language": "en-US"}, + "turn_detection": {"type": "server_vad"}, + } + assert [event["delta"] for event in events[3:5]] == ["four score", " and seven"] + assert events[6]["transcript"] == "Four score and seven" + assert events[6]["usage"] == {"type": "duration", "seconds": 2.0} + assert speech_client.requests[0].streaming_config.config.model == "chirp_3" + assert [bytes(request.audio) for request in speech_client.requests[1:]] == [b"\x00\x01" * 800, b"\x00\x01" * 800] 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 a5d7e47fb65..1262d124d89 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,6 +6,7 @@ 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, @@ -160,7 +161,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(MuseProtocolError, match=message): + with pytest.raises(RealtimeTranscriptionProtocolError, match=message): parse_session_update(_event("session.update", session={"type": "transcription", **session}), MUSE_MODEL) @@ -584,7 +585,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(MuseProtocolError, match=message): + with pytest.raises(RealtimeTranscriptionProtocolError, 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 new file mode 100644 index 00000000000..87ab3da37a5 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_backend.py @@ -0,0 +1,261 @@ +import json +from collections.abc import AsyncIterator, Sequence +from datetime import timedelta +from typing import Final + +import pytest +from google.cloud.speech_v2.types import ( + RecognitionResponseMetadata, + SpeechRecognitionAlternative, + StreamingRecognitionResult, + StreamingRecognizeRequest, + StreamingRecognizeResponse, +) +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_transformation import SpeechStreamingTarget + +TARGET: Final = SpeechStreamingTarget( + api_endpoint="us-speech.googleapis.com", + recognizer="projects/proj-1/locations/us/recognizers/_", + access_token="token", +) +CONFIGURE: Final = json.dumps( + {"kind": "configure", "model": "chirp_3", "language_codes": ["en-US"], "sample_rate_hertz": 16_000} +) +FINISH_TURN: Final = json.dumps({"kind": "finish_turn"}) +DISCARD_TURN: Final = json.dumps({"kind": "discard_turn"}) +ScriptItem = StreamingRecognizeResponse | Exception + + +def _response( + transcript: str | None, + *, + is_final: bool = False, + billed: float = 0.0, + event: str = "SPEECH_EVENT_TYPE_UNSPECIFIED", +) -> StreamingRecognizeResponse: + results = ( + [] + if transcript is None + else [ + StreamingRecognitionResult( + alternatives=[SpeechRecognitionAlternative(transcript=transcript)], is_final=is_final + ) + ] + ) + return StreamingRecognizeResponse( + results=results, + speech_event_type=event, + metadata=RecognitionResponseMetadata(total_billed_duration=timedelta(seconds=billed)), + ) + + +class _FakeTransport: + def __init__(self) -> None: + self.closed = False + + async def close(self) -> None: + self.closed = True + + +class _FakeSpeechClient: + def __init__(self, *scripts: Sequence[ScriptItem]) -> None: + self.transport: Final = _FakeTransport() + self.streams: Final[list[list[StreamingRecognizeRequest]]] = [] + self._scripts: Final = [list(script) for script in scripts] + + async def streaming_recognize( + self, requests: AsyncIterator[StreamingRecognizeRequest] | None = None + ) -> AsyncIterator[StreamingRecognizeResponse]: + assert requests is not None + script: Final = self._scripts.pop(0) if self._scripts else [] + received: Final[list[StreamingRecognizeRequest]] = [] + self.streams.append(received) + return self._respond(requests, script, received) + + async def _respond( + self, + requests: AsyncIterator[StreamingRecognizeRequest], + script: list[ScriptItem], + received: list[StreamingRecognizeRequest], + ) -> AsyncIterator[StreamingRecognizeResponse]: + async for request in requests: + received.append(request) + if request.audio and script: + yield self._next(script) + while script: + yield self._next(script) + + @staticmethod + def _next(script: list[ScriptItem]) -> StreamingRecognizeResponse: + item: Final = script.pop(0) + if isinstance(item, Exception): + raise item + return item + + +def _backend(client: _FakeSpeechClient, **kwargs: object) -> SpeechStreamingBackend: + return SpeechStreamingBackend(TARGET, client_factory=lambda target: client, **kwargs) + + +async def _recv(backend: SpeechStreamingBackend) -> dict[str, object]: + message: Final = await backend.recv() + assert isinstance(message, str) + return json.loads(message) + + +async def _configure(backend: SpeechStreamingBackend) -> None: + await backend.send(CONFIGURE) + assert await _recv(backend) == {"kind": "configured"} + + +def _audio(stream: list[StreamingRecognizeRequest]) -> list[bytes]: + return [bytes(request.audio) for request in stream[1:]] + + +@pytest.mark.asyncio +async def test_audio_streams_through_one_recognize_call_with_the_config_first(): + client = _FakeSpeechClient([_response("hello"), _response("hello world", is_final=True, billed=2.0)]) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x01\x02") + await backend.send(b"\x03\x04") + await backend.send(FINISH_TURN) + first, second, finished = [await _recv(backend) for _ in range(3)] + assert first == { + "kind": "response", + "speech_event": "none", + "results": [{"transcript": "hello", "is_final": False}], + "billed_seconds": 0.0, + } + assert second["results"] == [{"transcript": "hello world", "is_final": True}] + assert second["billed_seconds"] == 2.0 + assert finished == {"kind": "turn_finished"} + (requests,) = client.streams + assert requests[0].recognizer == TARGET.recognizer + config = requests[0].streaming_config + assert config.config.model == "chirp_3" + assert list(config.config.language_codes) == ["en-US"] + assert config.config.explicit_decoding_config.sample_rate_hertz == 16_000 + assert config.config.explicit_decoding_config.audio_channel_count == 1 + assert config.config.explicit_decoding_config.encoding.name == "LINEAR16" + assert config.streaming_features.interim_results + assert config.streaming_features.enable_voice_activity_events + assert _audio(requests) == [b"\x01\x02", b"\x03\x04"] + assert client.transport.closed + + +@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")]) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x00\x00") + await backend.send(b"\x00\x00") + begin, end = [await _recv(backend) for _ in range(2)] + assert (begin["speech_event"], begin["results"]) == ("begin", []) + assert end["speech_event"] == "end" + + +@pytest.mark.asyncio +async def test_audio_before_configure_is_rejected(): + backend = _backend(_FakeSpeechClient()) + with pytest.raises(RuntimeError, match="before the Speech-to-Text stream was configured"): + await backend.send(b"\x00\x00") + + +@pytest.mark.asyncio +async def test_stream_failure_closes_the_session_with_1011_and_the_reason(): + client = _FakeSpeechClient([PermissionError("IAM_PERMISSION_DENIED: speech.recognizers.recognize")]) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x00\x00") + with pytest.raises(ConnectionClosedError) as excinfo: + await backend.recv() + assert excinfo.value.rcvd is not None + assert excinfo.value.rcvd.code == 1011 + assert "IAM_PERMISSION_DENIED" in excinfo.value.rcvd.reason + assert client.transport.closed + + +@pytest.mark.asyncio +async def test_close_discards_the_open_turn_then_reports_a_normal_closure(): + client = _FakeSpeechClient([_response("hi")]) + backend = _backend(client) + await _configure(backend) + await backend.send(b"\x00\x00") + assert (await _recv(backend))["results"][0]["transcript"] == "hi" + await backend.close() + assert await _recv(backend) == {"kind": "turn_discarded"} + with pytest.raises(ConnectionClosedOK): + await backend.recv() + assert client.transport.closed + + +@pytest.mark.asyncio +async def test_turn_commands_without_audio_answer_immediately(): + backend = _backend(_FakeSpeechClient()) + await _configure(backend) + await backend.send(FINISH_TURN) + assert await _recv(backend) == {"kind": "turn_finished"} + await backend.send(DISCARD_TURN) + assert await _recv(backend) == {"kind": "turn_discarded"} + + +@pytest.mark.asyncio +async def test_discard_turn_cancels_the_open_stream_and_the_next_turn_starts_fresh(): + client = _FakeSpeechClient([_response("draft")], [_response("again", is_final=True)]) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert (await _recv(backend))["results"][0]["transcript"] == "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 [_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)]) + async with _backend(client) as backend: + await _configure(backend) + await backend.send(b"\x00\x00") + await backend.send(FINISH_TURN) + first = await _recv(backend) + assert await _recv(backend) == {"kind": "turn_finished"} + await backend.send(b"\x00\x00") + await backend.send(FINISH_TURN) + second = await _recv(backend) + assert await _recv(backend) == {"kind": "turn_finished"} + assert (first["billed_seconds"], second["billed_seconds"]) == (2.0, 5.0) + assert len(client.streams) == 2 + + +@pytest.mark.asyncio +async def test_streams_rotate_before_the_five_minute_limit_without_losing_audio(): + now = [0.0] + client = _FakeSpeechClient( + [_response("first"), _response("first half", is_final=True, billed=239.0)], + [_response("second")], + ) + 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" + now[0] = 239.0 + await backend.send(b"\x02\x02") + assert (await _recv(backend))["results"][0]["transcript"] == "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 + 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" 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 new file mode 100644 index 00000000000..e294c7cad36 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/audio_transcription/test_vertex_ai_realtime_transformation.py @@ -0,0 +1,355 @@ +import base64 +import json +from typing import Final +from unittest.mock import MagicMock + +import pytest + +from litellm.llms.base_llm.realtime.transformation import RealtimeBackend +from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import ( + MAX_AUDIO_MESSAGE_BYTES, + ChirpProtocolError, + ChirpSessionConfig, + SpeechStreamingTarget, + VertexChirpRealtimeConfig, + is_vertex_speech_to_text_model, + new_words, + parse_chirp_session_update, +) +from litellm.llms.vertex_ai.common_utils import VertexAIError +from litellm.types.llms.vertex_ai_speech_to_text import ( + VertexSpeechStreamingConfigured, + VertexSpeechStreamingResponse, + VertexSpeechStreamingResult, + VertexSpeechStreamingTurnDiscarded, + VertexSpeechStreamingTurnFinished, +) +from litellm.types.realtime import RealtimeResponseTransformInput + +MODEL: Final = "chirp_3" +EMPTY_TRANSFORM_INPUT: Final[RealtimeResponseTransformInput] = { + "session_configuration_request": None, + "current_output_item_id": None, + "current_response_id": None, + "current_delta_chunks": None, + "current_item_chunks": None, + "current_conversation_id": None, + "current_delta_type": None, +} +DELTA: Final = "conversation.item.input_audio_transcription.delta" +COMPLETED: Final = "conversation.item.input_audio_transcription.completed" + + +def _event(event_type: str, **fields: object) -> str: + return json.dumps({"type": event_type, **fields}) + + +def _ga_session_update( + rate: int = 24_000, turn_detection: str | None = "server_vad", language: str | None = "en", model: str = MODEL +) -> str: + transcription = {"model": model} if language is None else {"model": model, "language": language} + return _event( + "session.update", + session={ + "type": "transcription", + "audio": { + "input": { + "format": {"type": "audio/pcm", "rate": rate}, + "turn_detection": None if turn_detection is None else {"type": turn_detection}, + "transcription": transcription, + } + }, + }, + ) + + +def _config(location: str | None = "us") -> VertexChirpRealtimeConfig: + return VertexChirpRealtimeConfig(access_token="token", project="proj-1", location=location) + + +def _configured( + rate: int = 24_000, turn_detection: str | None = "server_vad", language: str | None = "en" +) -> VertexChirpRealtimeConfig: + config = _config() + config.transform_session_created_event(MODEL, "sess_1") + config.transform_realtime_request(_ga_session_update(rate, turn_detection, language), MODEL) + return config + + +def _backend_events(config: VertexChirpRealtimeConfig, frame: object) -> list[dict[str, object]]: + assert hasattr(frame, "model_dump_json") + response = config.transform_realtime_response(frame.model_dump_json(), MODEL, MagicMock(), EMPTY_TRANSFORM_INPUT)[ + "response" + ] + assert isinstance(response, list) + return response + + +def _response( + *results: tuple[str, bool], speech_event: str = "none", billed_seconds: float = 0.0 +) -> VertexSpeechStreamingResponse: + return VertexSpeechStreamingResponse( + speech_event=speech_event, + results=tuple(VertexSpeechStreamingResult(transcript=text, is_final=final) for text, final in results), + billed_seconds=billed_seconds, + ) + + +def _types(events: list[dict[str, object]]) -> list[object]: + return [event["type"] for event in events] + + +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)] + + +@pytest.mark.parametrize( + ("model", "expected"), + [ + ("vertex_ai/chirp_3", True), + ("chirp_3", True), + ("chirp_2", True), + ("gemini-live-2.5-flash", False), + ("vertex_ai/gemini-2.0-flash-live-preview-04-09", False), + ], +) +def test_is_vertex_speech_to_text_model(model: str, expected: bool): + assert is_vertex_speech_to_text_model(model) is expected + + +def test_ga_session_update_maps_to_a_speech_config(): + config = parse_chirp_session_update(_ga_session_update(16_000, "server_vad", "pt"), "vertex_ai/chirp_3") + assert config == ChirpSessionConfig(model=MODEL, language="pt-BR", sample_rate=16_000, server_vad=True) + assert json.loads(config.configure_command()) == { + "kind": "configure", + "model": MODEL, + "language_codes": ["pt-BR"], + "sample_rate_hertz": 16_000, + } + + +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}, + ), + MODEL, + ) + assert config == ChirpSessionConfig(model=MODEL, language=None, sample_rate=24_000, server_vad=False) + assert json.loads(config.configure_command())["language_codes"] == ["auto"] + + +@pytest.mark.parametrize( + ("payload", "message"), + [ + (_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"), + (_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"), + ], +) +def test_unsupported_session_settings_are_rejected(payload: str, message: str): + with pytest.raises(ChirpProtocolError, match=message): + parse_chirp_session_update(payload, MODEL) + + +def test_session_update_configures_once_and_later_updates_are_ignored(): + config = _config() + config.transform_session_created_event(MODEL, "sess_1") + first = _commands(config, _ga_session_update(16_000)) + assert first == [{"kind": "configure", "model": MODEL, "language_codes": ["en-US"], "sample_rate_hertz": 16_000}] + assert config.is_setup_message(first[0]) + assert _commands(config, _ga_session_update(8_000)) == [] + + +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) + with pytest.raises(ChirpProtocolError, match=r"session\.update must configure"): + config.transform_realtime_request(_event("input_audio_buffer.commit"), MODEL) + + +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] + assert b"".join(chunk for chunk in chunks if isinstance(chunk, bytes)) == audio + + +def test_commit_end_and_clear_map_to_turn_commands(): + config = _configured() + assert _commands(config, _event("input_audio_buffer.commit")) == [{"kind": "finish_turn"}] + assert _commands(config, _event("input_audio_buffer.end")) == [{"kind": "finish_turn"}] + assert _commands(config, _event("input_audio_buffer.clear")) == [{"kind": "discard_turn"}] + + +def test_unsupported_client_events_are_dropped(): + assert _commands(_configured(), _event("response.create")) == [] + + +def test_connect_announces_a_session_with_chirp_defaults(): + event = _config().transform_session_created_event(MODEL, "sess_1") + assert event["type"] == "session.created" + assert event["session"]["id"] == "sess_1" + assert event["session"]["audio"]["input"] == { + "format": {"type": "audio/pcm", "rate": 24_000}, + "transcription": {"model": MODEL}, + "turn_detection": {"type": "server_vad"}, + } + + +def test_configured_backend_reports_the_negotiated_session(): + config = _configured(rate=16_000, turn_detection=None, language="pt-BR") + events = _backend_events(config, VertexSpeechStreamingConfigured()) + assert _types(events) == ["session.created"] + session = events[0]["session"] + assert isinstance(session, dict) + assert session["id"] == "sess_1" + assert session["audio"]["input"] == { + "format": {"type": "audio/pcm", "rate": 16_000}, + "transcription": {"model": MODEL, "language": "pt-BR"}, + "turn_detection": None, + } + + +def test_backend_frames_before_session_update_are_an_error(): + config = _config() + config.transform_session_created_event(MODEL, "sess_1") + with pytest.raises(ChirpProtocolError, match=r"session\.update must configure"): + _backend_events(config, VertexSpeechStreamingConfigured()) + + +def test_server_vad_turn_streams_new_words_then_completes_with_usage(): + config = _configured() + assert _types(_backend_events(config, _response(speech_event="begin"))) == ["input_audio_buffer.speech_started"] + first = _backend_events(config, _response(("four score", False))) + assert [(event["type"], event["delta"]) for event in first] == [(DELTA, "four score")] + second = _backend_events(config, _response(("four score and seven", False))) + assert [event["delta"] for event in second] == [" and seven"] + final = _backend_events(config, _response(("Four score and seven years ago.", True), billed_seconds=3.5)) + assert _types(final) == [DELTA, "input_audio_buffer.speech_stopped", COMPLETED] + assert final[0]["delta"] == " years ago." + assert final[2]["transcript"] == "Four score and seven years ago." + assert final[2]["usage"] == {"type": "duration", "seconds": 3.5} + assert {event["item_id"] for event in (*first, *second, *final)} == {first[0]["item_id"]} + assert _backend_events(config, _response(speech_event="end")) == [] + + +def test_manual_turns_complete_on_commit_without_speech_events(): + config = _configured(turn_detection=None) + assert _backend_events(config, _response(speech_event="begin")) == [] + first = _backend_events(config, _response(("hello there", True), billed_seconds=1.25)) + assert [(event["type"], event["delta"]) for event in first] == [(DELTA, "hello there")] + second = _backend_events(config, _response(("world", True))) + assert [event["delta"] for event in second] == [" world"] + completed = _backend_events(config, VertexSpeechStreamingTurnFinished()) + assert _types(completed) == [COMPLETED] + assert completed[0]["transcript"] == "hello there world" + assert completed[0]["usage"] == {"type": "duration", "seconds": 1.25} + assert _backend_events(config, VertexSpeechStreamingTurnFinished()) == [] + + +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, VertexSpeechStreamingTurnFinished()) == [] + fresh = _backend_events(config, _response(("again", False))) + assert fresh[0]["delta"] == "again" + assert fresh[0]["item_id"] != draft[0]["item_id"] + + +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)) + second = _backend_events(config, _response(("two", True), billed_seconds=5.0)) + assert first[-1]["usage"] == {"type": "duration", "seconds": 2.0} + assert second[-1]["usage"] == {"type": "duration", "seconds": 3.0} + assert config.unbilled_usage_on_session_close(MODEL) is None + assert _backend_events(config, _response(billed_seconds=6.5)) == [] + assert config.unbilled_usage_on_session_close(MODEL) == {"type": "duration", "seconds": 1.5} + assert config.unbilled_usage_on_session_close(MODEL) is None + + +class _NullBackend: + async def __aenter__(self) -> "_NullBackend": + return self + + async def __aexit__(self, exc_type: object, exc_value: object, traceback: object) -> None: + return None + + async def send(self, message: str | bytes) -> None: + return None + + async def recv(self, decode: bool | None = None) -> str | bytes: + return "" + + async def close(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_open_backend_targets_the_regional_speech_endpoint(): + targets: list[SpeechStreamingTarget] = [] + + def factory(target: SpeechStreamingTarget) -> RealtimeBackend: + targets.append(target) + return _NullBackend() + + config = VertexChirpRealtimeConfig(access_token="token", project="proj-1", location=None, backend_factory=factory) + url = config.get_complete_url(None, "vertex_ai/chirp_3") + assert url == "us-speech.googleapis.com" + assert config.validate_environment({}, MODEL, "https://" + url) == {} + backend = await config.open_backend(url, {}) + assert isinstance(backend, _NullBackend) + assert targets == [ + SpeechStreamingTarget( + api_endpoint="us-speech.googleapis.com", + recognizer="projects/proj-1/locations/us/recognizers/_", + access_token="token", + ) + ] + + +@pytest.mark.parametrize( + ("location", "api_base", "endpoint"), + [ + ("global", None, "speech.googleapis.com"), + ("europe-west4", None, "europe-west4-speech.googleapis.com"), + ("us", "https://speech-proxy.internal:8443/v2", "speech-proxy.internal:8443"), + ], +) +def test_get_complete_url_honors_location_and_api_base(location: str, api_base: str | None, endpoint: str): + assert _config(location).get_complete_url(api_base, MODEL) == endpoint + + +def test_get_complete_url_rejects_non_speech_models(): + with pytest.raises(ValueError, match="Unsupported Speech-to-Text streaming model"): + _config().get_complete_url(None, "gemini-live-2.5-flash") + + +@pytest.mark.parametrize("location", ["bad loc", "../us"]) +def test_invalid_locations_are_rejected_up_front(location: str): + with pytest.raises(VertexAIError): + _config(location) + + +@pytest.mark.parametrize( + ("previous", "current", "delta"), + [ + ("", "hello", "hello"), + ("hello", "hello world", " world"), + ("hello", "Hello, world", " world"), + ("hello world", "hello world", ""), + ("hello there", "hello world", " world"), + ("hello world", "hello", ""), + ], +) +def test_new_words(previous: str, current: str, delta: str): + assert new_words(previous, current) == delta diff --git a/tests/test_litellm/realtime_api/test_main.py b/tests/test_litellm/realtime_api/test_main.py index 643e65af47c..86b25b2f9c8 100644 --- a/tests/test_litellm/realtime_api/test_main.py +++ b/tests/test_litellm/realtime_api/test_main.py @@ -499,3 +499,69 @@ async def test_arealtime_azure_env_beta_protocol_wins_over_a_ga_client(monkeypat assert await _azure_backend_url_dialed_for(_GA_CLIENT) == ( "wss://my-endpoint.openai.azure.com/openai/realtime?api-version=2024-10-01-preview&deployment=gpt-realtime" ) + + +async def _vertex_provider_config_for(monkeypatch, model: str, vertex_location: str | None): + from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig + from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig + + captured: dict[str, object] = {} + + def mock_get_llm_provider(model, api_base, api_key): + return model.removeprefix("vertex_ai/"), "vertex_ai", None, api_base + + async def mock_token_resolver(**kwargs): + return "access-token", kwargs["project_id"] + + async def mock_async_realtime(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(realtime_main, "get_llm_provider", mock_get_llm_provider) + monkeypatch.setattr(realtime_main, "vertex_access_token_resolver", mock_token_resolver) + monkeypatch.setattr(realtime_main.base_llm_http_handler, "async_realtime", mock_async_realtime) + monkeypatch.setattr(litellm, "vertex_location", None) + monkeypatch.delenv("VERTEXAI_LOCATION", raising=False) + await realtime_main._arealtime.__wrapped__( + model=model, + websocket=MagicMock(), + litellm_logging_obj=FakeLogging(), + query_params={"model": model, "intent": "transcription"}, + vertex_credentials="fake-credentials", + vertex_project="proj-1", + vertex_location=vertex_location, + ) + provider_config = captured["provider_config"] + assert isinstance(provider_config, (VertexAIRealtimeConfig, VertexChirpRealtimeConfig)) + return provider_config, captured["model"] + + +@pytest.mark.asyncio +async def test_arealtime_routes_chirp_models_to_the_speech_to_text_backend(monkeypatch): + from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig + + provider_config, model = await _vertex_provider_config_for(monkeypatch, "vertex_ai/chirp_3", None) + assert isinstance(provider_config, VertexChirpRealtimeConfig) + assert model == "chirp_3" + assert provider_config.get_complete_url(None, model) == "us-speech.googleapis.com" + assert provider_config.validate_environment({}, model, "https://us-speech.googleapis.com") == {} + + +@pytest.mark.asyncio +async def test_arealtime_routes_chirp_models_to_the_configured_speech_region(monkeypatch): + provider_config, model = await _vertex_provider_config_for(monkeypatch, "vertex_ai/chirp_3", "europe-west4") + assert provider_config.get_complete_url(None, model) == "europe-west4-speech.googleapis.com" + + +@pytest.mark.asyncio +async def test_arealtime_keeps_gemini_live_on_the_vertex_realtime_websocket(monkeypatch): + from litellm.llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig + + provider_config, model = await _vertex_provider_config_for(monkeypatch, "vertex_ai/gemini-live-2.5-flash", None) + assert isinstance(provider_config, VertexAIRealtimeConfig) + assert provider_config.get_complete_url(None, model).startswith("wss://us-central1-aiplatform.googleapis.com/") + + +@pytest.mark.asyncio +async def test_realtime_health_check_names_the_batch_mode_for_chirp_models(): + with pytest.raises(ValueError, match="mode audio_transcription"): + await realtime_main._realtime_health_check(model="chirp_3", custom_llm_provider="vertex_ai", api_key=None) diff --git a/uv.lock b/uv.lock index 35eaa20c39e..34c5ea4214d 100644 --- a/uv.lock +++ b/uv.lock @@ -10,7 +10,7 @@ resolution-markers = [ ] [options] -exclude-newer = "2026-09-14T20:32:38.482736111Z" +exclude-newer = "2026-09-15T00:05:13.895745Z" exclude-newer-span = "P3D" [manifest] @@ -2615,6 +2615,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/3d/f7/661d7a9023e877a226b5683429c3662f75a29ef45cb1464cf39adb689218/google_cloud_resource_manager-1.17.0-py3-none-any.whl", hash = "sha256:e479baf4b014a57f298e01b8279e3290b032e3476d69c8e5e1427af8f82739a5", size = 404403, upload-time = "2026-03-26T22:15:26.57Z" }, ] +[[package]] +name = "google-cloud-speech" +version = "2.40.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "google-api-core", version = "2.25.2", source = { registry = "https://pypi.org/simple" }, extra = ["grpc"], marker = "python_full_version >= '3.14'" }, + { name = "google-api-core", version = "2.30.3", source = { registry = "https://pypi.org/simple" }, extra = ["grpc"], marker = "python_full_version < '3.14'" }, + { name = "google-auth" }, + { name = "grpcio" }, + { name = "proto-plus" }, + { name = "protobuf" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/5a/c1/5dc9795314f4aefea0b01b02e9f5486a198341ecc15fe47f89a61c68df63/google_cloud_speech-2.40.0.tar.gz", hash = "sha256:e89e688e4ce0b926754038bf992d0d0f065c5f1c3503bb20e6c46d08b63658fc", size = 404366, upload-time = "2026-06-03T16:13:59.506Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cc/78/afeca8d597fab54bdd823f857aad15d6f9c4628ff3cb72aa237d01700721/google_cloud_speech-2.40.0-py3-none-any.whl", hash = "sha256:7cc0302b3b9ca33d2eae9669da94a44316601a240942895362ac70e765b9f39c", size = 345427, upload-time = "2026-06-03T16:12:40.909Z" }, +] + [[package]] name = "google-cloud-storage" version = "3.4.1" @@ -4566,6 +4583,7 @@ proxy-runtime = [ { name = "ddtrace" }, { name = "detect-secrets" }, { name = "google-cloud-aiplatform" }, + { name = "google-cloud-speech" }, { name = "google-genai" }, { name = "grpcio" }, { name = "langfuse" }, @@ -4593,6 +4611,9 @@ stt-nvidia-riva = [ { name = "nvidia-riva-client" }, { name = "soundfile" }, ] +stt-vertex-chirp = [ + { name = "google-cloud-speech" }, +] utils = [ { name = "numpydoc" }, ] @@ -4721,6 +4742,8 @@ requires-dist = [ { name = "google-cloud-aiplatform", marker = "extra == 'proxy-runtime'", specifier = ">=1.133.0,<2.0" }, { name = "google-cloud-iam", marker = "extra == 'extra-proxy'", specifier = ">=2.19.1,<3.0" }, { name = "google-cloud-kms", marker = "extra == 'extra-proxy'", specifier = ">=2.24.2,<3.0" }, + { name = "google-cloud-speech", marker = "extra == 'proxy-runtime'", specifier = ">=2.40.0,<3.0" }, + { name = "google-cloud-speech", marker = "extra == 'stt-vertex-chirp'", specifier = ">=2.40.0,<3.0" }, { name = "google-genai", marker = "extra == 'proxy-runtime'", specifier = ">=1.37.0,<2.0" }, { name = "granian", marker = "extra == 'proxy'", specifier = ">=2.7.4,<3.0" }, { name = "grpcio", marker = "extra == 'grpc'", specifier = "==1.78.0" }, @@ -4787,7 +4810,7 @@ requires-dist = [ { name = "uvloop", marker = "sys_platform != 'win32' and extra == 'proxy'", specifier = ">=0.22.1,<1.0" }, { name = "websockets", marker = "extra == 'proxy'", specifier = ">=15.0.1,<16.0" }, ] -provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "saml", "semantic-router", "mlflow", "grpc", "stt-nvidia-riva", "google", "bedrock-realtime", "proxy-runtime"] +provides-extras = ["proxy", "cli", "extra-proxy", "utils", "caching", "mcp", "saml", "semantic-router", "mlflow", "grpc", "stt-vertex-chirp", "stt-nvidia-riva", "google", "bedrock-realtime", "proxy-runtime"] [package.metadata.requires-dev] ci = [ From 7f9db61528bf3fd8a83b93851a3d3d5070868329 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 15:35:59 -0700 Subject: [PATCH 2/7] fix(vertex_ai): carry turns across stream rotation and route by model info Rotating the Speech-to-Text stream at 240 s no longer ends the active turn: the turn and its billed seconds continue on the new stream, forced at 280 s. Bound the request and event queues (64 and 256) so a slow peer applies backpressure instead of growing memory. Route a model to the Chirp realtime path from its cost-map entry (mode audio_transcription plus /v1/realtime) instead of a hardcoded name. Return on every branch of the recv and transform helpers (CodeQL mixed returns), have the shared protocol helper take the provider's error class so the Meta tests assert MuseProtocolError again, and pin google-cloud-speech in the ci group so unit shards import it. --- .../realtime/transcription_protocol.py | 84 ++++--- .../llms/base_llm/realtime/transformation.py | 3 +- litellm/llms/meta/realtime/transformation.py | 8 +- .../audio_transcription/realtime_backend.py | 209 +++++++++++------- .../realtime_transformation.py | 22 +- pyproject.toml | 1 + .../test_meta_realtime_transformation.py | 5 +- .../test_vertex_ai_realtime_backend.py | 162 ++++++++++++-- .../test_vertex_ai_realtime_transformation.py | 34 ++- uv.lock | 4 +- 10 files changed, 374 insertions(+), 158 deletions(-) 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" }, From 4356fc58d82a30e470ce25caa0949b002c51f3b4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:05:57 -0700 Subject: [PATCH 3/7] chore(proxy): restore the CI-generated lazy OpenAPI snapshot --- litellm/proxy/_lazy_openapi_snapshot.json | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 40b64160b71..213cd88b6ce 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -19616,7 +19616,7 @@ } } }, - "description": "\nUnified rate-limit error.\n\nEvery rate-limit condition surfaced by litellm \u2014 whether it originated from\nan upstream LLM provider, a vendor batch endpoint, or one of litellm's own\nproxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\nmax-iterations, etc.) \u2014 is raised as an instance of this class.\n\nThe :attr:`category` attribute lets callers distinguish the source. See\n:class:`RateLimitErrorCategory` for the available values.\n" + "description": "\n Unified rate-limit error.\n\n Every rate-limit condition surfaced by litellm \u2014 whether it originated from\n an upstream LLM provider, a vendor batch endpoint, or one of litellm's own\n proxy-side limiters (parallel-requests, dynamic-rate, batch-rate, budget,\n max-iterations, etc.) \u2014 is raised as an instance of this class.\n\n The :attr:`category` attribute lets callers distinguish the source. See\n :class:`RateLimitErrorCategory` for the available values.\n " }, "500": { "content": { From 3d7a771ea7d2a327a4cf09b5425fdfe4e3fbc69d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:53:06 -0700 Subject: [PATCH 4/7] fix(vertex_ai): apply finals before interims and refresh the token per stream --- .../audio_transcription/realtime_backend.py | 33 ++++++--- .../realtime_transformation.py | 30 ++++---- .../llms/vertex_ai/realtime/transformation.py | 23 +++++++ litellm/realtime_api/main.py | 35 ++++------ .../test_vertex_ai_realtime_backend.py | 68 ++++++++++++++++++- .../test_vertex_ai_realtime_transformation.py | 48 ++++++++++++- 6 files changed, 186 insertions(+), 51 deletions(-) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py index 87c72c193bc..0c16fea9e9f 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py @@ -80,7 +80,7 @@ class _Closed: pass -def open_speech_client(target: SpeechStreamingTarget) -> SpeechStreamingClient: +def open_speech_client(target: SpeechStreamingTarget, access_token: str) -> SpeechStreamingClient: try: from google.api_core.client_options import ClientOptions from google.cloud.speech_v2 import SpeechAsyncClient @@ -88,7 +88,7 @@ def open_speech_client(target: SpeechStreamingTarget) -> SpeechStreamingClient: except ImportError as e: raise ImportError(SPEECH_SDK_INSTALL_HINT) from e return SpeechAsyncClient( - credentials=Credentials(token=target.access_token), + credentials=Credentials(token=access_token), transport="grpc_asyncio", client_options=ClientOptions(api_endpoint=target.api_endpoint), ) @@ -157,6 +157,7 @@ class _RecognizeStream: self.speech_active: bool = False self.billed_seconds: float = 0.0 self._cancelled: bool = False + self._closed: bool = False self._task: asyncio.Task[None] | None = None async def send_audio(self, audio: bytes) -> None: @@ -170,8 +171,15 @@ class _RecognizeStream: if self._task is not None: self._task.cancel() + async def close(self) -> None: + if self._closed: + return + self._closed = True + await self._client.transport.close() + async def relay(self, outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", billed_before: float) -> float: if self._cancelled: + await self.close() return 0.0 task: Final = asyncio.create_task(self._forward(outbox, billed_before)) self._task = task @@ -181,6 +189,8 @@ class _RecognizeStream: task.cancel() await asyncio.wait((task,)) raise + finally: + await self.close() return self.billed_seconds async def _forward(self, outbox: "asyncio.Queue[str | _StreamFailure | _Closed]", billed_before: float) -> None: @@ -209,7 +219,7 @@ class SpeechStreamingBackend: self, target: SpeechStreamingTarget, *, - client_factory: Callable[[SpeechStreamingTarget], SpeechStreamingClient] = open_speech_client, + client_factory: Callable[[SpeechStreamingTarget, str], SpeechStreamingClient] = open_speech_client, clock: Callable[[], float] = time.monotonic, rotation_seconds: float = STREAM_ROTATION_SECONDS, rotation_deadline_seconds: float = STREAM_ROTATION_DEADLINE_SECONDS, @@ -222,7 +232,6 @@ class SpeechStreamingBackend: 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._turn: tuple[_RecognizeStream, ...] = () self._billed_before: float = 0.0 @@ -282,13 +291,16 @@ class SpeechStreamingBackend: 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() + await self._close_unrelayed_streams() if not self._outbox.full(): self._outbox.put_nowait(_Closed()) + async def _close_unrelayed_streams(self) -> None: + unrelayed: Final = tuple(self._links.get_nowait() for _ in range(self._links.qsize())) + for link in unrelayed: + if isinstance(link, _RecognizeStream): + await link.close() + async def _link(self, item: _RecognizeStream | str) -> None: if self._pump is None: self._pump = asyncio.create_task(self._pump_links()) @@ -333,10 +345,9 @@ class SpeechStreamingBackend: config: Final = self._config if config is None: raise RuntimeError("audio was sent before the Speech-to-Text stream was configured") - if self._client is None: - self._client = self._client_factory(self._target) + access_token: Final = await self._target.resolve_access_token() stream: Final = _RecognizeStream( - client=self._client, + client=self._client_factory(self._target, access_token), request_type=StreamingRecognizeRequest, first_request=StreamingRecognizeRequest(recognizer=self._target.recognizer, streaming_config=config), opened_at=self._clock(), diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py index 6ec7a21a134..dab2e980fd0 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_transformation.py @@ -1,4 +1,4 @@ -from collections.abc import Callable, Mapping +from collections.abc import Awaitable, Callable, Mapping from dataclasses import dataclass, replace from typing import Final @@ -6,7 +6,7 @@ from pydantic import JsonValue, TypeAdapter from typing_extensions import assert_never import litellm -from litellm import verbose_logger +from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.audio_utils.utils import normalize_transcription_language_to_bcp47 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -73,7 +73,7 @@ class ChirpProtocolError(RealtimeTranscriptionProtocolError): class SpeechStreamingTarget: api_endpoint: str recognizer: str - access_token: str + resolve_access_token: Callable[[], Awaitable[str]] @dataclass(frozen=True, slots=True) @@ -249,16 +249,14 @@ class ChirpEventTransformer: finals: Final = tuple( result.transcript.strip() for result in frame.results if result.is_final and result.transcript.strip() ) - begin_events: Final = self._begin() if frame.speech_event == "begin" or interim or finals else () - interim_events: Final = self._hypothesis(interim) if interim else () + begin_events: Final = self._begin() if frame.speech_event == "begin" else () final_events: Final = tuple(event for final in finals for event in self._final(final)) + interim_events: Final = self._hypothesis(interim) if interim else () end_events: Final = self._stop() if frame.speech_event == "end" else () - return (*begin_events, *interim_events, *final_events, *end_events) + return (*begin_events, *final_events, *interim_events, *end_events) def _begin(self) -> tuple[OpenAIRealtimeEvents, ...]: - if self._turn is None: - self._turn = _Turn(item_id=self._new_item_id()) - turn: Final = self._turn + turn: Final = self._require_turn() if turn.started_emitted or not self._require_config().server_vad: return () self._turn = replace(turn, started_emitted=True) @@ -272,21 +270,23 @@ class ChirpEventTransformer: return (speech_event("input_audio_buffer.speech_stopped", turn.item_id),) def _hypothesis(self, text: str) -> tuple[OpenAIRealtimeEvents, ...]: + begin_events: Final = self._begin() turn: Final = self._require_turn() hypothesis: Final = _join_transcript(turn.committed, text) delta: Final = new_words(turn.preview, hypothesis) self._turn = replace(turn, preview=hypothesis) - return (delta_event(turn.item_id, delta),) if delta else () + return (*begin_events, delta_event(turn.item_id, delta)) if delta else begin_events def _final(self, text: str) -> tuple[OpenAIRealtimeEvents, ...]: + begin_events: Final = self._begin() turn: Final = self._require_turn() committed: Final = _join_transcript(turn.committed, text) delta: Final = new_words(turn.preview, committed) self._turn = replace(turn, committed=committed, preview=committed) delta_events: Final[tuple[OpenAIRealtimeEvents, ...]] = (delta_event(turn.item_id, delta),) if delta else () if not self._require_config().server_vad: - return delta_events - return (*delta_events, *self._complete()) + return (*begin_events, *delta_events) + return (*begin_events, *delta_events, *self._complete()) def _finish_turn(self) -> tuple[OpenAIRealtimeEvents, ...]: if self._turn is None: @@ -326,12 +326,12 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): def __init__( self, *, - access_token: str, + resolve_access_token: Callable[[], Awaitable[str]], project: str, location: str | None, backend_factory: Callable[[SpeechStreamingTarget], RealtimeBackend] = _default_backend_factory, ) -> None: - self._access_token: Final = access_token + self._resolve_access_token: Final = resolve_access_token self._project: Final = validate_vertex_transcription_project_id(project) self._location: Final = validate_vertex_transcription_location(location, DEFAULT_SPEECH_TO_TEXT_LOCATION) self._backend_factory: Final = backend_factory @@ -357,7 +357,7 @@ class VertexChirpRealtimeConfig(BaseRealtimeConfig): SpeechStreamingTarget( api_endpoint=url, recognizer=f"projects/{self._project}/locations/{self._location}/recognizers/_", - access_token=self._access_token, + resolve_access_token=self._resolve_access_token, ) ) diff --git a/litellm/llms/vertex_ai/realtime/transformation.py b/litellm/llms/vertex_ai/realtime/transformation.py index fe59034c27b..9fed6d52f0e 100644 --- a/litellm/llms/vertex_ai/realtime/transformation.py +++ b/litellm/llms/vertex_ai/realtime/transformation.py @@ -12,10 +12,16 @@ Auth: OAuth2 Bearer token (not an API key). """ import json +from collections.abc import Awaitable, Callable from typing import Final from litellm import verbose_logger from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig +from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import ( + VertexChirpRealtimeConfig, + is_vertex_speech_to_text_model, +) +from litellm.llms.vertex_ai.vertex_llm_base import VertexBase class VertexAIRealtimeConfig(GeminiRealtimeConfig): @@ -232,3 +238,20 @@ class VertexAIRealtimeConfig(GeminiRealtimeConfig): return [] return super().transform_realtime_request(message, model, session_configuration_request) + + +def vertex_realtime_config( + model: str, + *, + access_token: str, + resolve_access_token: Callable[[], Awaitable[str]], + project: str, + location: str | None, +) -> VertexAIRealtimeConfig | VertexChirpRealtimeConfig: + if is_vertex_speech_to_text_model(model): + return VertexChirpRealtimeConfig(resolve_access_token=resolve_access_token, project=project, location=location) + return VertexAIRealtimeConfig( + access_token=access_token, + project=project, + location=VertexBase.get_vertex_region(vertex_region=location, model=model), + ) diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index aed7bf15bdc..0e83edab5e1 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -38,11 +38,8 @@ from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_pr from ..llms.bedrock.realtime.handler import BedrockRealtime from ..llms.custom_httpx.http_handler import get_shared_realtime_ssl_context from ..llms.openai.realtime.handler import OpenAIRealtime -from ..llms.vertex_ai.audio_transcription.realtime_transformation import ( - VertexChirpRealtimeConfig, - is_vertex_speech_to_text_model, -) -from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig +from ..llms.vertex_ai.audio_transcription.realtime_transformation import is_vertex_speech_to_text_model +from ..llms.vertex_ai.realtime.transformation import VertexAIRealtimeConfig, vertex_realtime_config from ..llms.vertex_ai.vertex_llm_base import VertexBase from ..llms.xai.realtime.handler import XAIRealtime from ..utils import client as wrapper_client @@ -555,9 +552,19 @@ async def _arealtime( timeout_seconds=REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS, ) - vertex_realtime_config: Final = _vertex_realtime_config( - model=model, + async def resolve_vertex_access_token() -> str: + refreshed_token, _ = await _resolve_vertex_access_token_bounded( + credentials=vertex_credentials, + project_id=resolved_project, + resolver=vertex_access_token_resolver, + timeout_seconds=REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS, + ) + return refreshed_token + + vertex_provider_config: Final = vertex_realtime_config( + model, access_token=access_token, + resolve_access_token=resolve_vertex_access_token, project=resolved_project, location=vertex_location, ) @@ -566,7 +573,7 @@ async def _arealtime( model=model, websocket=websocket, logging_obj=litellm_logging_obj, - provider_config=vertex_realtime_config, + provider_config=vertex_provider_config, api_base=dynamic_api_base or litellm_params.api_base, api_key=None, client=client, @@ -580,18 +587,6 @@ async def _arealtime( raise ValueError(f"Unsupported model: {model}") -def _vertex_realtime_config( - model: str, access_token: str, project: str, location: str | None -) -> VertexAIRealtimeConfig | VertexChirpRealtimeConfig: - if is_vertex_speech_to_text_model(model): - return VertexChirpRealtimeConfig(access_token=access_token, project=project, location=location) - return VertexAIRealtimeConfig( - access_token=access_token, - project=project, - location=vertex_llm_base.get_vertex_region(vertex_region=location, model=model), - ) - - def _is_transcription_only_realtime_model(model: str, custom_llm_provider: str) -> bool: try: model_info: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider) 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 49758e62415..d6f65c90806 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,7 @@ import asyncio import json from collections.abc import AsyncIterator, Sequence +from dataclasses import replace from datetime import timedelta from typing import Final @@ -17,10 +18,15 @@ from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK 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 + +async def _static_token() -> str: + return "token" + + TARGET: Final = SpeechStreamingTarget( api_endpoint="us-speech.googleapis.com", recognizer="projects/proj-1/locations/us/recognizers/_", - access_token="token", + resolve_access_token=_static_token, ) CONFIGURE: Final = json.dumps( {"kind": "configure", "model": "chirp_3", "language_codes": ["en-US"], "sample_rate_hertz": 16_000} @@ -101,7 +107,7 @@ class _FakeSpeechClient: def _backend(client: _FakeSpeechClient, **kwargs: object) -> SpeechStreamingBackend: - return SpeechStreamingBackend(TARGET, client_factory=lambda target: client, **kwargs) + return SpeechStreamingBackend(TARGET, client_factory=lambda target, access_token: client, **kwargs) async def _recv(backend: SpeechStreamingBackend) -> dict[str, object]: @@ -346,6 +352,64 @@ async def test_rotation_is_forced_at_the_deadline_during_continuous_speech(): assert [_audio(stream) for stream in client.streams] == [[b"\x01\x01", b"\x02\x02"], [b"\x03\x03"]] +@pytest.mark.asyncio +async def test_every_stream_opens_its_own_client_with_a_freshly_resolved_token(): + now = [0.0] + tokens = iter(("token-1", "token-2")) + seen_tokens: list[str] = [] + clients = [_FakeSpeechClient([_response("first")]), _FakeSpeechClient([_response("second")])] + unopened = iter(clients) + + async def resolve_access_token() -> str: + return next(tokens) + + def open_client(target: SpeechStreamingTarget, access_token: str) -> _FakeSpeechClient: + seen_tokens.append(access_token) + return next(unopened) + + backend = SpeechStreamingBackend( + replace(TARGET, resolve_access_token=resolve_access_token), + client_factory=open_client, + clock=lambda: now[0], + rotation_seconds=240.0, + ) + async with backend: + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "first" + now[0] = 240.0 + await backend.send(b"\x02\x02") + assert await _transcript(backend) == "second" + assert clients[0].transport.closed + assert not clients[1].transport.closed + assert seen_tokens == ["token-1", "token-2"] + assert [len(client.streams) for client in clients] == [1, 1] + assert clients[1].transport.closed + + +@pytest.mark.asyncio +async def test_close_releases_a_rotated_stream_that_never_started_relaying(): + now = [0.0] + hold = asyncio.Event() + clients = [_FakeSpeechClient([_response("first"), hold]), _FakeSpeechClient([_response("never")])] + unopened = iter(clients) + backend = SpeechStreamingBackend( + TARGET, + client_factory=lambda target, access_token: next(unopened), + clock=lambda: now[0], + rotation_seconds=240.0, + ) + await _configure(backend) + await backend.send(b"\x01\x01") + assert await _transcript(backend) == "first" + now[0] = 240.0 + await backend.send(b"\x02\x02") + await asyncio.sleep(0) + assert clients[1].streams == [] + await backend.close() + assert [client.transport.closed for client in clients] == [True, True] + + @pytest.mark.asyncio async def test_discard_turn_cancels_every_stream_of_the_turn(): now = [0.0] 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 fcbfad8bf12..719dd621c82 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 @@ -63,8 +63,12 @@ def _ga_session_update( ) +async def _token() -> str: + return "token" + + def _config(location: str | None = "us") -> VertexChirpRealtimeConfig: - return VertexChirpRealtimeConfig(access_token="token", project="proj-1", location=location) + return VertexChirpRealtimeConfig(resolve_access_token=_token, project="proj-1", location=location) def _configured( @@ -261,6 +265,41 @@ def test_server_vad_turn_streams_new_words_then_completes_with_usage(): assert _backend_events(config, _response(speech_event="end")) == [] +def test_server_vad_final_result_completes_before_the_interim_that_follows_it(): + config = _configured() + _backend_events(config, _response(speech_event="begin")) + events = _backend_events(config, _response(("four score", True), ("and seven", False))) + assert _types(events) == [ + DELTA, + "input_audio_buffer.speech_stopped", + COMPLETED, + "input_audio_buffer.speech_started", + DELTA, + ] + assert events[2]["transcript"] == "four score" + assert events[4]["delta"] == "and seven" + assert events[4]["item_id"] != events[2]["item_id"] + assert events[4]["item_id"] == events[3]["item_id"] + finished = _backend_events(config, _response(("and seven years", True))) + assert [(event["type"], event.get("delta", event.get("transcript"))) for event in finished] == [ + (DELTA, " years"), + ("input_audio_buffer.speech_stopped", None), + (COMPLETED, "and seven years"), + ] + assert {event["item_id"] for event in finished} == {events[4]["item_id"]} + + +def test_manual_turn_keeps_the_interim_that_follows_a_final_in_the_same_frame(): + config = _configured(turn_detection=None) + first = _backend_events(config, _response(("four score", True), ("and seven", False))) + assert [(event["type"], event["delta"]) for event in first] == [(DELTA, "four score"), (DELTA, " and seven")] + second = _backend_events(config, _response(("and seven years", True))) + assert [event["delta"] for event in second] == [" years"] + completed = _backend_events(config, VertexSpeechStreamingTurnFinished()) + assert [(event["type"], event["transcript"]) for event in completed] == [(COMPLETED, "four score and seven years")] + assert {event["item_id"] for event in (*first, *second, *completed)} == {first[0]["item_id"]} + + def test_manual_turns_complete_on_commit_without_speech_events(): config = _configured(turn_detection=None) assert _backend_events(config, _response(speech_event="begin")) == [] @@ -322,7 +361,9 @@ async def test_open_backend_targets_the_regional_speech_endpoint(): targets.append(target) return _NullBackend() - config = VertexChirpRealtimeConfig(access_token="token", project="proj-1", location=None, backend_factory=factory) + config = VertexChirpRealtimeConfig( + resolve_access_token=_token, project="proj-1", location=None, backend_factory=factory + ) url = config.get_complete_url(None, "vertex_ai/chirp_3") assert url == "us-speech.googleapis.com" assert config.validate_environment({}, MODEL, "https://" + url) == {} @@ -332,9 +373,10 @@ async def test_open_backend_targets_the_regional_speech_endpoint(): SpeechStreamingTarget( api_endpoint="us-speech.googleapis.com", recognizer="projects/proj-1/locations/us/recognizers/_", - access_token="token", + resolve_access_token=_token, ) ] + assert await targets[0].resolve_access_token() == "token" @pytest.mark.parametrize( 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 5/7] 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)) From 3911d62bbeee55f940ac4294275ada2c5580bf88 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 17:44:28 -0700 Subject: [PATCH 6/7] fix(vertex_ai): prune a discarded turn's id once its marker is delivered --- .../audio_transcription/realtime_backend.py | 27 ++++++++++++++----- .../test_vertex_ai_realtime_backend.py | 1 + 2 files changed, 21 insertions(+), 7 deletions(-) diff --git a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py index 4c8338c027e..859a883463c 100644 --- a/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py +++ b/litellm/llms/vertex_ai/audio_transcription/realtime_backend.py @@ -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)) 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 15601c5ca6c..d5e88706e23 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 @@ -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}] From 17c519c40a3c98671eac5966eaed3ab9cb8b2261 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 18 Sep 2026 18:09:33 -0700 Subject: [PATCH 7/7] test(custom_httpx): pass the token resolver and two-argument client factory in the realtime bridge test --- .../llms/custom_httpx/test_llm_http_handler.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py index 1f71ffd43f6..82220d53375 100644 --- a/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/test_litellm/llms/custom_httpx/test_llm_http_handler.py @@ -3780,11 +3780,17 @@ async def test_async_realtime_bridges_a_transcription_session_through_the_provid yield script.pop(0) speech_client = FakeSpeechClient() + + async def resolve_access_token() -> str: + return "token" + provider_config = VertexChirpRealtimeConfig( - access_token="token", + resolve_access_token=resolve_access_token, project="proj-1", location="us", - backend_factory=lambda target: SpeechStreamingBackend(target, client_factory=lambda target: speech_client), + backend_factory=lambda target: SpeechStreamingBackend( + target, client_factory=lambda target, access_token: speech_client + ), ) audio = base64.b64encode(b"\x00\x01" * 800).decode() client_ws = _ScriptedClientWebSocket(