mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(proxy): close streamed speech when the route fails before responding
A response headers hook, or any other step between aspeech returning an unread audio stream and the route returning its StreamingResponse, could raise and leave that stream open. The upstream connection and the deployment's max_parallel_requests slot then stayed held, so later clips got 429s. The route now closes the audio on any error before it returns
This commit is contained in:
mateo-berri
parent
5fa3c58bac
commit
94fd515f91
2 changed files with 119 additions and 56 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue