mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
Merge pull request #41064 from BerriAI/litellm_bedrock_realtime_propagate_provider_failures
fix(bedrock/realtime): propagate deferred Nova Sonic stream failures to the router
This commit is contained in:
Yassin Kortam
GitHub
commit
ee03bad8c6
3 changed files with 526 additions and 166 deletions
|
|
@ -314,6 +314,10 @@ REALTIME_CREDENTIAL_RESOLUTION_TIMEOUT_SECONDS: Final = float(
|
|||
# RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
|
||||
|
||||
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update"
|
||||
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed"
|
||||
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure"
|
||||
|
||||
# SSL/TLS cipher configuration for faster handshakes
|
||||
# Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones
|
||||
# This balances performance with broad compatibility
|
||||
|
|
|
|||
|
|
@ -7,13 +7,20 @@ This uses aws_sdk_bedrock_runtime for bidirectional streaming with Nova Sonic.
|
|||
import asyncio
|
||||
import contextlib
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Mapping
|
||||
from typing import Final, Protocol
|
||||
from collections.abc import AsyncIterator, Mapping, MutableMapping
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final, NoReturn, Protocol
|
||||
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm._logging import _redact_string, verbose_proxy_logger
|
||||
from litellm.constants import (
|
||||
BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY,
|
||||
BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY,
|
||||
BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY,
|
||||
)
|
||||
from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
|
|
@ -28,6 +35,32 @@ from .transformation import BedrockRealtimeConfig
|
|||
_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({})
|
||||
|
||||
_BEDROCK_STREAM_ERROR_STATUS: Final[Mapping[str, int]] = MappingProxyType(
|
||||
{
|
||||
"AccessDeniedException": 403,
|
||||
"ConflictException": 400,
|
||||
"InternalServerException": 500,
|
||||
"ModelErrorException": 424,
|
||||
"ModelNotReadyException": 429,
|
||||
"ModelStreamErrorException": 424,
|
||||
"ModelTimeoutException": 408,
|
||||
"ResourceNotFoundException": 404,
|
||||
"ServiceQuotaExceededException": 400,
|
||||
"ServiceUnavailableException": 503,
|
||||
"ThrottlingException": 429,
|
||||
"ValidationException": 400,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _as_bedrock_error(error: BaseException) -> BaseException:
|
||||
status_code: Final = _BEDROCK_STREAM_ERROR_STATUS.get(type(error).__name__)
|
||||
if status_code is None:
|
||||
return error
|
||||
return BedrockError(status_code=status_code, message=f"{type(error).__name__}: {error}")
|
||||
|
||||
|
||||
def _json_dict(value: JsonValue) -> dict[str, JsonValue]:
|
||||
return value if isinstance(value, dict) else {}
|
||||
|
|
@ -51,6 +84,8 @@ def _should_log_event(openai_message: Mapping[str, object]) -> bool:
|
|||
class RealtimeClientWebSocket(Protocol):
|
||||
"""The client-facing websocket surface the realtime bridge talks to."""
|
||||
|
||||
scope: MutableMapping[str, object] # mutable-ok: the ASGI scope is the per-connection state store
|
||||
|
||||
async def receive_text(self) -> str: ...
|
||||
|
||||
async def send_text(self, data: str) -> None: ...
|
||||
|
|
@ -85,6 +120,81 @@ class BedrockBidirectionalStream(Protocol):
|
|||
async def await_output(self) -> tuple[object, BedrockOutputStream]: ...
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BridgeOutcome:
|
||||
logged_events: tuple[OpenAIRealtimeEvents, ...]
|
||||
provider_failure: BaseException | None
|
||||
client_disconnected: bool
|
||||
|
||||
|
||||
async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: str | None) -> AsyncIterator[str]:
|
||||
if initial_message is not None:
|
||||
yield initial_message
|
||||
while True:
|
||||
try:
|
||||
yield await client_ws.receive_text()
|
||||
except Exception as e: # noqa: BLE001 # any receive failure means the client is gone
|
||||
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
|
||||
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)
|
||||
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:
|
||||
scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = _redact_string(str(error))
|
||||
raise error from failure
|
||||
|
||||
|
||||
def _parse_client_message(message: str) -> Mapping[str, JsonValue]:
|
||||
try:
|
||||
return _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message))
|
||||
except ValueError:
|
||||
return _EMPTY_JSON_OBJECT
|
||||
|
||||
|
||||
async def _ack_session_update(
|
||||
client_ws: RealtimeClientWebSocket,
|
||||
bedrock_stream: BedrockBidirectionalStream,
|
||||
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
|
||||
if logging_obj is None:
|
||||
return True
|
||||
requested_modalities: Final = _CLIENT_MODALITIES_ADAPTER.validate_python(
|
||||
_json_dict(parsed_client_message.get("session")).get("modalities")
|
||||
)
|
||||
try:
|
||||
await client_ws.send_text(
|
||||
json.dumps(transformation_config.session_updated_event(model, logging_obj, requested_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)
|
||||
return False
|
||||
return True
|
||||
|
||||
|
||||
class BedrockRealtime(BaseAWSLLM):
|
||||
"""Handler for Bedrock Nova Sonic realtime speech-to-speech API."""
|
||||
|
||||
|
|
@ -132,6 +242,8 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
except ImportError:
|
||||
raise ImportError("Missing aws_sdk_bedrock_runtime. Install with: pip install aws-sdk-bedrock-runtime")
|
||||
|
||||
pending_session_update: Final = _pending_session_update(websocket.scope)
|
||||
|
||||
# Get AWS region
|
||||
if aws_region_name is None:
|
||||
optional_params: Final = {
|
||||
|
|
@ -190,90 +302,105 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
|
||||
transformation_config: Final = BedrockRealtimeConfig()
|
||||
|
||||
try:
|
||||
# Initialize the bidirectional stream
|
||||
bedrock_stream: Final = await open_bidirectional_stream()
|
||||
bedrock_stream: Final = await open_bidirectional_stream()
|
||||
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established")
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bidirectional stream established")
|
||||
|
||||
if pending_session_update 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")
|
||||
|
||||
# Track state for transformation
|
||||
session_state: Final[RealtimeResponseTransformInput] = {
|
||||
"current_output_item_id": None,
|
||||
"current_response_id": None,
|
||||
"current_conversation_id": None,
|
||||
"current_delta_chunks": None,
|
||||
"current_item_chunks": None,
|
||||
"current_delta_type": None,
|
||||
"session_configuration_request": None,
|
||||
}
|
||||
# Track state for transformation
|
||||
session_state: Final[RealtimeResponseTransformInput] = {
|
||||
"current_output_item_id": None,
|
||||
"current_response_id": None,
|
||||
"current_conversation_id": None,
|
||||
"current_delta_chunks": None,
|
||||
"current_item_chunks": None,
|
||||
"current_delta_type": None,
|
||||
"session_configuration_request": None,
|
||||
}
|
||||
|
||||
# Create tasks for bidirectional forwarding
|
||||
client_to_bedrock_task: Final = asyncio.create_task(
|
||||
self._forward_client_to_bedrock(
|
||||
websocket,
|
||||
bedrock_stream,
|
||||
transformation_config,
|
||||
model,
|
||||
session_state,
|
||||
logging_obj,
|
||||
outcome: Final = await self._bridge(
|
||||
websocket,
|
||||
bedrock_stream,
|
||||
transformation_config,
|
||||
model,
|
||||
session_state,
|
||||
logging_obj,
|
||||
initial_message=pending_session_update,
|
||||
)
|
||||
|
||||
logged_events: Final = (
|
||||
*outcome.logged_events,
|
||||
*(
|
||||
leftover_event
|
||||
for leftover_event in transformation_config.leftover_usage_done_events()
|
||||
if _should_log_event(leftover_event)
|
||||
),
|
||||
)
|
||||
if logged_events:
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
logging_obj.dispatch_success_handlers(
|
||||
list(logged_events), # mutable-ok: realtime spend logging requires a list result
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
)
|
||||
|
||||
async def forward_bedrock_and_collect_logged_events() -> tuple[OpenAIRealtimeEvents, ...]:
|
||||
return tuple(
|
||||
[
|
||||
event
|
||||
async for event in self._forward_bedrock_to_client(
|
||||
bedrock_stream,
|
||||
websocket,
|
||||
transformation_config,
|
||||
model,
|
||||
logging_obj,
|
||||
session_state,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
bedrock_to_client_task: Final = asyncio.create_task(forward_bedrock_and_collect_logged_events())
|
||||
|
||||
# Wait for both tasks to complete
|
||||
await asyncio.gather(
|
||||
client_to_bedrock_task,
|
||||
bedrock_to_client_task,
|
||||
return_exceptions=True,
|
||||
if outcome.provider_failure is None:
|
||||
return
|
||||
if outcome.client_disconnected:
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock Realtime: stream failed after the client disconnected: %s", outcome.provider_failure
|
||||
)
|
||||
return
|
||||
_raise_provider_failure(websocket.scope, outcome.provider_failure)
|
||||
|
||||
forwarded_logged_events: Final = (
|
||||
bedrock_to_client_task.result()
|
||||
if not bedrock_to_client_task.cancelled() and bedrock_to_client_task.exception() is None
|
||||
else ()
|
||||
)
|
||||
logged_events: Final = (
|
||||
*forwarded_logged_events,
|
||||
*(
|
||||
leftover_event
|
||||
for leftover_event in transformation_config.leftover_usage_done_events()
|
||||
if _should_log_event(leftover_event)
|
||||
),
|
||||
)
|
||||
if logged_events:
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
logging_obj.dispatch_success_handlers(
|
||||
list(logged_events), # mutable-ok: realtime spend logging requires a list result
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
)
|
||||
async def _bridge(
|
||||
self,
|
||||
websocket: RealtimeClientWebSocket,
|
||||
bedrock_stream: BedrockBidirectionalStream,
|
||||
transformation_config: BedrockRealtimeConfig,
|
||||
model: str,
|
||||
session_state: RealtimeResponseTransformInput,
|
||||
logging_obj: LiteLLMLogging,
|
||||
initial_message: str | None,
|
||||
) -> _BridgeOutcome:
|
||||
"""Run both forwarding directions until the client leaves or either side fails."""
|
||||
logged: Final[list[OpenAIRealtimeEvents]] = [] # mutable-ok: events forwarded before a failure are still spend
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error in BedrockRealtime.async_realtime: %s", e)
|
||||
try:
|
||||
await websocket.close(code=1011, reason=_redact_string(f"Internal error: {e}"))
|
||||
except Exception:
|
||||
pass
|
||||
raise
|
||||
async def collect_logged_events() -> None:
|
||||
async for event in self._forward_bedrock_to_client(
|
||||
bedrock_stream, websocket, transformation_config, model, logging_obj, session_state
|
||||
):
|
||||
logged.append(event)
|
||||
|
||||
client_task: Final = asyncio.create_task(
|
||||
self._forward_client_to_bedrock(
|
||||
websocket, bedrock_stream, transformation_config, model, session_state, logging_obj, initial_message
|
||||
)
|
||||
)
|
||||
bedrock_task: Final = asyncio.create_task(collect_logged_events())
|
||||
|
||||
await asyncio.wait((client_task, bedrock_task), return_when=asyncio.FIRST_EXCEPTION)
|
||||
client_disconnected: Final = (
|
||||
client_task.done() and not client_task.cancelled() and client_task.exception() is None
|
||||
)
|
||||
client_task.cancel()
|
||||
bedrock_task.cancel()
|
||||
client_outcome, bedrock_outcome = await asyncio.gather(client_task, bedrock_task, return_exceptions=True)
|
||||
|
||||
return _BridgeOutcome(
|
||||
logged_events=tuple(logged),
|
||||
provider_failure=(
|
||||
client_outcome
|
||||
if isinstance(client_outcome, Exception)
|
||||
else bedrock_outcome
|
||||
if isinstance(bedrock_outcome, Exception)
|
||||
else None
|
||||
),
|
||||
client_disconnected=client_disconnected,
|
||||
)
|
||||
|
||||
async def _forward_client_to_bedrock(
|
||||
self,
|
||||
|
|
@ -283,8 +410,12 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
model: str,
|
||||
session_state: RealtimeResponseTransformInput,
|
||||
logging_obj: LiteLLMLogging | None = None,
|
||||
):
|
||||
"""Forward messages from client WebSocket to Bedrock stream."""
|
||||
initial_message: str | None = None,
|
||||
) -> None:
|
||||
"""Forward messages from client WebSocket to Bedrock stream.
|
||||
|
||||
Returns once the client is gone; provider failures (input stream or readiness) propagate to the caller.
|
||||
"""
|
||||
from aws_sdk_bedrock_runtime.models import (
|
||||
BidirectionalInputPayloadPart,
|
||||
InvokeModelWithBidirectionalStreamInputChunk,
|
||||
|
|
@ -299,41 +430,28 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200])
|
||||
|
||||
try:
|
||||
while True:
|
||||
# Receive message from client
|
||||
message = await client_ws.receive_text()
|
||||
async for message in _client_messages(client_ws, initial_message):
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Received from client: %s", message[:200])
|
||||
parsed_client_message = _parse_client_message(message)
|
||||
is_session_update = _json_str(parsed_client_message.get("type")) == "session.update"
|
||||
if is_session_update:
|
||||
client_ws.scope[BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY] = (
|
||||
message # rebind-ok: scope outlives the attempt
|
||||
)
|
||||
|
||||
# Transform OpenAI format to Bedrock format
|
||||
transformed_messages = transformation_config.transform_realtime_request(
|
||||
message=message,
|
||||
model=model,
|
||||
session_configuration_request=session_state.get("session_configuration_request"),
|
||||
)
|
||||
|
||||
# Send transformed messages to Bedrock
|
||||
for bedrock_message in transformed_messages:
|
||||
await send_to_bedrock(bedrock_message)
|
||||
|
||||
if logging_obj is not None:
|
||||
client_message_type: str | None = None
|
||||
requested_modalities: list[str] | None = None
|
||||
with contextlib.suppress(Exception):
|
||||
parsed_client_message = _json_dict(_CLIENT_MESSAGE_ADAPTER.validate_json(message))
|
||||
client_message_type = _json_str(parsed_client_message.get("type"))
|
||||
if client_message_type == "session.update":
|
||||
requested_modalities = _CLIENT_MODALITIES_ADAPTER.validate_python(
|
||||
_json_dict(parsed_client_message.get("session")).get("modalities")
|
||||
)
|
||||
if client_message_type == "session.update":
|
||||
await client_ws.send_text(
|
||||
json.dumps(
|
||||
transformation_config.session_updated_event(model, logging_obj, requested_modalities)
|
||||
)
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("Client to Bedrock forwarding ended: %s", e, exc_info=True)
|
||||
if is_session_update and not await _ack_session_update(
|
||||
client_ws, bedrock_stream, transformation_config, model, logging_obj, parsed_client_message
|
||||
):
|
||||
break
|
||||
finally:
|
||||
for close_message in transformation_config.session_close_messages():
|
||||
with contextlib.suppress(Exception):
|
||||
await send_to_bedrock(close_message)
|
||||
|
|
@ -349,68 +467,71 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
logging_obj: LiteLLMLogging,
|
||||
session_state: RealtimeResponseTransformInput,
|
||||
) -> AsyncIterator[OpenAIRealtimeEvents]:
|
||||
"""Forward messages from Bedrock to the client, yielding the ones to record for spend logging."""
|
||||
try:
|
||||
while True:
|
||||
# Receive from Bedrock
|
||||
output = await bedrock_stream.await_output()
|
||||
result = await output[1].receive()
|
||||
"""Forward messages from Bedrock to the client, yielding the ones to record for spend logging.
|
||||
|
||||
if result is None:
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended")
|
||||
break
|
||||
Provider failures propagate to the caller; the client websocket is only closed on a normal stream end.
|
||||
"""
|
||||
|
||||
payload_bytes = result.value.bytes_ if result.value else None
|
||||
if payload_bytes:
|
||||
bedrock_response = payload_bytes.decode("utf-8")
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200])
|
||||
|
||||
# Transform Bedrock format to OpenAI format
|
||||
realtime_response_transform_input: RealtimeResponseTransformInput = {
|
||||
"current_output_item_id": session_state.get("current_output_item_id"),
|
||||
"current_response_id": session_state.get("current_response_id"),
|
||||
"current_conversation_id": session_state.get("current_conversation_id"),
|
||||
"current_delta_chunks": session_state.get("current_delta_chunks"),
|
||||
"current_item_chunks": session_state.get("current_item_chunks"),
|
||||
"current_delta_type": session_state.get("current_delta_type"),
|
||||
"session_configuration_request": session_state.get("session_configuration_request"),
|
||||
}
|
||||
|
||||
transformed_response = transformation_config.transform_realtime_response(
|
||||
message=bedrock_response,
|
||||
model=model,
|
||||
logging_obj=logging_obj,
|
||||
realtime_response_transform_input=realtime_response_transform_input,
|
||||
)
|
||||
|
||||
# Update session state
|
||||
session_state.update(
|
||||
{
|
||||
"current_output_item_id": transformed_response.get("current_output_item_id"),
|
||||
"current_response_id": transformed_response.get("current_response_id"),
|
||||
"current_conversation_id": transformed_response.get("current_conversation_id"),
|
||||
"current_delta_chunks": transformed_response.get("current_delta_chunks"),
|
||||
"current_item_chunks": transformed_response.get("current_item_chunks"),
|
||||
"current_delta_type": transformed_response.get("current_delta_type"),
|
||||
"session_configuration_request": transformed_response.get("session_configuration_request"),
|
||||
}
|
||||
)
|
||||
|
||||
# Send transformed messages to client
|
||||
response_value = transformed_response["response"]
|
||||
openai_messages = response_value if isinstance(response_value, list) else (response_value,)
|
||||
for openai_message in openai_messages:
|
||||
message_json = json.dumps(openai_message)
|
||||
await client_ws.send_text(message_json)
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200])
|
||||
if _should_log_event(openai_message):
|
||||
yield openai_message
|
||||
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True)
|
||||
finally:
|
||||
# Close the client WebSocket
|
||||
async def send_to_client(message_json: str) -> bool:
|
||||
try:
|
||||
await client_ws.close()
|
||||
except Exception:
|
||||
pass
|
||||
await client_ws.send_text(message_json)
|
||||
except Exception as e: # noqa: BLE001 # any send failure means the client is gone
|
||||
verbose_proxy_logger.debug("Bedrock to client forwarding ended: %s", e, exc_info=True)
|
||||
return False
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Sent to client: %s", message_json[:200])
|
||||
return True
|
||||
|
||||
output: Final = await bedrock_stream.await_output()
|
||||
while True:
|
||||
result = await output[1].receive()
|
||||
|
||||
if result is None:
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended")
|
||||
with contextlib.suppress(Exception):
|
||||
await client_ws.close()
|
||||
return
|
||||
|
||||
payload_bytes = result.value.bytes_ if result.value else None
|
||||
if payload_bytes:
|
||||
bedrock_response = payload_bytes.decode("utf-8")
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Received from Bedrock: %s", bedrock_response[:200])
|
||||
|
||||
# Transform Bedrock format to OpenAI format
|
||||
realtime_response_transform_input: RealtimeResponseTransformInput = {
|
||||
"current_output_item_id": session_state.get("current_output_item_id"),
|
||||
"current_response_id": session_state.get("current_response_id"),
|
||||
"current_conversation_id": session_state.get("current_conversation_id"),
|
||||
"current_delta_chunks": session_state.get("current_delta_chunks"),
|
||||
"current_item_chunks": session_state.get("current_item_chunks"),
|
||||
"current_delta_type": session_state.get("current_delta_type"),
|
||||
"session_configuration_request": session_state.get("session_configuration_request"),
|
||||
}
|
||||
|
||||
transformed_response = transformation_config.transform_realtime_response(
|
||||
message=bedrock_response,
|
||||
model=model,
|
||||
logging_obj=logging_obj,
|
||||
realtime_response_transform_input=realtime_response_transform_input,
|
||||
)
|
||||
|
||||
# Update session state
|
||||
session_state.update(
|
||||
{
|
||||
"current_output_item_id": transformed_response.get("current_output_item_id"),
|
||||
"current_response_id": transformed_response.get("current_response_id"),
|
||||
"current_conversation_id": transformed_response.get("current_conversation_id"),
|
||||
"current_delta_chunks": transformed_response.get("current_delta_chunks"),
|
||||
"current_item_chunks": transformed_response.get("current_item_chunks"),
|
||||
"current_delta_type": transformed_response.get("current_delta_type"),
|
||||
"session_configuration_request": transformed_response.get("session_configuration_request"),
|
||||
}
|
||||
)
|
||||
|
||||
# Send transformed messages to client
|
||||
response_value = transformed_response["response"]
|
||||
openai_messages = response_value if isinstance(response_value, list) else (response_value,)
|
||||
for openai_message in openai_messages:
|
||||
if not await send_to_client(json.dumps(openai_message)):
|
||||
return
|
||||
if _should_log_event(openai_message):
|
||||
yield openai_message
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
|
|
@ -51,6 +52,27 @@ class FakeBedrockStream:
|
|||
def __init__(self, input_stream=None):
|
||||
self.input_stream = input_stream if input_stream is not None else FakeInputStream()
|
||||
|
||||
async def await_output(self):
|
||||
return (None, EndedBedrockReceiver())
|
||||
|
||||
|
||||
class ServiceUnavailableException(Exception):
|
||||
"""Named like the modeled AWS SDK error so the handler maps it to HTTP 503"""
|
||||
|
||||
|
||||
class ModelStreamErrorException(Exception):
|
||||
"""Named like the modeled AWS SDK error so the handler maps it to HTTP 424"""
|
||||
|
||||
|
||||
class UnavailableBedrockStream:
|
||||
"""Lazy duplex stream whose HTTP response only fails once the output is awaited"""
|
||||
|
||||
def __init__(self):
|
||||
self.input_stream = FakeInputStream()
|
||||
|
||||
async def await_output(self):
|
||||
raise ServiceUnavailableException("fault injected: Bedrock realtime unavailable")
|
||||
|
||||
|
||||
class FakeLogging:
|
||||
def __init__(self, trace_id="trace-nova-sonic"):
|
||||
|
|
@ -61,6 +83,7 @@ class DisconnectingClientWS:
|
|||
def __init__(self, messages):
|
||||
self._messages = list(messages)
|
||||
self.sent_to_client = []
|
||||
self.scope = {}
|
||||
|
||||
async def receive_text(self):
|
||||
if self._messages:
|
||||
|
|
@ -93,6 +116,7 @@ class RealtimeClientWS:
|
|||
def __init__(self):
|
||||
self.closed = False
|
||||
self.sent_to_client = []
|
||||
self.scope = {}
|
||||
|
||||
async def receive_text(self):
|
||||
raise RuntimeError("client disconnected")
|
||||
|
|
@ -104,6 +128,25 @@ class RealtimeClientWS:
|
|||
self.closed = True
|
||||
|
||||
|
||||
class ConnectedClientWS(RealtimeClientWS):
|
||||
"""Client that sends its scripted messages and then stays connected until the server closes it"""
|
||||
|
||||
def __init__(self, messages):
|
||||
super().__init__()
|
||||
self._messages = list(messages)
|
||||
self._closed_event = asyncio.Event()
|
||||
|
||||
async def receive_text(self):
|
||||
if self._messages:
|
||||
return self._messages.pop(0)
|
||||
await self._closed_event.wait()
|
||||
raise RuntimeError("client disconnected")
|
||||
|
||||
async def close(self, code=None, reason=None):
|
||||
self.closed = True
|
||||
self._closed_event.set()
|
||||
|
||||
|
||||
class ScriptedBedrockReceiver:
|
||||
def __init__(self, payloads):
|
||||
self._payloads = list(payloads)
|
||||
|
|
@ -115,10 +158,48 @@ class ScriptedBedrockReceiver:
|
|||
return SimpleNamespace(value=SimpleNamespace(bytes_=payload.encode("utf-8")))
|
||||
|
||||
|
||||
class ScriptedBedrockStream:
|
||||
class BreakingBedrockReceiver(ScriptedBedrockReceiver):
|
||||
"""Delivers its payloads, then the provider stream breaks instead of ending normally"""
|
||||
|
||||
async def receive(self):
|
||||
if not self._payloads:
|
||||
await asyncio.sleep(0)
|
||||
raise ModelStreamErrorException("Nova Sonic stream broke")
|
||||
return await super().receive()
|
||||
|
||||
|
||||
class DrainedThenOpenBedrockReceiver(ScriptedBedrockReceiver):
|
||||
"""Delivers its payloads, flags `drained`, then stays open like a live Nova Sonic turn"""
|
||||
|
||||
def __init__(self, payloads):
|
||||
super().__init__(payloads)
|
||||
self.drained = asyncio.Event()
|
||||
|
||||
async def receive(self):
|
||||
if not self._payloads:
|
||||
self.drained.set()
|
||||
await asyncio.Event().wait()
|
||||
return await super().receive()
|
||||
|
||||
|
||||
class ResetOnAudioInputStream(FakeInputStream):
|
||||
"""Accepts session setup, then the provider resets the input side once the first response was delivered"""
|
||||
|
||||
def __init__(self, drained):
|
||||
super().__init__()
|
||||
self._drained = drained
|
||||
|
||||
async def send(self, event):
|
||||
if "audioInput" in json.loads(event.value.bytes_.decode("utf-8")).get("event", {}):
|
||||
await self._drained.wait()
|
||||
raise RuntimeError("bedrock input stream reset")
|
||||
self.sent.append(event)
|
||||
|
||||
|
||||
class ScriptedBedrockStream:
|
||||
def __init__(self, payloads, receiver_type=ScriptedBedrockReceiver):
|
||||
self.input_stream = FakeInputStream()
|
||||
self._receiver = ScriptedBedrockReceiver(payloads)
|
||||
self._receiver = receiver_type(payloads)
|
||||
|
||||
async def await_output(self):
|
||||
return (None, self._receiver)
|
||||
|
|
@ -163,6 +244,11 @@ def stub_aws_sdk_client(monkeypatch):
|
|||
|
||||
async def invoke_model_with_bidirectional_stream(self, operation_input):
|
||||
captured["operation_input"] = operation_input
|
||||
if captured.get("streams"):
|
||||
stream = captured["streams"].pop(0)
|
||||
if isinstance(stream, Exception):
|
||||
raise stream
|
||||
return stream
|
||||
return ScriptedBedrockStream(captured.get("scripted_payloads", []))
|
||||
|
||||
package = types.ModuleType("aws_sdk_bedrock_runtime")
|
||||
|
|
@ -263,7 +349,8 @@ class TestBedrockRealtimeHandler:
|
|||
[json.dumps({"type": "session.update", "session": {"instructions": "You are helpful."}})]
|
||||
)
|
||||
|
||||
await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {})
|
||||
with pytest.raises(RuntimeError, match="bedrock send failed"):
|
||||
await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {})
|
||||
|
||||
assert stream.input_stream.closed
|
||||
|
||||
|
|
@ -464,6 +551,154 @@ class TestBedrockRealtimeSessionLifecycle:
|
|||
assert client_ws.sent_to_client == []
|
||||
|
||||
|
||||
class TestBedrockRealtimeProviderFailurePropagation:
|
||||
"""Deferred Nova Sonic failures must escape async_realtime so the router can fall back / cool down (LIT-6484)"""
|
||||
|
||||
SESSION_UPDATE = json.dumps({"type": "session.update", "session": {"instructions": "hi", "modalities": ["text"]}})
|
||||
AWS_PARAMS = {"aws_region_name": "us-east-1", "aws_access_key_id": "k", "aws_secret_access_key": "s"}
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_readiness_failure_escapes_and_fallback_replays_session_update(self, stub_aws_sdk_client):
|
||||
handler = BedrockRealtime()
|
||||
websocket = ConnectedClientWS([self.SESSION_UPDATE])
|
||||
healthy_stream = ScriptedBedrockStream([])
|
||||
eager_failure = ServiceUnavailableException("fault injected before the stream was returned")
|
||||
stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), eager_failure, 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
|
||||
assert [json.loads(m)["type"] for m in websocket.sent_to_client] == ["session.created"]
|
||||
assert not websocket.closed, "the proxy route owns the client-facing error event and 1011 close"
|
||||
|
||||
with pytest.raises(ServiceUnavailableException):
|
||||
await handler.async_realtime(
|
||||
model="amazon.nova-sonic-v1:0", websocket=websocket, logging_obj=FakeLogging(), **self.AWS_PARAMS
|
||||
)
|
||||
|
||||
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", "session.updated"]
|
||||
replayed = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in healthy_stream.input_stream.sent]
|
||||
assert [next(iter(event["event"])) for event in replayed][:2] == ["sessionStart", "promptStart"]
|
||||
assert websocket.closed
|
||||
|
||||
TEXT_TURN = (
|
||||
json.dumps({"event": {"contentStart": {"role": "ASSISTANT", "type": "TEXT"}}}),
|
||||
json.dumps({"event": {"textOutput": {"content": "Hi"}}}),
|
||||
json.dumps({"event": {"contentEnd": {"stopReason": "END_TURN"}}}),
|
||||
)
|
||||
|
||||
@pytest.fixture
|
||||
def spend_dispatch(self, monkeypatch):
|
||||
import litellm.llms.bedrock.realtime.handler as handler_module
|
||||
|
||||
dispatched = {}
|
||||
|
||||
class RecordingLogging(FakeLogging):
|
||||
async def dispatch_success_handlers(self, result=None, prefer_async_handlers=False, **kwargs):
|
||||
dispatched["events"] = result
|
||||
|
||||
class RecordingLoggingWorker:
|
||||
def ensure_initialized_and_enqueue(self, coro):
|
||||
dispatched["coro"] = coro
|
||||
|
||||
monkeypatch.setattr(handler_module, "GLOBAL_LOGGING_WORKER", RecordingLoggingWorker())
|
||||
dispatched["logging_obj"] = RecordingLogging()
|
||||
return dispatched
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mid_stream_failure_escapes_keeps_partial_spend_and_blocks_replay(
|
||||
self, stub_aws_sdk_client, spend_dispatch
|
||||
):
|
||||
handler = BedrockRealtime()
|
||||
websocket = ConnectedClientWS([self.SESSION_UPDATE])
|
||||
stream = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=BreakingBedrockReceiver)
|
||||
stub_aws_sdk_client["streams"] = [stream]
|
||||
|
||||
with pytest.raises(BedrockError) as failure:
|
||||
await handler.async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=websocket,
|
||||
logging_obj=spend_dispatch["logging_obj"],
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
|
||||
assert failure.value.status_code == 424
|
||||
await spend_dispatch["coro"]
|
||||
assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"]
|
||||
assert "response.done" in [json.loads(m)["type"] for m in websocket.sent_to_client]
|
||||
flushed = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in stream.input_stream.sent]
|
||||
assert [next(iter(event["event"])) for event in flushed][-2:] == ["promptEnd", "sessionEnd"]
|
||||
assert stream.input_stream.closed
|
||||
|
||||
with pytest.raises(BedrockError) as replay:
|
||||
await handler.async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=websocket,
|
||||
logging_obj=spend_dispatch["logging_obj"],
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
|
||||
assert replay.value.status_code == 400, "a committed session must not be silently restarted on a fallback"
|
||||
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"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_input_side_failure_keeps_spend_for_responses_already_delivered(
|
||||
self, stub_aws_sdk_client, spend_dispatch
|
||||
):
|
||||
receiver = DrainedThenOpenBedrockReceiver(self.TEXT_TURN)
|
||||
stream = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=lambda _payloads: receiver)
|
||||
stream.input_stream = ResetOnAudioInputStream(receiver.drained)
|
||||
stub_aws_sdk_client["streams"] = [stream]
|
||||
websocket = ConnectedClientWS(
|
||||
[self.SESSION_UPDATE, json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"})]
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="bedrock input stream reset"):
|
||||
await BedrockRealtime().async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=websocket,
|
||||
logging_obj=spend_dispatch["logging_obj"],
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
|
||||
assert "response.done" in [json.loads(m)["type"] for m in websocket.sent_to_client]
|
||||
await spend_dispatch["coro"]
|
||||
assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stream_failure_after_client_disconnect_is_not_a_provider_failure(self, stub_aws_sdk_client):
|
||||
stream = ScriptedBedrockStream([], receiver_type=BreakingBedrockReceiver)
|
||||
stub_aws_sdk_client["streams"] = [stream]
|
||||
|
||||
await BedrockRealtime().async_realtime(
|
||||
model="amazon.nova-sonic-v1:0", websocket=RealtimeClientWS(), logging_obj=FakeLogging(), **self.AWS_PARAMS
|
||||
)
|
||||
|
||||
assert stream.input_stream.closed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_session_updated_is_not_sent_before_bedrock_is_ready(self, stub_aws_models):
|
||||
handler = BedrockRealtime()
|
||||
stream = UnavailableBedrockStream()
|
||||
client_ws = DisconnectingClientWS([self.SESSION_UPDATE])
|
||||
|
||||
with pytest.raises(ServiceUnavailableException):
|
||||
await handler._forward_client_to_bedrock(
|
||||
client_ws, stream, BedrockRealtimeConfig(), "amazon.nova-sonic-v1:0", {}, FakeLogging()
|
||||
)
|
||||
|
||||
assert client_ws.sent_to_client == []
|
||||
assert stream.input_stream.closed
|
||||
|
||||
|
||||
class TestBedrockRealtimeAwsAuth:
|
||||
"""AWS auth params passed via litellm_params must reach the Smithy client config (LIT-3923 regression)"""
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue