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