mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
Merge pull request #41721 from BerriAI/litellm_vertex_chirp3_streaming_stt
feat(vertex_ai): stream Chirp speech-to-text over /v1/realtime
This commit is contained in:
commit
9486caf584
21 changed files with 2636 additions and 208 deletions
|
|
@ -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 = "",
|
||||
|
|
|
|||
275
litellm/llms/base_llm/realtime/transcription_protocol.py
Normal file
275
litellm/llms/base_llm/realtime/transcription_protocol.py
Normal file
|
|
@ -0,0 +1,275 @@
|
|||
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")
|
||||
|
||||
|
||||
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 error("invalid JSON object") from None
|
||||
if not isinstance(value, dict):
|
||||
raise error("message must be a JSON object")
|
||||
return value
|
||||
|
||||
|
||||
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 error(f"{name} must be an object")
|
||||
return value
|
||||
|
||||
|
||||
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 error(f"{name} must be a string")
|
||||
return value
|
||||
|
||||
|
||||
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 error(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, error: ProtocolErrorType = RealtimeTranscriptionProtocolError
|
||||
) -> TranscriptionSessionUpdate:
|
||||
message: Final = json_object(payload, error)
|
||||
if message.get("type") not in SESSION_UPDATE_EVENT_TYPES:
|
||||
raise error("expected session.update")
|
||||
session: Final = json_mapping(message.get("session"), "session", error)
|
||||
if not session:
|
||||
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 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", 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", 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], 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 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", error),
|
||||
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", error)
|
||||
return TranscriptionAudioFormat(
|
||||
layout="ga",
|
||||
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,
|
||||
error: ProtocolErrorType = RealtimeTranscriptionProtocolError,
|
||||
) -> bytes:
|
||||
if not isinstance(audio, str):
|
||||
raise error("Audio must be a base64 string")
|
||||
if max_encoded_bytes is not None and len(audio) > max_encoded_bytes:
|
||||
raise error("Audio append exceeds the four-second backlog limit")
|
||||
try:
|
||||
decoded: Final = base64.b64decode(audio, validate=True)
|
||||
except (binascii.Error, ValueError):
|
||||
raise error("Audio must be valid base64") from None
|
||||
if len(decoded) % 2:
|
||||
raise error("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
|
||||
|
|
@ -1,8 +1,10 @@
|
|||
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
|
||||
|
||||
import httpx
|
||||
from typing_extensions import Self
|
||||
|
||||
from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents
|
||||
from litellm.types.realtime import (
|
||||
|
|
@ -21,6 +23,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 +97,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,
|
||||
|
|
|
|||
|
|
@ -6287,7 +6287,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:
|
||||
|
|
|
|||
|
|
@ -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, MuseProtocolError)
|
||||
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, MuseProtocolError)
|
||||
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, MuseProtocolError)
|
||||
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, MuseProtocolError)
|
||||
buffered: Final = self._pending_audio + audio
|
||||
packet_end: Final = len(buffered) - len(buffered) % config.packet_bytes
|
||||
self._pending_audio = buffered[packet_end:]
|
||||
|
|
|
|||
418
litellm/llms/vertex_ai/audio_transcription/realtime_backend.py
Normal file
418
litellm/llms/vertex_ai/audio_transcription/realtime_backend.py
Normal file
|
|
@ -0,0 +1,418 @@
|
|||
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
|
||||
|
||||
from pydantic import TypeAdapter
|
||||
from typing_extensions import Self, assert_never
|
||||
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
|
||||
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()
|
||||
_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
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _TurnResult:
|
||||
turn: int
|
||||
event: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _TurnDiscarded:
|
||||
turn: int
|
||||
|
||||
|
||||
@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:
|
||||
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=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()
|
||||
|
||||
|
||||
def _normal_closure() -> ConnectionClosedOK:
|
||||
return ConnectionClosedOK(rcvd=Close(1000, ""), sent=None)
|
||||
|
||||
|
||||
class _RecognizeStream:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
client: SpeechStreamingClient,
|
||||
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
|
||||
)
|
||||
self._requests.put_nowait(first_request)
|
||||
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:
|
||||
await self._requests.put(self._request_type(audio=audio))
|
||||
|
||||
async def half_close(self) -> None:
|
||||
await self._requests.put(None)
|
||||
|
||||
def cancel(self) -> None:
|
||||
self._cancelled = True
|
||||
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[_OutboxItem], 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
|
||||
try:
|
||||
await asyncio.wait((task,))
|
||||
except asyncio.CancelledError:
|
||||
task.cancel()
|
||||
await asyncio.wait((task,))
|
||||
raise
|
||||
finally:
|
||||
await self.close()
|
||||
return self.billed_seconds
|
||||
|
||||
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(
|
||||
_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}"))
|
||||
|
||||
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:
|
||||
yield request
|
||||
|
||||
|
||||
_Link = _RecognizeStream | str | _TurnDiscarded
|
||||
|
||||
|
||||
class SpeechStreamingBackend:
|
||||
def __init__(
|
||||
self,
|
||||
target: SpeechStreamingTarget,
|
||||
*,
|
||||
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,
|
||||
) -> None:
|
||||
self._target: Final = target
|
||||
self._client_factory: Final = client_factory
|
||||
self._clock: Final = clock
|
||||
self._rotation_seconds: Final = rotation_seconds
|
||||
self._rotation_deadline_seconds: Final = rotation_deadline_seconds
|
||||
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
|
||||
|
||||
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 self._closed:
|
||||
raise _normal_closure()
|
||||
if isinstance(message, bytes):
|
||||
await self._send_audio(message)
|
||||
return
|
||||
command: Final = _COMMAND_ADAPTER.validate_json(message)
|
||||
match command:
|
||||
case VertexSpeechStreamingConfigure():
|
||||
self._config = _streaming_config(command)
|
||||
await self._link(_CONFIGURED_EVENT)
|
||||
case VertexSpeechStreamingFinishTurn():
|
||||
await self._finish_turn()
|
||||
case VertexSpeechStreamingDiscardTurn():
|
||||
await self._discard_turn()
|
||||
case _:
|
||||
assert_never(command)
|
||||
|
||||
async def recv(self, decode: bool | None = None) -> str | bytes:
|
||||
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(
|
||||
rcvd=Close(STREAM_FAILURE_CLOSE_CODE, item.reason[:_CLOSE_REASON_MAX_CHARS]), sent=None
|
||||
)
|
||||
case _Closed():
|
||||
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 _:
|
||||
assert_never(item)
|
||||
|
||||
async def close(self) -> None:
|
||||
if self._closed:
|
||||
return
|
||||
self._closed = True
|
||||
self._turn = ()
|
||||
pump: Final = self._pump
|
||||
if pump is not None:
|
||||
pump.cancel()
|
||||
await asyncio.wait((pump,))
|
||||
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: _Link) -> None:
|
||||
if self._pump is None:
|
||||
self._pump = asyncio.create_task(self._pump_links())
|
||||
await self._links.put(item)
|
||||
|
||||
async def _pump_links(self) -> None:
|
||||
while True:
|
||||
await self._relay(await self._links.get())
|
||||
|
||||
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(
|
||||
_TurnDiscardedEvent(
|
||||
turn=link.turn,
|
||||
event=VertexSpeechStreamingTurnDiscarded(billed_seconds=self._billed_before).model_dump_json(),
|
||||
)
|
||||
)
|
||||
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
|
||||
if config is None:
|
||||
raise RuntimeError("audio was sent before the Speech-to-Text stream was configured")
|
||||
access_token: Final = await self._target.resolve_access_token()
|
||||
stream: Final = _RecognizeStream(
|
||||
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(),
|
||||
turn=self._turn_index,
|
||||
)
|
||||
await self._link(stream)
|
||||
return stream
|
||||
|
||||
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)
|
||||
|
||||
async def _discard_turn(self) -> None:
|
||||
streams: Final = self._turn
|
||||
turn: Final = self._turn_index
|
||||
self._turn = ()
|
||||
self._discarded_turns |= {turn}
|
||||
self._turn_index += 1
|
||||
for stream in streams:
|
||||
stream.cancel()
|
||||
await self._link(_TurnDiscarded(turn=turn))
|
||||
|
|
@ -0,0 +1,446 @@
|
|||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from dataclasses import dataclass, replace
|
||||
from typing import Final
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
from typing_extensions import assert_never
|
||||
|
||||
import litellm
|
||||
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
|
||||
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_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()
|
||||
_DISCARD_TURN_COMMAND: Final = VertexSpeechStreamingDiscardTurn().model_dump_json()
|
||||
|
||||
|
||||
class ChirpProtocolError(RealtimeTranscriptionProtocolError):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SpeechStreamingTarget:
|
||||
api_endpoint: str
|
||||
recognizer: str
|
||||
resolve_access_token: Callable[[], Awaitable[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:
|
||||
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:
|
||||
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, 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:
|
||||
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._billed_seconds = max(self._billed_seconds, frame.billed_seconds)
|
||||
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)
|
||||
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" 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, *final_events, *interim_events, *end_events)
|
||||
|
||||
def _begin(self) -> tuple[OpenAIRealtimeEvents, ...]:
|
||||
turn: Final = self._require_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, ...]:
|
||||
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 (*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 (*begin_events, *delta_events)
|
||||
return (*begin_events, *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,
|
||||
*,
|
||||
resolve_access_token: Callable[[], Awaitable[str]],
|
||||
project: str,
|
||||
location: str | None,
|
||||
backend_factory: Callable[[SpeechStreamingTarget], RealtimeBackend] = _default_backend_factory,
|
||||
) -> None:
|
||||
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
|
||||
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/_",
|
||||
resolve_access_token=self._resolve_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, ChirpProtocolError)
|
||||
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"), error=ChirpProtocolError)
|
||||
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]
|
||||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -48905,7 +48905,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": {
|
||||
|
|
|
|||
|
|
@ -38,7 +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.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
|
||||
|
|
@ -541,8 +542,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,
|
||||
|
|
@ -553,17 +552,28 @@ async def _arealtime(
|
|||
timeout_seconds=REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS,
|
||||
)
|
||||
|
||||
vertex_realtime_config: Final = VertexAIRealtimeConfig(
|
||||
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=resolved_location,
|
||||
location=vertex_location,
|
||||
)
|
||||
|
||||
await base_llm_http_handler.async_realtime(
|
||||
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,
|
||||
|
|
@ -684,6 +694,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),
|
||||
|
|
|
|||
|
|
@ -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,66 @@ 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"
|
||||
billed_seconds: float
|
||||
|
||||
|
||||
VertexSpeechStreamingEventUnion = (
|
||||
VertexSpeechStreamingResponse
|
||||
| VertexSpeechStreamingConfigured
|
||||
| VertexSpeechStreamingTurnFinished
|
||||
| VertexSpeechStreamingTurnDiscarded
|
||||
)
|
||||
VertexSpeechStreamingEvent = Annotated[VertexSpeechStreamingEventUnion, Field(discriminator="kind")]
|
||||
|
|
|
|||
|
|
@ -48905,7 +48905,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": {
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
@ -270,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",
|
||||
|
|
|
|||
0
tests/test_litellm/llms/base_llm/realtime/__init__.py
Normal file
0
tests/test_litellm/llms/base_llm/realtime/__init__.py
Normal file
|
|
@ -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})
|
||||
|
|
@ -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,149 @@ 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()
|
||||
|
||||
async def resolve_access_token() -> str:
|
||||
return "token"
|
||||
|
||||
provider_config = VertexChirpRealtimeConfig(
|
||||
resolve_access_token=resolve_access_token,
|
||||
project="proj-1",
|
||||
location="us",
|
||||
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(
|
||||
[
|
||||
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]
|
||||
|
|
|
|||
|
|
@ -0,0 +1,491 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Callable, Sequence
|
||||
from dataclasses import replace
|
||||
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 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/_",
|
||||
resolve_access_token=_static_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 | asyncio.Event
|
||||
|
||||
|
||||
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 await self._next(script)
|
||||
while script:
|
||||
yield await self._next(script)
|
||||
|
||||
@staticmethod
|
||||
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
|
||||
|
||||
|
||||
def _backend(client: _FakeSpeechClient, **kwargs: object) -> SpeechStreamingBackend:
|
||||
return SpeechStreamingBackend(TARGET, client_factory=lambda target, access_token: client, **kwargs)
|
||||
|
||||
|
||||
async def _recv(backend: SpeechStreamingBackend) -> dict[str, object]:
|
||||
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"}
|
||||
|
||||
|
||||
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:]]
|
||||
|
||||
|
||||
@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_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 _transcript(backend) == "hi"
|
||||
await backend.close()
|
||||
with pytest.raises(ConnectionClosedOK):
|
||||
await backend.recv()
|
||||
with pytest.raises(ConnectionClosedOK):
|
||||
await backend.send(b"\x00\x00")
|
||||
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", "billed_seconds": 0.0}
|
||||
|
||||
|
||||
@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 _transcript(backend) == "draft"
|
||||
await backend.send(DISCARD_TURN)
|
||||
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}
|
||||
assert backend._discarded_turns == frozenset()
|
||||
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(
|
||||
[_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_ending_the_turn():
|
||||
now = [0.0]
|
||||
client = _FakeSpeechClient(
|
||||
[_response("first"), _response("first half", is_final=True, billed=239.0)],
|
||||
[_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 _transcript(backend) == "first"
|
||||
now[0] = 239.0
|
||||
await backend.send(b"\x02\x02")
|
||||
assert await _transcript(backend) == "first half"
|
||||
now[0] = 240.0
|
||||
await backend.send(b"\x03\x03")
|
||||
second = await _recv(backend)
|
||||
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_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]
|
||||
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", "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"]]
|
||||
|
||||
|
||||
@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"
|
||||
|
|
@ -0,0 +1,424 @@
|
|||
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,
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
async def _token() -> str:
|
||||
return "token"
|
||||
|
||||
|
||||
def _config(location: str | None = "us") -> VertexChirpRealtimeConfig:
|
||||
return VertexChirpRealtimeConfig(resolve_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", 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):
|
||||
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_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")) == []
|
||||
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(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))
|
||||
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(
|
||||
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) == {}
|
||||
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/_",
|
||||
resolve_access_token=_token,
|
||||
)
|
||||
]
|
||||
assert await targets[0].resolve_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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
27
uv.lock
generated
27
uv.lock
generated
|
|
@ -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" },
|
||||
]
|
||||
|
|
@ -4607,6 +4628,7 @@ ci = [
|
|||
{ name = "blockbuster" },
|
||||
{ name = "claude-agent-sdk" },
|
||||
{ name = "detect-secrets" },
|
||||
{ name = "google-cloud-speech" },
|
||||
{ name = "google-generativeai" },
|
||||
{ name = "jsonlines" },
|
||||
{ name = "langchain" },
|
||||
|
|
@ -4721,6 +4743,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 +4811,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 = [
|
||||
|
|
@ -4799,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" },
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue