diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index d23396422d1..e6c6722e21b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -13740,64 +13740,68 @@ async def audio_speech( ### ALERTING ### asyncio.create_task(proxy_logging_obj.update_request_status(litellm_call_id=litellm_call_id, status="success")) - ### RESPONSE HEADERS ### - hidden_params: Final = getattr(response, "_hidden_params", {}) or {} - model_id: Final = hidden_params.get("model_id", None) or "" - cache_key: Final = hidden_params.get("cache_key", None) or "" - api_base: Final = hidden_params.get("api_base", None) or "" - response_cost: Final = hidden_params.get("response_cost", None) or "" - response_call_id: Final = hidden_params.get("litellm_call_id", None) or "" + try: + ### RESPONSE HEADERS ### + hidden_params: Final = getattr(response, "_hidden_params", {}) or {} + model_id: Final = hidden_params.get("model_id", None) or "" + cache_key: Final = hidden_params.get("cache_key", None) or "" + api_base: Final = hidden_params.get("api_base", None) or "" + response_cost: Final = hidden_params.get("response_cost", None) or "" + response_call_id: Final = hidden_params.get("litellm_call_id", None) or "" - custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( - user_api_key_dict=user_api_key_dict, - model_id=model_id, - cache_key=cache_key, - api_base=api_base, - version=version, - response_cost=response_cost, - model_region=getattr(user_api_key_dict, "allowed_model_region", ""), - fastest_response_batch_completion=None, - call_id=response_call_id, - request_data=data, - hidden_params=hidden_params, - ) + custom_headers: Final = ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + model_id=model_id, + cache_key=cache_key, + api_base=api_base, + version=version, + response_cost=response_cost, + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + fastest_response_batch_completion=None, + call_id=response_call_id, + request_data=data, + hidden_params=hidden_params, + ) - # Call response headers hook (matches audio_transcription behavior) - callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook( - data=data, - user_api_key_dict=user_api_key_dict, - response=response, - request_headers=dict(request.headers), - ) - if callback_headers: - custom_headers.update(callback_headers) + # Call response headers hook (matches audio_transcription behavior) + callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook( + data=data, + user_api_key_dict=user_api_key_dict, + response=response, + request_headers=dict(request.headers), + ) + if callback_headers: + custom_headers.update(callback_headers) - requested_format: Final = data.get("response_format") - upstream_content_type: Final = ( - response.content_type if isinstance(response, HttpxBinaryResponseContent) else None - ) - media_type: Final = resolve_speech_media_type( - upstream_content_type=upstream_content_type, - response_format=requested_format if isinstance(requested_format, str) else None, - ) + requested_format: Final = data.get("response_format") + upstream_content_type: Final = ( + response.content_type if isinstance(response, HttpxBinaryResponseContent) else None + ) + media_type: Final = resolve_speech_media_type( + upstream_content_type=upstream_content_type, + response_format=requested_format if isinstance(requested_format, str) else None, + ) - audio: Final = cast( # cast-ok: the speech routes return this, and resolve_speech_media_type above already reads it as such - HttpxBinaryResponseContent, response - ) - return StreamingResponse( - _audio_speech_chunk_generator( - audio, - on_upstream_error=partial( - _log_audio_speech_stream_failure, - request_data=data, - user_api_key_dict=user_api_key_dict, - response_cost=float(hidden_params.get("response_cost") or 0.0), + audio: Final = cast( # cast-ok: the speech routes return this, and resolve_speech_media_type above already reads it as such + HttpxBinaryResponseContent, response + ) + return StreamingResponse( + _audio_speech_chunk_generator( + audio, + on_upstream_error=partial( + _log_audio_speech_stream_failure, + request_data=data, + user_api_key_dict=user_api_key_dict, + response_cost=float(hidden_params.get("response_cost") or 0.0), + ), ), - ), - media_type=media_type, - headers=custom_headers, - background=BackgroundTask(_close_audio_speech_response, audio), - ) + media_type=media_type, + headers=custom_headers, + background=BackgroundTask(_close_audio_speech_response, audio), + ) + except BaseException: + await _close_audio_speech_response(response) + raise except Exception as e: await proxy_logging_obj.post_call_failure_hook( diff --git a/tests/unit/proxy/proxy_server/test_routes_audio.py b/tests/unit/proxy/proxy_server/test_routes_audio.py index b3ed1235500..9e160399c78 100644 --- a/tests/unit/proxy/proxy_server/test_routes_audio.py +++ b/tests/unit/proxy/proxy_server/test_routes_audio.py @@ -16,7 +16,7 @@ import io import json import math import wave -from collections.abc import AsyncIterator, Callable, Mapping +from collections.abc import AsyncIterator, Awaitable, Callable, Mapping from contextlib import AbstractContextManager from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -498,7 +498,11 @@ async def _serve_until_disconnect(response: StreamingResponse) -> None: await response({"type": "http", "asgi": {"spec_version": "2.4"}}, _receive, _send) -def _route_speech_to(monkeypatch: pytest.MonkeyPatch, response: HttpxBinaryResponseContent) -> AsyncMock: +def _route_speech_to( + monkeypatch: pytest.MonkeyPatch, + response: HttpxBinaryResponseContent, + response_headers_hook: Callable[..., Awaitable[dict[str, str]]] | None = None, +) -> AsyncMock: failure_hook: Final = AsyncMock() monkeypatch.setattr(proxy_server, "llm_router", MagicMock()) monkeypatch.setattr( @@ -507,7 +511,7 @@ def _route_speech_to(monkeypatch: pytest.MonkeyPatch, response: HttpxBinaryRespo MagicMock( pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]), post_call_failure_hook=failure_hook, - post_call_response_headers_hook=AsyncMock(return_value={}), + post_call_response_headers_hook=response_headers_hook or AsyncMock(return_value={}), update_request_status=AsyncMock(), ), ) @@ -586,6 +590,61 @@ async def test_audio_speech_route_releases_an_upstream_whose_body_is_never_read( assert stream.closed is True +def _elevenlabs_client(stream: httpx.AsyncByteStream) -> litellm.AsyncHTTPHandler: + return litellm.AsyncHTTPHandler( + transport=httpx.MockTransport( + lambda _: httpx.Response(200, headers={"content-type": "audio/mpeg"}, stream=stream) + ) + ) + + +async def _headers_from_the_audio_body(*, response: HttpxBinaryResponseContent, **_: object) -> dict[str, str]: + return {"x-audio-bytes": str(len(response.content))} + + +@pytest.mark.asyncio +async def test_audio_speech_route_frees_the_deployment_slot_when_a_response_hook_fails( + monkeypatch: pytest.MonkeyPatch, +): + """A response hook that fails before any audio is sent must not keep the deployment's only + slot, or every later clip on max_parallel_requests: 1 is turned away.""" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = Router( + model_list=[ + { + "model_name": "tts", + "litellm_params": { + "model": "elevenlabs/eleven_multilingual_v2", + "api_key": "fake-key", + "max_parallel_requests": 1, + }, + } + ] + ) + first_upstream: Final = _RecordingAudioStream([b"audio"], []) + first_clip: Final = await router.aspeech( + model="tts", input="hello", voice="voice-id", client=_elevenlabs_client(first_upstream), stream=True + ) + _route_speech_to(monkeypatch, first_clip, response_headers_hook=_headers_from_the_audio_body) + + with pytest.raises(ProxyException): + await proxy_server.audio_speech( + request=_speech_request({"model": "tts", "input": "hello", "voice": "voice-id", "stream": True}), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + next_clip: Final = await router.aspeech( + model="tts", + input="hello", + voice="voice-id", + client=_elevenlabs_client(_RecordingAudioStream([b"audio"], [])), + stream=True, + ) + await next_clip.aclose() + assert first_upstream.closed is True + + class _BreaksMidClipAudioStream(httpx.AsyncByteStream): async def __aiter__(self): yield b"a" * AUDIO_SPEECH_CHUNK_SIZE