mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(bedrock)!: fail a Nova Sonic realtime session whose Bedrock stream drops instead of closing 1000 (#49170)
* fix(bedrock): fail a Nova Sonic realtime session whose Bedrock stream drops instead of closing 1000 The Bedrock SDK ends the output stream the same way for a dropped connection as for a session that sent sessionEnd, so the handler treated a dropped stream as a clean finish. The client got a 1000 close with no error event and the router logged 200 OK with no retry An output end now counts as clean only after the client left and the handler began closing the Bedrock input. An earlier end raises a 503 BedrockError, so the client gets an error event and a 1011 close, the failure handlers run, partial spend is still logged, and a retry of a committed session is still refused * fix(bedrock): fail a dropped Nova Sonic session right away and log it as failed When Bedrock's output side ends mid-session, by a drop or a modeled stream error, the client side no longer spends the close timeout sending close messages into the dead input stream, so the error reaches the app right after the drop. A session the app ends itself still sends Bedrock its close messages and closes the input The replay refusal now leads with the committed cause, so the error event and the cut 1011 close reason both start with it. It stays a 400 and is tagged as a failure outside the selected deployment, so router retries and fallbacks of a dropped session no longer add extra failure counts A clean realtime session now counts once toward its deployment's cooldown success counter, so cooldown reads the real failure rate instead of 100 percent. Realtime tpm and rpm counting is unchanged The spend row of a session the provider ended now reads failure with its error information, while the events already forwarded are still billed. A session the app ended keeps logging success * fix(bedrock): agree on a dropped realtime session's failure across loggers and the spend row A realtime session's success event now reports status failure, with the same error_information the spend row stores, when the session recorded a provider failure in its metadata. Usage already streamed stays billed and the router still counts one cooldown failure The spend row builder now runs every row's error_information through sanitize_error_information_for_spend_logs (DB size cap plus prompt redaction), so the realtime failure rows get the same treatment as the proxy failure hook, which no longer sanitizes on its own to avoid truncating twice Drop the selection-failure tag on the committed-session refusal, since the shared logging object already keeps a retried refusal out of cooldown, and pin the router's realtime success counting with an OpenAI realtime deployment * fix(bedrock): count realtime router successes only for sessions that logged one, and log first-turn drops The router counts a realtime success only when the session's success handlers ran, so a refused handshake or an upstream close before any frame no longer offsets its failure. The success event's payload status stays success again, so each logger gets one failure record per dropped session, while the spend row still reads failure from the metadata the Bedrock handler records. A committed Bedrock session that fails before any logged event now gets one zero-cost Logs row. After an output failure the input stream is closed in the background, so the SDK body writer finishes without delaying the client's error. Every input send failure reports the same 503 stream-ended cause as an output drop * fix(realtime): count a router success only for a session that ended cleanly A realtime session the upstream closes with 1011 or 1006 after its first frame now marks itself as ended abnormally where that close is seen, in the shared relay and in the OpenAI Live handler, so the router no longer counts it as a cooldown success. The success callbacks still run, so the proxy keeps the budget and parallel-slot settlement it had A Bedrock session that fails before Bedrock sends any usage now logs no success event, as on base. The spend-row failure fields go into a new metadata mapping instead of the request's own dict, validated with a pydantic TypeAdapter instead of an unchecked TypeIs guard. The handler tests drop a wall-clock bound and type their helpers * fix(bedrock): log a dropped realtime session's usage from its own logging state The failure event replaced the shared standard_logging_object while the queued success dispatch was still running its callbacks, so loggers received $0 and 0 tokens for a billed dropped session. * fix(proxy): settle a rate-limit reservation once when a request logs a failure and a success A dropped realtime session fires a failure event and then a success event for the same call. The v3 limiter refunded the TPM reservation on the failure and released it again on the success, so the key's TPM counter went negative. The success path now skips a reservation the failure already released and marks it settled for a later failure. The realtime handler test callbacks are typed * fix(rate-limit): settle each token reservation once when a request logs both a failure and a success A dropped realtime session logs a success for its usage and then a failure. After a success the limiter now clears every reserved amount (TPM, PTU, ITPM, OTPM), so a later failure refunds nothing twice. A new reservation written in the same context after a released one now starts clean, so reserve, fail, reserve, succeed charges the second request's usage once instead of 78 TPM or 28 ITPM for 27 and 20 used * fix(rate-limit): take a token hold before awaiting so concurrent success and failure callbacks settle it once * fix(rate-limit): track the ITPM/OTPM hold's release apart from TPM so a nested pre-call keeps the request's project IO charge * test(rate-limit): pin the split hold-release flags across nested MCP pre-calls * fix(realtime): count a relay that fails on its own as an abnormal end * test(realtime): pin the relay failure exit with a client frame the provider rejects --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
norkelbyte
GitHub
parent
5b546259d5
commit
864b98acec
15 changed files with 1551 additions and 118 deletions
|
|
@ -451,6 +451,7 @@ OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES: Final = frozenset(
|
|||
{"response.completed", "response.incomplete", "response.failed"}
|
||||
)
|
||||
REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged"
|
||||
REALTIME_SESSION_ENDED_ABNORMALLY_KEY: Final = "realtime_session_ended_abnormally"
|
||||
|
||||
# SSL/TLS cipher configuration for faster handshakes
|
||||
# Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones
|
||||
|
|
|
|||
|
|
@ -21,6 +21,7 @@ from litellm.constants import (
|
|||
REALTIME_LAZY_SETUP_CONTENT_GRACE_SECONDS,
|
||||
REALTIME_LAZY_SETUP_DEADLINE_SECONDS,
|
||||
REALTIME_METERED_USAGE_KEY,
|
||||
REALTIME_SESSION_ENDED_ABNORMALLY_KEY,
|
||||
REALTIME_SESSION_FAILURE_LOGGED_KEY,
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
)
|
||||
|
|
@ -65,6 +66,10 @@ class BackendClose:
|
|||
code: int
|
||||
reason: str
|
||||
|
||||
@property
|
||||
def ended_cleanly(self) -> bool:
|
||||
return self.code == 1000
|
||||
|
||||
@property
|
||||
def message(self) -> str:
|
||||
if not self.reason:
|
||||
|
|
@ -75,6 +80,7 @@ class BackendClose:
|
|||
class ClientLoopExit(Enum):
|
||||
CLIENT_DISCONNECTED = auto()
|
||||
BACKEND_CLOSED = auto()
|
||||
RELAY_FAILED = auto()
|
||||
|
||||
|
||||
def backend_close_from(error: "ConnectionClosed") -> BackendClose:
|
||||
|
|
@ -2526,7 +2532,7 @@ class RealTimeStreaming:
|
|||
return ClientLoopExit.BACKEND_CLOSED
|
||||
except Exception as e:
|
||||
verbose_logger.debug("Error in client ack messages: %s", e)
|
||||
return ClientLoopExit.CLIENT_DISCONNECTED
|
||||
return ClientLoopExit.RELAY_FAILED
|
||||
|
||||
async def bidirectional_forward(self) -> None:
|
||||
if self._is_translation_session and translation_guardrails_apply(self.request_data):
|
||||
|
|
@ -2539,9 +2545,15 @@ class RealTimeStreaming:
|
|||
client_task: Final = asyncio.create_task(self.client_ack_messages())
|
||||
try:
|
||||
await asyncio.wait((forward_task, client_task), return_when=asyncio.FIRST_COMPLETED)
|
||||
if client_task.done() and client_task.result() is ClientLoopExit.CLIENT_DISCONNECTED:
|
||||
client_exit: Final = client_task.result() if client_task.done() else None
|
||||
if client_exit is ClientLoopExit.RELAY_FAILED and self.logging_obj:
|
||||
self.logging_obj.model_call_details[REALTIME_SESSION_ENDED_ABNORMALLY_KEY] = True
|
||||
if client_exit in (ClientLoopExit.CLIENT_DISCONNECTED, ClientLoopExit.RELAY_FAILED):
|
||||
return
|
||||
await self._close_client(await forward_task)
|
||||
close: Final = await forward_task
|
||||
if not close.ended_cleanly and self.logging_obj:
|
||||
self.logging_obj.model_call_details[REALTIME_SESSION_ENDED_ABNORMALLY_KEY] = True
|
||||
await self._close_client(close)
|
||||
finally:
|
||||
forward_task.cancel()
|
||||
client_task.cancel()
|
||||
|
|
@ -2551,7 +2563,7 @@ class RealTimeStreaming:
|
|||
redacted_message: Final = redact_internal_details_from_client_message(close.message)
|
||||
redacted_reason: Final = redact_internal_details_from_client_message(close.reason)
|
||||
try:
|
||||
if close.code != 1000:
|
||||
if not close.ended_cleanly:
|
||||
await self.websocket.send_text(realtime_error_event(redacted_message, error_type="server_error"))
|
||||
await self.websocket.close(
|
||||
code=client_close_code(close.code),
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import contextlib
|
|||
import importlib.metadata
|
||||
import json
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping, MutableMapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from types import MappingProxyType
|
||||
from typing import Final, NoReturn, Protocol, runtime_checkable
|
||||
|
||||
|
|
@ -34,12 +34,15 @@ from litellm.llms.base_llm.realtime.transformation import RealtimeSessionUpdateR
|
|||
from litellm.types.llms.bedrock import AwsAuthParams
|
||||
from litellm.types.llms.openai import OpenAIRealtimeEvents
|
||||
from litellm.types.realtime import RealtimeResponseTransformInput
|
||||
from litellm.types.utils import StandardLoggingPayloadErrorInformation
|
||||
|
||||
from ..base_aws_llm import BaseAWSLLM, run_aws_signing
|
||||
from ..common_utils import BedrockError
|
||||
from .transformation import BedrockRealtimeConfig, requested_modalities
|
||||
|
||||
_CLIENT_MESSAGE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
|
||||
_REQUEST_MAPPING_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None)
|
||||
_SPEND_LOG_METADATA_KEYS: Final = ("litellm_metadata", "metadata")
|
||||
|
||||
_EMPTY_JSON_OBJECT: Final[Mapping[str, JsonValue]] = MappingProxyType({})
|
||||
|
||||
|
|
@ -61,6 +64,10 @@ _BEDROCK_STREAM_ERROR_STATUS: Final[Mapping[str, int]] = MappingProxyType(
|
|||
)
|
||||
|
||||
|
||||
def _stream_ended_before_session_closed() -> BedrockError:
|
||||
return BedrockError(status_code=503, message="Bedrock realtime stream ended before the session closed")
|
||||
|
||||
|
||||
def _as_bedrock_error(error: BaseException) -> BaseException:
|
||||
status_code: Final = _BEDROCK_STREAM_ERROR_STATUS.get(type(error).__name__)
|
||||
if status_code is None:
|
||||
|
|
@ -165,6 +172,12 @@ class _BridgeOutcome:
|
|||
client_disconnected: bool
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class SessionEndSignals:
|
||||
input_closing: asyncio.Event = field(default_factory=asyncio.Event)
|
||||
output_failed: asyncio.Event = field(default_factory=asyncio.Event)
|
||||
|
||||
|
||||
async def _client_messages(client_ws: RealtimeClientWebSocket, initial_message: str | None) -> AsyncIterator[str]:
|
||||
if initial_message is not None:
|
||||
yield initial_message
|
||||
|
|
@ -188,8 +201,8 @@ def raise_if_session_committed(scope: Mapping[str, object] | None) -> None:
|
|||
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 "")
|
||||
(f"{committed_failure}. " if committed_failure else "")
|
||||
+ "Bedrock realtime session already committed to a provider stream; it cannot be replayed"
|
||||
),
|
||||
)
|
||||
|
||||
|
|
@ -201,11 +214,48 @@ def _pending_opening_message(scope: Mapping[str, object]) -> str | None:
|
|||
return pending if isinstance(pending, str) else None
|
||||
|
||||
|
||||
def _spend_log_metadata(litellm_params: Mapping[str, object]) -> tuple[str, Mapping[str, object]] | None:
|
||||
return next(
|
||||
(
|
||||
(key, metadata)
|
||||
for key in _SPEND_LOG_METADATA_KEYS
|
||||
if (metadata := _REQUEST_MAPPING_ADAPTER.validate_python(litellm_params.get(key)))
|
||||
),
|
||||
None,
|
||||
)
|
||||
|
||||
|
||||
def _record_failure_on_spend_log(logging_obj: LiteLLMLogging, error: BaseException) -> None:
|
||||
"""The spend row of a session the provider ended reads failure while it still bills the forwarded events."""
|
||||
litellm_params: Final = _REQUEST_MAPPING_ADAPTER.validate_python(
|
||||
logging_obj.model_call_details.get("litellm_params")
|
||||
)
|
||||
if litellm_params is None:
|
||||
return
|
||||
spend_log_metadata: Final = _spend_log_metadata(litellm_params)
|
||||
if spend_log_metadata is None:
|
||||
return
|
||||
status_code: Final[object] = getattr(error, "status_code", None)
|
||||
error_information: Final[StandardLoggingPayloadErrorInformation] = {
|
||||
"error_code": "" if status_code is None else str(status_code),
|
||||
"error_class": type(error).__name__,
|
||||
"error_message": redact_string(str(error)),
|
||||
"llm_provider": "bedrock",
|
||||
}
|
||||
metadata_key, metadata = spend_log_metadata
|
||||
logging_obj.model_call_details["litellm_params"] = { # rebind-ok: the spend row reads the logging object
|
||||
**litellm_params,
|
||||
metadata_key: {**metadata, "status": "failure", "error_information": error_information},
|
||||
}
|
||||
|
||||
|
||||
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 _is_session_committed(scope):
|
||||
scope[BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY] = redact_string(str(error))
|
||||
if error is failure:
|
||||
raise failure
|
||||
raise error from failure
|
||||
|
||||
|
||||
|
|
@ -267,6 +317,25 @@ async def _reject_session_update(
|
|||
return True
|
||||
|
||||
|
||||
async def _close_bedrock_input(
|
||||
bedrock_stream: BedrockBidirectionalStream,
|
||||
close_messages: Sequence[str],
|
||||
send_to_bedrock: Callable[[str], Awaitable[None]],
|
||||
timeout: float,
|
||||
) -> None:
|
||||
async def flush_and_close() -> None:
|
||||
for close_message in close_messages:
|
||||
with contextlib.suppress(Exception):
|
||||
await send_to_bedrock(close_message)
|
||||
with contextlib.suppress(Exception):
|
||||
await bedrock_stream.input_stream.close()
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(flush_and_close(), timeout)
|
||||
except (asyncio.TimeoutError, TimeoutError):
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: timed out closing Bedrock input stream after %.1fs", timeout)
|
||||
|
||||
|
||||
class BedrockRealtime(BaseAWSLLM):
|
||||
"""Handler for Bedrock Nova Sonic realtime speech-to-speech API."""
|
||||
|
||||
|
|
@ -278,6 +347,18 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
super().__init__()
|
||||
self._sdk_version_lookup: Final = sdk_version_lookup
|
||||
self._close_timeout: Final = close_timeout
|
||||
self._failed_input_closes: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong refs to running closes
|
||||
|
||||
def _close_failed_input_in_background(
|
||||
self,
|
||||
bedrock_stream: BedrockBidirectionalStream,
|
||||
send_to_bedrock: Callable[[str], Awaitable[None]],
|
||||
) -> None:
|
||||
close_task: Final = asyncio.create_task(
|
||||
_close_bedrock_input(bedrock_stream, (), send_to_bedrock, self._close_timeout)
|
||||
)
|
||||
self._failed_input_closes.add(close_task)
|
||||
close_task.add_done_callback(self._failed_input_closes.discard)
|
||||
|
||||
async def async_realtime(
|
||||
self,
|
||||
|
|
@ -433,23 +514,26 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
if _should_log_event(leftover_event)
|
||||
),
|
||||
)
|
||||
if outcome.provider_failure is not None and outcome.client_disconnected:
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock Realtime: stream failed after the client disconnected: %s", outcome.provider_failure
|
||||
)
|
||||
provider_failure: Final = None if outcome.client_disconnected else outcome.provider_failure
|
||||
if logged_events:
|
||||
if provider_failure is not None:
|
||||
_record_failure_on_spend_log(logging_obj, _as_bedrock_error(provider_failure))
|
||||
usage_logging_obj: Final = logging_obj if provider_failure is None else logging_obj.attempt_scoped_copy()
|
||||
GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(
|
||||
logging_obj.dispatch_success_handlers(
|
||||
usage_logging_obj.dispatch_success_handlers(
|
||||
list(logged_events),
|
||||
prefer_async_handlers=True,
|
||||
)
|
||||
)
|
||||
logging_obj.model_call_details[REALTIME_SESSION_SUCCESS_LOGGED_KEY] = True
|
||||
|
||||
if outcome.provider_failure is None:
|
||||
if 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)
|
||||
_raise_provider_failure(websocket.scope, provider_failure)
|
||||
|
||||
async def _bridge(
|
||||
self,
|
||||
|
|
@ -463,16 +547,34 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
) -> _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
|
||||
signals: Final = SessionEndSignals()
|
||||
|
||||
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)
|
||||
try:
|
||||
async for event in self._forward_bedrock_to_client(
|
||||
bedrock_stream,
|
||||
websocket,
|
||||
transformation_config,
|
||||
model,
|
||||
logging_obj,
|
||||
session_state,
|
||||
signals=signals,
|
||||
):
|
||||
logged.append(event)
|
||||
except Exception:
|
||||
signals.output_failed.set()
|
||||
raise
|
||||
|
||||
client_task: Final = asyncio.create_task(
|
||||
self._forward_client_to_bedrock(
|
||||
websocket, bedrock_stream, transformation_config, model, session_state, logging_obj, initial_message
|
||||
websocket,
|
||||
bedrock_stream,
|
||||
transformation_config,
|
||||
model,
|
||||
session_state,
|
||||
logging_obj,
|
||||
initial_message,
|
||||
signals=signals,
|
||||
)
|
||||
)
|
||||
bedrock_task: Final = asyncio.create_task(collect_logged_events())
|
||||
|
|
@ -506,6 +608,8 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
session_state: RealtimeResponseTransformInput,
|
||||
logging_obj: LiteLLMLogging | None = None,
|
||||
initial_message: str | None = None,
|
||||
*,
|
||||
signals: SessionEndSignals,
|
||||
) -> None:
|
||||
"""Forward messages from client WebSocket to Bedrock stream.
|
||||
|
||||
|
|
@ -521,7 +625,10 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
|
||||
async def send_to_bedrock(bedrock_message: str) -> None:
|
||||
event: Final = build_input_chunk(bedrock_message.encode("utf-8"))
|
||||
await bedrock_stream.input_stream.send(event)
|
||||
try:
|
||||
await bedrock_stream.input_stream.send(event)
|
||||
except Exception as e:
|
||||
raise _stream_ended_before_session_closed() from e
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Sent to Bedrock: %s", bedrock_message[:200])
|
||||
|
||||
try:
|
||||
|
|
@ -553,20 +660,13 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
client_ws, transformation_config, model, logging_obj, parsed_client_message
|
||||
):
|
||||
break
|
||||
signals.input_closing.set()
|
||||
finally:
|
||||
|
||||
async def close_bedrock_input() -> None:
|
||||
for close_message in transformation_config.session_close_messages():
|
||||
with contextlib.suppress(Exception):
|
||||
await send_to_bedrock(close_message)
|
||||
with contextlib.suppress(Exception):
|
||||
await bedrock_stream.input_stream.close()
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(close_bedrock_input(), self._close_timeout)
|
||||
except (asyncio.TimeoutError, TimeoutError):
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock Realtime: timed out closing Bedrock input stream after %.1fs", self._close_timeout
|
||||
if signals.output_failed.is_set():
|
||||
self._close_failed_input_in_background(bedrock_stream, send_to_bedrock)
|
||||
else:
|
||||
await _close_bedrock_input(
|
||||
bedrock_stream, transformation_config.session_close_messages(), send_to_bedrock, self._close_timeout
|
||||
)
|
||||
|
||||
async def _forward_bedrock_to_client(
|
||||
|
|
@ -577,6 +677,8 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
model: str,
|
||||
logging_obj: LiteLLMLogging,
|
||||
session_state: RealtimeResponseTransformInput,
|
||||
*,
|
||||
signals: SessionEndSignals,
|
||||
) -> AsyncIterator[OpenAIRealtimeEvents]:
|
||||
"""Forward messages from Bedrock to the client, yielding the ones to record for spend logging.
|
||||
|
||||
|
|
@ -597,6 +699,8 @@ class BedrockRealtime(BaseAWSLLM):
|
|||
result = await output[1].receive()
|
||||
|
||||
if result is None:
|
||||
if not signals.input_closing.is_set():
|
||||
raise _stream_ended_before_session_closed()
|
||||
verbose_proxy_logger.debug("Bedrock Realtime: Bedrock stream ended")
|
||||
with contextlib.suppress(Exception):
|
||||
await client_ws.close()
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from litellm._logging import redact_secrets
|
|||
from litellm.constants import (
|
||||
OPENAI_LIVE_SESSION_CLOSE_TIMEOUT_SECONDS,
|
||||
OPENAI_LIVE_TERMINAL_RESPONSE_EVENT_TYPES,
|
||||
REALTIME_SESSION_ENDED_ABNORMALLY_KEY,
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES,
|
||||
)
|
||||
|
|
@ -316,12 +317,15 @@ async def _forward_live_backend_frames(
|
|||
backend_websocket: LiveBackendWebSocket,
|
||||
client_gone: asyncio.Event,
|
||||
event_retention: _LiveEventRetention,
|
||||
upstream_ended_abnormally: asyncio.Event,
|
||||
) -> None:
|
||||
from starlette.websockets import WebSocketDisconnect
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
|
||||
event_loop: Final = asyncio.get_running_loop()
|
||||
async for message, close_deadline in _live_backend_messages(websocket, backend_websocket):
|
||||
async for message, close_deadline in _live_backend_messages(
|
||||
websocket, backend_websocket, upstream_ended_abnormally
|
||||
):
|
||||
if (event := _metering_event(message)) is not None:
|
||||
event_retention.retain(event)
|
||||
if client_gone.is_set():
|
||||
|
|
@ -341,6 +345,7 @@ async def _forward_live_backend_frames(
|
|||
async def _live_backend_messages(
|
||||
websocket: LiveClientWebSocket,
|
||||
backend_websocket: LiveBackendWebSocket,
|
||||
upstream_ended_abnormally: asyncio.Event,
|
||||
) -> AsyncIterator[tuple[str, float | None]]:
|
||||
from starlette.websockets import WebSocketDisconnect
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
|
|
@ -350,6 +355,8 @@ async def _live_backend_messages(
|
|||
yield message, None
|
||||
except ConnectionClosed as error:
|
||||
upstream_close: Final = backend_close_from(error)
|
||||
if not upstream_close.ended_cleanly:
|
||||
upstream_ended_abnormally.set()
|
||||
with suppress(WebSocketDisconnect, RuntimeError, ConnectionClosed, OSError):
|
||||
await websocket.close(code=upstream_close.code, reason=upstream_close.reason)
|
||||
return
|
||||
|
|
@ -434,6 +441,7 @@ class OpenAILiveSessions(OpenAIRealtime):
|
|||
"complete_input_dict": {"session_start": rewritten_start},
|
||||
}
|
||||
event_retention: Final = _LiveEventRetention()
|
||||
session_ended_abnormally: Final = asyncio.Event()
|
||||
try:
|
||||
logging_obj.pre_call(input=None, api_key=api_key, additional_args=pre_call_args)
|
||||
async with self._websocket_connector(
|
||||
|
|
@ -448,7 +456,11 @@ class OpenAILiveSessions(OpenAIRealtime):
|
|||
websocket=websocket,
|
||||
backend_websocket=backend_websocket,
|
||||
event_retention=event_retention,
|
||||
upstream_ended_abnormally=session_ended_abnormally,
|
||||
)
|
||||
except Exception:
|
||||
session_ended_abnormally.set()
|
||||
raise
|
||||
finally:
|
||||
session_end_time: Final = self._clock()
|
||||
retained_events: Final = event_retention.results()
|
||||
|
|
@ -463,6 +475,10 @@ class OpenAILiveSessions(OpenAIRealtime):
|
|||
logging_obj.model_call_details[ # rebind-ok: flags the shared logging object as already logged
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY
|
||||
] = True
|
||||
if session_ended_abnormally.is_set():
|
||||
logging_obj.model_call_details[ # rebind-ok: flags the shared logging object for the router
|
||||
REALTIME_SESSION_ENDED_ABNORMALLY_KEY
|
||||
] = True
|
||||
except InvalidStatus as error:
|
||||
await close_after_upstream_handshake_refusal(websocket, error.response.status_code)
|
||||
except Exception as error: # noqa: BLE001 # any relay failure must still close the client socket with a redacted reason
|
||||
|
|
@ -484,6 +500,7 @@ class OpenAILiveSessions(OpenAIRealtime):
|
|||
websocket: LiveClientWebSocket,
|
||||
backend_websocket: LiveBackendWebSocket,
|
||||
event_retention: _LiveEventRetention,
|
||||
upstream_ended_abnormally: asyncio.Event,
|
||||
) -> None:
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
|
||||
|
|
@ -503,6 +520,7 @@ class OpenAILiveSessions(OpenAIRealtime):
|
|||
backend_websocket=backend_websocket,
|
||||
client_gone=client_gone,
|
||||
event_retention=event_retention,
|
||||
upstream_ended_abnormally=upstream_ended_abnormally,
|
||||
)
|
||||
)
|
||||
client_gone_task: Final = asyncio.create_task(client_gone.wait())
|
||||
|
|
|
|||
|
|
@ -709,7 +709,7 @@ class RequestRateLimiterStash:
|
|||
A single mutable instance is shared by every context forked from the
|
||||
request task (the SDK call, streaming generators, and the logging worker's
|
||||
captured context all see the same object), which is what makes the
|
||||
``reservation_released`` flag and ``parallel_slot`` clearing effective
|
||||
released flags and ``parallel_slot`` clearing effective
|
||||
across sibling callbacks: the first release wins, later callbacks observe
|
||||
the cleared state.
|
||||
|
||||
|
|
@ -746,9 +746,32 @@ class RequestRateLimiterStash:
|
|||
batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None
|
||||
batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = ()
|
||||
reservation_released: bool = False
|
||||
io_reservation_released: bool = False
|
||||
tpm_limited_tags: frozenset[str] = field(default_factory=frozenset)
|
||||
socket_agent_windows: tuple[RateLimitDescriptor, ...] = ()
|
||||
|
||||
def clear_reserved_tokens(self) -> None:
|
||||
self.reserved_tokens = 0
|
||||
self.ptu_reserved_tokens = 0
|
||||
self.itpm_reserved_tokens = 0
|
||||
self.otpm_reserved_tokens = 0
|
||||
|
||||
def mark_reservation_released(self) -> None:
|
||||
self.reservation_released = True
|
||||
self.io_reservation_released = True
|
||||
|
||||
def discard_released_tpm_reservation(self) -> None:
|
||||
if self.reservation_released:
|
||||
self.reserved_tokens = 0
|
||||
self.ptu_reserved_tokens = 0
|
||||
self.reservation_released = False
|
||||
|
||||
def discard_released_io_reservation(self) -> None:
|
||||
if self.io_reservation_released:
|
||||
self.itpm_reserved_tokens = 0
|
||||
self.otpm_reserved_tokens = 0
|
||||
self.io_reservation_released = False
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class CounterRefund:
|
||||
|
|
@ -4238,6 +4261,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
requested_model=requested_model,
|
||||
)
|
||||
|
||||
stash.discard_released_io_reservation()
|
||||
if itpm_reserved > 0:
|
||||
itpm_scopes: Final = tuple(
|
||||
(d["key"], d["value"]) for d in io_token_descriptors if d["key"] == PROJECT_ITPM_DESCRIPTOR_KEY
|
||||
|
|
@ -4597,6 +4621,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
# incremented so post-call reconciliation only applies
|
||||
# the (actual - reserved) delta to those — unreserved
|
||||
# scopes get charged the full actual usage instead.
|
||||
stash.discard_released_tpm_reservation()
|
||||
stash.reserved_tokens = estimated_tokens
|
||||
stash.ptu_reserved_tokens = ptu_estimated_tokens
|
||||
stash.ptu_reserved_window_identities = frozenset(
|
||||
|
|
@ -5140,7 +5165,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
self._response_usage(response_obj)
|
||||
) or self._aggregate_only_total_tokens(self._response_usage(kwargs.get("combined_usage_object")))
|
||||
|
||||
if not response_usage[2] and not combined_usage[2] and aggregate_total <= 0 and not stash.reservation_released:
|
||||
if (
|
||||
not response_usage[2]
|
||||
and not combined_usage[2]
|
||||
and aggregate_total <= 0
|
||||
and not stash.io_reservation_released
|
||||
):
|
||||
return ()
|
||||
resolved_usage: Final = (
|
||||
response_usage
|
||||
|
|
@ -5153,25 +5183,25 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
billable_input, completion_tokens, _ = resolved_usage
|
||||
|
||||
if stash.reservation_released or (
|
||||
if stash.io_reservation_released or (
|
||||
not stash.itpm_reserved_window_identities and not stash.otpm_reserved_window_identities
|
||||
):
|
||||
return self._build_reservation_aware_tpm_ops(
|
||||
targets=tuple(stash.itpm_reserved_scopes),
|
||||
reserved_scopes=frozenset() if stash.reservation_released else stash.itpm_reserved_scopes,
|
||||
reserved_scopes=frozenset() if stash.io_reservation_released else stash.itpm_reserved_scopes,
|
||||
actual_tokens=billable_input,
|
||||
reserved_tokens=0 if stash.reservation_released else itpm_reserved,
|
||||
reserved_tokens=0 if stash.io_reservation_released else itpm_reserved,
|
||||
) + self._build_reservation_aware_tpm_ops(
|
||||
targets=tuple(stash.otpm_reserved_scopes),
|
||||
reserved_scopes=frozenset() if stash.reservation_released else stash.otpm_reserved_scopes,
|
||||
reserved_scopes=frozenset() if stash.io_reservation_released else stash.otpm_reserved_scopes,
|
||||
actual_tokens=completion_tokens,
|
||||
reserved_tokens=0 if stash.reservation_released else otpm_reserved,
|
||||
reserved_tokens=0 if stash.io_reservation_released else otpm_reserved,
|
||||
)
|
||||
|
||||
itpm_ops: Final[Sequence[ReservationAwareIncrementOperation]] = (
|
||||
self._build_project_reservation_ops(
|
||||
targets=tuple(stash.itpm_reserved_scopes),
|
||||
reserved_scopes=frozenset() if stash.reservation_released else stash.itpm_reserved_scopes,
|
||||
reserved_scopes=frozenset() if stash.io_reservation_released else stash.itpm_reserved_scopes,
|
||||
actual_tokens=billable_input,
|
||||
reserved_tokens=itpm_reserved,
|
||||
reservation_window_identities=stash.itpm_reserved_window_identities,
|
||||
|
|
@ -5182,7 +5212,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
otpm_ops: Final[Sequence[ReservationAwareIncrementOperation]] = (
|
||||
self._build_project_reservation_ops(
|
||||
targets=tuple(stash.otpm_reserved_scopes),
|
||||
reserved_scopes=frozenset() if stash.reservation_released else stash.otpm_reserved_scopes,
|
||||
reserved_scopes=frozenset() if stash.io_reservation_released else stash.otpm_reserved_scopes,
|
||||
actual_tokens=completion_tokens,
|
||||
reserved_tokens=otpm_reserved,
|
||||
reservation_window_identities=stash.otpm_reserved_window_identities,
|
||||
|
|
@ -5402,7 +5432,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
total_tokens = self._aggregate_only_total_tokens(usage=_usage)
|
||||
|
||||
stash: Final = get_request_stash_for_call(_call_id_from_callback_kwargs(kwargs))
|
||||
reserved_tokens: Final = stash.reserved_tokens if stash is not None else 0
|
||||
open_reservation: Final = stash if stash is not None and not stash.reservation_released else None
|
||||
reserved_tokens: Final = open_reservation.reserved_tokens if open_reservation is not None else 0
|
||||
reserved_model: Final = stash.reserved_model if stash is not None else None
|
||||
reserved_scopes: Final[frozenset[tuple[str, str]]] = stash.reserved_scopes if stash is not None else frozenset()
|
||||
# Reconciliation must target the same model-scoped counter that the
|
||||
|
|
@ -5464,7 +5495,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
reconcile_model=reconcile_model,
|
||||
reserved_scopes=reserved_scopes,
|
||||
reserved_ceiling=stash.ptu_ceiling if stash is not None else None,
|
||||
reserved_tokens=stash.ptu_reserved_tokens if stash is not None else 0,
|
||||
reserved_tokens=open_reservation.ptu_reserved_tokens if open_reservation is not None else 0,
|
||||
reserved_windows=stash.ptu_reserved_window_identities if stash is not None else frozenset(),
|
||||
total_tokens=0 if ptu_skips_socket_usage else total_tokens,
|
||||
)
|
||||
|
|
@ -5561,18 +5592,18 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
else None,
|
||||
)
|
||||
|
||||
await self._apply_usage_operations(
|
||||
self._build_success_event_pipeline_operations(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
rate_limit_type=rate_limit_type,
|
||||
),
|
||||
litellm_parent_otel_span,
|
||||
usage_operations: Final = self._build_success_event_pipeline_operations(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
rate_limit_type=rate_limit_type,
|
||||
)
|
||||
io_token_operations: Final = self._build_io_token_reservation_ops(
|
||||
kwargs=kwargs,
|
||||
response_obj=response_obj,
|
||||
)
|
||||
if stash is not None:
|
||||
stash.clear_reserved_tokens()
|
||||
await self._apply_usage_operations(usage_operations, litellm_parent_otel_span)
|
||||
if io_token_operations:
|
||||
if isinstance(io_token_operations, list):
|
||||
await self.async_increment_tokens_with_ttl_preservation(
|
||||
|
|
@ -5733,10 +5764,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
# already released it (proxy-level rejection that also bubbles up
|
||||
# here as an LLM-error callback). max_parallel_requests is its
|
||||
# own counter and is always decremented per call.
|
||||
reserved_tokens, itpm_reserved, otpm_reserved = (
|
||||
(0, 0, 0)
|
||||
if stash is None or stash.reservation_released
|
||||
else (stash.reserved_tokens, stash.itpm_reserved_tokens, stash.otpm_reserved_tokens)
|
||||
reserved_tokens: Final = 0 if stash is None or stash.reservation_released else stash.reserved_tokens
|
||||
itpm_reserved, otpm_reserved = (
|
||||
(0, 0)
|
||||
if stash is None or stash.io_reservation_released
|
||||
else (stash.itpm_reserved_tokens, stash.otpm_reserved_tokens)
|
||||
)
|
||||
tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(failed_call)
|
||||
|
||||
|
|
@ -5806,6 +5838,8 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
else ()
|
||||
)
|
||||
|
||||
if stash is not None and (reserved_tokens > 0 or itpm_reserved > 0 or otpm_reserved > 0):
|
||||
stash.mark_reservation_released()
|
||||
if pipeline_operations:
|
||||
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline_post_call(
|
||||
pipeline_operations, parent_otel_span=litellm_parent_otel_span
|
||||
|
|
@ -5820,8 +5854,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
pipeline_operations=project_operations,
|
||||
parent_otel_span=litellm_parent_otel_span,
|
||||
)
|
||||
if stash is not None and (reserved_tokens > 0 or itpm_reserved > 0 or otpm_reserved > 0):
|
||||
stash.reservation_released = True
|
||||
if stash is None and tpm_actual > 0 and _is_socket_turn(failed_call):
|
||||
await self._apply_usage_operations(
|
||||
self._build_success_event_pipeline_operations(
|
||||
|
|
@ -5927,7 +5959,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
|
||||
Idempotent: the slot release clears the stashed acquisition (and slot
|
||||
removal is a no-op ZREM on a second run), and the TPM/ITPM/OTPM
|
||||
refund is guarded by the stash's ``reservation_released`` flag — if
|
||||
refund is guarded by the stash's two released flags — if
|
||||
both this hook and async_log_failure_event end up running in the same
|
||||
flow, only the first release/refund applies. A mid-stream failure
|
||||
relayed here with recovered partial usage settles the reservation at
|
||||
|
|
@ -5953,11 +5985,9 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
)
|
||||
stash.batch_tpd_refund_ops = ()
|
||||
|
||||
if stash.reservation_released:
|
||||
return
|
||||
reserved_tokens: Final = stash.reserved_tokens
|
||||
itpm_reserved: Final = stash.itpm_reserved_tokens
|
||||
otpm_reserved: Final = stash.otpm_reserved_tokens
|
||||
reserved_tokens: Final = 0 if stash.reservation_released else stash.reserved_tokens
|
||||
itpm_reserved: Final = 0 if stash.io_reservation_released else stash.itpm_reserved_tokens
|
||||
otpm_reserved: Final = 0 if stash.io_reservation_released else stash.otpm_reserved_tokens
|
||||
if reserved_tokens <= 0 and itpm_reserved <= 0 and otpm_reserved <= 0:
|
||||
return
|
||||
tpm_actual, itpm_actual, otpm_actual = self._recovered_partial_usage_tokens(request_data)
|
||||
|
|
@ -6024,6 +6054,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
itpm_reserved,
|
||||
otpm_reserved,
|
||||
)
|
||||
stash.mark_reservation_released()
|
||||
if combined_ops:
|
||||
await self.internal_usage_cache.dual_cache.async_increment_cache_pipeline(
|
||||
increment_list=list(combined_ops),
|
||||
|
|
@ -6040,7 +6071,6 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger):
|
|||
pipeline_operations=project_ops,
|
||||
parent_otel_span=user_api_key_dict.parent_otel_span,
|
||||
)
|
||||
stash.reservation_released = True
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.exception("Error releasing TPM reservation on post-call failure: %s", e)
|
||||
return
|
||||
|
|
|
|||
|
|
@ -64,7 +64,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 #
|
|||
_sanitize_error_information_for_spend_logs, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
deployment_params_provider,
|
||||
get_request_model_access_groups,
|
||||
sanitize_error_information_for_spend_logs,
|
||||
sanitize_error_information_for_spend_logs, # pyright: ignore[reportUnusedImport] # backwards-compatible package export
|
||||
should_store_prompts_and_responses_in_spend_logs,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.usage_capture import CapturedUsage, capture_callback, captured_callback_kwargs
|
||||
|
|
@ -79,7 +79,6 @@ from litellm.types.utils import (
|
|||
CallTypes,
|
||||
LiteLLMBatch,
|
||||
StandardLoggingPayload,
|
||||
StandardLoggingPayloadErrorInformation,
|
||||
)
|
||||
from litellm.utils import get_end_user_id_for_cost_tracking
|
||||
|
||||
|
|
@ -492,13 +491,6 @@ class _ProxyDBLogger(CustomLogger):
|
|||
# ``.get("traceback")`` / truthy checks, and the TypedDict marks
|
||||
# the field as optional, so omitting is type-safe.
|
||||
_error_information.pop("traceback", None)
|
||||
# Strip echoed request input + apply DB-size cap before storing in
|
||||
# the spend-log metadata column (LIT-2992). Result is never None
|
||||
# here because the input above is constructed non-None.
|
||||
_error_information = cast(
|
||||
StandardLoggingPayloadErrorInformation,
|
||||
sanitize_error_information_for_spend_logs(_error_information, original_exception=original_exception),
|
||||
)
|
||||
_metadata["error_information"] = _error_information
|
||||
|
||||
_metadata = ( # rebind-ok: enriched identity is shared with the earlier immutable capture
|
||||
|
|
|
|||
|
|
@ -598,6 +598,7 @@ def get_logging_payload(
|
|||
kwargs = {}
|
||||
|
||||
rejected_as_unknown_model: Final = isinstance(response_obj, ProxyModelNotFoundError)
|
||||
failure_exception: Final = response_obj if isinstance(response_obj, BaseException) else None
|
||||
response_served_from_cache: Final = served_from_cache(response_obj)
|
||||
if response_obj is None:
|
||||
response_obj = {}
|
||||
|
|
@ -888,14 +889,11 @@ def get_logging_payload(
|
|||
additional_usage_values["cache_creation_input_tokens"] = cache_write_tokens
|
||||
clean_metadata["additional_usage_values"] = additional_usage_values
|
||||
drop_requester_ip_address: Final = requester_ip_address_logging_disabled()
|
||||
error_information: Final = clean_metadata["error_information"]
|
||||
if drop_requester_ip_address:
|
||||
clean_metadata["requester_ip_address"] = None
|
||||
if drop_requester_ip_address and isinstance(error_information, Mapping):
|
||||
clean_metadata["error_information"] = cast( # cast-ok: same keys, only the str values are masked
|
||||
StandardLoggingPayloadErrorInformation,
|
||||
_error_information_without_ip_addresses(_OBJECT_MAPPING.validate_python(error_information)),
|
||||
)
|
||||
clean_metadata["error_information"] = sanitize_error_information_for_spend_logs(
|
||||
clean_metadata["error_information"], original_exception=failure_exception
|
||||
)
|
||||
|
||||
if litellm.cache is None:
|
||||
cache_key = "Cache OFF"
|
||||
|
|
|
|||
|
|
@ -84,6 +84,8 @@ from litellm.constants import (
|
|||
MID_STREAM_FAILURE_LOG_WAIT_SECONDS,
|
||||
OUTPUT_TOKEN_CEILING_PARAMS,
|
||||
PROMPT_CACHE_PIN_CACHE_MAX_SIZE,
|
||||
REALTIME_SESSION_ENDED_ABNORMALLY_KEY,
|
||||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY,
|
||||
ROUTING_REQUEST_TAGS_METADATA_KEY,
|
||||
RUNTIME_UPDATABLE_ROUTER_SETTINGS,
|
||||
|
|
@ -672,6 +674,7 @@ _NO_SESSION_KWARGS: Final[Mapping[str, Mapping[str, object]]] = MappingProxyType
|
|||
_NO_DEFAULT_PARAMS: Final[Mapping[str, object]] = MappingProxyType({})
|
||||
_SESSION_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_CALL_KWARGS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_LITELLM_PARAMS_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_SILENT_MODEL_ADAPTER: Final = TypeAdapter(str | list[str])
|
||||
_RESOLVED_RETRY_POLICY_ADAPTER: Final = TypeAdapter(RetryPolicy | None)
|
||||
|
|
@ -6902,6 +6905,22 @@ class PythonCore(DeprecatedRouterMembers):
|
|||
raise safeguard_refusal_error(model=model, stop_details=refusal_details)
|
||||
|
||||
self.success_calls[model_name] += 1
|
||||
realtime_session_logging: Final = (
|
||||
_CALL_KWARGS_ADAPTER.validate_python(kwargs).get("litellm_logging_obj")
|
||||
if generic_function_name == "_arealtime"
|
||||
else None
|
||||
)
|
||||
realtime_session_deployment_id: Final = (
|
||||
_MODEL_INFO_ADAPTER.validate_python(kwargs["model_info"]).get("id")
|
||||
if isinstance(realtime_session_logging, LiteLLMLogging)
|
||||
and realtime_session_logging.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True
|
||||
and realtime_session_logging.model_call_details.get(REALTIME_SESSION_ENDED_ABNORMALLY_KEY) is not True
|
||||
else None
|
||||
)
|
||||
if realtime_session_deployment_id is not None:
|
||||
increment_deployment_successes_for_current_minute(
|
||||
ctx=self._core_ctx, deployment_id=str(realtime_session_deployment_id)
|
||||
)
|
||||
verbose_router_logger.info("ageneric_api_call_with_fallbacks(model=%s)\x1b[32m 200 OK\x1b[0m", model_name)
|
||||
|
||||
return response
|
||||
|
|
|
|||
|
|
@ -1,12 +1,16 @@
|
|||
import asyncio
|
||||
import copy
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from collections.abc import Mapping
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from typing import Final
|
||||
from typing import Final, Literal
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
|
||||
import litellm
|
||||
from litellm.constants import (
|
||||
|
|
@ -14,10 +18,15 @@ from litellm.constants import (
|
|||
REALTIME_SESSION_SUCCESS_LOGGED_KEY,
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.realtime_errors import websocket_close_reason
|
||||
from litellm.llms.bedrock.common_utils import BedrockError
|
||||
from litellm.llms.bedrock.realtime.handler import BedrockRealtime
|
||||
from litellm.llms.bedrock.realtime.handler import BedrockRealtime, SessionEndSignals
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
from litellm.llms.bedrock.realtime.transformation import BedrockRealtimeConfig
|
||||
from litellm.realtime_api import main as realtime_main
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.utils import StandardLoggingPayload, StandardLoggingPayloadErrorInformation
|
||||
from tests.unit.llms.bedrock.event_loop_probe import host_in_loop_probing_vault
|
||||
|
||||
|
||||
|
|
@ -36,6 +45,66 @@ def _isolate_host_aws_config(monkeypatch, tmp_path):
|
|||
monkeypatch.delenv(env_var, raising=False)
|
||||
|
||||
|
||||
def _closing_input() -> SessionEndSignals:
|
||||
signals: Final = SessionEndSignals()
|
||||
signals.input_closing.set()
|
||||
return signals
|
||||
|
||||
|
||||
_SPEND_ROW_METADATA_ADAPTER: Final = TypeAdapter(Mapping[str, object])
|
||||
_ERROR_INFORMATION_ADAPTER: Final = TypeAdapter(StandardLoggingPayloadErrorInformation | None)
|
||||
_BEDROCK_INPUT_CHUNK_ADAPTER: Final = TypeAdapter(Mapping[str, Mapping[str, object]])
|
||||
_STANDARD_LOGGING_PAYLOAD_ADAPTER: Final = TypeAdapter(StandardLoggingPayload)
|
||||
|
||||
|
||||
def _spend_row(
|
||||
logging_obj: "FakeLogging | LiteLLMLogging",
|
||||
) -> tuple[Literal["success", "failure"], StandardLoggingPayloadErrorInformation | None]:
|
||||
return _spend_row_of(logging_obj.model_call_details)
|
||||
|
||||
|
||||
def _spend_row_of(
|
||||
model_call_details: dict[str, object],
|
||||
) -> tuple[Literal["success", "failure"], StandardLoggingPayloadErrorInformation | None]:
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
row: Final = get_logging_payload(kwargs=model_call_details, response_obj=None, start_time=now, end_time=now)
|
||||
metadata: Final = _SPEND_ROW_METADATA_ADAPTER.validate_json(row["metadata"])
|
||||
return row["status"], _ERROR_INFORMATION_ADAPTER.validate_python(metadata.get("error_information"))
|
||||
|
||||
|
||||
class LoggedEventRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.events: Final[list[tuple[Literal["success", "failure"], dict[str, object]]]] = []
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.events.append(("success", dict(kwargs)))
|
||||
|
||||
async def async_log_failure_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.events.append(("failure", dict(kwargs)))
|
||||
|
||||
def payload(self, event: Literal["success", "failure"]) -> StandardLoggingPayload:
|
||||
(logged,) = (kwargs for logged_event, kwargs in self.events if logged_event == event)
|
||||
return _STANDARD_LOGGING_PAYLOAD_ADAPTER.validate_python(logged["standard_logging_object"])
|
||||
|
||||
|
||||
class HeldSuccessLogger(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.entered: Final = asyncio.Event()
|
||||
self.release: Final = asyncio.Event()
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.entered.set()
|
||||
await self.release.wait()
|
||||
|
||||
|
||||
class FakePayloadPart:
|
||||
def __init__(self, bytes_):
|
||||
self.bytes_ = bytes_
|
||||
|
|
@ -85,6 +154,19 @@ class HangingCloseInputStream(FakeInputStream):
|
|||
await asyncio.Future()
|
||||
|
||||
|
||||
class BodyWriterInputStream(FakeInputStream):
|
||||
"""Models the SDK input stream, whose request-body writer task only finishes once the stream is closed"""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self._close_requested = asyncio.Event()
|
||||
self.body_writer = asyncio.get_running_loop().create_task(self._close_requested.wait())
|
||||
|
||||
async def close(self):
|
||||
await super().close()
|
||||
self._close_requested.set()
|
||||
|
||||
|
||||
class FakeBedrockStream:
|
||||
def __init__(self, input_stream=None):
|
||||
self.input_stream = input_stream if input_stream is not None else FakeInputStream()
|
||||
|
|
@ -124,6 +206,11 @@ class FakeLogging:
|
|||
self.litellm_trace_id = trace_id
|
||||
self.model_call_details = {}
|
||||
|
||||
def attempt_scoped_copy(self):
|
||||
attempt = copy.copy(self)
|
||||
attempt.model_call_details = dict(self.model_call_details)
|
||||
return attempt
|
||||
|
||||
|
||||
class DisconnectingClientWS:
|
||||
def __init__(self, messages):
|
||||
|
|
@ -159,12 +246,15 @@ class EndedBedrockStream:
|
|||
|
||||
|
||||
class RealtimeClientWS:
|
||||
def __init__(self):
|
||||
def __init__(self, messages=()):
|
||||
self.closed = False
|
||||
self.sent_to_client = []
|
||||
self.scope = {}
|
||||
self._messages = list(messages)
|
||||
|
||||
async def receive_text(self):
|
||||
if self._messages:
|
||||
return self._messages.pop(0)
|
||||
raise RuntimeError("client disconnected")
|
||||
|
||||
async def send_text(self, message):
|
||||
|
|
@ -178,8 +268,7 @@ 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)
|
||||
super().__init__(messages)
|
||||
self._closed_event = asyncio.Event()
|
||||
|
||||
async def receive_text(self):
|
||||
|
|
@ -251,6 +340,77 @@ class ScriptedBedrockStream:
|
|||
return (None, self._receiver)
|
||||
|
||||
|
||||
def _sent_event_names(input_stream: FakeInputStream) -> list[str]:
|
||||
return [
|
||||
next(iter(_BEDROCK_INPUT_CHUNK_ADAPTER.validate_json(chunk.value.bytes_)["event"]))
|
||||
for chunk in input_stream.sent
|
||||
]
|
||||
|
||||
|
||||
class SessionEndingInputStream(FakeInputStream):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.session_ended = asyncio.Event()
|
||||
|
||||
async def send(self, event):
|
||||
await super().send(event)
|
||||
if "sessionEnd" in json.loads(event.value.bytes_.decode("utf-8")).get("event", {}):
|
||||
self.session_ended.set()
|
||||
|
||||
async def close(self):
|
||||
await super().close()
|
||||
self.session_ended.set()
|
||||
|
||||
|
||||
class SlowClosingInputStream(SessionEndingInputStream):
|
||||
async def close(self):
|
||||
await asyncio.Event().wait()
|
||||
|
||||
|
||||
class OutputEndsAfterSessionEndReceiver(ScriptedBedrockReceiver):
|
||||
def __init__(self, payloads, session_ended):
|
||||
super().__init__(payloads)
|
||||
self._session_ended = session_ended
|
||||
|
||||
async def receive(self):
|
||||
if not self._payloads:
|
||||
await self._session_ended.wait()
|
||||
return await super().receive()
|
||||
|
||||
|
||||
class SessionBoundBedrockStream:
|
||||
def __init__(self, payloads=(), input_stream_type=SessionEndingInputStream):
|
||||
self.input_stream = input_stream_type()
|
||||
self._receiver = OutputEndsAfterSessionEndReceiver(payloads, self.input_stream.session_ended)
|
||||
|
||||
async def await_output(self):
|
||||
return (None, self._receiver)
|
||||
|
||||
|
||||
class ConnectionResetOnAudioInputStream(FakeInputStream):
|
||||
def __init__(self, connection_reset):
|
||||
super().__init__()
|
||||
self._connection_reset = connection_reset
|
||||
|
||||
async def send(self, event):
|
||||
await asyncio.sleep(0)
|
||||
if self._connection_reset.is_set() or "audioInput" in json.loads(event.value.bytes_.decode("utf-8"))["event"]:
|
||||
self._connection_reset.set()
|
||||
raise RuntimeError("bedrock connection reset")
|
||||
self.sent.append(event)
|
||||
|
||||
|
||||
class OutputEndsAfterConnectionResetReceiver(ScriptedBedrockReceiver):
|
||||
def __init__(self, payloads, connection_reset):
|
||||
super().__init__(payloads)
|
||||
self._connection_reset = connection_reset
|
||||
|
||||
async def receive(self):
|
||||
if not self._payloads:
|
||||
await self._connection_reset.wait()
|
||||
return await super().receive()
|
||||
|
||||
|
||||
class FakeAWSCredentialsIdentity:
|
||||
def __init__(self, access_key_id, secret_access_key, session_token=None):
|
||||
self.access_key_id = access_key_id
|
||||
|
|
@ -352,7 +512,7 @@ def stub_aws_sdk_client(monkeypatch):
|
|||
raise stream
|
||||
captured["open_stream"] = stream
|
||||
return stream
|
||||
stream = ScriptedBedrockStream(captured.get("scripted_payloads", []))
|
||||
stream = SessionBoundBedrockStream(captured.get("scripted_payloads", []))
|
||||
captured["open_stream"] = stream
|
||||
return stream
|
||||
|
||||
|
|
@ -407,7 +567,9 @@ 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", {})
|
||||
await handler._forward_client_to_bedrock(
|
||||
client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, signals=SessionEndSignals()
|
||||
)
|
||||
|
||||
sent_events = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in stream.input_stream.sent]
|
||||
event_names = [next(iter(event["event"])) for event in sent_events]
|
||||
|
|
@ -422,7 +584,7 @@ class TestBedrockRealtimeHandler:
|
|||
stream = FakeBedrockStream()
|
||||
|
||||
await handler._forward_client_to_bedrock(
|
||||
DisconnectingClientWS([]), stream, config, "amazon.nova-sonic-v1:0", {}
|
||||
DisconnectingClientWS([]), stream, config, "amazon.nova-sonic-v1:0", {}, signals=SessionEndSignals()
|
||||
)
|
||||
|
||||
assert stream.input_stream.sent == []
|
||||
|
|
@ -437,9 +599,13 @@ class TestBedrockRealtimeHandler:
|
|||
[json.dumps({"type": "session.update", "session": {"instructions": "You are helpful."}})]
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="bedrock send failed"):
|
||||
await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {})
|
||||
with pytest.raises(BedrockError) as failure:
|
||||
await handler._forward_client_to_bedrock(
|
||||
client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, signals=SessionEndSignals()
|
||||
)
|
||||
|
||||
assert failure.value.status_code == 503
|
||||
assert str(failure.value.__cause__) == "bedrock send failed"
|
||||
assert stream.input_stream.closed
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -452,7 +618,10 @@ class TestBedrockRealtimeHandler:
|
|||
)
|
||||
|
||||
await asyncio.wait_for(
|
||||
handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}), 2
|
||||
handler._forward_client_to_bedrock(
|
||||
client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, signals=SessionEndSignals()
|
||||
),
|
||||
2,
|
||||
)
|
||||
|
||||
assert not stream.input_stream.closed
|
||||
|
|
@ -466,7 +635,9 @@ 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", {})
|
||||
await handler._forward_client_to_bedrock(
|
||||
client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, signals=SessionEndSignals()
|
||||
)
|
||||
|
||||
sent_events = [json.loads(chunk.value.bytes_.decode("utf-8")) for chunk in stream.input_stream.sent]
|
||||
event_names = [next(iter(event["event"])) for event in sent_events]
|
||||
|
|
@ -495,6 +666,7 @@ class TestBedrockRealtimeHandler:
|
|||
"amazon.nova-sonic-v1:0",
|
||||
FakeLogging(),
|
||||
{},
|
||||
signals=_closing_input(),
|
||||
)
|
||||
]
|
||||
|
||||
|
|
@ -526,6 +698,7 @@ class TestBedrockRealtimeHandler:
|
|||
"amazon.nova-sonic-v1:0",
|
||||
FakeLogging(),
|
||||
{},
|
||||
signals=_closing_input(),
|
||||
)
|
||||
]
|
||||
|
||||
|
|
@ -597,6 +770,7 @@ class TestBedrockRealtimeHandler:
|
|||
"amazon.nova-sonic-v1:0",
|
||||
MagicMock(),
|
||||
{},
|
||||
signals=_closing_input(),
|
||||
):
|
||||
pass
|
||||
|
||||
|
|
@ -656,7 +830,9 @@ class TestBedrockRealtimeSessionLifecycle:
|
|||
[json.dumps({"type": "session.update", "session": {"instructions": "hi", "modalities": ["text"]}})]
|
||||
)
|
||||
|
||||
await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging())
|
||||
await handler._forward_client_to_bedrock(
|
||||
client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging(), signals=SessionEndSignals()
|
||||
)
|
||||
|
||||
acked = [json.loads(message) for message in client_ws.sent_to_client]
|
||||
updated = [event for event in acked if event["type"] == "session.updated"]
|
||||
|
|
@ -670,7 +846,9 @@ class TestBedrockRealtimeSessionLifecycle:
|
|||
stream = FakeBedrockStream()
|
||||
client_ws = DisconnectingClientWS([json.dumps({"type": "session.update", "session": {"instructions": "hi"}})])
|
||||
|
||||
await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {})
|
||||
await handler._forward_client_to_bedrock(
|
||||
client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, signals=SessionEndSignals()
|
||||
)
|
||||
|
||||
assert client_ws.sent_to_client == []
|
||||
|
||||
|
|
@ -693,7 +871,9 @@ class TestBedrockRealtimeSessionLifecycle:
|
|||
stream = FakeBedrockStream()
|
||||
client_ws = DisconnectingClientWS([update([tool]), update([other_tool]), update([tool])])
|
||||
|
||||
await handler._forward_client_to_bedrock(client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging())
|
||||
await handler._forward_client_to_bedrock(
|
||||
client_ws, stream, config, "amazon.nova-sonic-v1:0", {}, FakeLogging(), signals=SessionEndSignals()
|
||||
)
|
||||
|
||||
acked: Final = [json.loads(message) for message in client_ws.sent_to_client]
|
||||
assert [event["type"] for event in acked] == ["session.updated", "error", "session.updated"]
|
||||
|
|
@ -720,10 +900,10 @@ class TestBedrockRealtimeSessionLifecycle:
|
|||
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([])
|
||||
stream: Final = SessionBoundBedrockStream()
|
||||
stub_aws_sdk_client["streams"] = [stream]
|
||||
user_text: Final = {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "hi"}]}
|
||||
websocket: Final = ConnectedClientWS(
|
||||
websocket: Final = RealtimeClientWS(
|
||||
[
|
||||
json.dumps({"type": "session.update", "session": session}),
|
||||
json.dumps({"type": "conversation.item.create", "item": user_text}),
|
||||
|
|
@ -761,8 +941,8 @@ class TestBedrockRealtimeProviderFailurePropagation:
|
|||
@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([])
|
||||
websocket = RealtimeClientWS([self.SESSION_UPDATE])
|
||||
healthy_stream = SessionBoundBedrockStream()
|
||||
eager_failure = ServiceUnavailableException("fault injected before the stream was returned")
|
||||
stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), eager_failure, healthy_stream]
|
||||
|
||||
|
|
@ -848,9 +1028,6 @@ class TestBedrockRealtimeProviderFailurePropagation:
|
|||
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(
|
||||
|
|
@ -877,8 +1054,8 @@ class TestBedrockRealtimeProviderFailurePropagation:
|
|||
@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([])
|
||||
websocket = RealtimeClientWS([self.USER_TEXT])
|
||||
healthy_stream = SessionBoundBedrockStream()
|
||||
stub_aws_sdk_client["streams"] = [UnavailableBedrockStream(), healthy_stream]
|
||||
|
||||
with pytest.raises(BedrockError) as failure:
|
||||
|
|
@ -907,7 +1084,7 @@ class TestBedrockRealtimeProviderFailurePropagation:
|
|||
[self.SESSION_UPDATE, json.dumps({"type": "input_audio_buffer.append", "audio": "AAAA"})]
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="bedrock input stream reset"):
|
||||
with pytest.raises(BedrockError) as failure:
|
||||
await BedrockRealtime().async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=websocket,
|
||||
|
|
@ -915,6 +1092,8 @@ class TestBedrockRealtimeProviderFailurePropagation:
|
|||
**self.AWS_PARAMS,
|
||||
)
|
||||
|
||||
assert failure.value.status_code == 503
|
||||
assert str(failure.value.__cause__) == "bedrock input stream reset", "the proxy log keeps the reset's cause"
|
||||
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"]
|
||||
|
|
@ -923,10 +1102,10 @@ class TestBedrockRealtimeProviderFailurePropagation:
|
|||
async def test_success_dispatch_stamps_the_ownership_marker_only_when_spend_was_logged(
|
||||
self, stub_aws_sdk_client, spend_dispatch
|
||||
):
|
||||
stub_aws_sdk_client["streams"] = [ScriptedBedrockStream(self.TEXT_TURN)]
|
||||
stub_aws_sdk_client["streams"] = [SessionBoundBedrockStream(self.TEXT_TURN)]
|
||||
await BedrockRealtime().async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=ConnectedClientWS([self.SESSION_UPDATE]),
|
||||
websocket=RealtimeClientWS([self.SESSION_UPDATE]),
|
||||
logging_obj=spend_dispatch["logging_obj"],
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
|
|
@ -935,15 +1114,315 @@ class TestBedrockRealtimeProviderFailurePropagation:
|
|||
assert spend_dispatch["logging_obj"].model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True
|
||||
|
||||
idle_logging = FakeLogging()
|
||||
stub_aws_sdk_client["streams"] = [ScriptedBedrockStream([])]
|
||||
stub_aws_sdk_client["streams"] = [SessionBoundBedrockStream()]
|
||||
await BedrockRealtime().async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=ConnectedClientWS([self.SESSION_UPDATE]),
|
||||
websocket=RealtimeClientWS([self.SESSION_UPDATE]),
|
||||
logging_obj=idle_logging,
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in idle_logging.model_call_details
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_stream_ending_mid_session_fails_the_session_and_keeps_partial_spend(
|
||||
self, stub_aws_sdk_client, spend_dispatch
|
||||
):
|
||||
handler: Final = BedrockRealtime()
|
||||
websocket: Final = ConnectedClientWS([self.SESSION_UPDATE])
|
||||
stream: Final = ScriptedBedrockStream(self.TEXT_TURN)
|
||||
stub_aws_sdk_client["streams"] = [stream]
|
||||
logging_obj: Final = spend_dispatch["logging_obj"]
|
||||
request_metadata: Final = {"user_api_key": "hashed-key"}
|
||||
logging_obj.model_call_details["litellm_params"] = {"litellm_metadata": request_metadata}
|
||||
|
||||
with pytest.raises(BedrockError) as failure:
|
||||
await handler.async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=websocket,
|
||||
logging_obj=logging_obj,
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
|
||||
assert failure.value.status_code == 503
|
||||
assert not websocket.closed, "the proxy route owns the client-facing error event and 1011 close"
|
||||
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"]
|
||||
status, error_information = _spend_row(logging_obj)
|
||||
assert status == "failure", "the Logs page must show the dropped session as failed while it bills its turn"
|
||||
assert error_information is not None
|
||||
assert error_information.get("error_code") == "503"
|
||||
assert error_information.get("error_message") == failure.value.message
|
||||
assert request_metadata == {"user_api_key": "hashed-key"}, "the request's own metadata stays as it was sent"
|
||||
|
||||
with pytest.raises(BedrockError) as replay:
|
||||
await handler.async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=websocket,
|
||||
logging_obj=logging_obj,
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
|
||||
assert replay.value.status_code == 400, "a router retry must not restart the dropped session"
|
||||
assert str(replay.value).startswith(failure.value.message), "the app's error event leads with the drop"
|
||||
assert websocket_close_reason(str(replay.value), fallback="Internal server error").startswith(
|
||||
failure.value.message
|
||||
), "the length-capped 1011 close reason must still carry the drop"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("receiver_type", "status_code"),
|
||||
[(ScriptedBedrockReceiver, 503), (BreakingBedrockReceiver, 424)],
|
||||
ids=["output-ends", "output-raises"],
|
||||
)
|
||||
async def test_a_failed_bedrock_output_reaches_the_app_without_waiting_on_the_input_close(
|
||||
self, stub_aws_sdk_client, spend_dispatch, receiver_type, status_code: int
|
||||
):
|
||||
stream: Final = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=receiver_type)
|
||||
stream.input_stream = HangingCloseInputStream()
|
||||
stub_aws_sdk_client["streams"] = [stream]
|
||||
|
||||
with pytest.raises(BedrockError) as failure:
|
||||
await asyncio.wait_for(
|
||||
BedrockRealtime(close_timeout=30).async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=ConnectedClientWS([self.SESSION_UPDATE]),
|
||||
logging_obj=spend_dispatch["logging_obj"],
|
||||
**self.AWS_PARAMS,
|
||||
),
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
assert failure.value.status_code == status_code
|
||||
assert stub_aws_sdk_client["client_closed"], "closing the SDK client still tears the dead stream down"
|
||||
await spend_dispatch["coro"]
|
||||
assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("receiver_type", "status_code"),
|
||||
[(ScriptedBedrockReceiver, 503), (BreakingBedrockReceiver, 424)],
|
||||
ids=["output-ends", "output-raises"],
|
||||
)
|
||||
async def test_a_failed_bedrock_output_leaves_no_pending_input_writer(
|
||||
self, stub_aws_sdk_client, spend_dispatch, receiver_type, status_code: int
|
||||
):
|
||||
stream: Final = ScriptedBedrockStream(self.TEXT_TURN, receiver_type=receiver_type)
|
||||
stream.input_stream = BodyWriterInputStream()
|
||||
stub_aws_sdk_client["streams"] = [stream]
|
||||
|
||||
with pytest.raises(BedrockError) as failure:
|
||||
await BedrockRealtime().async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=ConnectedClientWS([self.SESSION_UPDATE]),
|
||||
logging_obj=spend_dispatch["logging_obj"],
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
|
||||
assert failure.value.status_code == status_code
|
||||
await asyncio.wait_for(stream.input_stream.body_writer, timeout=1)
|
||||
await spend_dispatch["coro"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_dropped_session_logs_one_success_usage_record_and_a_failed_spend_row(
|
||||
self, stub_aws_sdk_client, spend_dispatch
|
||||
):
|
||||
stub_aws_sdk_client["streams"] = [ScriptedBedrockStream(self.TEXT_TURN)]
|
||||
recorder: Final = LoggedEventRecorder()
|
||||
logging_obj: Final = LiteLLMLogging(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="_arealtime",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="nova-sonic-drop",
|
||||
function_id="nova-sonic-drop",
|
||||
dynamic_async_success_callbacks=[recorder],
|
||||
)
|
||||
logging_obj.update_from_kwargs(
|
||||
kwargs={"litellm_metadata": {"user_api_key": "hashed-key"}},
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
optional_params={},
|
||||
litellm_params={"_arealtime": True},
|
||||
custom_llm_provider="bedrock",
|
||||
)
|
||||
|
||||
with pytest.raises(BedrockError) as failure:
|
||||
await BedrockRealtime().async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=ConnectedClientWS([self.SESSION_UPDATE]),
|
||||
logging_obj=logging_obj,
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
await spend_dispatch["coro"]
|
||||
|
||||
assert failure.value.status_code == 503
|
||||
((logged_event, usage_kwargs),) = recorder.events
|
||||
assert logged_event == "success"
|
||||
usage_record: Final = recorder.payload("success")
|
||||
assert usage_record["status"] == "success", "the router's failure event is each logger's one error record"
|
||||
assert [event["type"] for event in usage_record["response"]["results"]] == ["response.done"]
|
||||
assert logging_obj.model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True
|
||||
status, error_information = _spend_row_of(usage_kwargs)
|
||||
assert status == "failure"
|
||||
assert error_information is not None
|
||||
assert error_information.get("error_code") == "503"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_dropped_session_bills_its_usage_to_loggers_while_the_failure_event_runs(
|
||||
self, stub_aws_sdk_client, spend_dispatch
|
||||
):
|
||||
usage_event: Final = json.dumps(
|
||||
{"event": {"usageEvent": {"totalInputTokens": 1000, "totalOutputTokens": 50, "totalTokens": 1050}}}
|
||||
)
|
||||
stub_aws_sdk_client["streams"] = [ScriptedBedrockStream((*self.TEXT_TURN, usage_event))]
|
||||
held: Final = HeldSuccessLogger()
|
||||
recorder: Final = LoggedEventRecorder()
|
||||
logging_obj: Final = LiteLLMLogging(
|
||||
model="amazon.nova-2-5-sonic",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="_arealtime",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="nova-sonic-drop-order",
|
||||
function_id="nova-sonic-drop-order",
|
||||
dynamic_async_success_callbacks=[held, recorder],
|
||||
dynamic_async_failure_callbacks=[recorder],
|
||||
)
|
||||
logging_obj.update_from_kwargs(
|
||||
kwargs={"litellm_metadata": {"user_api_key": "hashed-key"}},
|
||||
model="amazon.nova-2-5-sonic",
|
||||
optional_params={},
|
||||
litellm_params={"_arealtime": True},
|
||||
custom_llm_provider="bedrock",
|
||||
)
|
||||
|
||||
with pytest.raises(BedrockError) as failure:
|
||||
await BedrockRealtime().async_realtime(
|
||||
model="amazon.nova-2-5-sonic",
|
||||
websocket=ConnectedClientWS([self.SESSION_UPDATE]),
|
||||
logging_obj=logging_obj,
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
success_dispatch: Final = asyncio.create_task(spend_dispatch["coro"])
|
||||
await asyncio.wait_for(held.entered.wait(), timeout=5)
|
||||
logging_obj.failure_handler(failure.value, "", logging_obj.start_time, datetime.now(), for_async_request=True)
|
||||
await logging_obj.async_failure_handler(failure.value, "", logging_obj.start_time, datetime.now())
|
||||
held.release.set()
|
||||
await success_dispatch
|
||||
|
||||
usage_record: Final = recorder.payload("success")
|
||||
assert recorder.payload("failure")["status"] == "failure"
|
||||
assert usage_record["status"] == "success"
|
||||
assert usage_record["total_tokens"] == 1050
|
||||
assert usage_record["response_cost"] > 0
|
||||
((_, usage_kwargs),) = (logged for logged in recorder.events if logged[0] == "success")
|
||||
status, _ = _spend_row_of(usage_kwargs)
|
||||
assert status == "failure"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_dropped_session_without_request_metadata_still_fails_and_bills_its_turn(
|
||||
self, stub_aws_sdk_client, spend_dispatch
|
||||
):
|
||||
stub_aws_sdk_client["streams"] = [ScriptedBedrockStream(self.TEXT_TURN)]
|
||||
logging_obj: Final = spend_dispatch["logging_obj"]
|
||||
logging_obj.model_call_details["litellm_params"] = {"litellm_metadata": None, "metadata": {}}
|
||||
|
||||
with pytest.raises(BedrockError) as failure:
|
||||
await BedrockRealtime().async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=ConnectedClientWS([self.SESSION_UPDATE]),
|
||||
logging_obj=logging_obj,
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
await spend_dispatch["coro"]
|
||||
|
||||
assert failure.value.status_code == 503
|
||||
assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("receiver_type", "status_code"),
|
||||
[(ScriptedBedrockReceiver, 503), (BreakingBedrockReceiver, 424)],
|
||||
ids=["output-ends", "output-raises"],
|
||||
)
|
||||
async def test_a_session_bedrock_fails_before_any_logged_event_dispatches_no_success(
|
||||
self, stub_aws_sdk_client, spend_dispatch, receiver_type, status_code: int
|
||||
):
|
||||
stub_aws_sdk_client["streams"] = [ScriptedBedrockStream((), receiver_type=receiver_type)]
|
||||
logging_obj: Final = spend_dispatch["logging_obj"]
|
||||
logging_obj.model_call_details["litellm_params"] = {"litellm_metadata": {"user_api_key": "hashed-key"}}
|
||||
|
||||
with pytest.raises(BedrockError) as failure:
|
||||
await BedrockRealtime().async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=ConnectedClientWS([self.SESSION_UPDATE]),
|
||||
logging_obj=logging_obj,
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
|
||||
assert failure.value.status_code == status_code
|
||||
assert "coro" not in spend_dispatch, "a failed session with no usage gets no success event, as on base"
|
||||
assert REALTIME_SESSION_SUCCESS_LOGGED_KEY not in logging_obj.model_call_details
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_output_ending_right_after_session_end_is_a_clean_finish(
|
||||
self, stub_aws_sdk_client, spend_dispatch
|
||||
):
|
||||
stream: Final = SessionBoundBedrockStream(self.TEXT_TURN, input_stream_type=SlowClosingInputStream)
|
||||
stub_aws_sdk_client["streams"] = [stream]
|
||||
websocket: Final = RealtimeClientWS([self.SESSION_UPDATE])
|
||||
spend_dispatch["logging_obj"].model_call_details["litellm_params"] = {
|
||||
"litellm_metadata": {"user_api_key": "hashed-key"}
|
||||
}
|
||||
|
||||
await asyncio.wait_for(
|
||||
BedrockRealtime(close_timeout=30).async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=websocket,
|
||||
logging_obj=spend_dispatch["logging_obj"],
|
||||
**self.AWS_PARAMS,
|
||||
),
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
assert _sent_event_names(stream.input_stream)[-1] == "sessionEnd"
|
||||
assert not stream.input_stream.closed, "the input close must still be pending when the output ends"
|
||||
assert websocket.closed
|
||||
await spend_dispatch["coro"]
|
||||
assert [event["type"] for event in spend_dispatch["events"]] == ["response.done"]
|
||||
assert spend_dispatch["logging_obj"].model_call_details.get(REALTIME_SESSION_SUCCESS_LOGGED_KEY) is True
|
||||
assert _spend_row(spend_dispatch["logging_obj"]) == ("success", None)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_output_ending_after_the_input_side_failed_still_fails_the_session(
|
||||
self, stub_aws_sdk_client, spend_dispatch
|
||||
):
|
||||
connection_reset: Final = asyncio.Event()
|
||||
stream: Final = ScriptedBedrockStream(
|
||||
self.TEXT_TURN,
|
||||
receiver_type=lambda payloads: OutputEndsAfterConnectionResetReceiver(payloads, connection_reset),
|
||||
)
|
||||
stream.input_stream = ConnectionResetOnAudioInputStream(connection_reset)
|
||||
stub_aws_sdk_client["streams"] = [stream]
|
||||
websocket: Final = ConnectedClientWS([self.SESSION_UPDATE, self.USER_AUDIO])
|
||||
|
||||
with pytest.raises(BedrockError) as failure:
|
||||
await BedrockRealtime().async_realtime(
|
||||
model="amazon.nova-sonic-v1:0",
|
||||
websocket=websocket,
|
||||
logging_obj=spend_dispatch["logging_obj"],
|
||||
**self.AWS_PARAMS,
|
||||
)
|
||||
|
||||
assert (failure.value.status_code, failure.value.message) == (
|
||||
503,
|
||||
"Bedrock realtime stream ended before the session closed",
|
||||
), "the app sees one cause whichever side of the stream reports the reset first"
|
||||
assert not websocket.closed, "the proxy route owns the client-facing error event and 1011 close"
|
||||
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)
|
||||
|
|
@ -982,7 +1461,13 @@ class TestBedrockRealtimeProviderFailurePropagation:
|
|||
|
||||
with pytest.raises(ServiceUnavailableException):
|
||||
await handler._forward_client_to_bedrock(
|
||||
client_ws, stream, BedrockRealtimeConfig(), "amazon.nova-sonic-v1:0", {}, FakeLogging()
|
||||
client_ws,
|
||||
stream,
|
||||
BedrockRealtimeConfig(),
|
||||
"amazon.nova-sonic-v1:0",
|
||||
{},
|
||||
FakeLogging(),
|
||||
signals=SessionEndSignals(),
|
||||
)
|
||||
|
||||
assert client_ws.sent_to_client == []
|
||||
|
|
@ -1128,7 +1613,6 @@ class TestBedrockRealtimeSdkLifecycle:
|
|||
)
|
||||
|
||||
assert stub_aws_sdk_client["client_closed"]
|
||||
assert stub_aws_sdk_client["input_closed_before_client_close"]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_client_without_close_completes_session(self, monkeypatch):
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
import asyncio
|
||||
import json
|
||||
from collections.abc import Coroutine
|
||||
from collections.abc import Callable, Coroutine
|
||||
from datetime import datetime
|
||||
from typing import Final, NoReturn, cast
|
||||
|
||||
|
|
@ -13,11 +13,13 @@ from websockets.frames import Close
|
|||
from websockets.http11 import Response
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
from litellm.constants import REALTIME_SESSION_SUCCESS_LOGGED_KEY, REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES
|
||||
from litellm.cost_calculator import handle_realtime_stream_cost_calculation
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.llms.openai.live import handler as live_handler
|
||||
from litellm.llms.openai.live.handler import LiveSessionStart, OpenAILiveSessions, parse_live_session_start
|
||||
from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute
|
||||
from litellm.types.utils import LiteLLMRealtimeStreamLoggingObject, Usage
|
||||
|
||||
_LIVE_EVENT_FRAME_ADAPTER: Final = TypeAdapter(dict[str, object])
|
||||
|
|
@ -1206,3 +1208,91 @@ async def test_async_live_session_reports_redacted_server_error_and_bounds_close
|
|||
finally:
|
||||
if worker.coroutine is not None:
|
||||
await worker.coroutine
|
||||
|
||||
|
||||
_LIVE_SESSION_STARTED: Final = '{"type":"session.started","session":{"model":"gpt-live-1"}}'
|
||||
_LIVE_SESSION_CLOSED: Final = '{"type":"session.closed","usage":{"seconds":4}}'
|
||||
|
||||
|
||||
def _upstream_closing_the_session() -> _FakeUpstream:
|
||||
return _FakeUpstream(client_frames=(), upstream_frames=(_LIVE_SESSION_STARTED, _LIVE_SESSION_CLOSED))
|
||||
|
||||
|
||||
def _upstream_failing_after_a_frame() -> _FakeUpstream:
|
||||
return _FakeUpstream(
|
||||
client_frames=(),
|
||||
upstream_frames=(_LIVE_SESSION_STARTED,),
|
||||
close_error=ConnectionClosedError(Close(code=1011, reason="upstream failed"), None),
|
||||
)
|
||||
|
||||
|
||||
def _upstream_dropping_after_a_frame() -> _FakeUpstream:
|
||||
return _FakeUpstream(
|
||||
client_frames=(), upstream_frames=(_LIVE_SESSION_STARTED,), close_error=ConnectionClosedError(None, None)
|
||||
)
|
||||
|
||||
|
||||
def _upstream_failing_a_client_send() -> _FakeUpstream:
|
||||
return _FakeUpstream(
|
||||
client_frames=(),
|
||||
upstream_frames=(_LIVE_SESSION_STARTED,),
|
||||
close_response=_LIVE_SESSION_CLOSED,
|
||||
fail_on_send_type="conversation.item.create",
|
||||
send_error=OSError("upstream socket write failed"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("upstream_factory", "counted_successes"),
|
||||
[
|
||||
(_upstream_closing_the_session, 1),
|
||||
(_upstream_failing_after_a_frame, 0),
|
||||
(_upstream_dropping_after_a_frame, 0),
|
||||
(_upstream_failing_a_client_send, 0),
|
||||
],
|
||||
ids=["session-closed", "upstream-1011-after-a-frame", "upstream-1006-after-a-frame", "upstream-send-failed"],
|
||||
)
|
||||
async def test_live_session_counts_a_router_success_only_when_it_ended_cleanly(
|
||||
upstream_factory: Callable[[], _FakeUpstream], counted_successes: int
|
||||
) -> None:
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{"model_name": "live-voice", "litellm_params": {"model": "openai/gpt-live-1"}, "model_info": {"id": "live-a"}}
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
worker: Final = _FakeLoggingWorker()
|
||||
|
||||
async def _arealtime(litellm_logging_obj: Logging, **kwargs: object) -> None:
|
||||
await OpenAILiveSessions(
|
||||
logging_worker=worker, websocket_connector=_FakeConnector(upstream_factory())
|
||||
).async_live_session(
|
||||
model="gpt-live-1",
|
||||
websocket=_FakeClientWebSocket(messages=('{"type":"conversation.item.create"}',)),
|
||||
logging_obj=litellm_logging_obj,
|
||||
session_start={"type": "session.start", "session": {"model": "gpt-live-1"}},
|
||||
api_base="https://api.openai.com/",
|
||||
api_key="live-test-key",
|
||||
)
|
||||
|
||||
await asyncio.wait_for(
|
||||
router._ageneric_api_call_with_fallbacks(
|
||||
model="live-voice",
|
||||
original_function=_arealtime,
|
||||
litellm_logging_obj=Logging(
|
||||
model="gpt-live-1",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="_arealtime",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="live-session",
|
||||
function_id="live-session",
|
||||
),
|
||||
),
|
||||
timeout=5,
|
||||
)
|
||||
|
||||
if worker.coroutine is not None:
|
||||
worker.coroutine.close()
|
||||
assert get_deployment_successes_for_current_minute(router, "live-a") == counted_successes
|
||||
|
|
|
|||
|
|
@ -3711,6 +3711,255 @@ async def test_failure_event_refunds_reservation_for_input_only_estimate_v3():
|
|||
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure_first", [True, False])
|
||||
@pytest.mark.parametrize(
|
||||
("response_obj", "used_tokens"),
|
||||
[
|
||||
([], 0),
|
||||
(ModelResponse(usage=Usage(prompt_tokens=20, completion_tokens=10, total_tokens=30)), 30),
|
||||
],
|
||||
)
|
||||
async def test_a_request_logging_a_failure_and_a_success_settles_its_reservation_once_v3(
|
||||
failure_first: bool, response_obj: object, used_tokens: int
|
||||
) -> None:
|
||||
api_key: Final = hash_token("sk-failure-and-success")
|
||||
local_cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
|
||||
tokens_key: Final = handler.create_rate_limit_keys(key="api_key", value=api_key, rate_limit_type="tokens")
|
||||
await _reserve_tpm_for_owner_call(handler, local_cache, api_key, "dropped-call")
|
||||
kwargs: Final = {
|
||||
"litellm_call_id": "dropped-call",
|
||||
"standard_logging_object": {"metadata": {"user_api_key_hash": api_key}},
|
||||
}
|
||||
now: Final = datetime.now()
|
||||
log_failure: Final = partial(
|
||||
handler.async_log_failure_event, kwargs=dict(kwargs), response_obj=None, start_time=now, end_time=now
|
||||
)
|
||||
log_success: Final = partial(
|
||||
handler.async_log_success_event, kwargs=dict(kwargs), response_obj=response_obj, start_time=now, end_time=now
|
||||
)
|
||||
|
||||
for log_event in (log_failure, log_success) if failure_first else (log_success, log_failure):
|
||||
await log_event()
|
||||
|
||||
assert int(await local_cache.async_get_cache(key=tokens_key) or 0) == used_tokens
|
||||
|
||||
|
||||
class _YieldingIncrementInMemoryCache(InMemoryCache):
|
||||
async def async_increment(self, key: str, value: float, **kwargs: object) -> float:
|
||||
await asyncio.sleep(0)
|
||||
return await super().async_increment(key, value, **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("success_starts_first", [True, False])
|
||||
@pytest.mark.parametrize("failure_path", ["failure_event", "post_call_failure_hook"])
|
||||
async def test_concurrent_failure_and_success_callbacks_settle_a_reservation_once_v3(
|
||||
success_starts_first: bool, failure_path: str
|
||||
) -> None:
|
||||
api_key: Final = hash_token("sk-concurrent-failure-and-success")
|
||||
model: Final = "gpt-4o-mini"
|
||||
project: Final = "proj-concurrent"
|
||||
local_cache: Final = DualCache(in_memory_cache=_YieldingIncrementInMemoryCache())
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
|
||||
caller: Final = UserAPIKeyAuth(
|
||||
api_key=api_key,
|
||||
tpm_limit=10_000,
|
||||
project_id=project,
|
||||
project_metadata={"model_itpm_limit": {model: 10_000}, "model_otpm_limit": {model: 10_000}},
|
||||
)
|
||||
counter_keys: Final = (
|
||||
handler.create_rate_limit_keys(key="api_key", value=api_key, rate_limit_type="tokens"),
|
||||
handler.create_rate_limit_keys(
|
||||
key="model_per_project_itpm", value=f"{project}:{model}", rate_limit_type="tokens"
|
||||
),
|
||||
handler.create_rate_limit_keys(
|
||||
key="model_per_project_otpm", value=f"{project}:{model}", rate_limit_type="tokens"
|
||||
),
|
||||
)
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=caller,
|
||||
cache=local_cache,
|
||||
data={
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 50,
|
||||
"litellm_call_id": "concurrent-call",
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
kwargs: Final = {
|
||||
"litellm_call_id": "concurrent-call",
|
||||
"model": model,
|
||||
"standard_logging_object": {"metadata": {"user_api_key_hash": api_key}},
|
||||
}
|
||||
now: Final = datetime.now()
|
||||
log_failure: Final = (
|
||||
partial(handler.async_log_failure_event, kwargs=dict(kwargs), response_obj=None, start_time=now, end_time=now)
|
||||
if failure_path == "failure_event"
|
||||
else partial(
|
||||
handler.async_post_call_failure_hook,
|
||||
request_data=dict(kwargs),
|
||||
original_exception=Exception("stream ended"),
|
||||
user_api_key_dict=caller,
|
||||
)
|
||||
)
|
||||
log_success: Final = partial(
|
||||
handler.async_log_success_event,
|
||||
kwargs=dict(kwargs),
|
||||
response_obj=ModelResponse(usage=Usage(prompt_tokens=20, completion_tokens=10, total_tokens=30)),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
)
|
||||
|
||||
await asyncio.gather(
|
||||
*((log_success(), log_failure()) if success_starts_first else (log_failure(), log_success()))
|
||||
)
|
||||
|
||||
assert tuple([int(await local_cache.async_get_cache(key=key) or 0) for key in counter_keys]) == (30, 20, 10)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("second_call_type", "second_model"),
|
||||
[("call_mcp_tool", "mcp-tool-call"), ("completion", "gpt-4o-mini")],
|
||||
)
|
||||
async def test_a_pre_call_after_a_failure_leaves_the_success_charging_project_io_once_v3(
|
||||
second_call_type: str, second_model: str
|
||||
) -> None:
|
||||
api_key: Final = hash_token("sk-retried-request-with-mcp-tool")
|
||||
model: Final = "gpt-4o-mini"
|
||||
project: Final = "proj-retried-mcp"
|
||||
local_cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
|
||||
caller: Final = UserAPIKeyAuth(
|
||||
api_key=api_key,
|
||||
tpm_limit=100_000,
|
||||
project_id=project,
|
||||
project_metadata={"model_itpm_limit": {model: 10_000}, "model_otpm_limit": {model: 10_000}},
|
||||
)
|
||||
request: Final = {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 50,
|
||||
"litellm_call_id": "retried-call-with-mcp-tool",
|
||||
}
|
||||
kwargs: Final = {
|
||||
"litellm_call_id": "retried-call-with-mcp-tool",
|
||||
"model": model,
|
||||
"standard_logging_object": {"metadata": {"user_api_key_hash": api_key}},
|
||||
}
|
||||
now: Final = datetime.now()
|
||||
itpm_key: Final = handler.create_rate_limit_keys(
|
||||
key="model_per_project_itpm", value=f"{project}:{model}", rate_limit_type="tokens"
|
||||
)
|
||||
otpm_key: Final = handler.create_rate_limit_keys(
|
||||
key="model_per_project_otpm", value=f"{project}:{model}", rate_limit_type="tokens"
|
||||
)
|
||||
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=caller, cache=local_cache, data=dict(request), call_type="completion"
|
||||
)
|
||||
await handler.async_log_failure_event(kwargs=dict(kwargs), response_obj=None, start_time=now, end_time=now)
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=caller,
|
||||
cache=local_cache,
|
||||
data={**request, "model": second_model},
|
||||
call_type=second_call_type,
|
||||
)
|
||||
await handler.async_log_success_event(
|
||||
kwargs=dict(kwargs),
|
||||
response_obj=ModelResponse(usage=Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27)),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
)
|
||||
|
||||
assert (
|
||||
int(await local_cache.async_get_cache(key=itpm_key) or 0),
|
||||
int(await local_cache.async_get_cache(key=otpm_key) or 0),
|
||||
) == (20, 7)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("steps", "key_tpm_limit", "expected_counters"),
|
||||
[
|
||||
(("failure", "mcp-tool-call", "failure"), 100_000, (0, 0, 0, 0)),
|
||||
(("failure", "tool-model-with-io-limit", "post_call_failure"), None, (0, 0, 0, 0)),
|
||||
(("failure", "mcp-tool-call", "post_call_failure", "success"), 100_000, (20, 7, 0, 0)),
|
||||
(("tool-model-over-io-limit", "success"), 100_000, (20, 7, 0, 0)),
|
||||
],
|
||||
)
|
||||
async def test_a_nested_mcp_pre_call_settles_each_project_io_hold_once_v3(
|
||||
steps: tuple[str, ...], key_tpm_limit: Optional[int], expected_counters: tuple[int, int, int, int]
|
||||
) -> None:
|
||||
api_key: Final = hash_token("sk-nested-mcp-hold-flags")
|
||||
model: Final = "gpt-4o-mini"
|
||||
tool_model: Final = "tool-model-with-io-limit"
|
||||
project: Final = "proj-nested-mcp-hold-flags"
|
||||
io_limits: Final = {model: 10_000, tool_model: 10_000, "tool-model-over-io-limit": 1}
|
||||
local_cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
|
||||
caller: Final = UserAPIKeyAuth(
|
||||
api_key=api_key,
|
||||
tpm_limit=key_tpm_limit,
|
||||
project_id=project,
|
||||
project_metadata={"model_itpm_limit": io_limits, "model_otpm_limit": io_limits},
|
||||
)
|
||||
request: Final = {
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 50,
|
||||
"litellm_call_id": "nested-mcp-hold-flags",
|
||||
}
|
||||
kwargs: Final = {
|
||||
"litellm_call_id": "nested-mcp-hold-flags",
|
||||
"model": model,
|
||||
"standard_logging_object": {"metadata": {"user_api_key_hash": api_key}},
|
||||
}
|
||||
now: Final = datetime.now()
|
||||
counter_keys: Final = tuple(
|
||||
handler.create_rate_limit_keys(
|
||||
key=f"model_per_project_{kind}", value=f"{project}:{scope}", rate_limit_type="tokens"
|
||||
)
|
||||
for scope, kind in ((model, "itpm"), (model, "otpm"), (tool_model, "itpm"), (tool_model, "otpm"))
|
||||
)
|
||||
mcp_pre_call: Final = partial(
|
||||
handler.async_pre_call_hook, user_api_key_dict=caller, cache=local_cache, call_type="call_mcp_tool"
|
||||
)
|
||||
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=caller, cache=local_cache, data=dict(request), call_type="completion"
|
||||
)
|
||||
for step in steps:
|
||||
match step:
|
||||
case "failure":
|
||||
await handler.async_log_failure_event(
|
||||
kwargs=dict(kwargs), response_obj=None, start_time=now, end_time=now
|
||||
)
|
||||
case "post_call_failure":
|
||||
await handler.async_post_call_failure_hook(
|
||||
request_data=dict(request),
|
||||
original_exception=HTTPException(status_code=429, detail="upstream rate limited"),
|
||||
user_api_key_dict=caller,
|
||||
)
|
||||
case "success":
|
||||
await handler.async_log_success_event(
|
||||
kwargs=dict(kwargs),
|
||||
response_obj=ModelResponse(usage=Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27)),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
)
|
||||
case "tool-model-over-io-limit":
|
||||
with pytest.raises(ProxyRateLimitError):
|
||||
await mcp_pre_call(data={**request, "model": step})
|
||||
case _:
|
||||
await mcp_pre_call(data={**request, "model": step})
|
||||
|
||||
assert tuple([int(await local_cache.async_get_cache(key=key) or 0) for key in counter_keys]) == expected_counters
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_post_call_failure_hook_settles_reservation_at_recovered_partial_usage_v3():
|
||||
"""
|
||||
|
|
@ -4198,6 +4447,125 @@ async def test_failure_event_settles_project_itpm_otpm_at_recovered_partial_usag
|
|||
assert int(await local_cache.async_get_cache(key=otpm_key) or 0) == 7
|
||||
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failure_first", [True, False])
|
||||
async def test_a_request_logging_a_failure_and_a_success_settles_its_project_itpm_otpm_once_v3(
|
||||
failure_first: bool,
|
||||
) -> None:
|
||||
api_key: Final = hash_token("sk-failure-and-success-project-io")
|
||||
local_cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key=api_key,
|
||||
project_id="proj-dropped",
|
||||
project_metadata={
|
||||
"model_itpm_limit": {"gpt-4o-mini": 10_000},
|
||||
"model_otpm_limit": {"gpt-4o-mini": 10_000},
|
||||
},
|
||||
),
|
||||
cache=local_cache,
|
||||
data={
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 50,
|
||||
"litellm_call_id": "dropped-project-call",
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
itpm_key: Final = handler.create_rate_limit_keys(
|
||||
key="model_per_project_itpm", value="proj-dropped:gpt-4o-mini", rate_limit_type="tokens"
|
||||
)
|
||||
otpm_key: Final = handler.create_rate_limit_keys(
|
||||
key="model_per_project_otpm", value="proj-dropped:gpt-4o-mini", rate_limit_type="tokens"
|
||||
)
|
||||
kwargs: Final = {
|
||||
"litellm_call_id": "dropped-project-call",
|
||||
"model": "gpt-4o-mini",
|
||||
"standard_logging_object": {"metadata": {"user_api_key_hash": api_key}},
|
||||
}
|
||||
now: Final = datetime.now()
|
||||
log_failure: Final = partial(
|
||||
handler.async_log_failure_event, kwargs=dict(kwargs), response_obj=None, start_time=now, end_time=now
|
||||
)
|
||||
log_success: Final = partial(
|
||||
handler.async_log_success_event,
|
||||
kwargs=dict(kwargs),
|
||||
response_obj=ModelResponse(usage=Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27)),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
)
|
||||
|
||||
for log_event in (log_failure, log_success) if failure_first else (log_success, log_failure):
|
||||
await log_event()
|
||||
|
||||
assert int(await local_cache.async_get_cache(key=itpm_key) or 0) == 20
|
||||
assert int(await local_cache.async_get_cache(key=otpm_key) or 0) == 7
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_a_reservation_made_after_a_released_one_in_the_same_context_settles_at_its_usage_v3() -> None:
|
||||
api_key: Final = hash_token("sk-second-reservation")
|
||||
local_cache: Final = DualCache()
|
||||
handler: Final = _PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(local_cache))
|
||||
user_api_key_dict: Final = UserAPIKeyAuth(
|
||||
api_key=api_key,
|
||||
tpm_limit=10_000,
|
||||
project_id="proj-second",
|
||||
project_metadata={
|
||||
"model_itpm_limit": {"gpt-4o-mini": 10_000},
|
||||
"model_otpm_limit": {"gpt-4o-mini": 10_000},
|
||||
},
|
||||
)
|
||||
now: Final = datetime.now()
|
||||
|
||||
async def reserve(call_id: str) -> None:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
cache=local_cache,
|
||||
data={
|
||||
"model": "gpt-4o-mini",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"max_tokens": 50,
|
||||
"litellm_call_id": call_id,
|
||||
},
|
||||
call_type="completion",
|
||||
)
|
||||
|
||||
def logged(call_id: str) -> dict[str, object]:
|
||||
return {
|
||||
"litellm_call_id": call_id,
|
||||
"model": "gpt-4o-mini",
|
||||
"standard_logging_object": {"metadata": {"user_api_key_hash": api_key}},
|
||||
}
|
||||
|
||||
await reserve("first-call")
|
||||
await handler.async_log_failure_event(kwargs=logged("first-call"), response_obj=None, start_time=now, end_time=now)
|
||||
await reserve("second-call")
|
||||
await handler.async_log_success_event(
|
||||
kwargs=logged("second-call"),
|
||||
response_obj=ModelResponse(usage=Usage(prompt_tokens=20, completion_tokens=7, total_tokens=27)),
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
)
|
||||
|
||||
counters: Final = {
|
||||
"tokens": ("api_key", api_key),
|
||||
"itpm": ("model_per_project_itpm", "proj-second:gpt-4o-mini"),
|
||||
"otpm": ("model_per_project_otpm", "proj-second:gpt-4o-mini"),
|
||||
}
|
||||
settled: Final = {
|
||||
name: int(
|
||||
await local_cache.async_get_cache(
|
||||
key=handler.create_rate_limit_keys(key=key, value=value, rate_limit_type="tokens")
|
||||
)
|
||||
or 0
|
||||
)
|
||||
for name, (key, value) in counters.items()
|
||||
}
|
||||
assert settled == {"tokens": 27, "itpm": 20, "otpm": 7}
|
||||
|
||||
# ----------------------- Per-MCP-server rate limiting (v3) -----------------------
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -3491,6 +3491,40 @@ async def test_sidecar_ignores_an_undecodable_event(): # test-quality-ok: a dis
|
|||
mock_proxy_logging.db_spend_update_writer.update_database.assert_not_awaited()
|
||||
|
||||
|
||||
def _spend_row_error_information(update_database_kwargs: Mapping[str, object]) -> dict[str, object]:
|
||||
row: Final = get_logging_payload(
|
||||
kwargs=cast(dict, update_database_kwargs["kwargs"]),
|
||||
response_obj=update_database_kwargs["completion_response"],
|
||||
start_time=cast(datetime, update_database_kwargs["start_time"]),
|
||||
end_time=cast(datetime, update_database_kwargs["end_time"]),
|
||||
)
|
||||
return json.loads(row["metadata"])["error_information"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_failure_hook_spend_row_truncates_a_long_provider_error_exactly_once(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setenv("MAX_STRING_LENGTH_PROMPT_IN_DB", "100")
|
||||
writer: Final = MagicMock(spec=DBSpendUpdateWriter)
|
||||
writer.update_database = AsyncMock()
|
||||
logger: Final = ProxyDBLogger(spend_writer=lambda: writer)
|
||||
provider_error: Final = "upstream failure " + "x" * 1000
|
||||
|
||||
await logger.async_post_call_failure_hook(
|
||||
request_data={"model": "gpt-4o", "messages": [{"role": "user", "content": "hi"}]},
|
||||
original_exception=litellm.InternalServerError(message=provider_error, llm_provider="openai", model="gpt-4o"),
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"),
|
||||
)
|
||||
|
||||
stored_message: Final = _spend_row_error_information(writer.update_database.call_args.kwargs)["error_message"]
|
||||
assert isinstance(stored_message, str)
|
||||
assert stored_message.count("skipped") == 1
|
||||
kept_chars: Final = len(stored_message.split("... (")[0]) + len(stored_message.rsplit(") ...", 1)[1])
|
||||
skipped_chars: Final = int(stored_message.split("skipped ", 1)[1].split(" chars", 1)[0])
|
||||
assert kept_chars + skipped_chars > len(provider_error)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_post_call_failure_hook_persists_no_raw_model_on_an_unknown_model_rejection():
|
||||
raw_model: Final = "opus-4.6 Please summarize my medical records\nPatient has diabetes"
|
||||
|
|
@ -3504,9 +3538,7 @@ async def test_async_post_call_failure_hook_persists_no_raw_model_on_an_unknown_
|
|||
user_api_key_dict=UserAPIKeyAuth(api_key="test_api_key"),
|
||||
)
|
||||
|
||||
error_information: Final = writer.update_database.call_args.kwargs["kwargs"]["litellm_params"]["metadata"][
|
||||
"error_information"
|
||||
]
|
||||
error_information: Final = _spend_row_error_information(writer.update_database.call_args.kwargs)
|
||||
assert "medical records" not in json.dumps(error_information)
|
||||
assert (
|
||||
error_information["error_message"]
|
||||
|
|
|
|||
|
|
@ -3060,7 +3060,7 @@ def test_retry_success_charges_released_project_io_reservations(
|
|||
stash.itpm_reserved_scopes = frozenset({itpm_scope})
|
||||
stash.otpm_reserved_tokens = 60
|
||||
stash.otpm_reserved_scopes = frozenset({otpm_scope})
|
||||
stash.reservation_released = True
|
||||
stash.mark_reservation_released()
|
||||
|
||||
operations = handler._build_io_token_reservation_ops(
|
||||
{},
|
||||
|
|
|
|||
|
|
@ -7486,6 +7486,46 @@ def test_failure_error_information_masks_ip_addresses_in_long_text_and_at_its_tr
|
|||
assert sanitized["error_code"] == "500"
|
||||
|
||||
|
||||
def test_spend_row_sanitizes_the_session_failure_a_realtime_success_row_carries(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.proxy.proxy_server import general_settings
|
||||
|
||||
monkeypatch.setitem(general_settings, "store_prompts_in_spend_logs", False)
|
||||
monkeypatch.setenv("MAX_STRING_LENGTH_PROMPT_IN_DB", "200")
|
||||
provider_error: Final = 'stream ended {"messages": [{"role": "user", "content": "private-prompt"}]} ' + "x" * 2000
|
||||
now: Final = datetime.datetime.now(timezone.utc)
|
||||
|
||||
row: Final = get_logging_payload(
|
||||
kwargs={
|
||||
"model": "bedrock/amazon.nova-sonic-v1:0",
|
||||
"call_type": "_arealtime",
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"user_api_key": "hashed-key",
|
||||
"status": "failure",
|
||||
"error_information": {
|
||||
"error_code": "503",
|
||||
"error_class": "BedrockError",
|
||||
"error_message": provider_error,
|
||||
"llm_provider": "bedrock",
|
||||
},
|
||||
}
|
||||
},
|
||||
},
|
||||
response_obj=None,
|
||||
start_time=now,
|
||||
end_time=now,
|
||||
)
|
||||
|
||||
stored: Final = json.loads(row["metadata"])["error_information"]
|
||||
assert row["status"] == "failure"
|
||||
assert "private-prompt" not in stored["error_message"]
|
||||
assert LITELLM_TRUNCATED_PAYLOAD_FIELD in stored["error_message"]
|
||||
assert len(stored["error_message"]) < len(provider_error)
|
||||
assert stored["error_class"] == "BedrockError"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"address", ["203.0.113.9", "2001:db8::9", "1:2:3:4:5:6:7:8", "1:2:3:4:5:6:7::", "::1:2:3:4:5:6:7", "fe80::1"]
|
||||
)
|
||||
|
|
|
|||
|
|
@ -3,8 +3,9 @@ import functools
|
|||
import gc
|
||||
import json
|
||||
import os
|
||||
import uuid
|
||||
import weakref
|
||||
from collections.abc import Awaitable, Callable, Iterator, Mapping
|
||||
from collections.abc import Awaitable, Callable, Coroutine, Iterator, Mapping
|
||||
from datetime import datetime, timezone
|
||||
from types import MethodType
|
||||
from typing import Final
|
||||
|
|
@ -13,6 +14,8 @@ from unittest.mock import AsyncMock, MagicMock, patch
|
|||
import httpx
|
||||
import pytest
|
||||
import respx
|
||||
from websockets.exceptions import ConnectionClosed
|
||||
from websockets.frames import Close
|
||||
|
||||
import litellm
|
||||
from litellm import Router
|
||||
|
|
@ -21,8 +24,14 @@ from litellm.caching.in_memory_cache import InMemoryCache
|
|||
from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging
|
||||
from litellm.litellm_core_utils.realtime_errors import close_after_upstream_handshake_refusal
|
||||
from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming
|
||||
from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper
|
||||
from litellm.llms.base_llm.realtime.transformation import BaseRealtimeConfig
|
||||
from litellm.llms.vertex_ai.audio_transcription.realtime_transformation import VertexChirpRealtimeConfig
|
||||
from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook
|
||||
from litellm.router_utils.router_callbacks.track_deployment_metrics import get_deployment_successes_for_current_minute
|
||||
from litellm.types.litellm_params import RoutingStrategyName
|
||||
from litellm.types.router import AlertingConfig, Deployment, DeploymentTypedDict, LiteLLM_Params, ModelInfo
|
||||
from litellm.types.utils import (
|
||||
|
|
@ -540,6 +549,242 @@ async def test_deployment_callback_on_success_counts_a_responses_websocket_turn(
|
|||
assert await router.cache.async_get_cache(key=tpm_key) == 17
|
||||
|
||||
|
||||
class _DiscardingLoggingWorker:
|
||||
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, object, None]) -> None:
|
||||
async_coroutine.close()
|
||||
|
||||
|
||||
async def _client_never_sends() -> str:
|
||||
await asyncio.Event().wait()
|
||||
raise AssertionError("unreachable")
|
||||
|
||||
|
||||
def _realtime_client_ws() -> MagicMock:
|
||||
client_ws: Final = MagicMock()
|
||||
client_ws.headers = {}
|
||||
client_ws.receive_text = AsyncMock(side_effect=_client_never_sends)
|
||||
client_ws.send_text = AsyncMock()
|
||||
client_ws.close = AsyncMock()
|
||||
return client_ws
|
||||
|
||||
|
||||
def _realtime_session_logging() -> LiteLLMLogging:
|
||||
return LiteLLMLogging(
|
||||
model="gpt-realtime-2.1",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="_arealtime",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id=str(uuid.uuid4()),
|
||||
function_id="realtime-session",
|
||||
)
|
||||
|
||||
|
||||
async def _relay_upstream(
|
||||
litellm_logging_obj: LiteLLMLogging,
|
||||
*upstream_frames: bytes | Exception,
|
||||
client_ws: MagicMock | None = None,
|
||||
provider_config: BaseRealtimeConfig | None = None,
|
||||
upstream_stays_open: bool = False,
|
||||
) -> None:
|
||||
backend_ws: Final = MagicMock()
|
||||
backend_ws.recv = AsyncMock(
|
||||
side_effect=_frames_then_silence(*upstream_frames) if upstream_stays_open else list(upstream_frames)
|
||||
)
|
||||
backend_ws.send = AsyncMock()
|
||||
relay: Final = RealTimeStreaming(
|
||||
client_ws or _realtime_client_ws(),
|
||||
backend_ws,
|
||||
litellm_logging_obj,
|
||||
provider_config=provider_config,
|
||||
model="gpt-realtime-2.1",
|
||||
logging_worker=_DiscardingLoggingWorker(),
|
||||
)
|
||||
await asyncio.wait_for(relay.bidirectional_forward(), timeout=2)
|
||||
|
||||
|
||||
_REALTIME_SESSION_CREATED: Final = json.dumps({"type": "session.created", "session": {"id": "sess_1"}}).encode()
|
||||
|
||||
|
||||
async def _clean_realtime_session(litellm_logging_obj: LiteLLMLogging) -> None:
|
||||
await _relay_upstream(litellm_logging_obj, _REALTIME_SESSION_CREATED, ConnectionClosed(Close(1000, ""), None))
|
||||
|
||||
|
||||
async def _upstream_internal_error_after_a_frame(litellm_logging_obj: LiteLLMLogging) -> None:
|
||||
await _relay_upstream(
|
||||
litellm_logging_obj, _REALTIME_SESSION_CREATED, ConnectionClosed(Close(1011, "upstream failed"), None)
|
||||
)
|
||||
|
||||
|
||||
async def _upstream_drop_without_close_frame_after_a_frame(litellm_logging_obj: LiteLLMLogging) -> None:
|
||||
await _relay_upstream(litellm_logging_obj, _REALTIME_SESSION_CREATED, ConnectionClosed(None, None))
|
||||
|
||||
|
||||
async def _upstream_close_before_any_frame(litellm_logging_obj: LiteLLMLogging) -> None:
|
||||
await _relay_upstream(litellm_logging_obj, ConnectionClosed(Close(1008, "model not found"), None))
|
||||
|
||||
|
||||
def _frames_then_silence(*frames: bytes | str) -> Callable[[], Awaitable[bytes | str]]:
|
||||
pending: Final = iter(frames)
|
||||
|
||||
async def _next_frame() -> bytes | str:
|
||||
frame: Final = next(pending, None)
|
||||
if frame is None:
|
||||
await asyncio.Event().wait()
|
||||
raise AssertionError("unreachable")
|
||||
return frame
|
||||
|
||||
return _next_frame
|
||||
|
||||
|
||||
def _client_ws_sending(*client_frames: str) -> MagicMock:
|
||||
client_ws: Final = _realtime_client_ws()
|
||||
client_ws.receive_text = AsyncMock(side_effect=_frames_then_silence(*client_frames))
|
||||
return client_ws
|
||||
|
||||
|
||||
async def _vertex_access_token() -> str:
|
||||
return "token"
|
||||
|
||||
|
||||
async def _client_frame_the_provider_rejects(litellm_logging_obj: LiteLLMLogging) -> None:
|
||||
deferred_setup_provider: Final = VertexChirpRealtimeConfig(
|
||||
resolve_access_token=_vertex_access_token, project="proj-1", location="us"
|
||||
)
|
||||
await _relay_upstream(
|
||||
litellm_logging_obj,
|
||||
client_ws=_client_ws_sending(
|
||||
json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"session": {"type": "transcription", "audio": {"input": {"format": {"type": "audio/pcmu"}}}},
|
||||
}
|
||||
)
|
||||
),
|
||||
provider_config=deferred_setup_provider,
|
||||
upstream_stays_open=True,
|
||||
)
|
||||
|
||||
|
||||
async def _client_hangs_up_after_a_frame(litellm_logging_obj: LiteLLMLogging) -> None:
|
||||
client_ws: Final = _realtime_client_ws()
|
||||
client_ws.receive_text = AsyncMock(side_effect=ConnectionClosed(Close(1000, ""), None))
|
||||
await _relay_upstream(litellm_logging_obj, _REALTIME_SESSION_CREATED, client_ws=client_ws, upstream_stays_open=True)
|
||||
|
||||
|
||||
async def _refused_handshake(litellm_logging_obj: LiteLLMLogging) -> None:
|
||||
await close_after_upstream_handshake_refusal(_realtime_client_ws(), 403)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("realtime_session", "counted_successes"),
|
||||
[
|
||||
(_clean_realtime_session, 1),
|
||||
(_upstream_close_before_any_frame, 0),
|
||||
(_refused_handshake, 0),
|
||||
(_upstream_internal_error_after_a_frame, 0),
|
||||
(_upstream_drop_without_close_frame_after_a_frame, 0),
|
||||
(_client_frame_the_provider_rejects, 0),
|
||||
(_client_hangs_up_after_a_frame, 1),
|
||||
],
|
||||
ids=[
|
||||
"clean-session",
|
||||
"upstream-close-before-any-frame",
|
||||
"refused-handshake",
|
||||
"upstream-1011-after-a-frame",
|
||||
"upstream-1006-after-a-frame",
|
||||
"client-frame-the-provider-rejects",
|
||||
"client-hangs-up-after-a-frame",
|
||||
],
|
||||
)
|
||||
async def test_realtime_session_counts_a_router_success_only_when_it_ended_without_a_failure(
|
||||
realtime_session: Callable[[LiteLLMLogging], Awaitable[None]], counted_successes: int
|
||||
) -> None:
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "realtime-voice",
|
||||
"litellm_params": {"model": "openai/gpt-realtime-2.1"},
|
||||
"model_info": {"id": "realtime-a"},
|
||||
}
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
|
||||
async def _arealtime(litellm_logging_obj: LiteLLMLogging, **kwargs: object) -> None:
|
||||
await realtime_session(litellm_logging_obj)
|
||||
|
||||
await router._ageneric_api_call_with_fallbacks(
|
||||
model="realtime-voice", original_function=_arealtime, litellm_logging_obj=_realtime_session_logging()
|
||||
)
|
||||
|
||||
assert get_deployment_successes_for_current_minute(router, "realtime-a") == counted_successes
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("realtime_model", ["bedrock/amazon.nova-sonic-v1:0", "openai/gpt-realtime-2.1"])
|
||||
@pytest.mark.parametrize(("clean_sessions", "cooled_down"), [(10, False), (3, True)])
|
||||
async def test_realtime_cooldown_reads_the_failure_rate_of_the_sessions_a_deployment_served(
|
||||
realtime_model: str, clean_sessions: int, cooled_down: bool
|
||||
) -> None:
|
||||
import time
|
||||
|
||||
provider, _, provider_model = realtime_model.partition("/")
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "realtime-voice",
|
||||
"litellm_params": {"model": realtime_model, "order": order},
|
||||
"model_info": {"id": deployment_id},
|
||||
}
|
||||
for order, deployment_id in ((1, "realtime-a"), (2, "realtime-b"))
|
||||
],
|
||||
num_retries=0,
|
||||
)
|
||||
drop: Final = litellm.ServiceUnavailableError(
|
||||
message="realtime provider stream ended before the session closed",
|
||||
llm_provider=provider,
|
||||
model=provider_model,
|
||||
)
|
||||
|
||||
def realtime_session(failure: Exception | None):
|
||||
async def _arealtime(litellm_logging_obj: LiteLLMLogging, **kwargs: object) -> None:
|
||||
if failure is not None:
|
||||
raise failure
|
||||
await _clean_realtime_session(litellm_logging_obj)
|
||||
|
||||
return _arealtime
|
||||
|
||||
for _ in range(clean_sessions):
|
||||
await router._ageneric_api_call_with_fallbacks(
|
||||
model="realtime-voice",
|
||||
original_function=realtime_session(None),
|
||||
litellm_logging_obj=_realtime_session_logging(),
|
||||
)
|
||||
|
||||
async def drop_session() -> bool:
|
||||
with pytest.raises(litellm.ServiceUnavailableError):
|
||||
await router._ageneric_api_call_with_fallbacks(
|
||||
model="realtime-voice",
|
||||
original_function=realtime_session(drop),
|
||||
litellm_logging_obj=_realtime_session_logging(),
|
||||
)
|
||||
return router.deployment_callback_on_failure(
|
||||
kwargs={
|
||||
"exception": drop,
|
||||
"litellm_params": {"metadata": {"model_group": "realtime-voice"}, "model_info": {"id": "realtime-a"}},
|
||||
},
|
||||
completion_response=None,
|
||||
start_time=time.time(),
|
||||
end_time=time.time(),
|
||||
)
|
||||
|
||||
drop_cooldowns: Final = [await drop_session() for _ in range(5)]
|
||||
|
||||
assert any(drop_cooldowns) is cooled_down
|
||||
|
||||
|
||||
async def _kwargs_routed_to(router: Router, model: str, function_name: str, **request: object) -> dict[str, object]:
|
||||
routed: Final[dict[str, object]] = {}
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue