From 3ef0b352559d768b38543c81c272ec1fd61d8973 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 17 Oct 2026 00:29:50 +0000 Subject: [PATCH] fix(bedrock)!: block realtime replay after an input-first start and reject malformed Nova Sonic session settings (#46936) * fix(bedrock): block realtime replay after input starts the session and reject non-string instructions and voice * test(bedrock): name the malformed realtime setting test for its string cases * fix(bedrock): replay the opening realtime message until Bedrock accepts the stream and refuse any fallback after commit Also rejects a malformed modalities, input_sample_rate_hertz or tools setting with invalid_type, dropped under drop_params. * fix(bedrock): replay a refused opening message on any fallback and reject OpenAI voices and non-string audio formats * fix(realtime): read the client scope only when the websocket has one and reject any-case OpenAI and empty Bedrock voices * test(bedrock): cover realtime replay, refusal and committed failure through a real proxy Fix the scripted Bedrock peer's port import and the nested tool call cell's stale expected events, so the Nova Sonic integration file collects again --- litellm/constants.py | 2 +- .../litellm_core_utils/realtime_streaming.py | 23 +- litellm/llms/bedrock/realtime/handler.py | 93 +++--- .../llms/bedrock/realtime/transformation.py | 123 ++++++-- litellm/realtime_api/main.py | 4 +- .../_support/bedrock_realtime_peer.py | 146 ++++++++-- ..._bedrock_nova_sonic_realtime_tool_calls.py | 269 +++++++++++++++++- .../test_realtime_streaming.py | 27 ++ .../realtime/test_bedrock_realtime_handler.py | 73 ++++- .../test_bedrock_realtime_transformation.py | 63 +++- 10 files changed, 726 insertions(+), 97 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 425fc22d340..f448fc351c3 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -391,7 +391,7 @@ DEEPGRAM_DEFAULT_API_BASE: Final = "https://api.deepgram.com/v1" NADIR_DEFAULT_API_BASE: Final = "https://api.getnadir.com/v1" DEEPGRAM_LISTEN_DEFAULT_MODEL: Final = "nova-3" -BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" +BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_opening_message" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure" BEDROCK_REALTIME_SDK_DISTRIBUTION: Final = "aws-sdk-bedrock-runtime" diff --git a/litellm/litellm_core_utils/realtime_streaming.py b/litellm/litellm_core_utils/realtime_streaming.py index e4eb9acc37d..927caec3350 100644 --- a/litellm/litellm_core_utils/realtime_streaming.py +++ b/litellm/litellm_core_utils/realtime_streaming.py @@ -1,11 +1,11 @@ import asyncio import json import traceback -from collections.abc import Callable, Coroutine, Mapping, Sequence +from collections.abc import Callable, Coroutine, Mapping, MutableMapping, Sequence from dataclasses import dataclass from enum import Enum, auto from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeAlias, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, NoReturn, Protocol, TypeAlias, TypedDict, cast, runtime_checkable from pydantic import TypeAdapter from typing_extensions import ReadOnly @@ -14,6 +14,7 @@ import litellm from litellm._logging import redact_internal_details_from_client_message, verbose_logger from litellm._uuid import uuid from litellm.constants import ( + BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, REALTIME_METERED_USAGE_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, @@ -100,6 +101,16 @@ class ScopedWebSocket(Protocol): def scope(self) -> _ASGIScope: ... +@runtime_checkable +class _ReplayScopeWebSocket(Protocol): + @property + def scope(self) -> MutableMapping[str, object]: ... + + +def client_websocket_scope(websocket: object) -> MutableMapping[str, object] | None: + return websocket.scope if isinstance(websocket, _ReplayScopeWebSocket) else None + + class _ClientWebSocket(ScopedWebSocket, Protocol): async def send_text(self, data: str) -> None: ... async def receive_text(self) -> str: ... @@ -233,6 +244,7 @@ class RealTimeStreaming: logging_worker: _LoggingWorker = GLOBAL_LOGGING_WORKER, ): self.websocket: _ClientWebSocket = websocket + self._replay_scope: Final = client_websocket_scope(websocket) self.backend_ws = backend_ws self.logging_obj = logging_obj self._logging_worker = logging_worker @@ -1451,6 +1463,13 @@ class RealTimeStreaming: return item async def _receive_client_message(self) -> str | None: + refused_opening_message: Final = ( + None + if self._replay_scope is None + else self._replay_scope.pop(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, None) + ) + if isinstance(refused_opening_message, str): + return refused_opening_message try: return await self.websocket.receive_text() except Exception as e: # noqa: BLE001 # whatever the client socket raises, the client is gone diff --git a/litellm/llms/bedrock/realtime/handler.py b/litellm/llms/bedrock/realtime/handler.py index c0a382b136c..3e7c2f7eee9 100644 --- a/litellm/llms/bedrock/realtime/handler.py +++ b/litellm/llms/bedrock/realtime/handler.py @@ -8,7 +8,7 @@ import asyncio import contextlib import importlib.metadata import json -from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence from dataclasses import dataclass from types import MappingProxyType from typing import Final, NoReturn, Protocol, runtime_checkable @@ -21,7 +21,7 @@ from litellm._uuid import uuid from litellm.constants import ( BEDROCK_REALTIME_CLOSE_TIMEOUT_SECONDS, BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY, - BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, + BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, BEDROCK_REALTIME_SDK_DISTRIBUTION, BEDROCK_REALTIME_SDK_SUPPORTED_RANGE, BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY, @@ -37,9 +37,8 @@ from litellm.types.realtime import RealtimeResponseTransformInput from ..base_aws_llm import BaseAWSLLM, run_aws_signing from ..common_utils import BedrockError -from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection +from .transformation import BedrockRealtimeConfig, BedrockSessionUpdateRejection, requested_modalities -_CLIENT_MODALITIES_ADAPTER: Final[TypeAdapter["list[str] | None"]] = TypeAdapter(list[str] | None) _CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({}) @@ -177,25 +176,35 @@ async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: return -def _pending_session_update(scope: Mapping[str, object]) -> str | None: - """A fallback attempt on the same websocket replays the session.update the failed attempt never acked.""" - if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True: - committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY) - raise BedrockError( - status_code=400, - message=( - "Bedrock realtime session already committed to a provider stream; it cannot be replayed" - + (f". The committed stream failed with: {committed_failure}" if committed_failure else "") - ), - ) - pending: Final = scope.get(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY) +def _is_session_committed(scope: Mapping[str, object]) -> bool: + return scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True + + +def raise_if_session_committed(scope: Mapping[str, object] | None) -> None: + """A fallback attempt on a websocket whose Bedrock stream already took input would silently start over.""" + if scope is None or not _is_session_committed(scope): + return + committed_failure: Final = scope.get(BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY) + raise BedrockError( + status_code=400, + message=( + "Bedrock realtime session already committed to a provider stream; it cannot be replayed" + + (f". The committed stream failed with: {committed_failure}" if committed_failure else "") + ), + ) + + +def _pending_opening_message(scope: Mapping[str, object]) -> str | None: + """A fallback attempt on the same websocket replays the opening message Bedrock never accepted a stream for.""" + raise_if_session_committed(scope) + pending: Final = scope.get(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY) return pending if isinstance(pending, str) else None def _raise_provider_failure(scope: MutableMapping[str, object], failure: BaseException) -> NoReturn: error: Final = _as_bedrock_error(failure) verbose_proxy_logger.error("Bedrock Realtime: provider stream failed: %s", redact_string(str(error))) - if scope.get(BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY) is True: + if _is_session_committed(scope): scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = redact_string(str(error)) raise error from failure @@ -207,26 +216,36 @@ def _parse_client_message(message: str) -> Mapping[str, JsonValue]: return _EMPTY_JSON_OBJECT -async def _ack_session_update( +async def _send_opening_message( client_ws: RealtimeClientWebSocket, bedrock_stream: BedrockBidirectionalStream, + message: str, + bedrock_messages: Sequence[str], + send_to_bedrock: Callable[[str], Awaitable[None]], +) -> None: + """Keep the opening message for a fallback to replay until Bedrock accepts the stream it opened.""" + client_ws.scope[BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY] = message # rebind-ok: scope outlives attempts + for bedrock_message in bedrock_messages: + await send_to_bedrock(bedrock_message) + await bedrock_stream.await_output() + client_ws.scope.pop(BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, None) + client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt + + +async def _ack_session_update( + client_ws: RealtimeClientWebSocket, transformation_config: BedrockRealtimeConfig, model: str, logging_obj: LiteLLMLogging | None, parsed_client_message: Mapping[str, JsonValue], ) -> bool: - """Ack the client's session.update once Bedrock accepted the stream; False means the client is gone.""" - await bedrock_stream.await_output() - client_ws.scope.pop(BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY, None) - client_ws.scope[BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY] = True # rebind-ok: scope outlives the attempt + """Ack the client's session.update; False means the client is gone.""" if logging_obj is None: return True - requested_modalities: Final = _CLIENT_MODALITIES_ADAPTER.validate_python( - _json_dict(parsed_client_message.get("session")).get("modalities") - ) + modalities: Final = requested_modalities(_json_dict(parsed_client_message.get("session"))) try: await client_ws.send_text( - json.dumps(transformation_config.session_updated_event(model, logging_obj, requested_modalities)) + json.dumps(transformation_config.session_updated_event(model, logging_obj, modalities)) ) except Exception as e: # noqa: BLE001 # any send failure means the client is gone verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True) @@ -314,7 +333,7 @@ class BedrockRealtime(BaseAWSLLM): except ImportError as e: raise _sdk_import_error(self._sdk_version_lookup(), e) from e - pending_session_update: Final = _pending_session_update(websocket.scope) + pending_opening_message: Final = _pending_opening_message(websocket.scope) # Get AWS region if aws_region_name is None: @@ -376,7 +395,7 @@ class BedrockRealtime(BaseAWSLLM): try: await self._run_session( - websocket, open_bidirectional_stream, model, logging_obj, pending_session_update, drop_params + websocket, open_bidirectional_stream, model, logging_obj, pending_opening_message, drop_params ) finally: await _close_bedrock_client(bedrock_client) @@ -387,7 +406,7 @@ class BedrockRealtime(BaseAWSLLM): open_bidirectional_stream: Callable[[], Awaitable[BedrockBidirectionalStream]], model: str, logging_obj: LiteLLMLogging, - pending_session_update: str | None, + pending_opening_message: str | None, drop_params: bool = False, ) -> None: transformation_config: Final = BedrockRealtimeConfig(drop_params=drop_params) @@ -396,7 +415,7 @@ class BedrockRealtime(BaseAWSLLM): verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established") - if pending_session_update is None: + if pending_opening_message is None: await websocket.send_text(json.dumps(transformation_config.session_created_event(model, logging_obj))) verbose_proxy_logger.debug("Bedrock Realtime: sent session.created to client on connect") @@ -418,7 +437,7 @@ class BedrockRealtime(BaseAWSLLM): model, session_state, logging_obj, - initial_message=pending_session_update, + initial_message=pending_opening_message, ) logged_events: Final = ( @@ -531,18 +550,22 @@ class BedrockRealtime(BaseAWSLLM): if not await _reject_session_update(client_ws, parsed_client_message, rejection): break continue - client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = message transformed_messages = transformation_config.transform_realtime_request( message=message, model=model, session_configuration_request=session_state.get("session_configuration_request"), ) - for bedrock_message in transformed_messages: - await send_to_bedrock(bedrock_message) + if transformed_messages and not _is_session_committed(client_ws.scope): + await _send_opening_message( + client_ws, bedrock_stream, message, transformed_messages, send_to_bedrock + ) + else: + for bedrock_message in transformed_messages: + await send_to_bedrock(bedrock_message) if is_session_update and not await _ack_session_update( - client_ws, bedrock_stream, transformation_config, model, logging_obj, parsed_client_message + client_ws, transformation_config, model, logging_obj, parsed_client_message ): break finally: diff --git a/litellm/llms/bedrock/realtime/transformation.py b/litellm/llms/bedrock/realtime/transformation.py index e8bc0482a28..9cb38822881 100644 --- a/litellm/llms/bedrock/realtime/transformation.py +++ b/litellm/llms/bedrock/realtime/transformation.py @@ -123,26 +123,81 @@ def _integer_setting(value: JsonValue) -> int | None: return None if number is None else int(number) -def _malformed_number_settings(session_config: Mapping[str, JsonValue]) -> tuple[str, ...]: +def requested_modalities(session_config: Mapping[str, JsonValue]) -> tuple[str, ...] | None: + modalities: Final = session_config.get("modalities") + if not isinstance(modalities, list): + return None + names: Final = tuple(modality for modality in modalities if isinstance(modality, str)) + return names if len(names) == len(modalities) else None + + +_OPENAI_VOICES: Final = frozenset( + ("alloy", "ash", "ballad", "coral", "echo", "sage", "shimmer", "verse", "marin", "cedar") +) + + +def _unplayable_voice(session_config: Mapping[str, JsonValue]) -> str | None: + voice: Final = session_config.get("voice") + if not isinstance(voice, str): + return None + return voice if not voice.strip() or voice.lower() in _OPENAI_VOICES else None + + +def _malformed_settings(session_config: Mapping[str, JsonValue]) -> tuple[tuple[str, str], ...]: + instructions: Final = session_config.get("instructions") + voice: Final = session_config.get("voice") + input_audio_format: Final = session_config.get("input_audio_format") + output_audio_format: Final = session_config.get("output_audio_format") + modalities: Final = session_config.get("modalities") temperature: Final = session_config.get("temperature") max_output_tokens: Final = session_config.get("max_response_output_tokens") + input_rate_hertz: Final = session_config.get("input_sample_rate_hertz") output_rate_hertz: Final = session_config.get("output_sample_rate_hertz") + tools: Final = session_config.get("tools") malformed: Final = ( - ("temperature", temperature is not None and _number_setting(temperature) is None), + ("instructions", "a string", instructions is not None and not isinstance(instructions, str)), + ("voice", "a string", voice is not None and not isinstance(voice, str)), + ("tools", "a list", tools is not None and not isinstance(tools, list)), + ( + "input_audio_format", + "a string", + input_audio_format is not None and not isinstance(input_audio_format, str), + ), + ( + "output_audio_format", + "a string", + output_audio_format is not None and not isinstance(output_audio_format, str), + ), + ( + "modalities", + "a list of strings", + modalities is not None and requested_modalities(session_config) is None, + ), + ("temperature", "a number", temperature is not None and _number_setting(temperature) is None), ( "max_response_output_tokens", + "a number or 'inf'", max_output_tokens is not None and max_output_tokens != "inf" and _integer_setting(max_output_tokens) is None, ), - ("output_sample_rate_hertz", output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None), + ( + "input_sample_rate_hertz", + "a number", + input_rate_hertz is not None and _integer_setting(input_rate_hertz) is None, + ), + ( + "output_sample_rate_hertz", + "a number", + output_rate_hertz is not None and _integer_setting(output_rate_hertz) is None, + ), ) - return tuple(name for name, is_malformed in malformed if is_malformed) + return tuple((name, expected) for name, expected, is_malformed in malformed if is_malformed) @dataclasses.dataclass(frozen=True, slots=True) class BedrockSessionUpdateRejection: - code: Literal["invalid_type", "session_update_not_applied"] + code: Literal["invalid_type", "invalid_value", "session_update_not_applied"] message: str param: str @@ -312,7 +367,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): return 8000 # G.711 typically uses 8kHz return 24000 if is_output else 16000 - def transform_session_update_event(self, json_message: dict) -> list[str]: + def transform_session_update_event(self, json_message: Mapping[str, JsonValue]) -> list[str]: """ Transform session.update event to Bedrock session configuration. @@ -325,7 +380,10 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): verbose_logger.debug("Handling session.update") messages: Final[list[str]] = [] - session_config: Final = json_message.get("session", {}) + requested_session: Final = json_message.get("session") + session_config: Final[Mapping[str, JsonValue]] = ( + requested_session if isinstance(requested_session, dict) else {} + ) locked_settings: Final = self._apply_locked_settings(session_config, self._started_settings) # Settings Bedrock locks at prompt start may only change before it @@ -336,13 +394,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self.output_sample_rate_hertz = locked_settings.output_sample_rate_hertz # Update audio input configuration from session if provided - if "input_audio_format" in session_config: - input_format: Final = session_config["input_audio_format"] + input_format: Final = session_config.get("input_audio_format") + if isinstance(input_format, str): self.input_sample_rate_hertz = self._map_audio_format_to_sample_rate(input_format, is_output=False) # Allow direct override of the input sample rate if provided (custom extension) - if "input_sample_rate_hertz" in session_config: - self.input_sample_rate_hertz = session_config["input_sample_rate_hertz"] + requested_input_rate: Final = _integer_setting(session_config.get("input_sample_rate_hertz")) + if requested_input_rate is not None: + self.input_sample_rate_hertz = requested_input_rate if self.prompt_started: return messages @@ -352,7 +411,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): messages.append(json.dumps(session_start)) # Send prompt start - tools: Final[Sequence[JsonValue] | None] = session_config.get("tools") + tools: Final = session_config.get("tools") prompt_start_config: Final[Mapping[str, object]] = { "promptName": self.prompt_name, "textOutputConfiguration": {"mediaType": self.text_media_type}, @@ -370,7 +429,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): "toolUseOutputConfiguration": {"mediaType": "application/json"}, "toolConfiguration": {"tools": self._transform_tools_to_bedrock_format(tools)}, } - if tools + if isinstance(tools, list) and tools else {} ), } @@ -445,12 +504,14 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): if isinstance(requested_tools, list) else base.tools_json ) - system_prompt: Final = base.system_prompt if instructions is None else str(instructions) - voice_id: Final = base.voice_id if voice is None else str(voice) + system_prompt: Final = instructions if isinstance(instructions, str) else base.system_prompt + voice_id: Final = ( + voice if isinstance(voice, str) and _unplayable_voice(session_config) is None else base.voice_id + ) format_sample_rate: Final = ( - base.output_sample_rate_hertz - if output_audio_format is None - else self._map_audio_format_to_sample_rate(str(output_audio_format), is_output=True) + self._map_audio_format_to_sample_rate(output_audio_format, is_output=True) + if isinstance(output_audio_format, str) + else base.output_sample_rate_hertz ) requested_output_rate: Final = _integer_setting(session_config.get("output_sample_rate_hertz")) output_sample_rate_hertz: Final = ( @@ -476,16 +537,26 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): session_config: Final = json_message.get("session") if not isinstance(session_config, Mapping): return None - malformed: Final = () if self.drop_params else _malformed_number_settings(session_config) + malformed: Final = () if self.drop_params else _malformed_settings(session_config) if malformed: + name, expected = malformed[0] return BedrockSessionUpdateRejection( code="invalid_type", message=( - f"Invalid type for 'session.{malformed[0]}': Bedrock Nova Sonic needs a number for temperature " - "and output_sample_rate_hertz, and a number or 'inf' for max_response_output_tokens. " + f"Invalid type for 'session.{name}': Bedrock Nova Sonic needs {expected} there. " "Set drop_params to drop it instead" ), - param=f"session.{malformed[0]}", + param=f"session.{name}", + ) + unplayable_voice: Final = None if self.drop_params else _unplayable_voice(session_config) + if unplayable_voice is not None: + return BedrockSessionUpdateRejection( + code="invalid_value", + message=( + f"Invalid value for 'session.voice': '{unplayable_voice}' is not a Bedrock Nova Sonic voice. " + "Set drop_params to drop it instead" + ), + param="session.voice", ) if not self.prompt_started: return None @@ -785,7 +856,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): List of Bedrock format messages (JSON strings) """ try: - json_message: Final[dict[str, object]] = json.loads(message) + json_message: Final[dict[str, JsonValue]] = json.loads(message) except json.JSONDecodeError: verbose_logger.warning("Invalid JSON message: %s", message[:200]) return [] @@ -816,11 +887,11 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self, model: str, logging_obj: LiteLLMLoggingObj, - modalities: list[str] | None = None, + modalities: Sequence[str] | None = None, ) -> OpenAIRealtimeStreamSession: session: Final = OpenAIRealtimeStreamSession( id=logging_obj.litellm_trace_id, - modalities=modalities if modalities is not None else ["text", "audio"], + modalities=list(modalities) if modalities is not None else ["text", "audio"], ) if model is not None and isinstance(model, str): session["model"] = model @@ -842,7 +913,7 @@ class BedrockRealtimeConfig(BaseRealtimeConfig): self, model: str, logging_obj: LiteLLMLoggingObj, - modalities: list[str] | None = None, + modalities: Sequence[str] | None = None, ) -> OpenAIRealtimeStreamSessionEvents: """Build the OpenAI session.updated ack reflecting the client's requested modalities.""" return OpenAIRealtimeStreamSessionEvents( diff --git a/litellm/realtime_api/main.py b/litellm/realtime_api/main.py index 4acce4de586..96b4b081c05 100644 --- a/litellm/realtime_api/main.py +++ b/litellm/realtime_api/main.py @@ -37,9 +37,10 @@ from ..litellm_core_utils.core_helpers import normalize_drop_params from ..litellm_core_utils.credential_accessor import CredentialAccessor from ..litellm_core_utils.get_litellm_params import get_litellm_params from ..litellm_core_utils.litellm_logging import Logging as LiteLLMLogging +from ..litellm_core_utils.realtime_streaming import client_websocket_scope from ..llms.azure.common_utils import get_azure_ad_token from ..llms.azure.realtime.handler import AzureOpenAIRealtime, azure_realtime_protocol_for_client -from ..llms.bedrock.realtime.handler import BedrockRealtime +from ..llms.bedrock.realtime.handler import BedrockRealtime, raise_if_session_committed from ..llms.custom_httpx.http_handler import realtime_ssl_for_url from ..llms.openai.live.handler import OpenAILiveSessions from ..llms.openai.realtime.handler import OpenAIRealtime @@ -471,6 +472,7 @@ async def _arealtime( For PROXY use only. """ + raise_if_session_committed(client_websocket_scope(websocket)) headers = cast(dict | None, kwargs.get("headers")) extra_headers: Final = cast(dict | None, kwargs.get("extra_headers")) if headers is None: diff --git a/tests/integration/_support/bedrock_realtime_peer.py b/tests/integration/_support/bedrock_realtime_peer.py index d30bf1b92f5..ff618a00136 100644 --- a/tests/integration/_support/bedrock_realtime_peer.py +++ b/tests/integration/_support/bedrock_realtime_peer.py @@ -14,6 +14,12 @@ text message (the injected trigger-audio block is USER/AUDIO, so it never trips ``tool_result`` fires on the Nth inbound ``toolResult`` event, which the proxy sends inside a TOOL contentStart/toolResult/contentEnd triple whose ``toolResultInputConfiguration`` ``toolUseId`` names the call being answered. + +A ``Refusal`` peer withholds each stream's initial response until that stream's opening arrived +(its SYSTEM prompt block or its first USER text block closed), then answers the HTTP error Bedrock +sends for an unknown model id. A ``StreamFailure`` peer accepts the stream, plays its turns, and +then sends an event-stream exception frame and ends the stream, as Bedrock does for a mid-call +model error. ``streams()`` groups the inbound events per opened stream. """ from __future__ import annotations @@ -26,22 +32,24 @@ import uuid from collections.abc import Iterator, Mapping from contextlib import contextmanager, suppress from dataclasses import dataclass -from itertools import chain +from itertools import chain, count from pathlib import Path from queue import SimpleQueue -from typing import Final, Literal +from typing import Final, Literal, TypeVar import h2.config import h2.connection import h2.events import h2.exceptions from integration._support.async_server import Handler, LoopbackServer -from integration._support.process import _free_port +from integration._support.ports import allocate_port from integration._support.tls import server_context from integration._support.upstream import aws_event_stream_frame from pydantic import JsonValue, TypeAdapter BedrockEvent = Mapping[str, JsonValue] +OpeningEnd = Literal["system_prompt", "user_input"] +Recorded = TypeVar("Recorded") _JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) @@ -177,6 +185,20 @@ class Turn: events: tuple[BedrockEvent, ...] +@dataclass(frozen=True, slots=True) +class Refusal: + after: OpeningEnd + status: int = 400 + error_type: str = "ValidationException" + message: str = "The provided model identifier is invalid." + + +@dataclass(frozen=True, slots=True) +class StreamFailure: + exception_type: str = "modelStreamErrorException" + message: str = "The system encountered an unexpected error during processing. Try your request again." + + def _event_body(event: BedrockEvent) -> tuple[str, Mapping[str, JsonValue]]: body: Final = event["event"] assert isinstance(body, Mapping) and len(body) == 1, event @@ -214,6 +236,17 @@ def _frame_payload(raw: bytes) -> bytes: return raw[12 + header_length : -4] +def _exception_frame(failure: StreamFailure) -> bytes: + return aws_event_stream_frame( + { + ":message-type": "exception", + ":exception-type": failure.exception_type, + ":content-type": "application/json", + }, + json.dumps({"message": failure.message}).encode(), + ) + + def _decode_inbound(chunk_payload: bytes) -> BedrockEvent | None: """One signed chunk frame's payload -> the inner chunk's {"bytes": ...} event.""" inner_payload: Final = _frame_payload(chunk_payload) @@ -239,6 +272,7 @@ class _Inbound: self.tool_use_ids: Final[dict[str, str]] = {} self.tool_results_seen = 0 self.user_text_closed = False + self.system_text_closed = False def record(self, event: BedrockEvent) -> None: self.events.put(event) @@ -253,6 +287,8 @@ class _Inbound: role_type: Final = self.open_blocks.pop(content_name, ("", "")) if role_type == ("USER", "TEXT"): self.user_text_closed = True + elif role_type == ("SYSTEM", "TEXT"): + self.system_text_closed = True elif name == "toolResult": self.tool_results_seen += 1 @@ -283,10 +319,35 @@ class _Script: return turn.events +def _opening_arrived(after: OpeningEnd, inbound: _Inbound) -> bool: + match after: + case "system_prompt": + return inbound.system_text_closed + case "user_input": + return inbound.user_text_closed + + class _PeerState: - def __init__(self, turns: tuple[Turn, ...]) -> None: + def __init__(self, turns: tuple[Turn, ...], refusal: Refusal | None, failure: StreamFailure | None) -> None: self.inbound: Final = _Inbound() self.script: Final = _Script(turns) + self.refusal: Final = refusal + self.failure: Final = failure + self.ordinals: Final = count() + self.opened: Final = SimpleQueue[int]() + self.stream_events: Final = SimpleQueue[tuple[int, BedrockEvent]]() + + +def _drain(queue: SimpleQueue[Recorded], log: list[Recorded]) -> tuple[Recorded, ...]: + while not queue.empty(): + log.append(queue.get()) # drained log, not a built value + return tuple(log) + + +@dataclass(frozen=True, slots=True) +class _Stream: + ordinal: int + inbound: _Inbound @dataclass(frozen=True, slots=True) @@ -294,11 +355,16 @@ class BedrockRealtimePeer: api_base: str _state: _PeerState _drained: list[BedrockEvent] + _opened: list[int] + _stream_events: list[tuple[int, BedrockEvent]] def received(self) -> tuple[BedrockEvent, ...]: - while not self._state.inbound.events.empty(): - self._drained.append(self._state.inbound.events.get()) # drained log, not a built value - return tuple(self._drained) + return _drain(self._state.inbound.events, self._drained) + + def streams(self) -> tuple[tuple[BedrockEvent, ...], ...]: + opened: Final = _drain(self._state.opened, self._opened) + recorded: Final = _drain(self._state.stream_events, self._stream_events) + return tuple(tuple(event for ordinal, event in recorded if ordinal == stream) for stream in opened) def tool_specs(self) -> tuple[Mapping[str, JsonValue], ...]: return tuple(chain.from_iterable(_tool_specs(event) for event in self.received())) @@ -344,6 +410,7 @@ class _H2Session: self._writer: Final = writer self._h2: Final = h2.connection.H2Connection(config=h2.config.H2Configuration(client_side=False)) self._buffers: Final[dict[int, bytes]] = {} + self._streams: Final[dict[int, _Stream]] = {} self._ended: set[int] = set() self._closed = False @@ -367,10 +434,14 @@ class _H2Session: def _apply(self, event: h2.events.Event) -> None: if isinstance(event, h2.events.RequestReceived): - self._h2.send_headers( - event.stream_id, - [(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")], - ) + ordinal: Final = next(self._state.ordinals) + self._streams[event.stream_id] = _Stream(ordinal, _Inbound()) + self._state.opened.put(ordinal) + if self._state.refusal is None: + self._h2.send_headers( + event.stream_id, + [(":status", "200"), ("content-type", "application/vnd.amazon.eventstream")], + ) elif isinstance(event, h2.events.DataReceived): self._h2.acknowledge_received_data(event.flow_controlled_length, event.stream_id) self._buffers[event.stream_id] = self._buffers.get(event.stream_id, b"") + event.data @@ -380,15 +451,44 @@ class _H2Session: decoded: Final = _decode_inbound(payload) if decoded is None: continue - self._state.inbound.record(decoded) - if event.stream_id not in self._ended: - for outbound in self._state.script.ready_events(self._state.inbound, decoded): - self._h2.send_data(event.stream_id, _outbound_chunk(outbound)) + self._record(event.stream_id, decoded) elif isinstance(event, (h2.events.StreamEnded, h2.events.StreamReset)): self._close_stream(event.stream_id) elif isinstance(event, h2.events.ConnectionTerminated): self._writer.close() + def _record(self, stream_id: int, decoded: BedrockEvent) -> None: + stream: Final = self._streams[stream_id] + self._state.inbound.record(decoded) + stream.inbound.record(decoded) + self._state.stream_events.put((stream.ordinal, decoded)) + if stream_id in self._ended: + return + refusal: Final = self._state.refusal + if refusal is not None: + if _opening_arrived(refusal.after, stream.inbound): + self._refuse(stream_id, refusal) + return + outbound_events: Final = self._state.script.ready_events(self._state.inbound, decoded) + for outbound in outbound_events: + self._h2.send_data(stream_id, _outbound_chunk(outbound)) + failure: Final = self._state.failure + if outbound_events and failure is not None and self._state.script.exhausted: + self._h2.send_data(stream_id, _exception_frame(failure)) + self._end_stream(stream_id) + + def _refuse(self, stream_id: int, refusal: Refusal) -> None: + self._h2.send_headers( + stream_id, + [ + (":status", str(refusal.status)), + ("content-type", "application/json"), + ("x-amzn-errortype", refusal.error_type), + ], + ) + self._h2.send_data(stream_id, json.dumps({"message": refusal.message}).encode()) + self._end_stream(stream_id) + def _close_stream(self, stream_id: int) -> None: self._end_stream(stream_id) self._closed = True @@ -411,19 +511,27 @@ def _session_handler(state: _PeerState) -> Handler: @contextmanager -def bedrock_realtime_peer(turns: tuple[Turn, ...], cert: tuple[Path, Path]) -> Iterator[BedrockRealtimePeer]: +def bedrock_realtime_peer( + turns: tuple[Turn, ...], + cert: tuple[Path, Path], + *, + refusal: Refusal | None = None, + failure: StreamFailure | None = None, +) -> Iterator[BedrockRealtimePeer]: """Serve scripted `turns` on a loopback TLS h2 endpoint; yield the peer. `cert` is the self-signed pair the owned proxy already trusts via SSL_CERT_FILE. SigV4 headers are ignored; the stream closes cleanly when the client ends it. """ - port: Final = _free_port() + port: Final = allocate_port() context: Final = server_context(cert[0], cert[1]) context.set_alpn_protocols(["h2"]) - state: Final = _PeerState(turns) + state: Final = _PeerState(turns, refusal, failure) server: Final = LoopbackServer(port, _session_handler(state), ssl_context=context) server.start() try: - yield BedrockRealtimePeer(api_base=f"https://localhost:{port}", _state=state, _drained=[]) + yield BedrockRealtimePeer( + api_base=f"https://localhost:{port}", _state=state, _drained=[], _opened=[], _stream_events=[] + ) finally: server.stop() diff --git a/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py index ccef839a163..b36da940a42 100644 --- a/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py +++ b/tests/integration/providers/test_bedrock_nova_sonic_realtime_tool_calls.py @@ -8,6 +8,11 @@ dials: ``AWSCRTHTTPClient`` -> TLS ALPN h2 -> invoke-with-bidirectional-stream) sides: the OpenAI realtime events the websocket client receives, and the toolSpec / toolResult documents the proxy forwarded to the peer. +The fallback cells cover a websocket that outlives one provider attempt. A refused stream keeps the +client's opening message for the fallback, which gets it once whether it dials Bedrock again or +an OpenAI realtime deployment, and once Bedrock accepted a stream no later attempt may open a second +provider session on that websocket. + The file boots its own proxy instead of the shared gateway so it can pass ``SSL_CERT_FILE`` pointing at the peer's self-signed cert, which is the only way the CRT h2 client trusts it. """ @@ -21,12 +26,19 @@ import uuid from collections.abc import AsyncIterator, Callable, Iterator, Mapping from contextlib import asynccontextmanager from pathlib import Path +from types import MappingProxyType from typing import Final +import httpx import pytest import websockets +import yaml from integration._support.bedrock_realtime_peer import ( + BedrockEvent, BedrockRealtimePeer, + OpeningEnd, + Refusal, + StreamFailure, Turn, assistant_text, bedrock_realtime_peer, @@ -43,12 +55,15 @@ from integration._support.client import ( object_value, string_value, ) -from integration._support.process import owned_proxy +from integration._support.process import graceful_stop_seconds, owned_proxy from integration._support.tls import write_self_signed_cert +from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import RealtimeResponse from pydantic import JsonValue from websockets.asyncio.client import ClientConnection +from websockets.exceptions import ConnectionClosed -pytestmark: Final = pytest.mark.timeout(180) +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) DEPLOYMENT: Final = "bedrock/amazon.nova-2-sonic-v1:0" WEATHER_PARAMETERS: Final = {"type": "object", "properties": {"location": {"type": "string"}}, "required": ["location"]} @@ -60,6 +75,30 @@ WEATHER_FUNCTION: Final = { NESTED_TOOL: Final = {"type": "function", "function": WEATHER_FUNCTION} FLAT_TOOL: Final = {"type": "function", **WEATHER_FUNCTION} TOOL_RESULT_OUTPUT: Final = '{"temperature_f":72,"conditions":"sunny"}' +QUESTION: Final = "What is the weather in New York right now?" +ANSWER: Final = "It is 72 and sunny" +INSTRUCTIONS: Final = "Answer every weather question in one short sentence." +OPENAI_VOICE_MODEL: Final = "gpt-realtime-2.1" +BEDROCK_REPLAY_GROUPS: Final = MappingProxyType( + { + "system_prompt": ("sonic-refused-session-update-first", "sonic-replay-session-update-first"), + "user_input": ("sonic-refused-text-first", "sonic-replay-text-first"), + } +) +OPENAI_REPLAY_GROUPS: Final = ("sonic-refused-openai-fallback", "openai-replay-target") +OPENAI_COMMITTED_GROUPS: Final = ("sonic-committed-openai-fallback", "openai-committed-target") +FALLBACK_GROUPS: Final = (*BEDROCK_REPLAY_GROUPS.values(), OPENAI_REPLAY_GROUPS, OPENAI_COMMITTED_GROUPS) +SCRIPTED_DONE: Final[dict[str, JsonValue]] = { + "type": "response.done", + "event_id": "evt_$UNIQUE_ID", + "response": { + "id": "resp_$UNIQUE_ID", + "object": "realtime.response", + "status": "completed", + "output": [], + "usage": {"total_tokens": 0, "input_tokens": 0, "output_tokens": 0}, + }, +} @pytest.fixture(scope="module") @@ -67,15 +106,35 @@ def cert(tmp_path_factory: pytest.TempPathFactory) -> tuple[Path, Path]: return write_self_signed_cert(tmp_path_factory.mktemp("bedrock-sonic-cert")) +def _fallback_config(directory: Path) -> Path: + stock: Final = JSON_OBJECT.validate_python( + yaml.safe_load((Path(__file__).resolve().parents[1] / "proxy_config.yaml").read_text()) + ) + config: Final = { + **stock, + "router_settings": { + **object_value(stock["router_settings"]), + "fallbacks": [{source: [target]} for source, target in FALLBACK_GROUPS], + }, + } + path: Final = directory / "bedrock-sonic-fallbacks.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + @pytest.fixture(scope="module") def candidate(tmp_path_factory: pytest.TempPathFactory, cert: tuple[Path, Path]) -> Iterator[Gateway]: + directory: Final = tmp_path_factory.mktemp("bedrock-sonic") with gateway_from_environment() as base: - with owned_proxy(base, tmp_path_factory.mktemp("bedrock-sonic"), {"SSL_CERT_FILE": str(cert[0])}) as proxy: + with owned_proxy(base, directory, {"SSL_CERT_FILE": str(cert[0])}, config=_fallback_config(directory)) as proxy: yield proxy def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) -> str: - name: Final = f"bedrock-sonic-{uuid.uuid4().hex[:12]}" + return _bedrock_deployment(gateway, scenario, peer, f"bedrock-sonic-{uuid.uuid4().hex[:12]}") + + +def _bedrock_deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer, name: str) -> str: created: Final = gateway.post( "/model/new", { @@ -94,6 +153,44 @@ def _deployment(gateway: Gateway, scenario: Scenario, peer: BedrockRealtimePeer) return name +def _openai_deployment(gateway: Gateway, scenario: Scenario, name: str) -> str: + scenario_id: Final = f"bedrock-sonic-fallback-{uuid.uuid4().hex[:12]}" + handle: Final = register_scenario( + scenario_id, RealtimeResponse(content_type="application/x-realtime", events=(SCRIPTED_DONE,)) + ) + scenario.cleanups.callback(delete_scenario, handle) + created: Final = gateway.post( + "/model/new", + { + "model_name": name, + "litellm_params": { + "model": f"openai/{OPENAI_VOICE_MODEL}", + "api_key": scenario_id, + "api_base": f"{gateway.upstream_url}/v1", + }, + "model_info": {"mode": "realtime"}, + }, + ) + scenario.cleanups.callback(scenario.delete_model, string_value(object_value(created["model_info"])["id"])) + return scenario_id + + +def _upstream_requests(gateway: Gateway, scenario_id: str) -> tuple[dict[str, JsonValue], ...]: + with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream: + observations: Final = JSON_OBJECT.validate_json(upstream.get("/__observations").content) + requests: Final = observations["requests"] + assert isinstance(requests, list), observations + return tuple( + request + for request in map(object_value, requests) + if request["authorization"] == f"Bearer {scenario_id}" or request["api_key"] == scenario_id + ) + + +def _upstream_frames(requests: tuple[dict[str, JsonValue], ...]) -> tuple[dict[str, JsonValue], ...]: + return tuple(object_value(request["body"]) for request in requests if request["method"] == "WEBSOCKET_FRAME") + + def _session_update(tool: Mapping[str, JsonValue]) -> str: return json.dumps( { @@ -113,12 +210,16 @@ def _user_text() -> str: "item": { "type": "message", "role": "user", - "content": [{"type": "input_text", "text": "What is the weather in New York right now?"}], + "content": [{"type": "input_text", "text": QUESTION}], }, } ) +def _instructions_update() -> str: + return json.dumps({"type": "session.update", "session": {"instructions": INSTRUCTIONS}}) + + _RESPONSE_CREATE: Final = json.dumps({"type": "response.create"}) @@ -207,6 +308,7 @@ def test_nested_tool_call_before_assistant_block(candidate: Gateway, cert: tuple assert _tail_types(events) == ( "response.created", "response.output_item.added", + "response.function_call_arguments.delta", "response.function_call_arguments.done", "response.output_item.done", "response.done", @@ -376,3 +478,160 @@ def test_malformed_tool_content(candidate: Gateway, cert: tuple[Path, Path], con events: Final = asyncio.run(_one_call_session(candidate, model, key, NESTED_TOOL)) (call,) = _function_calls(events) assert string_value(call["arguments"]) == (content if content is not None else "{}"), events + + +def _of_type(events: tuple[Event, ...], event_type: str) -> tuple[Event, ...]: + return tuple(event for event in events if event.get("type") == event_type) + + +def _bedrock_fields(events: tuple[BedrockEvent, ...], name: str) -> tuple[dict[str, JsonValue], ...]: + bodies: Final = tuple(object_value(event["event"]) for event in events) + return tuple(object_value(body[name]) for body in bodies if name in body) + + +def _text_inputs(events: tuple[BedrockEvent, ...], role: str) -> tuple[str, ...]: + roles: Final = MappingProxyType( + {string_value(start["contentName"]): start.get("role") for start in _bedrock_fields(events, "contentStart")} + ) + return tuple( + string_value(text["content"]) + for text in _bedrock_fields(events, "textInput") + if roles.get(string_value(text["contentName"])) == role + ) + + +def _answer_turn() -> Turn: + return Turn("user_input", assistant_text(ANSWER) + (usage(30, 12),)) + + +def _answer_text(events: tuple[Event, ...]) -> str: + return "".join(string_value(event["delta"]) for event in _of_type(events, "response.text.delta")) + + +async def _opened_with(socket: ClientConnection, opening: tuple[str, ...], stop: frozenset[str]) -> tuple[Event, ...]: + created: Final = JSON_OBJECT.validate_json(await socket.recv()) + for message in opening: + await socket.send(message) + return (created,) + await _events_until(socket, lambda event, _events: event.get("type") in stop) + + +async def _session( + gateway: Gateway, model: str, key: str, opening: tuple[str, ...], stop: frozenset[str] +) -> tuple[Event, ...]: + async with _realtime(gateway, model, key) as socket: + return await _opened_with(socket, opening, stop) + + +async def _session_until_closed( + gateway: Gateway, model: str, key: str, opening: tuple[str, ...] +) -> tuple[tuple[Event, ...], int | None]: + async with _realtime(gateway, model, key) as socket: + events: Final = await _opened_with(socket, opening, frozenset(("error", "session.created"))) + if events[-1].get("type") != "error": + return events, None + try: + await asyncio.wait_for(socket.recv(), 30) + except ConnectionClosed: + return events, socket.close_code + return events, None + + +@pytest.mark.parametrize("after", ["system_prompt", "user_input"], ids=["session-update-first", "text-first"]) +def test_refused_opening_is_replayed_on_the_bedrock_fallback( + candidate: Gateway, cert: tuple[Path, Path], after: OpeningEnd +) -> None: + source, target = BEDROCK_REPLAY_GROUPS[after] + update_first: Final = after == "system_prompt" + opening: Final = ((_instructions_update(),) if update_first else ()) + (_user_text(), _RESPONSE_CREATE) + opening_role, opening_text = ("SYSTEM", INSTRUCTIONS) if update_first else ("USER", QUESTION) + with ( + bedrock_realtime_peer((), cert, refusal=Refusal(after)) as refusing, + bedrock_realtime_peer((_answer_turn(),), cert) as healthy, + candidate.scenario() as scenario, + ): + _bedrock_deployment(candidate, scenario, refusing, source) + _bedrock_deployment(candidate, scenario, healthy, target) + events: Final = asyncio.run( + _session( + candidate, source, scenario.key(), opening, frozenset(("response.done", "error", "session.created")) + ) + ) + assert len(_of_type(events, "session.created")) == 1, events + assert len(_of_type(events, "session.updated")) == int(update_first), events + assert _of_type(events, "error") == (), events + assert len(_response_dones(events)) == 1, events + assert _answer_text(events) == ANSWER, events + (replayed,) = healthy.streams() + assert _text_inputs(replayed, "USER") == (QUESTION,), replayed + assert _text_inputs(replayed, "SYSTEM").count(INSTRUCTIONS) == int(update_first), replayed + refused: Final = refusing.streams() + assert refused and all(opening_text in _text_inputs(stream, opening_role) for stream in refused), refused + + +def test_refused_text_opening_is_replayed_on_the_openai_fallback(candidate: Gateway, cert: tuple[Path, Path]) -> None: + source, target = OPENAI_REPLAY_GROUPS + with ( + bedrock_realtime_peer((), cert, refusal=Refusal("user_input")) as refusing, + candidate.scenario() as scenario, + ): + _bedrock_deployment(candidate, scenario, refusing, source) + scenario_id: Final = _openai_deployment(candidate, scenario, target) + events: Final = asyncio.run( + _session( + candidate, + source, + scenario.key(), + (_user_text(), _RESPONSE_CREATE), + frozenset(("response.done", "error")), + ) + ) + assert _of_type(events, "error") == (), events + assert len(_response_dones(events)) == 1, events + requests: Final = _upstream_requests(candidate, scenario_id) + assert len(tuple(request for request in requests if request["method"] == "WEBSOCKET")) == 1, requests + assert _upstream_frames(requests) == (json.loads(_user_text()), json.loads(_RESPONSE_CREATE)), requests + + +def test_committed_session_failure_closes_the_socket_without_another_provider_stream( + candidate: Gateway, cert: tuple[Path, Path] +) -> None: + source, target = OPENAI_COMMITTED_GROUPS + with ( + bedrock_realtime_peer((_answer_turn(),), cert, failure=StreamFailure()) as committed, + candidate.scenario() as scenario, + ): + _bedrock_deployment(candidate, scenario, committed, source) + scenario_id: Final = _openai_deployment(candidate, scenario, target) + events, close_code = asyncio.run( + _session_until_closed(candidate, source, scenario.key(), (_user_text(), _RESPONSE_CREATE)) + ) + assert len(_of_type(events, "session.created")) == 1, events + assert len(_response_dones(events)) == 1, events + assert _answer_text(events) == ANSWER, events + assert string_value(events[-1]["type"]) == "error", events + assert close_code == 1011, events + assert _upstream_requests(candidate, scenario_id) == () + assert len(committed.streams()) == 1 + + +def test_voice_object_is_refused_and_never_reaches_bedrock(candidate: Gateway, cert: tuple[Path, Path]) -> None: + voice_update: Final = json.dumps({"type": "session.update", "session": {"voice": {"id": "voice_123"}}}) + with bedrock_realtime_peer((_answer_turn(),), cert) as peer, candidate.scenario() as scenario: + model: Final = _deployment(candidate, scenario, peer) + events: Final = asyncio.run( + _session( + candidate, + model, + scenario.key(), + (voice_update, _user_text(), _RESPONSE_CREATE), + frozenset(("response.done", "session.created")), + ) + ) + errors: Final = tuple(object_value(event["error"]) for event in _of_type(events, "error")) + assert tuple((error.get("code"), error.get("param")) for error in errors) == ( + ("invalid_type", "session.voice"), + ), events + assert _of_type(events, "session.updated") == (), events + assert len(_of_type(events, "session.created")) == 1, events + assert _answer_text(events) == ANSWER, events + assert "voice_123" not in json.dumps(peer.received()) diff --git a/tests/unit/litellm_core_utils/test_realtime_streaming.py b/tests/unit/litellm_core_utils/test_realtime_streaming.py index 9b2c42450ae..7d85da3f756 100644 --- a/tests/unit/litellm_core_utils/test_realtime_streaming.py +++ b/tests/unit/litellm_core_utils/test_realtime_streaming.py @@ -13,6 +13,7 @@ from websockets.frames import Close import litellm from litellm.constants import ( + BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY, REALTIME_METERED_USAGE_KEY, REALTIME_SESSION_FAILURE_LOGGED_KEY, REALTIME_SESSION_SUCCESS_LOGGED_KEY, @@ -543,6 +544,32 @@ async def test_backend_to_client_stores_normalized_events_for_logging(): assert streaming.messages[0]["response"]["usage"]["total_tokens"] == 0 +@pytest.mark.asyncio +async def test_client_ack_messages_replays_the_opening_message_a_refused_bedrock_stream_left(): + opening_message: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "What is 17 * 3?"}]}, + } + ) + next_message: Final = json.dumps({"type": "response.create"}) + client_ws = MagicMock() + client_ws.scope = {"headers": [], BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY: opening_message} + client_ws.receive_text = AsyncMock(side_effect=[next_message, Exception("connection closed")]) + backend_ws = MagicMock() + backend_ws.send = AsyncMock() + logging_obj: Final = _mock_logging() + logging_obj.pre_call = MagicMock() + streaming = RealTimeStreaming(client_ws, backend_ws, logging_obj) + + await streaming.client_ack_messages() + + sent_types: Final = [json.loads(call.args[0])["type"] for call in backend_ws.send.call_args_list] + assert sent_types == ["conversation.item.create", "response.create"] + assert json.loads(backend_ws.send.call_args_list[0].args[0])["item"]["content"][0]["text"] == "What is 17 * 3?" + assert BEDROCK_REALTIME_PENDING_OPENING_MESSAGE_SCOPE_KEY not in client_ws.scope + + @pytest.mark.asyncio async def test_client_ack_messages_keeps_beta_session_shape_for_beta_clients(): client_ws = MagicMock() diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py index 8f95a502137..a181c5a5425 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_handler.py @@ -17,6 +17,7 @@ from litellm.constants import ( from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.realtime.handler import BedrockRealtime from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig +from litellm.realtime_api import main as realtime_main @pytest.fixture(autouse=True) @@ -109,6 +110,14 @@ class UnavailableBedrockStream: raise ServiceUnavailableException("fault injected: Bedrock realtime unavailable") +class RecordingRealtimeHandler: + def __init__(self): + self.calls = 0 + + async def async_realtime(self, **kwargs): + self.calls += 1 + + class FakeLogging: def __init__(self, trace_id="trace-nova-sonic"): self.litellm_trace_id = trace_id @@ -678,18 +687,23 @@ class TestBedrockRealtimeSessionLifecycle: @pytest.mark.asyncio @pytest.mark.parametrize( - "drop_params,expected_errors", - [(False, [("invalid_type", "session.temperature")]), (True, [])], + "session,drop_params,expected_errors", + [ + ({"temperature": "0.2", "modalities": ["text"]}, False, [("invalid_type", "session.temperature")]), + ({"temperature": "0.2", "modalities": ["text"]}, True, []), + ({"modalities": "text"}, False, [("invalid_type", "session.modalities")]), + ({"modalities": "text"}, True, []), + ], ) - async def test_malformed_number_setting_is_rejected_or_dropped_and_the_session_still_answers( - self, stub_aws_sdk_client, drop_params, expected_errors + async def test_malformed_setting_is_rejected_or_dropped_and_the_session_still_answers( + self, stub_aws_sdk_client, session, drop_params, expected_errors ): stream: Final = ScriptedBedrockStream([]) stub_aws_sdk_client["streams"] = [stream] user_text: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]} websocket: Final = ConnectedClientWS( [ - json.dumps({"type": "session.update", "session": {"temperature": "0.2", "modalities": ["text"]}}), + json.dumps({"type": "session.update", "session": session}), json.dumps({"type": "conversation.item.create", "item": user_text}), ] ) @@ -777,12 +791,26 @@ class TestBedrockRealtimeProviderFailurePropagation: dispatched["logging_obj"] = RecordingLogging() return dispatched + USER_TEXT: Final = json.dumps( + { + "type": "conversation.item.create", + "item": {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}, + } + ) + REJECTED_SESSION_UPDATE: Final = json.dumps({"type": "session.update", "session": {"temperature": "0.2"}}) + USER_AUDIO: Final = json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"}) + @pytest.mark.asyncio + @pytest.mark.parametrize( + "client_messages", + [[SESSION_UPDATE], [USER_TEXT], [USER_AUDIO], [REJECTED_SESSION_UPDATE, USER_TEXT]], + ids=["session-update-first", "text-first", "audio-first", "input-after-rejected-update"], + ) async def test_mid_stream_failure_escapes_keeps_partial_spend_and_blocks_replay( - self, stub_aws_sdk_client, spend_dispatch + self, stub_aws_sdk_client, spend_dispatch, monkeypatch, client_messages: list[str] ): handler = BedrockRealtime() - websocket = ConnectedClientWS([self.SESSION_UPDATE]) + websocket = ConnectedClientWS(client_messages) stream = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=BreakingBedrockReceiver) stub_aws_sdk_client["streams"] = [stream] @@ -814,6 +842,37 @@ class TestBedrockRealtimeProviderFailurePropagation: assert not litellm.should_retry(replay.value.status_code), "the router must not retry the replay refusal" assert "Nova Sonic stream broke" in replay.value.message, "the router surfaces the last attempt's error" + other_provider: Final = RecordingRealtimeHandler() + monkeypatch.setattr(realtime_main, "openai_realtime", other_provider) + with pytest.raises(BedrockError) as fallback: + await realtime_main._arealtime.__wrapped__( + model="openai/gpt-realtime", websocket=websocket, litellm_logging_obj=FakeLogging() + ) + + assert fallback.value.status_code == 400, "a fallback to another provider must not start a second session" + assert other_provider.calls == 0 + + @pytest.mark.asyncio + async def test_readiness_failure_on_an_input_first_start_replays_the_input(self, stub_aws_sdk_client): + handler = BedrockRealtime() + websocket = ConnectedClientWS([self.USER_TEXT]) + healthy_stream = ScriptedBedrockStream([]) + stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), healthy_stream] + + with pytest.raises(BedrockError) as failure: + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + assert failure.value.status_code == 503 + + await handler.async_realtime( + model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS + ) + + assert [json.loads(m)["type"] for m in websocket.sent_to_client] == ["session.created"] + replayed = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in healthy_stream.input_stream.sent] + assert any(event["event"].get("textInput", {}).get("content") == "hi" for event in replayed) + @pytest.mark.asyncio async def test_input_side_failure_keeps_spend_for_responses_already_delivered( self, stub_aws_sdk_client, spend_dispatch diff --git a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py index 1cc135a8671..3e10c3bcbdc 100644 --- a/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py +++ b/tests/unit/llms/bedrock/realtime/test_bedrock_realtime_transformation.py @@ -1655,9 +1655,17 @@ class TestBedrockRealtimeUserEventsAndUsage: ({"max_response_output_tokens": float("inf")}, "session.max_response_output_tokens"), ({"max_response_output_tokens": True}, "session.max_response_output_tokens"), ({"output_sample_rate_hertz": "16000"}, "session.output_sample_rate_hertz"), + ({"instructions": {"role": "pirate"}}, "session.instructions"), + ({"voice": {"id": "voice_123"}}, "session.voice"), + ({"modalities": "text"}, "session.modalities"), + ({"modalities": ["text", 1]}, "session.modalities"), + ({"input_sample_rate_hertz": "16000"}, "session.input_sample_rate_hertz"), + ({"tools": {"type": "function", "name": "lookup"}}, "session.tools"), + ({"input_audio_format": {"type": "audio/pcm"}}, "session.input_audio_format"), + ({"output_audio_format": {"type": "audio/pcmu"}}, "session.output_audio_format"), ], ) - def test_malformed_number_setting_is_rejected_without_drop_params(self, session, param): + def test_malformed_setting_is_rejected_without_drop_params(self, session, param): rejection: Final = BedrockRealtimeConfig().session_update_rejection( {"type": "session.update", "session": session} ) @@ -1677,6 +1685,59 @@ class TestBedrockRealtimeUserEventsAndUsage: inference: Final = self._inference_configuration(config, session) assert (inference["temperature"], inference["maxTokens"]) == (0.7, 1024) + def test_non_number_input_sample_rate_is_dropped_with_drop_params(self): + config: Final = BedrockRealtimeConfig(drop_params=True) + session: Final = {"input_sample_rate_hertz": "8000"} + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + config.transform_session_update_event({"type": "session.update", "session": session}) + assert config.input_sample_rate_hertz == BedrockRealtimeConfig().input_sample_rate_hertz + + def test_non_string_instructions_and_voice_are_dropped_with_drop_params(self): + config: Final = BedrockRealtimeConfig(drop_params=True) + session: Final = {"instructions": {"role": "pirate"}, "voice": {"id": "voice_123"}} + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + sent: Final = self._voice_and_system_prompt(config, session) + assert sent == self._voice_and_system_prompt(BedrockRealtimeConfig(), {}) + + @pytest.mark.parametrize("voice", ["alloy", "marin", "Alloy", "", " "]) + def test_openai_or_empty_voice_is_rejected_without_drop_params(self, voice): + rejection: Final = BedrockRealtimeConfig().session_update_rejection( + {"type": "session.update", "session": {"voice": voice}} + ) + assert rejection is not None + assert (rejection.code, rejection.param) == ("invalid_value", "session.voice") + assert f"'{voice}'" in rejection.message + + def test_openai_voice_is_dropped_with_drop_params(self): + config: Final = BedrockRealtimeConfig(drop_params=True) + session: Final = {"voice": "alloy", "instructions": "Answer as a pirate."} + assert config.session_update_rejection({"type": "session.update", "session": session}) is None + voice, system_prompt = self._voice_and_system_prompt(config, session) + assert voice == self._voice_and_system_prompt(BedrockRealtimeConfig(), {})[0] + assert system_prompt == "Answer as a pirate." + + def test_nova_voice_is_applied(self): + session: Final = {"voice": "tiffany"} + assert BedrockRealtimeConfig().session_update_rejection({"type": "session.update", "session": session}) is None + assert self._voice_and_system_prompt(BedrockRealtimeConfig(), session)[0] == "tiffany" + + def test_g711_input_audio_format_sets_the_input_sample_rate(self): + config: Final = BedrockRealtimeConfig() + config.transform_session_update_event({"type": "session.update", "session": {"input_audio_format": "g711_ulaw"}}) + assert config.input_sample_rate_hertz == 8000 + + @staticmethod + def _voice_and_system_prompt(config: BedrockRealtimeConfig, session: dict[str, object]) -> tuple[object, object]: + events: Final = [ + json.loads(message)["event"] + for message in config.transform_session_update_event({"type": "session.update", "session": session}) + ] + voice: Final = next( + e["promptStart"]["audioOutputConfiguration"]["voiceId"] for e in events if "promptStart" in e + ) + system_prompt: Final = next(e["textInput"]["content"] for e in events if "textInput" in e) + return voice, system_prompt + def test_null_number_settings_count_as_unset(self): session: Final = {"temperature": None, "max_response_output_tokens": None} config: Final = BedrockRealtimeConfig()