diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 650dcae3239..09652dfca2a 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -139,6 +139,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, LITELLM_PASS_THROUGH_ENDPOINT_MARKER, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, + WEBSOCKET_CLIENT_TEXT_COUNT_KEY, EndpointType, PassthroughStandardLoggingPayload, ) @@ -2489,13 +2490,14 @@ def _with_trace_context(headers: Mapping[str, str], parent_span: object) -> dict async def websocket_passthrough_request( websocket: WebSocket, target: str, - custom_headers: dict, + custom_headers: Mapping[str, str], user_api_key_dict: UserAPIKeyAuth, forward_headers: bool | None = False, endpoint: str | None = None, cost_per_request: float | None = None, accept_websocket: bool = True, setup_model_rewriter: Callable[[str], str] | None = None, + count_client_text: Callable[[str | bytes], int] | None = None, ): """ WebSocket passthrough request handler. @@ -2509,6 +2511,9 @@ async def websocket_passthrough_request( endpoint: The endpoint path (for logging purposes) cost_per_request: Optional field - cost per request to the target endpoint setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream + count_client_text: Optional read-only count of each client frame, text or binary, after it reaches the + upstream, such as billable characters. Only the sum reaches the success handler, under WEBSOCKET_CLIENT_TEXT_COUNT_KEY + in request_body, so routes whose client text must not be stored can still bill it """ from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.proxy_server import proxy_config, proxy_logging_obj @@ -2519,6 +2524,7 @@ async def websocket_passthrough_request( # Initialize tracking variables start_time: Final = datetime.now() websocket_messages: Final[list[dict[str, object]]] = [] + client_text_count: int = 0 # rebind-ok: forward_client_to_upstream adds each client frame's count litellm_call_id: Final = str(uuid.uuid4()) verbose_proxy_logger.info("WebSocket passthrough (%s): Starting WebSocket connection to %s", endpoint, target) @@ -2637,6 +2643,7 @@ async def websocket_passthrough_request( async def forward_client_to_upstream() -> None: """Forward messages from client to upstream WebSocket""" + nonlocal client_text_count try: while True: message = await websocket.receive() @@ -2706,8 +2713,12 @@ async def websocket_passthrough_request( # Not a JSON message or doesn't contain setup data await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter)) + if count_client_text is not None: + client_text_count += count_client_text(text_data) elif bytes_data is not None: await upstream_ws.send(bytes_data) + if count_client_text is not None: + client_text_count += count_client_text(bytes_data) except asyncio.CancelledError: raise except Exception: @@ -2844,7 +2855,9 @@ async def websocket_passthrough_request( end_time=end_time, logging_obj=logging_obj, cache_hit=False, - request_body={}, + request_body={} + if count_client_text is None + else {WEBSOCKET_CLIENT_TEXT_COUNT_KEY: client_text_count}, **success_kwargs, ) ) diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 0d99cd54a6b..40a78d79194 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -27,6 +27,8 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__ SETTLES_BACKGROUND_RESPONSE_KWARG: Final = "settles_background_response" +WEBSOCKET_CLIENT_TEXT_COUNT_KEY: Final = "litellm_websocket_client_text_count" + class BackgroundResponseSettlement(str, Enum): SETTLES = "settles" diff --git a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py index b0504af1b5e..a797f8084b9 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_pass_through_endpoints.py @@ -66,6 +66,7 @@ from litellm.types.guardrails import GuardrailEventHooks from litellm.types.passthrough_endpoints.pass_through_endpoints import ( LITELLM_PASS_THROUGH_DEPLOYMENT_MODEL_INFO_STATE_KEY, LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY, + WEBSOCKET_CLIENT_TEXT_COUNT_KEY, EndpointType, PassthroughStandardLoggingPayload, ) @@ -7808,6 +7809,151 @@ async def _track_cost_for_passthrough_kwargs(kwargs: dict) -> AsyncMock: return increment_spend_counters +DEEPGRAM_SPEAK_TARGET = "wss://api.deepgram.com/v1/speak?model=aura-2-thalia-en" +SPEAK_AUDIO_CHUNK = bytes(range(256)) * 8 +SPOKEN_TEXT = "The quick brown fox jumps over the lazy dog." +SPEAK_FRAME = json.dumps({"type": "Speak", "text": SPOKEN_TEXT}) +FLUSH_FRAME = json.dumps({"type": "Flush"}) +FLUSHED_FRAME = json.dumps({"type": "Flushed", "sequence_id": 0}) +SPEAK_CLOSE_FRAME = json.dumps({"type": "Close"}) +CLIENT_BINARY_FRAME = b"\x00\xffclient-binary" +SPEAK_SESSION_FRAMES = ({"text": SPEAK_FRAME}, {"bytes": CLIENT_BINARY_FRAME}, {"text": FLUSH_FRAME}) +SPEAK_SESSION_CLIENT_LENGTH = sum(map(len, (SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME))) +RELAY_HANG_GUARD_SECONDS = 5.0 + + +class SpeakingUpstreamWebSocket: + """Answers each Speak with audio as soon as it arrives and ends the session on Close, the way Deepgram speak + does, so a relay that held audio back until the end would never see the Close it waits for""" + + def __init__(self, close_code: int = 1000, close_reason: str = ""): + self.received: list[str | bytes] = [] + self.close = AsyncMock() + self._close_code = close_code + self._close_reason = close_reason + self._replies: asyncio.Queue[str | bytes | None] = asyncio.Queue() + + async def send(self, message: str | bytes) -> None: + self.received.append(message) + match json.loads(message).get("type") if isinstance(message, str) else None: + case "Speak": + self._replies.put_nowait(SPEAK_AUDIO_CHUNK) + case "Flush": + self._replies.put_nowait(FLUSHED_FRAME) + case "Close": + self._replies.put_nowait(None) + + async def recv(self, decode: bool | None = None) -> str | bytes: + from websockets.exceptions import ConnectionClosedError, ConnectionClosedOK + from websockets.frames import Close + + reply = await self._replies.get() + if reply is not None: + return reply + close = Close(self._close_code, self._close_reason) + closed = ConnectionClosedOK if self._close_code == 1000 else ConnectionClosedError + raise closed(rcvd=close, sent=close, rcvd_then_sent=True) + + +def _speak_client(frames: tuple[dict[str, str | bytes], ...] = SPEAK_SESSION_FRAMES): + """A client socket that sends ``frames``, then Close once the first audio has reached it""" + first_audio = asyncio.Event() + + async def script(): + for frame in frames: + yield {"type": "websocket.receive", **frame} + await first_audio.wait() + yield {"type": "websocket.receive", "text": SPEAK_CLOSE_FRAME} + await asyncio.Event().wait() + + steps = script() + + async def receive(): + return await anext(steps) + + websocket = _client_websocket(receive) + websocket.send_bytes = AsyncMock(side_effect=lambda data: first_audio.set()) + return websocket + + +async def _relay_deepgram_speak(upstream_ws, websocket, count_client_text): + """Runs the generic relay the way the Deepgram speak route does and returns the success handler mock""" + with ( + _patched_websocket_passthrough_environment(upstream_ws), + patch( # test-quality-ok: pass_through_endpoint_logging is a module global read inside websocket_passthrough_request; there is no injection seam + "litellm.proxy.pass_through_endpoints.pass_through_endpoints." + "pass_through_endpoint_logging.pass_through_async_success_handler", + new=AsyncMock(), + ) as success_handler, + ): + await asyncio.wait_for( + websocket_passthrough_request( + websocket=websocket, + target=DEEPGRAM_SPEAK_TARGET, + custom_headers={"Authorization": "Token dg-provider-key"}, + user_api_key_dict=UserAPIKeyAuth(), + forward_headers=False, + endpoint="/deepgram/v1/speak", + accept_websocket=False, + count_client_text=count_client_text, + ), + timeout=RELAY_HANG_GUARD_SECONDS, + ) + return success_handler + + +@pytest.mark.asyncio +async def test_websocket_passthrough_counts_client_text_without_changing_what_is_relayed(): + """Audio reaches the client while the session is open, every client frame reaches the upstream unchanged, and + the success handler gets only the summed count of the text and binary frames, never the frames""" + upstream_ws = SpeakingUpstreamWebSocket() + websocket = _speak_client() + + success_handler = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len) + + assert upstream_ws.received == [SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME] + assert [call.args[0] for call in websocket.send_bytes.await_args_list] == [SPEAK_AUDIO_CHUNK] + assert [call.args[0] for call in websocket.send_text.await_args_list] == [FLUSHED_FRAME] + assert success_handler.call_args.kwargs["request_body"] == { + WEBSOCKET_CLIENT_TEXT_COUNT_KEY: SPEAK_SESSION_CLIENT_LENGTH + } + websocket.close.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_websocket_passthrough_without_a_counter_reports_an_empty_request_body(): + """Listen, realtime and every other caller pass no counter, so their success handler sees what it saw before""" + upstream_ws = SpeakingUpstreamWebSocket() + + success_handler = await _relay_deepgram_speak(upstream_ws, _speak_client(), count_client_text=None) + + assert upstream_ws.received == [SPEAK_FRAME, CLIENT_BINARY_FRAME, FLUSH_FRAME, SPEAK_CLOSE_FRAME] + assert success_handler.call_args.kwargs["request_body"] == {} + assert success_handler.call_args.kwargs["response_body"] == [json.loads(FLUSHED_FRAME)] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("close_code", "close_reason"), + [ + pytest.param(1008, "DATA-0000: invalid text", id="policy"), + pytest.param(1011, "NET-0001: internal", id="internal"), + ], +) +async def test_websocket_passthrough_counts_client_text_when_the_upstream_ends_with_an_error_close( + close_code: int, close_reason: str +): + upstream_ws = SpeakingUpstreamWebSocket(close_code=close_code, close_reason=close_reason) + websocket = _speak_client() + + success_handler = await _relay_deepgram_speak(upstream_ws, websocket, count_client_text=len) + + websocket.close.assert_awaited_once_with(code=close_code, reason=close_reason) + assert success_handler.call_args.kwargs["request_body"] == { + WEBSOCKET_CLIENT_TEXT_COUNT_KEY: SPEAK_SESSION_CLIENT_LENGTH + } + + @pytest.mark.asyncio async def test_passthrough_success_reconciles_budget_reservation(): """