From cae5a9bdd6c559e538a412fc5594fd93fdb81a3f Mon Sep 17 00:00:00 2001 From: Gabriel Creti Date: Tue, 13 Oct 2026 16:53:21 +0200 Subject: [PATCH 01/20] feat(tts): stream audio from the provider as it arrives on /v1/audio/speech The proxy read the whole audio body from the provider before sending the first byte to the client, so time-to-first-byte was the full synthesis time. Pass stream=True to get the bytes forwarded as they arrive. The response type does not change: it is still HttpxBinaryResponseContent, only unread, so every existing caller, logging path and spend hook works the same. Providers opt in with supports_streaming on their config, the others ignore the flag and buffer as before. --- .../text_to_speech/transformation.py | 4 + litellm/llms/azure/azure.py | 35 +++--- .../azure/text_to_speech/transformation.py | 4 + .../base_llm/text_to_speech/transformation.py | 4 +- litellm/llms/custom_httpx/llm_http_handler.py | 11 ++ .../text_to_speech/transformation.py | 2 + litellm/llms/openai/openai.py | 36 +++--- litellm/main.py | 6 + litellm/proxy/proxy_server.py | 7 +- litellm/utils.py | 16 +++ .../custom_httpx/test_llm_http_handler.py | 98 ++++++++++++++++ .../llms/openai/test_openai_audio_speech.py | 105 ++++++++++++++++++ .../proxy/proxy_server/test_routes_audio.py | 50 +++++++++ tests/unit/test_utils.py | 4 + 14 files changed, 351 insertions(+), 31 deletions(-) create mode 100644 tests/unit/llms/openai/test_openai_audio_speech.py diff --git a/litellm/llms/aws_polly/text_to_speech/transformation.py b/litellm/llms/aws_polly/text_to_speech/transformation.py index 01fec02c310..debb19b4a7f 100644 --- a/litellm/llms/aws_polly/text_to_speech/transformation.py +++ b/litellm/llms/aws_polly/text_to_speech/transformation.py @@ -35,6 +35,8 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): Reference: https://docs.aws.amazon.com/polly/latest/dg/API_SynthesizeSpeech.html """ + supports_streaming = True + def __init__(self): BaseTextToSpeechConfig.__init__(self) BaseAWSLLM.__init__(self) @@ -82,6 +84,7 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): aspeech: bool, api_base: str | None, api_key: str | None, + stream: bool = False, **kwargs: object, ) -> Union[ "HttpxBinaryResponseContent", @@ -128,6 +131,7 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): extra_headers=extra_headers, client=None, _is_async=aspeech, + stream=stream, ) return response diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index d396e1f4867..2090d4d76de 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -2,7 +2,7 @@ import asyncio import json import time from collections.abc import Callable, Coroutine -from typing import Final +from typing import Final, cast import httpx from openai import ( @@ -13,6 +13,7 @@ from openai import ( BadRequestError, OpenAI, ) +from openai.types.audio.speech_create_params import SpeechCreateParams import litellm from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES @@ -1380,6 +1381,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): aspeech: bool | None = None, client=None, litellm_params: dict | None = None, + stream: bool = False, ) -> HttpxBinaryResponseContent: max_retries = optional_params.pop("max_retries", 2) @@ -1399,6 +1401,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): logging_obj=logging_obj, client=client, litellm_params=litellm_params, + stream=stream, ) azure_client: Final[AzureOpenAI] = self.get_azure_openai_client( @@ -1420,13 +1423,15 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): }, ) - response: Final = azure_client.audio.speech.create( - model=model, - voice=voice, - input=input, - **optional_params, + params: Final = cast( # cast-ok: the provider-mapped body is this SDK call's own params shape + SpeechCreateParams, {"model": model, "voice": voice, "input": input, **optional_params} ) - return HttpxBinaryResponseContent(response=response.response) + raw_response: Final = ( + azure_client.audio.speech.with_streaming_response.create(**params).__enter__().http_response + if stream + else azure_client.audio.speech.create(**params).response + ) + return HttpxBinaryResponseContent(response=raw_response) async def async_audio_speech( self, @@ -1444,6 +1449,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): logging_obj: LiteLLMLoggingObj, client=None, litellm_params: dict | None = None, + stream: bool = False, ) -> HttpxBinaryResponseContent: azure_client: Final[AsyncAzureOpenAI] = self.get_azure_openai_client( api_base=api_base, @@ -1464,14 +1470,15 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): }, ) - azure_response: Final = await azure_client.audio.speech.create( - model=model, - voice=voice, - input=input, - **optional_params, + params: Final = cast( # cast-ok: the provider-mapped body is this SDK call's own params shape + SpeechCreateParams, {"model": model, "voice": voice, "input": input, **optional_params} ) - - return HttpxBinaryResponseContent(response=azure_response.response) + raw_response: Final = ( + (await azure_client.audio.speech.with_streaming_response.create(**params).__aenter__()).http_response + if stream + else (await azure_client.audio.speech.create(**params)).response + ) + return HttpxBinaryResponseContent(response=raw_response) def get_headers( self, diff --git a/litellm/llms/azure/text_to_speech/transformation.py b/litellm/llms/azure/text_to_speech/transformation.py index eed7a3178ca..9720bcc31ba 100644 --- a/litellm/llms/azure/text_to_speech/transformation.py +++ b/litellm/llms/azure/text_to_speech/transformation.py @@ -33,6 +33,8 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): Reference: https://learn.microsoft.com/en-us/azure/ai-services/speech-service/rest-text-to-speech """ + supports_streaming = True + # Azure endpoint domains DEFAULT_VOICE = "en-US-AriaNeural" COGNITIVE_SERVICES_DOMAIN = "api.cognitive.microsoft.com" @@ -73,6 +75,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): aspeech: bool, api_base: str | None, api_key: str | None, + stream: bool = False, **kwargs: object, ) -> Union[ "HttpxBinaryResponseContent", @@ -129,6 +132,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig): extra_headers=extra_headers, client=None, _is_async=aspeech, + stream=stream, ) return response diff --git a/litellm/llms/base_llm/text_to_speech/transformation.py b/litellm/llms/base_llm/text_to_speech/transformation.py index fb85cdb7687..aed733658ab 100644 --- a/litellm/llms/base_llm/text_to_speech/transformation.py +++ b/litellm/llms/base_llm/text_to_speech/transformation.py @@ -1,6 +1,6 @@ import types from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any, TypedDict +from typing import TYPE_CHECKING, Any, ClassVar, TypedDict import httpx @@ -35,6 +35,8 @@ class TextToSpeechRequestData(TypedDict, total=False): class BaseTextToSpeechConfig(ABC): + supports_streaming: ClassVar[bool] = False + def __init__(self): pass diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 134ab96ab56..29f4eb5f85f 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -12106,6 +12106,7 @@ class BaseLLMHTTPHandler: extra_headers: dict[str, object] | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, _is_async: bool = False, + stream: bool = False, ) -> Union[ "HttpxBinaryResponseContent", Coroutine[object, object, "HttpxBinaryResponseContent"], @@ -12127,6 +12128,7 @@ class BaseLLMHTTPHandler: extra_headers=extra_headers, timeout=timeout, client=client if isinstance(client, AsyncHTTPHandler) else None, + stream=stream, ) if client is None or not isinstance(client, HTTPHandler): @@ -12174,6 +12176,8 @@ class BaseLLMHTTPHandler: }, ) + should_stream: Final = stream and text_to_speech_provider_config.supports_streaming + try: # Determine request body type and send appropriately if "dict_body" in request_data: @@ -12182,6 +12186,7 @@ class BaseLLMHTTPHandler: headers=headers, json=request_data["dict_body"], timeout=timeout, + stream=should_stream, ) elif "ssml_body" in request_data: response = sync_httpx_client.post( @@ -12189,6 +12194,7 @@ class BaseLLMHTTPHandler: headers=headers, data=request_data["ssml_body"], timeout=timeout, + stream=should_stream, ) else: raise ValueError( @@ -12222,6 +12228,7 @@ class BaseLLMHTTPHandler: timeout: float | httpx.Timeout, extra_headers: dict[str, object] | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, + stream: bool = False, ) -> "HttpxBinaryResponseContent": """ Async version of the text-to-speech handler. @@ -12275,6 +12282,8 @@ class BaseLLMHTTPHandler: }, ) + should_stream: Final = stream and text_to_speech_provider_config.supports_streaming + try: # Determine request body type and send appropriately if "dict_body" in request_data: @@ -12283,6 +12292,7 @@ class BaseLLMHTTPHandler: headers=headers, json=request_data["dict_body"], timeout=timeout, + stream=should_stream, ) elif "ssml_body" in request_data: response = await async_httpx_client.post( @@ -12290,6 +12300,7 @@ class BaseLLMHTTPHandler: headers=headers, data=request_data["ssml_body"], timeout=timeout, + stream=should_stream, ) else: raise ValueError( diff --git a/litellm/llms/elevenlabs/text_to_speech/transformation.py b/litellm/llms/elevenlabs/text_to_speech/transformation.py index eb93543df46..991bc002e16 100644 --- a/litellm/llms/elevenlabs/text_to_speech/transformation.py +++ b/litellm/llms/elevenlabs/text_to_speech/transformation.py @@ -37,6 +37,8 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): Reference: https://elevenlabs.io/docs/api-reference/text-to-speech/convert """ + supports_streaming = True + TTS_BASE_URL = "https://api.elevenlabs.io" TTS_ENDPOINT_PATH = "/v1/text-to-speech" DEFAULT_OUTPUT_FORMAT = "pcm_44100" diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 3a12326d433..4dcf9263a52 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -17,6 +17,7 @@ from openai._constants import RAW_RESPONSE_HEADER from openai._legacy_response import LegacyAPIResponse from openai._types import RequestOptions from openai.types import CreateEmbeddingResponse +from openai.types.audio.speech_create_params import SpeechCreateParams from openai.types.beta.assistant_deleted import AssistantDeleted from openai.types.file_deleted import FileDeleted from pydantic import BaseModel, TypeAdapter @@ -1618,6 +1619,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): aspeech: bool | None = None, client=None, shared_session: Optional["ClientSession"] = None, + stream: bool = False, ) -> HttpxBinaryResponseContent: if aspeech is not None and aspeech is True: return self.async_audio_speech( @@ -1634,6 +1636,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj=logging_obj, client=client, shared_session=shared_session, + stream=stream, ) openai_client: Final = self._get_openai_client( @@ -1656,14 +1659,16 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): }, ) - response: Final = sync_client.audio.speech.create( - model=model, - voice=voice, - input=input, - **optional_params, + params: Final = cast( # cast-ok: the provider-mapped body is this SDK call's own params shape + SpeechCreateParams, {"model": model, "voice": voice, "input": input, **optional_params} ) - speech_response: Final = HttpxBinaryResponseContent(response=response.response) - set_provider_response_headers_in_hidden_params(speech_response, response.response.headers) + raw_response: Final = ( + sync_client.audio.speech.with_streaming_response.create(**params).__enter__().http_response + if stream + else sync_client.audio.speech.create(**params).response + ) + speech_response: Final = HttpxBinaryResponseContent(response=raw_response) + set_provider_response_headers_in_hidden_params(speech_response, raw_response.headers) return speech_response async def async_audio_speech( @@ -1681,6 +1686,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj: LiteLLMLoggingObj, client=None, shared_session: Optional["ClientSession"] = None, + stream: bool = False, ) -> HttpxBinaryResponseContent: openai_client: Final = cast( AsyncOpenAI, @@ -1704,14 +1710,16 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): }, ) - response: Final = await openai_client.audio.speech.create( - model=model, - voice=voice, - input=input, - **optional_params, + params: Final = cast( # cast-ok: the provider-mapped body is this SDK call's own params shape + SpeechCreateParams, {"model": model, "voice": voice, "input": input, **optional_params} ) - speech_response: Final = HttpxBinaryResponseContent(response=response.response) - set_provider_response_headers_in_hidden_params(speech_response, response.response.headers) + raw_response: Final = ( + (await openai_client.audio.speech.with_streaming_response.create(**params).__aenter__()).http_response + if stream + else (await openai_client.audio.speech.create(**params)).response + ) + speech_response: Final = HttpxBinaryResponseContent(response=raw_response) + set_provider_response_headers_in_hidden_params(speech_response, raw_response.headers) return speech_response diff --git a/litellm/main.py b/litellm/main.py index 474d82e8ef1..a759e57ab42 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8273,6 +8273,7 @@ def speech( headers: dict | None = None, custom_llm_provider: str | None = None, aspeech: bool | None = None, + stream: bool = False, base_url: str | None = None, **kwargs, ) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]: @@ -8410,6 +8411,7 @@ def speech( client=client, # pass AsyncOpenAI, OpenAI client aspeech=aspeech, shared_session=shared_session, + stream=stream, ) elif custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS: # Check if this is Azure Speech Service (Cognitive Services TTS) @@ -8442,6 +8444,7 @@ def speech( aspeech=aspeech or False, api_base=api_base, api_key=api_key, + stream=stream, **kwargs, ) else: @@ -8489,6 +8492,7 @@ def speech( client=client, # pass AsyncOpenAI, OpenAI client aspeech=aspeech, litellm_params=litellm_params_dict, + stream=stream, ) elif custom_llm_provider == "elevenlabs": from litellm.llms.elevenlabs.text_to_speech.transformation import ( @@ -8533,6 +8537,7 @@ def speech( extra_headers=extra_headers, client=client, _is_async=aspeech or False, + stream=stream, ) elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "vertex_ai_beta": from litellm.llms.vertex_ai.text_to_speech.transformation import ( @@ -8731,6 +8736,7 @@ def speech( aspeech=aspeech or False, api_base=api_base, api_key=api_key, + stream=stream, **kwargs, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c2ca045236c..cac9487a489 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12834,8 +12834,11 @@ async def _audio_speech_chunk_generator( # too large: latency is low, but memory usage is high # 8192 is a good compromise _generator: Final = await _response.aiter_bytes(chunk_size=AUDIO_SPEECH_CHUNK_SIZE) - async for chunk in _generator: - yield chunk + try: + async for chunk in _generator: + yield chunk + finally: + await _response.aclose() @router.post( diff --git a/litellm/utils.py b/litellm/utils.py index 3f10d2a3ebd..d2e84b36e4b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2376,6 +2376,16 @@ _STREAMING_CALL_TYPES: Final = frozenset( ) +_SPEECH_CALL_TYPES: Final = frozenset( + { + CallTypes.speech, + CallTypes.aspeech, + CallTypes.speech.value, + CallTypes.aspeech.value, + } +) + + def _is_streaming_request( kwargs: dict[str, object], call_type: CallTypes | str, @@ -2385,7 +2395,13 @@ def _is_streaming_request( Returns True if: - if "stream=True" in kwargs (litellm chat completion, litellm text completion, litellm messages) - if call_type is generate_content_stream or agenerate_content_stream (litellm google genai) + + Speech is excluded: "stream=True" there asks for the audio bytes to be forwarded as they + arrive, the return value is still a single response object, so it is logged like any other + non-streaming call. """ + if call_type in _SPEECH_CALL_TYPES: + return False if "stream" in kwargs and kwargs["stream"] is True: return True return call_type in _STREAMING_CALL_TYPES diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 347840b5c5d..9ff458b0986 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -4788,6 +4788,104 @@ async def test_async_text_to_speech_handler_records_upstream_response_headers(): _assert_upstream_headers_recorded(response) +class _RecordingSyncAudioStream(httpx.SyncByteStream): + def __init__(self, chunks: list[bytes], produced: list[bytes]): + self._chunks = chunks + self._produced = produced + + def __iter__(self): + for chunk in self._chunks: + self._produced.append(chunk) + yield chunk + + +def _sync_audio_chunk_transport(produced: list[bytes]) -> httpx.MockTransport: + def _respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "audio/mpeg"}, + stream=_RecordingSyncAudioStream([b"one", b"two", b"three"], produced), + ) + + return httpx.MockTransport(_respond) + + +class _StreamingTextToSpeechConfig(_HeaderTextToSpeechConfig): + supports_streaming = True + + +class _RecordingAudioStream(httpx.AsyncByteStream): + def __init__(self, chunks: list[bytes], produced: list[bytes]): + self._chunks = chunks + self._produced = produced + + async def __aiter__(self): + for chunk in self._chunks: + self._produced.append(chunk) + yield chunk + + +def _audio_chunk_transport(produced: list[bytes]) -> httpx.MockTransport: + def _respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "audio/mpeg"}, + stream=_RecordingAudioStream([b"one", b"two", b"three"], produced), + ) + + return httpx.MockTransport(_respond) + + +@pytest.mark.asyncio +async def test_async_text_to_speech_handler_yields_chunks_before_upstream_finishes(): + produced: Final[list[bytes]] = [] + client: Final = AsyncHTTPHandler(transport=_audio_chunk_transport(produced)) + + response: Final = await BaseLLMHTTPHandler().async_text_to_speech_handler( + client=client, + stream=True, + **{ + **_text_to_speech_call_kwargs(), + "text_to_speech_provider_config": _StreamingTextToSpeechConfig(), + }, + ) + audio: Final = await response.aiter_bytes() + + assert await anext(audio) == b"one" + assert produced == [b"one"] + + +@pytest.mark.asyncio +async def test_async_text_to_speech_handler_buffers_when_provider_does_not_support_streaming(): + produced: Final[list[bytes]] = [] + client: Final = AsyncHTTPHandler(transport=_audio_chunk_transport(produced)) + + response: Final = await BaseLLMHTTPHandler().async_text_to_speech_handler( + client=client, stream=True, **_text_to_speech_call_kwargs() + ) + + assert produced == [b"one", b"two", b"three"] + assert response.content == b"onetwothree" + + +def test_text_to_speech_handler_yields_chunks_before_upstream_finishes(): + produced: Final[list[bytes]] = [] + client: Final = HTTPHandler() + client.client = httpx.Client(transport=_sync_audio_chunk_transport(produced)) + + response: Final = BaseLLMHTTPHandler().text_to_speech_handler( + client=client, + stream=True, + **{ + **_text_to_speech_call_kwargs(), + "text_to_speech_provider_config": _StreamingTextToSpeechConfig(), + }, + ) + + assert next(response.iter_bytes()) == b"one" + assert produced == [b"one"] + + async def _get_by_id_with_upstream(handler_name: str, upstream_response: httpx.Response) -> object: async_client: Final = AsyncHTTPHandler() await async_client.close() diff --git a/tests/unit/llms/openai/test_openai_audio_speech.py b/tests/unit/llms/openai/test_openai_audio_speech.py new file mode 100644 index 00000000000..27ae0ffd543 --- /dev/null +++ b/tests/unit/llms/openai/test_openai_audio_speech.py @@ -0,0 +1,105 @@ +"""Streaming behaviour of the OpenAI and Azure OpenAI ``/audio/speech`` paths.""" + +from __future__ import annotations + +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest +from openai import AsyncAzureOpenAI, AsyncOpenAI + +from litellm.llms.azure.azure import AzureChatCompletion +from litellm.llms.openai.openai import OpenAIChatCompletion + + +class _RecordingAudioStream(httpx.AsyncByteStream): + def __init__(self, chunks: list[bytes], produced: list[bytes]): + self._chunks = chunks + self._produced = produced + + async def __aiter__(self): + for chunk in self._chunks: + self._produced.append(chunk) + yield chunk + + +def _audio_client(produced: list[bytes]) -> httpx.AsyncClient: + def _respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "audio/mpeg"}, + stream=_RecordingAudioStream([b"one", b"two", b"three"], produced), + ) + + return httpx.AsyncClient(transport=httpx.MockTransport(_respond)) + + +async def _openai_speech(produced: list[bytes], *, stream: bool): + return await OpenAIChatCompletion().async_audio_speech( + model="tts-1", + input="hello", + voice="alloy", + optional_params={}, + api_key="fake-key", + api_base=None, + organization=None, + project=None, + max_retries=0, + timeout=10.0, + logging_obj=Mock(), + client=AsyncOpenAI(api_key="fake-key", http_client=_audio_client(produced)), + stream=stream, + ) + + +@pytest.mark.asyncio +async def test_openai_audio_speech_yields_chunks_before_upstream_finishes(): + produced: Final[list[bytes]] = [] + + response: Final = await _openai_speech(produced, stream=True) + audio: Final = await response.aiter_bytes() + + assert await anext(audio) == b"one" + assert produced == [b"one"] + + +@pytest.mark.asyncio +async def test_openai_audio_speech_buffers_without_stream(): + produced: Final[list[bytes]] = [] + + response: Final = await _openai_speech(produced, stream=False) + + assert produced == [b"one", b"two", b"three"] + assert response.content == b"onetwothree" + + +@pytest.mark.asyncio +async def test_azure_audio_speech_yields_chunks_before_upstream_finishes(): + produced: Final[list[bytes]] = [] + + response: Final = await AzureChatCompletion().async_audio_speech( + model="tts-1", + input="hello", + voice="alloy", + optional_params={}, + api_key="fake-key", + api_base="https://fake.openai.azure.com", + api_version="2024-02-01", + azure_ad_token=None, + azure_ad_token_provider=None, + max_retries=0, + timeout=10.0, + logging_obj=Mock(), + client=AsyncAzureOpenAI( + api_key="fake-key", + api_version="2024-02-01", + azure_endpoint="https://fake.openai.azure.com", + http_client=_audio_client(produced), + ), + stream=True, + ) + audio: Final = await response.aiter_bytes() + + assert await anext(audio) == b"one" + assert produced == [b"one"] diff --git a/tests/unit/proxy/proxy_server/test_routes_audio.py b/tests/unit/proxy/proxy_server/test_routes_audio.py index de76c7257cf..cb186537de7 100644 --- a/tests/unit/proxy/proxy_server/test_routes_audio.py +++ b/tests/unit/proxy/proxy_server/test_routes_audio.py @@ -15,6 +15,7 @@ from unittest.mock import AsyncMock, MagicMock import httpx import pytest +from litellm.constants import AUDIO_SPEECH_CHUNK_SIZE from litellm.proxy import proxy_server from litellm.types.llms.openai import HttpxBinaryResponseContent @@ -254,3 +255,52 @@ def test_audio_transcription_error(client, auth_as, patched_transcription_error, response = client.post(path, files=files, data=data) assert response.status_code == 500 assert len(response.content) > 0 + + +class _RecordingAudioStream(httpx.AsyncByteStream): + def __init__(self, chunks: list[bytes], produced: list[bytes]): + self._chunks = chunks + self._produced = produced + self.closed = False + + async def __aiter__(self): + for chunk in self._chunks: + self._produced.append(chunk) + yield chunk + + async def aclose(self) -> None: + self.closed = True + + +@pytest.mark.asyncio +async def test_audio_speech_chunk_generator_forwards_an_unread_upstream_stream(): + """A streaming upstream response reaches the client chunk by chunk, not buffered whole.""" + produced: list[bytes] = [] + first, second = b"a" * AUDIO_SPEECH_CHUNK_SIZE, b"b" * AUDIO_SPEECH_CHUNK_SIZE + stream = _RecordingAudioStream([first, second], produced) + response = HttpxBinaryResponseContent( + httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}, stream=stream) + ) + + chunks = proxy_server._audio_speech_chunk_generator(response) + + assert await anext(chunks) == first + assert produced == [first] + assert b"".join([chunk async for chunk in chunks]) == second + assert stream.closed is True + + +@pytest.mark.asyncio +async def test_audio_speech_chunk_generator_closes_a_stream_abandoned_midway(): + """An abandoned stream still releases its connection, a client disconnect cannot leak one.""" + first = b"a" * AUDIO_SPEECH_CHUNK_SIZE + stream = _RecordingAudioStream([first, b"b" * AUDIO_SPEECH_CHUNK_SIZE], []) + response = HttpxBinaryResponseContent( + httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}, stream=stream) + ) + + chunks = proxy_server._audio_speech_chunk_generator(response) + assert await anext(chunks) == first + await chunks.aclose() + + assert stream.closed is True diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 6ad4f75f97f..9b54c4783c1 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -3642,6 +3642,10 @@ class TestIsStreamingRequest: def test_non_streaming_call_type_enum(self): assert _is_streaming_request(kwargs={}, call_type=CallTypes.acompletion) is False + @pytest.mark.parametrize("call_type", [CallTypes.aspeech, CallTypes.aspeech.value, CallTypes.speech]) + def test_speech_with_stream_true_is_not_a_streaming_request(self, call_type): + assert _is_streaming_request(kwargs={"stream": True}, call_type=call_type) is False + def test_stream_true_overrides_non_streaming_call_type(self): assert _is_streaming_request(kwargs={"stream": True}, call_type=CallTypes.acompletion) is True From 539ea3abff194804d9b50ec0113e1db442825c36 Mon Sep 17 00:00:00 2001 From: Gabriel Creti Date: Tue, 13 Oct 2026 18:04:27 +0200 Subject: [PATCH 02/20] fix(tts): keep streamed speech inside its concurrency and cleanup guarantees Streaming the audio out moved two things that used to be covered by the buffered response, and the review on #46222 caught both. A deployment's max_parallel_requests counted syntheses that had started rather than syntheses in flight, because aspeech now returns at the response headers with the audio still arriving. The slot is held until the response is closed, the way the chat streaming path already does it. A client that walks away mid-clip leaves the body unread, and httpx only releases the connection when it is read to the end. The body iterator closes it on the way out, and the route arms the same close as a background task for a response whose body is never read at all. These are two different paths on purpose: Starlette raises ClientDisconnect out of the response before background tasks run, so neither one covers both cases. Also pulls the speech body cast into one helper beside speech_request_body, which drops a banned typing.cast import. --- litellm/litellm_core_utils/logging_utils.py | 12 ++ litellm/llms/azure/azure.py | 17 ++- litellm/llms/openai/openai.py | 15 +- litellm/proxy/proxy_server.py | 16 ++- litellm/router.py | 32 ++++- .../proxy/proxy_server/test_routes_audio.py | 136 ++++++++++++++++-- .../test_router_max_parallel_requests.py | 61 ++++++++ 7 files changed, 261 insertions(+), 28 deletions(-) diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 5ac6b5435dc..66ed0706a8b 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -20,6 +20,7 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + from openai.types.audio.speech_create_params import SpeechCreateParams from opentelemetry.trace import Span as _Span from litellm import ModelResponse as _ModelResponse @@ -313,6 +314,17 @@ def _set_duration_in_model_call_details( verbose_logger.warning("Error setting `llm_api_duration_ms`: %s", e) +def speech_create_params( + model: str, voice: str, input: str, optional_params: Mapping[str, object] +) -> "SpeechCreateParams": + """The provider-mapped speech body, in the shape the OpenAI and Azure SDKs take.""" + from openai.types.audio.speech_create_params import SpeechCreateParams + + return cast( # cast-ok: map_openai_params already mapped these to this SDK call's params + SpeechCreateParams, {"model": model, "voice": voice, "input": input, **optional_params} + ) + + def speech_request_body(model: str, voice: str, optional_params: Mapping[str, object]) -> Mapping[str, object]: """Speech request body for telemetry, without the caller headers the provider SDKs take as request kwargs rather than body fields.""" diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 2090d4d76de..0bc0730e581 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -2,7 +2,7 @@ import asyncio import json import time from collections.abc import Callable, Coroutine -from typing import Final, cast +from typing import Final import httpx from openai import ( @@ -13,12 +13,15 @@ from openai import ( BadRequestError, OpenAI, ) -from openai.types.audio.speech_create_params import SpeechCreateParams import litellm from litellm.constants import AZURE_OPERATION_POLLING_TIMEOUT, DEFAULT_MAX_RETRIES from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing +from litellm.litellm_core_utils.logging_utils import ( + speech_create_params, + speech_request_body, + track_llm_api_timing, +) from litellm.litellm_core_utils.url_utils import SSRFError, assert_same_origin from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -1423,9 +1426,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): }, ) - params: Final = cast( # cast-ok: the provider-mapped body is this SDK call's own params shape - SpeechCreateParams, {"model": model, "voice": voice, "input": input, **optional_params} - ) + params: Final = speech_create_params(model, voice, input, optional_params) raw_response: Final = ( azure_client.audio.speech.with_streaming_response.create(**params).__enter__().http_response if stream @@ -1470,9 +1471,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): }, ) - params: Final = cast( # cast-ok: the provider-mapped body is this SDK call's own params shape - SpeechCreateParams, {"model": model, "voice": voice, "input": input, **optional_params} - ) + params: Final = speech_create_params(model, voice, input, optional_params) raw_response: Final = ( (await azure_client.audio.speech.with_streaming_response.create(**params).__aenter__()).http_response if stream diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 4dcf9263a52..f7becfff592 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -17,7 +17,6 @@ from openai._constants import RAW_RESPONSE_HEADER from openai._legacy_response import LegacyAPIResponse from openai._types import RequestOptions from openai.types import CreateEmbeddingResponse -from openai.types.audio.speech_create_params import SpeechCreateParams from openai.types.beta.assistant_deleted import AssistantDeleted from openai.types.file_deleted import FileDeleted from pydantic import BaseModel, TypeAdapter @@ -30,7 +29,11 @@ from litellm.constants import DEFAULT_MAX_RETRIES from litellm.files.types import FileContentStreamingResult from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.litellm_core_utils.logging_utils import speech_request_body, track_llm_api_timing +from litellm.litellm_core_utils.logging_utils import ( + speech_create_params, + speech_request_body, + track_llm_api_timing, +) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException from litellm.llms.bedrock.chat.invoke_handler import MockResponseIterator @@ -1659,9 +1662,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): }, ) - params: Final = cast( # cast-ok: the provider-mapped body is this SDK call's own params shape - SpeechCreateParams, {"model": model, "voice": voice, "input": input, **optional_params} - ) + params: Final = speech_create_params(model, voice, input, optional_params) raw_response: Final = ( sync_client.audio.speech.with_streaming_response.create(**params).__enter__().http_response if stream @@ -1710,9 +1711,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): }, ) - params: Final = cast( # cast-ok: the provider-mapped body is this SDK call's own params shape - SpeechCreateParams, {"model": model, "voice": voice, "input": input, **optional_params} - ) + params: Final = speech_create_params(model, voice, input, optional_params) raw_response: Final = ( (await openai_client.audio.speech.with_streaming_response.create(**params).__aenter__()).http_response if stream diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cac9487a489..34ecab1be70 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -765,6 +765,7 @@ except ImportError: build_billing_metrics_recorder = None shutdown_billing_metrics_recorder = None from fastapi.exception_handlers import http_exception_handler +from starlette.background import BackgroundTask from starlette.exceptions import HTTPException as StarletteHTTPException from starlette.websockets import WebSocketState @@ -12838,6 +12839,15 @@ async def _audio_speech_chunk_generator( async for chunk in _generator: yield chunk finally: + # Reading the body to the end closes the upstream, but a client that walks away + # mid-clip never gets there. Starlette raises ClientDisconnect out of the response + # before its background task runs, so this is the path that covers a disconnect. + await _close_audio_speech_response(_response) + + +async def _close_audio_speech_response(_response: HttpxBinaryResponseContent) -> None: + """Release the upstream connection holding this response's audio.""" + with contextlib.suppress(Exception): await _response.aclose() @@ -12943,10 +12953,14 @@ async def audio_speech( 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(response), + _audio_speech_chunk_generator(audio), media_type=media_type, headers=custom_headers, + background=BackgroundTask(_close_audio_speech_response, audio), ) except Exception as e: diff --git a/litellm/router.py b/litellm/router.py index cdede46ef5c..7d323f4fe32 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -505,6 +505,29 @@ _FALLBACK_HOP_ADAPTER: Final = TypeAdapter(Mapping[str, object]) _DEPLOYMENT_SELECTED_EVENT: Final = "litellm.request.deployment_selected" +async def _hold_slot_until_closed(response: object, slot: contextlib.AsyncExitStack) -> None: + """Keep a deployment slot held until a streamed binary response is closed. + + A buffered response already has its body in hand, so its slot is released right + away. A streamed one is still pulling audio from the provider, so the slot has to + survive until whoever consumes it closes the response. + """ + raw: Final = getattr(response, "response", None) + if not isinstance(raw, httpx.Response) or raw.is_stream_consumed: + await slot.aclose() + return + + close: Final = raw.aclose + + async def _aclose_and_release() -> None: + try: + await close() + finally: + await slot.aclose() + + raw.aclose = _aclose_and_release # pyright: ignore[reportAttributeAccessIssue] # wrapping the bound method keeps every caller's close path releasing the slot + + def _is_fallback_hop(request_kwargs: Mapping[str, object]) -> bool: fallback_depth: Final = request_kwargs.get("fallback_depth") return isinstance(fallback_depth, int) and fallback_depth > 0 @@ -5024,7 +5047,10 @@ class Router: ) self.total_calls[model_name] += 1 - async with self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span): + async with contextlib.AsyncExitStack() as deployment_slot: + await deployment_slot.enter_async_context( + self._deployment_slot(deployment=deployment, kwargs=kwargs, parent_otel_span=parent_otel_span) + ) response = await litellm.aspeech( **{ **data, @@ -5034,6 +5060,10 @@ class Router: **kwargs, } ) + # A streamed response returns at the headers with its audio still arriving, so + # the slot has to outlive this block or the deployment's max_parallel_requests + # would only bound how many syntheses can *start*. + await _hold_slot_until_closed(response, deployment_slot.pop_all()) self.success_calls[model_name] += 1 verbose_router_logger.info("litellm.aspeech(model=%s)\x1b[32m 200 OK\x1b[0m", model_name) diff --git a/tests/unit/proxy/proxy_server/test_routes_audio.py b/tests/unit/proxy/proxy_server/test_routes_audio.py index cb186537de7..052556d9ad9 100644 --- a/tests/unit/proxy/proxy_server/test_routes_audio.py +++ b/tests/unit/proxy/proxy_server/test_routes_audio.py @@ -9,14 +9,26 @@ Pins (PR2): from __future__ import annotations +import asyncio +import contextlib +import gc import io +from typing import Final + +import orjson from unittest.mock import AsyncMock, MagicMock import httpx import pytest +from starlette.background import BackgroundTask +from starlette.responses import StreamingResponse + from litellm.constants import AUDIO_SPEECH_CHUNK_SIZE +from fastapi import Request, Response + from litellm.proxy import proxy_server +from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.openai import HttpxBinaryResponseContent @@ -290,17 +302,123 @@ async def test_audio_speech_chunk_generator_forwards_an_unread_upstream_stream() assert stream.closed is True -@pytest.mark.asyncio -async def test_audio_speech_chunk_generator_closes_a_stream_abandoned_midway(): - """An abandoned stream still releases its connection, a client disconnect cannot leak one.""" - first = b"a" * AUDIO_SPEECH_CHUNK_SIZE - stream = _RecordingAudioStream([first, b"b" * AUDIO_SPEECH_CHUNK_SIZE], []) - response = HttpxBinaryResponseContent( +def _speech_request(payload: dict) -> Request: + body: Final = orjson.dumps(payload) + + async def _receive(): + return {"type": "http.request", "body": body, "more_body": False} + + return Request( + { + "type": "http", + "method": "POST", + "path": "/v1/audio/speech", + "headers": [(b"content-type", b"application/json")], + "query_string": b"", + }, + receive=_receive, + ) + + +def _audio_response(stream: _RecordingAudioStream) -> HttpxBinaryResponseContent: + return HttpxBinaryResponseContent( httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}, stream=stream) ) - chunks = proxy_server._audio_speech_chunk_generator(response) - assert await anext(chunks) == first - await chunks.aclose() + +async def _serve_until_disconnect(response: StreamingResponse) -> None: + """Drive the response through one ASGI cycle where the client drops at the first chunk.""" + + async def _send(message): + if message["type"] == "http.response.body": + raise OSError("client gone") + + async def _receive(): + return {"type": "http.disconnect"} + + with contextlib.suppress(Exception): + await response({"type": "http", "asgi": {"spec_version": "2.4"}}, _receive, _send) + + +@pytest.fixture +def patched_streaming_speech(monkeypatch): + """Route /v1/audio/speech to an upstream whose audio body is left unread.""" + stream = _RecordingAudioStream([b"a" * AUDIO_SPEECH_CHUNK_SIZE], []) + monkeypatch.setattr(proxy_server, "llm_router", MagicMock()) + monkeypatch.setattr( + proxy_server, + "proxy_logging_obj", + MagicMock( + pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]), + post_call_failure_hook=AsyncMock(), + post_call_response_headers_hook=AsyncMock(return_value={}), + update_request_status=AsyncMock(), + ), + ) + + async def _add_data(data, **kwargs): + return data + + monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data) + + async def _llm_call(): + return _audio_response(stream) + + async def _fake_route_request(*args, **kwargs): + return _llm_call() + + monkeypatch.setattr(proxy_server, "route_request", _fake_route_request) + yield stream + + +@pytest.mark.asyncio +async def test_audio_speech_releases_its_upstream_when_the_client_walks_away(): + """A client that drops mid-clip still has its upstream connection released. Starlette + raises ClientDisconnect out of the response before any background task runs, so the + body iterator's own cleanup is what has to cover this.""" + stream = _RecordingAudioStream([b"a" * AUDIO_SPEECH_CHUNK_SIZE] * 3, []) + response = _audio_response(stream) + + await _serve_until_disconnect( + StreamingResponse( + proxy_server._audio_speech_chunk_generator(response), + media_type="audio/mpeg", + background=BackgroundTask(proxy_server._close_audio_speech_response, response), + ) + ) + # The abandoned iterator is suspended mid-yield; it runs its cleanup once dropped. + gc.collect() + await asyncio.sleep(0.1) + + assert stream.closed is True + + +@pytest.mark.asyncio +async def test_audio_speech_route_releases_an_upstream_whose_body_is_never_read( + monkeypatch, patched_streaming_speech +): + """A response the server never iterates leaves the body iterator's cleanup unrun and + out of reach of GC, so the route itself has to arm the release.""" + stream = patched_streaming_speech + captured: list[StreamingResponse] = [] + + original = proxy_server.StreamingResponse + + def _capture(*args, **kwargs): + response = original(*args, **kwargs) + captured.append(response) + return response + + monkeypatch.setattr(proxy_server, "StreamingResponse", _capture) + + request = _speech_request({"model": "tts-1", "input": "Hi", "voice": "alloy", "stream": True}) + await proxy_server.audio_speech( + request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test") + ) + + assert len(captured) == 1 + assert stream.closed is False + + await captured[0].background() assert stream.closed is True diff --git a/tests/unit/test_router/test_router_max_parallel_requests.py b/tests/unit/test_router/test_router_max_parallel_requests.py index fce273cd1ab..0f3ecd134b2 100644 --- a/tests/unit/test_router/test_router_max_parallel_requests.py +++ b/tests/unit/test_router/test_router_max_parallel_requests.py @@ -1,7 +1,10 @@ from __future__ import annotations +import asyncio + from typing import Final +import httpx import litellm import pytest from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit @@ -139,3 +142,61 @@ async def test_max_parallel_requests_tpm_rate_limiting_base_case() -> None: with pytest.raises(litellm.RateLimitError): await router.acompletion(**request) + + +class _OneChunkAudioStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield b"audio" + + async def aclose(self) -> None: + return None + + +def _speech_router() -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "tts", + "litellm_params": { + "model": "elevenlabs/eleven_multilingual_v2", + "api_key": "fake-key", + "max_parallel_requests": 1, + }, + } + ] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [True, False], ids=["streamed", "buffered"]) +async def test_aspeech_holds_its_deployment_slot_until_the_audio_is_released( + monkeypatch: pytest.MonkeyPatch, stream: bool +) -> None: + """A streamed synthesis occupies its deployment's only slot for as long as audio is + still arriving, so a second caller is turned away rather than doubling up on it.""" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _speech_router() + + def _client() -> litellm.AsyncHTTPHandler: + return litellm.AsyncHTTPHandler( + transport=httpx.MockTransport( + lambda _: httpx.Response(200, headers={"content-type": "audio/mpeg"}, stream=_OneChunkAudioStream()) + ) + ) + + response: Final = await router.aspeech( + model="tts", input="hello", voice="voice-id", client=_client(), stream=stream + ) + + async def _second_caller() -> str: + try: + await router.aspeech(model="tts", input="hello", voice="voice-id", client=_client(), stream=stream) + except litellm.RateLimitError: + return "refused" + return "admitted" + + while_audio_open: Final = await _second_caller() + await response.aclose() + await asyncio.sleep(0) + + assert (while_audio_open, await _second_caller()) == ("refused" if stream else "admitted", "admitted") From 318554085601a1d55986c1b452b6c57c9ec9b711 Mon Sep 17 00:00:00 2001 From: Gabriel Creti Date: Wed, 14 Oct 2026 10:01:23 +0200 Subject: [PATCH 03/20] test(tts): cover the slot-holding helper directly router_code_coverage.py wants every function in router.py called by name from a unit test, and _hold_slot_until_closed was only reached through router.aspeech(). Test it on its own instead of adding it to the ignore list: which responses keep their slot is the part worth pinning, so the cases are buffered body, non-binary response and unread body. --- .../test_router_max_parallel_requests.py | 38 +++++++++++++++++++ 1 file changed, 38 insertions(+) diff --git a/tests/unit/test_router/test_router_max_parallel_requests.py b/tests/unit/test_router/test_router_max_parallel_requests.py index 0f3ecd134b2..866a0843592 100644 --- a/tests/unit/test_router/test_router_max_parallel_requests.py +++ b/tests/unit/test_router/test_router_max_parallel_requests.py @@ -1,13 +1,16 @@ from __future__ import annotations import asyncio +import contextlib from typing import Final import httpx import litellm import pytest +from litellm.router import _hold_slot_until_closed from litellm.router_utils.client_initalization_utils import MaxParallelRequestsLimit +from litellm.types.llms.openai import HttpxBinaryResponseContent from litellm.utils import calculate_max_parallel_requests @@ -200,3 +203,38 @@ async def test_aspeech_holds_its_deployment_slot_until_the_audio_is_released( await asyncio.sleep(0) assert (while_audio_open, await _second_caller()) == ("refused" if stream else "admitted", "admitted") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("response", "releases_immediately"), + [ + pytest.param(HttpxBinaryResponseContent(httpx.Response(200, content=b"audio")), True, id="buffered-body"), + pytest.param("not-a-binary-response", True, id="non-binary-response"), + pytest.param( + HttpxBinaryResponseContent(httpx.Response(200, stream=_OneChunkAudioStream())), False, id="unread-body" + ), + ], +) +async def test_hold_slot_until_closed_keys_off_whether_the_body_is_still_arriving( + response: object, releases_immediately: bool +) -> None: + """Only a response whose body has yet to be read keeps its slot, everything else frees it.""" + released: Final[list[str]] = [] + + @contextlib.asynccontextmanager + async def _slot(): + try: + yield + finally: + released.append("released") + + stack: Final = contextlib.AsyncExitStack() + await stack.enter_async_context(_slot()) + + await _hold_slot_until_closed(response, stack) + + assert bool(released) is releases_immediately + if not releases_immediately: + await response.aclose() + assert released == ["released"] From 1dc337b68099a0e905a253f3bda41f54e2e74011 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 11:44:34 +0000 Subject: [PATCH 04/20] fix(tts): stream ElevenLabs speech from its /stream endpoint ElevenLabs' convert endpoint synthesizes the whole clip before it sends the first byte, so stream: true still waited for full synthesis. Requests with stream: true now call /v1/text-to-speech/{voice_id}/stream. --- .../text_to_speech/transformation.py | 5 ++- litellm/main.py | 2 ++ ...levenlabs_text_to_speech_transformation.py | 31 +++++++++++++++++++ 3 files changed, 37 insertions(+), 1 deletion(-) diff --git a/litellm/llms/elevenlabs/text_to_speech/transformation.py b/litellm/llms/elevenlabs/text_to_speech/transformation.py index 991bc002e16..001415d8366 100644 --- a/litellm/llms/elevenlabs/text_to_speech/transformation.py +++ b/litellm/llms/elevenlabs/text_to_speech/transformation.py @@ -41,6 +41,7 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): TTS_BASE_URL = "https://api.elevenlabs.io" TTS_ENDPOINT_PATH = "/v1/text-to-speech" + TTS_STREAM_PATH_SUFFIX = "/stream" DEFAULT_OUTPUT_FORMAT = "pcm_44100" VOICE_MAPPINGS = { "alloy": "21m00Tcm4TlvDq8ikWAM", # Rachel @@ -65,6 +66,7 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): ELEVENLABS_QUERY_PARAMS_KEY = "__elevenlabs_query_params__" ELEVENLABS_VOICE_ID_KEY = "__elevenlabs_voice_id__" + ELEVENLABS_STREAM_KEY = "__elevenlabs_stream__" def get_supported_openai_params(self, model: str) -> list: """ @@ -307,7 +309,8 @@ class ElevenLabsTextToSpeechConfig(BaseTextToSpeechConfig): raise ValueError("ElevenLabs voice_id is required. Pass `voice` when calling `litellm.speech()`.") encoded_voice_id: Final = encode_url_path_segment(voice_id, field_name="voice_id") - url = f"{base_url}{self.TTS_ENDPOINT_PATH}/{encoded_voice_id}" + stream_suffix: Final = self.TTS_STREAM_PATH_SUFFIX if self.ELEVENLABS_STREAM_KEY in litellm_params else "" + url = f"{base_url}{self.TTS_ENDPOINT_PATH}/{encoded_voice_id}{stream_suffix}" query_params: Final = litellm_params.get(self.ELEVENLABS_QUERY_PARAMS_KEY, {}) if query_params: diff --git a/litellm/main.py b/litellm/main.py index a759e57ab42..3bcf6c71903 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8518,6 +8518,8 @@ def speech( litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_QUERY_PARAMS_KEY] = query_params litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY] = voice_id + if stream: + litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_STREAM_KEY] = True if api_base is not None: litellm_params_dict["api_base"] = api_base diff --git a/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py b/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py index d05371d7df9..775b47657be 100644 --- a/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py +++ b/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py @@ -55,3 +55,34 @@ def test_speech_keeps_an_internal_prefixed_kwarg_out_of_the_elevenlabs_request(r sent: Final = json.loads(respx_mock.calls[0].request.content) assert "_litellm_undeclared_sentinel" not in sent, sent assert sent["text"] == "hi" + + +@pytest.mark.parametrize( + ("stream", "expected_path"), + [ + (True, "/v1/text-to-speech/21m00Tcm4TlvDq8ikWAM/stream"), + (False, "/v1/text-to-speech/21m00Tcm4TlvDq8ikWAM"), + ], +) +def test_speech_calls_the_elevenlabs_streaming_endpoint_only_when_stream_is_requested( + respx_mock: respx.MockRouter, stream: bool, expected_path: str +) -> None: + api_base: Final = "http://localhost:12346" + respx_mock.post(url__regex=rf"{api_base}/v1/text-to-speech/.*").mock( + return_value=httpx.Response(status_code=200, content=b"audio", headers={"content-type": "audio/mpeg"}) + ) + + response: Final = litellm.speech( + model="elevenlabs/eleven_flash_v2_5", + input="hi", + voice="21m00Tcm4TlvDq8ikWAM", + api_base=api_base, + api_key="fake_elevenlabs_api_key", + response_format="mp3", + stream=stream, + ) + + assert response.read() == b"audio" + sent_url: Final = respx_mock.calls[0].request.url + assert sent_url.path == expected_path + assert sent_url.params["output_format"] == "mp3_44100_128" From 6f3810bb1b7cb3943a5e16dd3175e1af629897a8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:03:34 +0000 Subject: [PATCH 05/20] fix(tts): log streamed speech like buffered speech A deployment-hook sync copied stream=True onto the logging object for speech, so the success handler waited for an assembled chat stream that never comes and streamed speech never reached spend logs or callbacks. --- litellm/utils.py | 5 ++--- tests/unit/test_utils.py | 41 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index d2e84b36e4b..3b94ba0c953 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2030,9 +2030,8 @@ def client(original_function): kwargs = modified_kwargs # Sync logging_obj.stream after deployment hooks (they may convert it). - _hook_stream: Final = kwargs.get("stream") - if _hook_stream is not None and logging_obj.stream != _hook_stream: - logging_obj.stream = _hook_stream + if kwargs.get("stream") is not None: + logging_obj.stream = _is_streaming_request(kwargs=kwargs, call_type=call_type) kwargs["litellm_logging_obj"] = logging_obj ## LOAD CREDENTIALS diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 9b54c4783c1..fe2d0ce332a 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -3650,6 +3650,47 @@ class TestIsStreamingRequest: assert _is_streaming_request(kwargs={"stream": True}, call_type=CallTypes.acompletion) is True +class _SpeechSuccessLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged_payloads: Final[asyncio.Queue[dict[str, object]]] = asyncio.Queue() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + await self.logged_payloads.put(kwargs["standard_logging_object"]) + + +@pytest.mark.usefixtures("local_model_cost_map") +async def test_aspeech_logs_streamed_audio_like_buffered_audio( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + api_base: Final = "http://localhost:12348" + respx_mock.post(url__regex=rf"{api_base}/v1/text-to-speech/.*").mock( + return_value=httpx.Response(status_code=200, content=b"audio", headers={"content-type": "audio/mpeg"}) + ) + speech_logger: Final = _SpeechSuccessLogger() + monkeypatch.setattr(litellm, "callbacks", [speech_logger]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + + async def logged_speech(stream: bool) -> dict[str, object]: + response: Final = await litellm.aspeech( + model="elevenlabs/eleven_multilingual_v2", + input="hello there", + voice="21m00Tcm4TlvDq8ikWAM", + api_base=api_base, + api_key="fake_elevenlabs_api_key", + stream=stream, + ) + assert await response.aread() == b"audio" + return await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=5) + + buffered: Final = await logged_speech(stream=False) + streamed: Final = await logged_speech(stream=True) + + assert streamed["status"] == "success" + assert streamed["call_type"] == buffered["call_type"] + assert streamed["response_cost"] == buffered["response_cost"] + + class TestCallbackAsyncSyncSeparation: """Test that LoggingCallbackManager auto-routes async callbacks to async lists.""" From 1ca97abce382d5951f7cacf3ffff58cc3721fbce Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 12:56:24 +0000 Subject: [PATCH 06/20] fix(tts): forward streamed speech audio as it arrives --- litellm/proxy/proxy_server.py | 9 ++++----- tests/unit/proxy/proxy_server/test_routes_audio.py | 5 +++-- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 34ecab1be70..6b1274f3d62 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12830,11 +12830,10 @@ async def moderations( async def _audio_speech_chunk_generator( _response: HttpxBinaryResponseContent, ) -> AsyncGenerator[bytes, None]: - # chunk_size has a big impact on latency, it can't be too small or too large - # too small: latency is high - # too large: latency is low, but memory usage is high - # 8192 is a good compromise - _generator: Final = await _response.aiter_bytes(chunk_size=AUDIO_SPEECH_CHUNK_SIZE) + upstream_still_streaming: Final = not _response.response.is_stream_consumed + _generator: Final = await _response.aiter_bytes( + chunk_size=None if upstream_still_streaming else AUDIO_SPEECH_CHUNK_SIZE + ) try: async for chunk in _generator: yield chunk diff --git a/tests/unit/proxy/proxy_server/test_routes_audio.py b/tests/unit/proxy/proxy_server/test_routes_audio.py index 052556d9ad9..a1bac615e38 100644 --- a/tests/unit/proxy/proxy_server/test_routes_audio.py +++ b/tests/unit/proxy/proxy_server/test_routes_audio.py @@ -286,9 +286,10 @@ class _RecordingAudioStream(httpx.AsyncByteStream): @pytest.mark.asyncio async def test_audio_speech_chunk_generator_forwards_an_unread_upstream_stream(): - """A streaming upstream response reaches the client chunk by chunk, not buffered whole.""" + """A streaming upstream response reaches the client piece by piece as it arrives, with no + wait for a full AUDIO_SPEECH_CHUNK_SIZE before the first audio goes out.""" produced: list[bytes] = [] - first, second = b"a" * AUDIO_SPEECH_CHUNK_SIZE, b"b" * AUDIO_SPEECH_CHUNK_SIZE + first, second = b"a" * (AUDIO_SPEECH_CHUNK_SIZE // 8), b"b" * AUDIO_SPEECH_CHUNK_SIZE stream = _RecordingAudioStream([first, second], produced) response = HttpxBinaryResponseContent( httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}, stream=stream) From 020af3cf0973d51c125f3dbf760f6a5fcc832adb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 14:24:03 +0000 Subject: [PATCH 07/20] fix(tts): log streamed speech when the audio ends Streamed speech now logs success when the caller reads the audio to the end or closes it, and logs a failure when the upstream breaks mid-clip. Key max_parallel_requests slots stay held for the whole stream, the logged duration and bytes match what was sent, and the router slot is no longer pinned by a pending logging task. --- .../litellm_core_utils/on_end_byte_stream.py | 59 +++++++ litellm/router.py | 4 + litellm/types/llms/openai.py | 3 + litellm/utils.py | 105 +++++++++-- tests/unit/test_utils.py | 163 ++++++++++++++++-- 5 files changed, 307 insertions(+), 27 deletions(-) create mode 100644 litellm/litellm_core_utils/on_end_byte_stream.py diff --git a/litellm/litellm_core_utils/on_end_byte_stream.py b/litellm/litellm_core_utils/on_end_byte_stream.py new file mode 100644 index 00000000000..11a7df99374 --- /dev/null +++ b/litellm/litellm_core_utils/on_end_byte_stream.py @@ -0,0 +1,59 @@ +from collections.abc import AsyncIterator, Awaitable, Callable, Iterator +from typing import Final + +import httpx + + +class OnEndAsyncByteStream(httpx.AsyncByteStream): + def __init__(self, stream: httpx.AsyncByteStream, on_end: Callable[[Exception | None], Awaitable[None]]) -> None: + self._stream: Final = stream + self._on_end: Final = on_end + self._ended = False + + async def __aiter__(self) -> AsyncIterator[bytes]: + try: + async for chunk in self._stream: + yield chunk + except Exception as error: + await self._end(error) + raise + await self._end(None) + + async def aclose(self) -> None: + try: + await self._stream.aclose() + finally: + await self._end(None) + + async def _end(self, error: Exception | None) -> None: + if self._ended: + return + self._ended = True + await self._on_end(error) + + +class OnEndSyncByteStream(httpx.SyncByteStream): + def __init__(self, stream: httpx.SyncByteStream, on_end: Callable[[Exception | None], None]) -> None: + self._stream: Final = stream + self._on_end: Final = on_end + self._ended = False + + def __iter__(self) -> Iterator[bytes]: + try: + yield from self._stream + except Exception as error: + self._end(error) + raise + self._end(None) + + def close(self) -> None: + try: + self._stream.close() + finally: + self._end(None) + + def _end(self, error: Exception | None) -> None: + if self._ended: + return + self._ended = True + self._on_end(error) diff --git a/litellm/router.py b/litellm/router.py index 7d323f4fe32..815a6cc7dc5 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -4977,6 +4977,10 @@ class Router: async def aspeech(self, model: str, input: str, voice: str | None = None, **kwargs): """ + With stream=True the response comes back at the provider's headers and keeps its + deployment slot until it is read to the end or closed, so close a response you stop + reading early. + Example Usage: ``` diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 26f65fee4fe..d582c3a6aa7 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -158,6 +158,9 @@ class HttpxBinaryResponseContent(_HttpxBinaryResponseContent, Generic[_ResponseT return self.response.num_bytes_downloaded return len(content) + def unread_response(self) -> httpx.Response | None: + return None if self.response.is_stream_consumed else self.response + def set_response_cost(self, response_cost: float | None) -> None: if response_cost is None: self._hidden_params.pop("response_cost", None) diff --git a/litellm/utils.py b/litellm/utils.py index 3b94ba0c953..4ebf2e457ee 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -90,6 +90,7 @@ from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) +from litellm.litellm_core_utils.on_end_byte_stream import OnEndAsyncByteStream, OnEndSyncByteStream from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens from litellm.rust_bridge import tokenizer as tokenizer_dispatch @@ -221,6 +222,7 @@ from litellm.types.llms.openai import ( ChatCompletionNamedToolChoiceParam, ChatCompletionToolParam, ChatCompletionToolParamFunctionChunk, + HttpxBinaryResponseContent, OpenAITextCompletionUserMessage, OpenAIWebSearchOptions, ) @@ -290,7 +292,7 @@ except (ImportError, AttributeError, TypeError): # Convert to str (if necessary) claude_json_str = json.dumps(json_data) import importlib.metadata -from collections.abc import AsyncIterator, Callable, Collection, Iterable, Iterator, Mapping, Sequence +from collections.abc import AsyncIterator, Awaitable, Callable, Collection, Iterable, Iterator, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeVar, cast, runtime_checkable from typing_extensions import assert_never @@ -1370,6 +1372,62 @@ def _dispatch_success_logging( ) +def _unread_speech_upstream(result: object, call_type: str) -> httpx.Response | None: + if call_type not in _SPEECH_CALL_TYPES or not isinstance(result, HttpxBinaryResponseContent): + return None + return result.unread_response() + + +def _async_speech_end_logger( + logging_obj: LiteLLMLoggingObject, + result: object, + start_time: datetime.datetime, + is_completion_with_fallbacks: bool, + is_litellm_internal_call: bool, +) -> Callable[[Exception | None], Awaitable[None]]: + async def log_speech_end(error: Exception | None) -> None: + end_time: Final = datetime.datetime.now() # noqa: DTZ005 # matches the naive start_time the wrapper took + if error is None: + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + is_litellm_internal_call=is_litellm_internal_call, + ) + return + if is_litellm_internal_call: + return + traceback_exception: Final = "".join(traceback.format_exception(error)) + unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) + logging_obj.failure_handler(error, traceback_exception, start_time, end_time) + await logging_obj.async_failure_handler(error, traceback_exception, start_time, end_time) + + return log_speech_end + + +def _sync_speech_end_logger( + logging_obj: LiteLLMLoggingObject, + result: object, + start_time: datetime.datetime, + is_litellm_internal_call: bool, + executor: BoundedLoggingThreadPoolExecutor, +) -> Callable[[Exception | None], None]: + ctx: Final = contextvars.copy_context() + + def log_speech_end(error: Exception | None) -> None: + end_time: Final = datetime.datetime.now() # noqa: DTZ005 # matches the naive start_time the wrapper took + if error is not None: + logging_obj.failure_handler(error, "".join(traceback.format_exception(error)), start_time, end_time) + return + if is_litellm_internal_call: + return + executor.submit(ctx.run, logging_obj.success_handler, result, start_time, end_time) + + return log_speech_end + + def _schedule_async_success_logging( logging_obj: LiteLLMLoggingObject, result: object, @@ -1906,11 +1964,23 @@ def client(original_function): # LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated verbose_logger.info("Wrapper: Completed Call, calling success_handler") - if not is_internal_call.get(): + executor: Final[BoundedLoggingThreadPoolExecutor] = getattr(sys.modules[__name__], "executor") + unread_speech: Final = _unread_speech_upstream(result, call_type) + if unread_speech is not None and isinstance(unread_speech.stream, httpx.SyncByteStream): + unread_speech.stream = OnEndSyncByteStream( + unread_speech.stream, + _sync_speech_end_logger( + logging_obj=logging_obj, + result=result, + start_time=start_time, + is_litellm_internal_call=is_internal_call.get(), + executor=executor, + ), + ) + elif not is_internal_call.get(): # Copy the current context to propagate it to the background thread # This is essential for OpenTelemetry span context propagation ctx: Final = contextvars.copy_context() - executor: Final[BoundedLoggingThreadPoolExecutor] = getattr(sys.modules[__name__], "executor") executor.submit( ctx.run, logging_obj.success_handler, @@ -2216,14 +2286,27 @@ def client(original_function): start_time=start_time, end_time=end_time, ) - _dispatch_success_logging( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - is_litellm_internal_call=_is_litellm_internal_call, - ) + unread_speech: Final = _unread_speech_upstream(result, call_type) + if unread_speech is not None and isinstance(unread_speech.stream, httpx.AsyncByteStream): + unread_speech.stream = OnEndAsyncByteStream( + unread_speech.stream, + _async_speech_end_logger( + logging_obj=logging_obj, + result=result, + start_time=start_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + is_litellm_internal_call=_is_litellm_internal_call, + ), + ) + else: + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + is_litellm_internal_call=_is_litellm_internal_call, + ) return result except Exception as e: diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index fe2d0ce332a..becdd4fa4fc 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -10,10 +10,12 @@ import logging import os import queue import threading -from collections.abc import Callable, Iterator, Mapping +import time +from collections.abc import AsyncIterator, Callable, Iterator, Mapping from concurrent.futures import Future, ThreadPoolExecutor from datetime import datetime, timedelta, timezone from pathlib import Path, PurePath +from types import MappingProxyType from typing import Final, cast from unittest import mock from unittest.mock import AsyncMock, MagicMock, Mock, patch @@ -3650,36 +3652,100 @@ class TestIsStreamingRequest: assert _is_streaming_request(kwargs={"stream": True}, call_type=CallTypes.acompletion) is True -class _SpeechSuccessLogger(CustomLogger): +class _SpeechLogger(CustomLogger): def __init__(self) -> None: super().__init__() self.logged_payloads: Final[asyncio.Queue[dict[str, object]]] = asyncio.Queue() + self.failed_payloads: Final[asyncio.Queue[dict[str, object]]] = asyncio.Queue() async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: await self.logged_payloads.put(kwargs["standard_logging_object"]) + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time) -> None: + await self.failed_payloads.put(kwargs["standard_logging_object"]) + + +class _SyncSpeechLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged_payloads: Final[queue.Queue[dict[str, object]]] = queue.Queue() + + def log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + self.logged_payloads.put(kwargs["standard_logging_object"]) + + +class _PausedAsyncAudio(httpx.AsyncByteStream): + def __init__(self, first: bytes, rest: bytes | Exception, pause_seconds: float) -> None: + self._first: Final = first + self._rest: Final = rest + self._pause_seconds: Final = pause_seconds + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield self._first + await asyncio.sleep(self._pause_seconds) + if isinstance(self._rest, Exception): + raise self._rest + yield self._rest + + +class _PausedSyncAudio(httpx.SyncByteStream): + def __init__(self, first: bytes, rest: bytes, pause_seconds: float) -> None: + self._first: Final = first + self._rest: Final = rest + self._pause_seconds: Final = pause_seconds + + def __iter__(self) -> Iterator[bytes]: + yield self._first + time.sleep(self._pause_seconds) + yield self._rest + + +_SPEECH_API_BASE: Final = "http://localhost:12348" +_SPEECH_ARGS: Final = MappingProxyType( + { + "model": "elevenlabs/eleven_multilingual_v2", + "input": "hello there", + "voice": "21m00Tcm4TlvDq8ikWAM", + "api_base": _SPEECH_API_BASE, + "api_key": "fake_elevenlabs_api_key", + } +) + + +def _mock_elevenlabs_audio(respx_mock: respx.MockRouter, audio: httpx.AsyncByteStream | httpx.SyncByteStream) -> None: + respx_mock.post(url__regex=rf"{_SPEECH_API_BASE}/v1/text-to-speech/.*").mock( + return_value=httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}, stream=audio) + ) + + +def _audio_summary(num_bytes: int) -> dict[str, object]: + return {"object": "binary", "content_type": "audio/mpeg", "num_bytes": num_bytes} + + +def _logged_seconds(payload: dict[str, object]) -> float: + start_time, end_time = payload["startTime"], payload["endTime"] + assert isinstance(start_time, float) and isinstance(end_time, float) + return end_time - start_time + + +@pytest.fixture +def speech_logger(monkeypatch: pytest.MonkeyPatch) -> _SpeechLogger: + logger: Final = _SpeechLogger() + monkeypatch.setattr(litellm, "callbacks", [logger]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + return logger + @pytest.mark.usefixtures("local_model_cost_map") async def test_aspeech_logs_streamed_audio_like_buffered_audio( - respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch + respx_mock: respx.MockRouter, speech_logger: _SpeechLogger ) -> None: - api_base: Final = "http://localhost:12348" - respx_mock.post(url__regex=rf"{api_base}/v1/text-to-speech/.*").mock( + respx_mock.post(url__regex=rf"{_SPEECH_API_BASE}/v1/text-to-speech/.*").mock( return_value=httpx.Response(status_code=200, content=b"audio", headers={"content-type": "audio/mpeg"}) ) - speech_logger: Final = _SpeechSuccessLogger() - monkeypatch.setattr(litellm, "callbacks", [speech_logger]) - monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) async def logged_speech(stream: bool) -> dict[str, object]: - response: Final = await litellm.aspeech( - model="elevenlabs/eleven_multilingual_v2", - input="hello there", - voice="21m00Tcm4TlvDq8ikWAM", - api_base=api_base, - api_key="fake_elevenlabs_api_key", - stream=stream, - ) + response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=stream) assert await response.aread() == b"audio" return await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=5) @@ -3691,6 +3757,71 @@ async def test_aspeech_logs_streamed_audio_like_buffered_audio( assert streamed["response_cost"] == buffered["response_cost"] +@pytest.mark.usefixtures("local_model_cost_map") +async def test_aspeech_stream_is_logged_when_the_audio_ends( + respx_mock: respx.MockRouter, speech_logger: _SpeechLogger +) -> None: + _mock_elevenlabs_audio(respx_mock, _PausedAsyncAudio(b"first", b"rest", pause_seconds=0.3)) + + response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=True) + assert await response.aread() == b"firstrest" + logged: Final = await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=5) + + assert logged["response"] == _audio_summary(len(b"firstrest")) + assert _logged_seconds(logged) >= 0.3 + + +@pytest.mark.usefixtures("local_model_cost_map") +async def test_aspeech_stream_logs_a_failure_when_the_upstream_breaks_mid_clip( + respx_mock: respx.MockRouter, speech_logger: _SpeechLogger +) -> None: + _mock_elevenlabs_audio( + respx_mock, _PausedAsyncAudio(b"first", httpx.ReadError("upstream reset"), pause_seconds=0) + ) + + response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=True) + with pytest.raises(httpx.ReadError): + await response.aread() + await response.aclose() + failed: Final = await asyncio.wait_for(speech_logger.failed_payloads.get(), timeout=5) + + assert failed["status"] == "failure" + with pytest.raises(TimeoutError): + await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=0.5) + + +@pytest.mark.usefixtures("local_model_cost_map") +async def test_aspeech_stream_closed_mid_clip_is_logged_with_the_audio_sent( + respx_mock: respx.MockRouter, speech_logger: _SpeechLogger +) -> None: + _mock_elevenlabs_audio(respx_mock, _PausedAsyncAudio(b"first", b"never sent", pause_seconds=60)) + + response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=True) + audio: Final = await response.aiter_bytes() + assert await anext(audio) == b"first" + await response.aclose() + logged: Final = await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=5) + + assert logged["status"] == "success" + assert logged["response"] == _audio_summary(len(b"first")) + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_speech_stream_is_logged_when_the_audio_ends( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + _mock_elevenlabs_audio(respx_mock, _PausedSyncAudio(b"first", b"rest", pause_seconds=0.3)) + speech_logger: Final = _SyncSpeechLogger() + monkeypatch.setattr(litellm, "callbacks", [speech_logger]) + + response: Final = litellm.speech(**_SPEECH_ARGS, stream=True) + assert response.read() == b"firstrest" + logged: Final = speech_logger.logged_payloads.get(timeout=5) + + assert logged["response"] == _audio_summary(len(b"firstrest")) + assert _logged_seconds(logged) >= 0.3 + + class TestCallbackAsyncSyncSeparation: """Test that LoggingCallbackManager auto-routes async callbacks to async lists.""" From 7f09947a6535bf810072a16a80d1c768e0bdaf74 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 16:48:37 +0000 Subject: [PATCH 08/20] fix(tts): bill and map streamed speech failures, defer logging only when streaming A streamed clip that breaks mid-way is now logged as a failure at the clip's cost, since speech is priced per input character and the provider bills the whole request. Transport errors map to litellm.Timeout or litellm.APIConnectionError, and the proxy runs post_call_failure_hook with the clip's cost. Success logging waits for the stream end only when the caller asked for stream=True. --- .../litellm_core_utils/on_end_byte_stream.py | 24 ++++++++- litellm/proxy/proxy_server.py | 32 ++++++++++-- litellm/utils.py | 37 +++++++++++-- .../proxy/proxy_server/test_routes_audio.py | 52 ++++++++++++++++--- tests/unit/test_utils.py | 50 +++++++++++++----- 5 files changed, 164 insertions(+), 31 deletions(-) diff --git a/litellm/litellm_core_utils/on_end_byte_stream.py b/litellm/litellm_core_utils/on_end_byte_stream.py index 11a7df99374..1030b06fd7c 100644 --- a/litellm/litellm_core_utils/on_end_byte_stream.py +++ b/litellm/litellm_core_utils/on_end_byte_stream.py @@ -5,8 +5,14 @@ import httpx class OnEndAsyncByteStream(httpx.AsyncByteStream): - def __init__(self, stream: httpx.AsyncByteStream, on_end: Callable[[Exception | None], Awaitable[None]]) -> None: + def __init__( + self, + stream: httpx.AsyncByteStream, + map_error: Callable[[httpx.TransportError], Exception], + on_end: Callable[[Exception | None], Awaitable[None]], + ) -> None: self._stream: Final = stream + self._map_error: Final = map_error self._on_end: Final = on_end self._ended = False @@ -14,6 +20,10 @@ class OnEndAsyncByteStream(httpx.AsyncByteStream): try: async for chunk in self._stream: yield chunk + except httpx.TransportError as error: + mapped: Final = self._map_error(error) + await self._end(mapped) + raise mapped from error except Exception as error: await self._end(error) raise @@ -33,14 +43,24 @@ class OnEndAsyncByteStream(httpx.AsyncByteStream): class OnEndSyncByteStream(httpx.SyncByteStream): - def __init__(self, stream: httpx.SyncByteStream, on_end: Callable[[Exception | None], None]) -> None: + def __init__( + self, + stream: httpx.SyncByteStream, + map_error: Callable[[httpx.TransportError], Exception], + on_end: Callable[[Exception | None], None], + ) -> None: self._stream: Final = stream + self._map_error: Final = map_error self._on_end: Final = on_end self._ended = False def __iter__(self) -> Iterator[bytes]: try: yield from self._stream + except httpx.TransportError as error: + mapped: Final = self._map_error(error) + self._end(mapped) + raise mapped from error except Exception as error: self._end(error) raise diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 6b1274f3d62..cf9ca1fa810 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12829,6 +12829,7 @@ async def moderations( async def _audio_speech_chunk_generator( _response: HttpxBinaryResponseContent, + on_upstream_error: Callable[[Exception], Awaitable[None]], ) -> AsyncGenerator[bytes, None]: upstream_still_streaming: Final = not _response.response.is_stream_consumed _generator: Final = await _response.aiter_bytes( @@ -12837,13 +12838,26 @@ async def _audio_speech_chunk_generator( try: async for chunk in _generator: yield chunk + except Exception as error: + await on_upstream_error(error) + raise finally: - # Reading the body to the end closes the upstream, but a client that walks away - # mid-clip never gets there. Starlette raises ClientDisconnect out of the response - # before its background task runs, so this is the path that covers a disconnect. await _close_audio_speech_response(_response) +async def _log_audio_speech_stream_failure( + error: Exception, + request_data: Mapping[str, object], + user_api_key_dict: UserAPIKeyAuth, + response_cost: float, +) -> None: + await proxy_logging_obj.post_call_failure_hook( + user_api_key_dict=user_api_key_dict, + original_exception=error, + request_data={**request_data, "combined_usage_object": litellm.Usage(), "response_cost": response_cost}, + ) + + async def _close_audio_speech_response(_response: HttpxBinaryResponseContent) -> None: """Release the upstream connection holding this response's audio.""" with contextlib.suppress(Exception): @@ -12872,7 +12886,7 @@ async def audio_speech( """ global proxy_logging_obj litellm_call_id: Final = resolve_litellm_call_id(request.headers.get("x-litellm-call-id")) - data: dict = {"litellm_call_id": litellm_call_id} + data: dict[str, object] = {"litellm_call_id": litellm_call_id} try: # Use orjson to parse JSON data, orjson speeds up requests significantly body: Final = await request.body() @@ -12956,7 +12970,15 @@ async def audio_speech( HttpxBinaryResponseContent, response ) return StreamingResponse( - _audio_speech_chunk_generator(audio), + _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), diff --git a/litellm/utils.py b/litellm/utils.py index 4ebf2e457ee..b31467f34c4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1372,12 +1372,33 @@ def _dispatch_success_logging( ) -def _unread_speech_upstream(result: object, call_type: str) -> httpx.Response | None: - if call_type not in _SPEECH_CALL_TYPES or not isinstance(result, HttpxBinaryResponseContent): +def _unread_speech_upstream(result: object, call_type: str, stream_requested: bool) -> httpx.Response | None: + if ( + not stream_requested + or call_type not in _SPEECH_CALL_TYPES + or not isinstance(result, HttpxBinaryResponseContent) + ): return None return result.unread_response() +def _bill_failed_speech_stream(logging_obj: LiteLLMLoggingObject, result: object) -> None: + logging_obj.record_partial_usage_for_failure(Usage(), logging_obj.response_cost_calculator(result=result) or 0.0) + + +def _speech_stream_error_mapper( + model: str | None, logging_obj: LiteLLMLoggingObject +) -> Callable[[httpx.TransportError], Exception]: + llm_provider: Final = str(logging_obj.model_call_details.get("custom_llm_provider") or "") + + def map_error(error: httpx.TransportError) -> Exception: + if isinstance(error, httpx.TimeoutException): + return litellm.Timeout(message=str(error), model=model or "", llm_provider=llm_provider) + return litellm.APIConnectionError(message=str(error), llm_provider=llm_provider, model=model or "") + + return map_error + + def _async_speech_end_logger( logging_obj: LiteLLMLoggingObject, result: object, @@ -1400,6 +1421,7 @@ def _async_speech_end_logger( if is_litellm_internal_call: return traceback_exception: Final = "".join(traceback.format_exception(error)) + _bill_failed_speech_stream(logging_obj, result) unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) logging_obj.failure_handler(error, traceback_exception, start_time, end_time) await logging_obj.async_failure_handler(error, traceback_exception, start_time, end_time) @@ -1419,6 +1441,7 @@ def _sync_speech_end_logger( def log_speech_end(error: Exception | None) -> None: end_time: Final = datetime.datetime.now() # noqa: DTZ005 # matches the naive start_time the wrapper took if error is not None: + _bill_failed_speech_stream(logging_obj, result) logging_obj.failure_handler(error, "".join(traceback.format_exception(error)), start_time, end_time) return if is_litellm_internal_call: @@ -1965,10 +1988,13 @@ def client(original_function): # LOG SUCCESS - handle streaming success logging in the _next_ object, remove `handle_success` once it's deprecated verbose_logger.info("Wrapper: Completed Call, calling success_handler") executor: Final[BoundedLoggingThreadPoolExecutor] = getattr(sys.modules[__name__], "executor") - unread_speech: Final = _unread_speech_upstream(result, call_type) + unread_speech: Final = _unread_speech_upstream( + result, call_type, "stream" in kwargs and kwargs["stream"] is True + ) if unread_speech is not None and isinstance(unread_speech.stream, httpx.SyncByteStream): unread_speech.stream = OnEndSyncByteStream( unread_speech.stream, + _speech_stream_error_mapper(model, logging_obj), _sync_speech_end_logger( logging_obj=logging_obj, result=result, @@ -2286,10 +2312,13 @@ def client(original_function): start_time=start_time, end_time=end_time, ) - unread_speech: Final = _unread_speech_upstream(result, call_type) + unread_speech: Final = _unread_speech_upstream( + result, call_type, "stream" in kwargs and kwargs["stream"] is True + ) if unread_speech is not None and isinstance(unread_speech.stream, httpx.AsyncByteStream): unread_speech.stream = OnEndAsyncByteStream( unread_speech.stream, + _speech_stream_error_mapper(model, logging_obj), _async_speech_end_logger( logging_obj=logging_obj, result=result, diff --git a/tests/unit/proxy/proxy_server/test_routes_audio.py b/tests/unit/proxy/proxy_server/test_routes_audio.py index a1bac615e38..dbb1b935887 100644 --- a/tests/unit/proxy/proxy_server/test_routes_audio.py +++ b/tests/unit/proxy/proxy_server/test_routes_audio.py @@ -295,7 +295,7 @@ async def test_audio_speech_chunk_generator_forwards_an_unread_upstream_stream() httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}, stream=stream) ) - chunks = proxy_server._audio_speech_chunk_generator(response) + chunks = proxy_server._audio_speech_chunk_generator(response, on_upstream_error=AsyncMock()) assert await anext(chunks) == first assert produced == [first] @@ -341,17 +341,15 @@ async def _serve_until_disconnect(response: StreamingResponse) -> None: await response({"type": "http", "asgi": {"spec_version": "2.4"}}, _receive, _send) -@pytest.fixture -def patched_streaming_speech(monkeypatch): - """Route /v1/audio/speech to an upstream whose audio body is left unread.""" - stream = _RecordingAudioStream([b"a" * AUDIO_SPEECH_CHUNK_SIZE], []) +def _route_speech_to(monkeypatch: pytest.MonkeyPatch, response: HttpxBinaryResponseContent) -> AsyncMock: + failure_hook: Final = AsyncMock() monkeypatch.setattr(proxy_server, "llm_router", MagicMock()) monkeypatch.setattr( proxy_server, "proxy_logging_obj", MagicMock( pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]), - post_call_failure_hook=AsyncMock(), + post_call_failure_hook=failure_hook, post_call_response_headers_hook=AsyncMock(return_value={}), update_request_status=AsyncMock(), ), @@ -363,12 +361,20 @@ def patched_streaming_speech(monkeypatch): monkeypatch.setattr(proxy_server, "add_litellm_data_to_request", _add_data) async def _llm_call(): - return _audio_response(stream) + return response async def _fake_route_request(*args, **kwargs): return _llm_call() monkeypatch.setattr(proxy_server, "route_request", _fake_route_request) + return failure_hook + + +@pytest.fixture +def patched_streaming_speech(monkeypatch): + """Route /v1/audio/speech to an upstream whose audio body is left unread.""" + stream = _RecordingAudioStream([b"a" * AUDIO_SPEECH_CHUNK_SIZE], []) + _route_speech_to(monkeypatch, _audio_response(stream)) yield stream @@ -382,7 +388,7 @@ async def test_audio_speech_releases_its_upstream_when_the_client_walks_away(): await _serve_until_disconnect( StreamingResponse( - proxy_server._audio_speech_chunk_generator(response), + proxy_server._audio_speech_chunk_generator(response, on_upstream_error=AsyncMock()), media_type="audio/mpeg", background=BackgroundTask(proxy_server._close_audio_speech_response, response), ) @@ -423,3 +429,33 @@ async def test_audio_speech_route_releases_an_upstream_whose_body_is_never_read( await captured[0].background() assert stream.closed is True + + +class _BreaksMidClipAudioStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield b"a" * AUDIO_SPEECH_CHUNK_SIZE + raise httpx.ReadError("upstream reset") + + +@pytest.mark.asyncio +async def test_audio_speech_bills_a_clip_whose_upstream_breaks_mid_stream(monkeypatch: pytest.MonkeyPatch): + response: Final = HttpxBinaryResponseContent( + httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}, stream=_BreaksMidClipAudioStream()) + ) + response.hidden_params = {"response_cost": 0.25} + failure_hook: Final = _route_speech_to(monkeypatch, response) + + request: Final = _speech_request({"model": "tts-1", "input": "Hi", "voice": "alloy", "stream": True}) + streaming: Final = await proxy_server.audio_speech( + request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test") + ) + assert isinstance(streaming, StreamingResponse) + body: Final = aiter(streaming.body_iterator) + assert await anext(body) == b"a" * AUDIO_SPEECH_CHUNK_SIZE + with pytest.raises(httpx.ReadError): + await anext(body) + + failure_hook.assert_awaited_once() + assert failure_hook.await_args is not None + assert isinstance(failure_hook.await_args.kwargs["original_exception"], httpx.ReadError) + assert failure_hook.await_args.kwargs["request_data"]["response_cost"] == 0.25 diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index becdd4fa4fc..e2401e3bab7 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -58,7 +58,7 @@ from litellm.proxy.utils import is_valid_api_key from litellm.responses.main import aresponses, responses from litellm.types.caching import CachingSupportedCallTypes from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY, STREAM_OPTIONS_STASH_KEYS -from litellm.types.llms.openai import ResponsesAPIResponse +from litellm.types.llms.openai import HttpxBinaryResponseContent, ResponsesAPIResponse from litellm.types.router import CredentialLiteLLMParams, GenericLiteLLMParams, LiteLLM_Params from litellm.types.utils import ( ADDRESSED_RESPONSE_ID_FIELD, @@ -3712,6 +3712,10 @@ _SPEECH_ARGS: Final = MappingProxyType( ) +async def _aspeech(stream: bool) -> HttpxBinaryResponseContent: + return await litellm.aspeech(**_SPEECH_ARGS, stream=stream) + + def _mock_elevenlabs_audio(respx_mock: respx.MockRouter, audio: httpx.AsyncByteStream | httpx.SyncByteStream) -> None: respx_mock.post(url__regex=rf"{_SPEECH_API_BASE}/v1/text-to-speech/.*").mock( return_value=httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}, stream=audio) @@ -3745,7 +3749,7 @@ async def test_aspeech_logs_streamed_audio_like_buffered_audio( ) async def logged_speech(stream: bool) -> dict[str, object]: - response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=stream) + response: Final = await _aspeech(stream=stream) assert await response.aread() == b"audio" return await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=5) @@ -3763,7 +3767,7 @@ async def test_aspeech_stream_is_logged_when_the_audio_ends( ) -> None: _mock_elevenlabs_audio(respx_mock, _PausedAsyncAudio(b"first", b"rest", pause_seconds=0.3)) - response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=True) + response: Final = await _aspeech(stream=True) assert await response.aread() == b"firstrest" logged: Final = await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=5) @@ -3772,20 +3776,42 @@ async def test_aspeech_stream_is_logged_when_the_audio_ends( @pytest.mark.usefixtures("local_model_cost_map") -async def test_aspeech_stream_logs_a_failure_when_the_upstream_breaks_mid_clip( - respx_mock: respx.MockRouter, speech_logger: _SpeechLogger +@pytest.mark.parametrize( + ("upstream_error", "raised"), + [ + (httpx.ReadError("upstream reset"), litellm.APIConnectionError), + (httpx.ReadTimeout("upstream stalled"), litellm.Timeout), + ], +) +async def test_aspeech_stream_broken_mid_clip_is_billed_as_a_failure( + respx_mock: respx.MockRouter, + speech_logger: _SpeechLogger, + upstream_error: httpx.TransportError, + raised: type[Exception], ) -> None: - _mock_elevenlabs_audio( - respx_mock, _PausedAsyncAudio(b"first", httpx.ReadError("upstream reset"), pause_seconds=0) + respx_mock.post(url__regex=rf"{_SPEECH_API_BASE}/v1/text-to-speech/.*").mock( + side_effect=[ + httpx.Response( + status_code=200, + headers={"content-type": "audio/mpeg"}, + stream=_PausedAsyncAudio(b"first", rest, pause_seconds=0), + ) + for rest in (b"rest", upstream_error) + ] ) - response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=True) - with pytest.raises(httpx.ReadError): - await response.aread() - await response.aclose() + finished: Final = await _aspeech(stream=True) + await finished.aread() + logged: Final = await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=5) + broken: Final = await _aspeech(stream=True) + with pytest.raises(raised): + await broken.aread() + await broken.aclose() failed: Final = await asyncio.wait_for(speech_logger.failed_payloads.get(), timeout=5) assert failed["status"] == "failure" + assert isinstance(logged["response_cost"], float) and logged["response_cost"] > 0 + assert failed["response_cost"] == logged["response_cost"] with pytest.raises(TimeoutError): await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=0.5) @@ -3796,7 +3822,7 @@ async def test_aspeech_stream_closed_mid_clip_is_logged_with_the_audio_sent( ) -> None: _mock_elevenlabs_audio(respx_mock, _PausedAsyncAudio(b"first", b"never sent", pause_seconds=60)) - response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=True) + response: Final = await _aspeech(stream=True) audio: Final = await response.aiter_bytes() assert await anext(audio) == b"first" await response.aclose() From ab48060051ffe3a54105ce04aef313480f837a42 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 18:58:39 +0000 Subject: [PATCH 09/20] fix(tts): stream only on a literal true, mark hang-ups, and cut off without a fake DB outage --- .../litellm_core_utils/on_end_byte_stream.py | 49 ++++++++---- litellm/main.py | 13 +-- litellm/proxy/proxy_server.py | 23 ++++-- litellm/utils.py | 80 +++++++++++++------ .../proxy/proxy_server/test_routes_audio.py | 4 +- tests/unit/test_utils.py | 24 ++++++ 6 files changed, 140 insertions(+), 53 deletions(-) diff --git a/litellm/litellm_core_utils/on_end_byte_stream.py b/litellm/litellm_core_utils/on_end_byte_stream.py index 1030b06fd7c..c8b2285cdd8 100644 --- a/litellm/litellm_core_utils/on_end_byte_stream.py +++ b/litellm/litellm_core_utils/on_end_byte_stream.py @@ -1,15 +1,34 @@ from collections.abc import AsyncIterator, Awaitable, Callable, Iterator -from typing import Final +from dataclasses import dataclass +from typing import Final, TypeAlias import httpx +@dataclass(frozen=True, slots=True) +class StreamCompleted: + pass + + +@dataclass(frozen=True, slots=True) +class StreamClosedEarly: + pass + + +@dataclass(frozen=True, slots=True) +class StreamFailed: + error: Exception + + +StreamEnd: TypeAlias = StreamCompleted | StreamClosedEarly | StreamFailed + + class OnEndAsyncByteStream(httpx.AsyncByteStream): def __init__( self, stream: httpx.AsyncByteStream, map_error: Callable[[httpx.TransportError], Exception], - on_end: Callable[[Exception | None], Awaitable[None]], + on_end: Callable[[StreamEnd], Awaitable[None]], ) -> None: self._stream: Final = stream self._map_error: Final = map_error @@ -22,24 +41,24 @@ class OnEndAsyncByteStream(httpx.AsyncByteStream): yield chunk except httpx.TransportError as error: mapped: Final = self._map_error(error) - await self._end(mapped) + await self._end(StreamFailed(mapped)) raise mapped from error except Exception as error: - await self._end(error) + await self._end(StreamFailed(error)) raise - await self._end(None) + await self._end(StreamCompleted()) async def aclose(self) -> None: try: await self._stream.aclose() finally: - await self._end(None) + await self._end(StreamClosedEarly()) - async def _end(self, error: Exception | None) -> None: + async def _end(self, end: StreamEnd) -> None: if self._ended: return self._ended = True - await self._on_end(error) + await self._on_end(end) class OnEndSyncByteStream(httpx.SyncByteStream): @@ -47,7 +66,7 @@ class OnEndSyncByteStream(httpx.SyncByteStream): self, stream: httpx.SyncByteStream, map_error: Callable[[httpx.TransportError], Exception], - on_end: Callable[[Exception | None], None], + on_end: Callable[[StreamEnd], None], ) -> None: self._stream: Final = stream self._map_error: Final = map_error @@ -59,21 +78,21 @@ class OnEndSyncByteStream(httpx.SyncByteStream): yield from self._stream except httpx.TransportError as error: mapped: Final = self._map_error(error) - self._end(mapped) + self._end(StreamFailed(mapped)) raise mapped from error except Exception as error: - self._end(error) + self._end(StreamFailed(error)) raise - self._end(None) + self._end(StreamCompleted()) def close(self) -> None: try: self._stream.close() finally: - self._end(None) + self._end(StreamClosedEarly()) - def _end(self, error: Exception | None) -> None: + def _end(self, end: StreamEnd) -> None: if self._ended: return self._ended = True - self._on_end(error) + self._on_end(end) diff --git a/litellm/main.py b/litellm/main.py index 3bcf6c71903..16b759f6edc 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8287,6 +8287,7 @@ def speech( model=model, custom_llm_provider=custom_llm_provider, api_base=api_base or base_url ) kwargs.pop("tags", []) + stream_audio: Final = stream is True optional_params = {} if response_format is not None: @@ -8411,7 +8412,7 @@ def speech( client=client, # pass AsyncOpenAI, OpenAI client aspeech=aspeech, shared_session=shared_session, - stream=stream, + stream=stream_audio, ) elif custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS: # Check if this is Azure Speech Service (Cognitive Services TTS) @@ -8444,7 +8445,7 @@ def speech( aspeech=aspeech or False, api_base=api_base, api_key=api_key, - stream=stream, + stream=stream_audio, **kwargs, ) else: @@ -8492,7 +8493,7 @@ def speech( client=client, # pass AsyncOpenAI, OpenAI client aspeech=aspeech, litellm_params=litellm_params_dict, - stream=stream, + stream=stream_audio, ) elif custom_llm_provider == "elevenlabs": from litellm.llms.elevenlabs.text_to_speech.transformation import ( @@ -8518,7 +8519,7 @@ def speech( litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_QUERY_PARAMS_KEY] = query_params litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_VOICE_ID_KEY] = voice_id - if stream: + if stream_audio: litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_STREAM_KEY] = True if api_base is not None: @@ -8539,7 +8540,7 @@ def speech( extra_headers=extra_headers, client=client, _is_async=aspeech or False, - stream=stream, + stream=stream_audio, ) elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "vertex_ai_beta": from litellm.llms.vertex_ai.text_to_speech.transformation import ( @@ -8738,7 +8739,7 @@ def speech( aspeech=aspeech or False, api_base=api_base, api_key=api_key, - stream=stream, + stream=stream_audio, **kwargs, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cf9ca1fa810..91c49b636b6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -12827,6 +12827,18 @@ async def moderations( ) +class AudioSpeechStreamCutOff(Exception): + status_code: Final = status.HTTP_502_BAD_GATEWAY + + +async def _audio_chunks_or_error(chunks: AsyncIterator[bytes]) -> AsyncGenerator[bytes | Exception, None]: + try: + async for chunk in chunks: + yield chunk + except Exception as error: # noqa: BLE001 # handed to the caller, which logs it and cuts the response off + yield error + + async def _audio_speech_chunk_generator( _response: HttpxBinaryResponseContent, on_upstream_error: Callable[[Exception], Awaitable[None]], @@ -12836,11 +12848,12 @@ async def _audio_speech_chunk_generator( chunk_size=None if upstream_still_streaming else AUDIO_SPEECH_CHUNK_SIZE ) try: - async for chunk in _generator: - yield chunk - except Exception as error: - await on_upstream_error(error) - raise + async with contextlib.aclosing(_audio_chunks_or_error(_generator)) as chunks: + async for chunk in chunks: + if isinstance(chunk, Exception): + await on_upstream_error(chunk) + raise AudioSpeechStreamCutOff(f"The provider stopped sending audio mid-clip: {chunk}") + yield chunk finally: await _close_audio_speech_response(_response) diff --git a/litellm/utils.py b/litellm/utils.py index b31467f34c4..ba9b5bcc4f3 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -90,7 +90,14 @@ from litellm.litellm_core_utils.fallback_generalizations import ( match_capability_generalizations, match_fill_missing_generalizations, ) -from litellm.litellm_core_utils.on_end_byte_stream import OnEndAsyncByteStream, OnEndSyncByteStream +from litellm.litellm_core_utils.on_end_byte_stream import ( + OnEndAsyncByteStream, + OnEndSyncByteStream, + StreamClosedEarly, + StreamCompleted, + StreamEnd, + StreamFailed, +) from litellm.litellm_core_utils.sensitive_data_masker import redact_credentials_in_payload from litellm.litellm_core_utils.tokenizer import strip_special_tokens from litellm.rust_bridge import tokenizer as tokenizer_dispatch @@ -1399,32 +1406,48 @@ def _speech_stream_error_mapper( return map_error +def _mark_speech_stream_closed_early(logging_obj: LiteLLMLoggingObject) -> None: + litellm_params: Final = logging_obj.model_call_details.get("litellm_params") + if not isinstance(litellm_params, dict): + return + metadata: Final = litellm_params.get("metadata") + if isinstance(metadata, dict): + metadata["client_disconnected"] = True + return + litellm_params["metadata"] = {"client_disconnected": True} + + def _async_speech_end_logger( logging_obj: LiteLLMLoggingObject, result: object, start_time: datetime.datetime, is_completion_with_fallbacks: bool, is_litellm_internal_call: bool, -) -> Callable[[Exception | None], Awaitable[None]]: - async def log_speech_end(error: Exception | None) -> None: +) -> Callable[[StreamEnd], Awaitable[None]]: + async def log_speech_end(end: StreamEnd) -> None: end_time: Final = datetime.datetime.now() # noqa: DTZ005 # matches the naive start_time the wrapper took - if error is None: - _dispatch_success_logging( - logging_obj=logging_obj, - result=result, - start_time=start_time, - end_time=end_time, - is_completion_with_fallbacks=is_completion_with_fallbacks, - is_litellm_internal_call=is_litellm_internal_call, - ) - return - if is_litellm_internal_call: - return - traceback_exception: Final = "".join(traceback.format_exception(error)) - _bill_failed_speech_stream(logging_obj, result) - unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) - logging_obj.failure_handler(error, traceback_exception, start_time, end_time) - await logging_obj.async_failure_handler(error, traceback_exception, start_time, end_time) + match end: + case StreamFailed(error=error): + if is_litellm_internal_call: + return + traceback_exception: Final = "".join(traceback.format_exception(error)) + _bill_failed_speech_stream(logging_obj, result) + unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) + logging_obj.failure_handler(error, traceback_exception, start_time, end_time) + await logging_obj.async_failure_handler(error, traceback_exception, start_time, end_time) + return + case StreamClosedEarly(): + _mark_speech_stream_closed_early(logging_obj) + case StreamCompleted(): + pass + _dispatch_success_logging( + logging_obj=logging_obj, + result=result, + start_time=start_time, + end_time=end_time, + is_completion_with_fallbacks=is_completion_with_fallbacks, + is_litellm_internal_call=is_litellm_internal_call, + ) return log_speech_end @@ -1435,15 +1458,20 @@ def _sync_speech_end_logger( start_time: datetime.datetime, is_litellm_internal_call: bool, executor: BoundedLoggingThreadPoolExecutor, -) -> Callable[[Exception | None], None]: +) -> Callable[[StreamEnd], None]: ctx: Final = contextvars.copy_context() - def log_speech_end(error: Exception | None) -> None: + def log_speech_end(end: StreamEnd) -> None: end_time: Final = datetime.datetime.now() # noqa: DTZ005 # matches the naive start_time the wrapper took - if error is not None: - _bill_failed_speech_stream(logging_obj, result) - logging_obj.failure_handler(error, "".join(traceback.format_exception(error)), start_time, end_time) - return + match end: + case StreamFailed(error=error): + _bill_failed_speech_stream(logging_obj, result) + logging_obj.failure_handler(error, "".join(traceback.format_exception(error)), start_time, end_time) + return + case StreamClosedEarly(): + _mark_speech_stream_closed_early(logging_obj) + case StreamCompleted(): + pass if is_litellm_internal_call: return executor.submit(ctx.run, logging_obj.success_handler, result, start_time, end_time) diff --git a/tests/unit/proxy/proxy_server/test_routes_audio.py b/tests/unit/proxy/proxy_server/test_routes_audio.py index dbb1b935887..80f3a3f7505 100644 --- a/tests/unit/proxy/proxy_server/test_routes_audio.py +++ b/tests/unit/proxy/proxy_server/test_routes_audio.py @@ -29,6 +29,7 @@ from fastapi import Request, Response from litellm.proxy import proxy_server from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.types.llms.openai import HttpxBinaryResponseContent @@ -452,9 +453,10 @@ async def test_audio_speech_bills_a_clip_whose_upstream_breaks_mid_stream(monkey assert isinstance(streaming, StreamingResponse) body: Final = aiter(streaming.body_iterator) assert await anext(body) == b"a" * AUDIO_SPEECH_CHUNK_SIZE - with pytest.raises(httpx.ReadError): + with pytest.raises(proxy_server.AudioSpeechStreamCutOff) as cut_off: await anext(body) + assert PrismaDBExceptionHandler.is_database_service_unavailable_error_in_chain(cut_off.value) is False failure_hook.assert_awaited_once() assert failure_hook.await_args is not None assert isinstance(failure_hook.await_args.kwargs["original_exception"], httpx.ReadError) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index e2401e3bab7..e75a498e5db 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -3773,6 +3773,11 @@ async def test_aspeech_stream_is_logged_when_the_audio_ends( assert logged["response"] == _audio_summary(len(b"firstrest")) assert _logged_seconds(logged) >= 0.3 + assert logged["error_information"] != { + "error_code": "499", + "error_message": "Client disconnected the request", + "error_class": "ClientDisconnected", + } @pytest.mark.usefixtures("local_model_cost_map") @@ -3830,6 +3835,25 @@ async def test_aspeech_stream_closed_mid_clip_is_logged_with_the_audio_sent( assert logged["status"] == "success" assert logged["response"] == _audio_summary(len(b"first")) + assert logged["error_information"] == { + "error_code": "499", + "error_message": "Client disconnected the request", + "error_class": "ClientDisconnected", + } + + +@pytest.mark.usefixtures("local_model_cost_map") +@pytest.mark.parametrize("stream", [1, "true"]) +async def test_aspeech_stream_other_than_true_buffers_the_audio( + respx_mock: respx.MockRouter, speech_logger: _SpeechLogger, stream: object +) -> None: + _mock_elevenlabs_audio(respx_mock, _PausedAsyncAudio(b"first", b"rest", pause_seconds=0)) + + response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=stream) + logged: Final = await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=5) + + assert response.content == b"firstrest" + assert logged["response"] == _audio_summary(len(b"firstrest")) @pytest.mark.usefixtures("local_model_cost_map") From 05dc9bace8c3af4db78a53ee469f12ae65f4ef74 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 19:31:23 +0000 Subject: [PATCH 10/20] fix(tts): put the 499 hang-up marker on the spend row --- litellm/utils.py | 24 ++++++++++++++++++++---- tests/unit/test_utils.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 54 insertions(+), 4 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index ba9b5bcc4f3..b3ffa5597d1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -73,6 +73,7 @@ from litellm.constants import ( HF_CONFIG_FETCH_TIMEOUT_SECONDS, INITIAL_RETRY_DELAY, JITTER, + LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED, MAX_RETRY_DELAY, MAX_TOKEN_TRIMMING_ATTEMPTS, MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE, @@ -263,6 +264,7 @@ from litellm.types.utils import ( SandboxProviders, SearchProviders, SelectTokenizerResponse, + StandardLoggingPayloadErrorInformation, StreamingChoices, TextChoices, TextCompletionResponse, @@ -1406,15 +1408,29 @@ def _speech_stream_error_mapper( return map_error +_CLIENT_DISCONNECTED_ERROR_INFORMATION: Final[StandardLoggingPayloadErrorInformation] = { + "error_code": str(LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED), + "error_message": "Client disconnected the request", + "error_class": "ClientDisconnected", +} + + def _mark_speech_stream_closed_early(logging_obj: LiteLLMLoggingObject) -> None: litellm_params: Final = logging_obj.model_call_details.get("litellm_params") if not isinstance(litellm_params, dict): return - metadata: Final = litellm_params.get("metadata") - if isinstance(metadata, dict): - metadata["client_disconnected"] = True + request_metadata: Final = tuple( + metadata for metadata in map(litellm_params.get, ("metadata", "litellm_metadata")) if isinstance(metadata, dict) + ) + if not request_metadata: + litellm_params["metadata"] = { + "client_disconnected": True, + "error_information": dict(_CLIENT_DISCONNECTED_ERROR_INFORMATION), + } return - litellm_params["metadata"] = {"client_disconnected": True} + for metadata in request_metadata: + metadata["client_disconnected"] = True + metadata["error_information"] = dict(_CLIENT_DISCONNECTED_ERROR_INFORMATION) def _async_speech_end_logger( diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index e75a498e5db..da1f51faf3f 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -54,6 +54,8 @@ from litellm.litellm_core_utils.thread_pool_executor import executor as logging_ from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, headers from litellm.llms.openai_like.json_loader import JSONProviderRegistry +from litellm.proxy._types import SpendLogsPayload +from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.proxy.utils import is_valid_api_key from litellm.responses.main import aresponses, responses from litellm.types.caching import CachingSupportedCallTypes @@ -3842,6 +3844,38 @@ async def test_aspeech_stream_closed_mid_clip_is_logged_with_the_audio_sent( } +class _SpendRowLogger(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.rows: Final[asyncio.Queue[SpendLogsPayload]] = asyncio.Queue() + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time) -> None: + await self.rows.put(get_logging_payload(kwargs, response_obj, start_time, end_time)) + + +@pytest.mark.usefixtures("local_model_cost_map") +async def test_aspeech_stream_closed_mid_clip_marks_the_spend_row_as_a_client_disconnect( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + spend_rows: Final = _SpendRowLogger() + monkeypatch.setattr(litellm, "callbacks", [spend_rows]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + _mock_elevenlabs_audio(respx_mock, _PausedAsyncAudio(b"first", b"never sent", pause_seconds=60)) + + response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=True, metadata={"user_api_key": "hashed-key"}) + audio: Final = await response.aiter_bytes() + assert await anext(audio) == b"first" + await response.aclose() + row: Final = await asyncio.wait_for(spend_rows.rows.get(), timeout=5) + + assert row["status"] == "success" + assert json.loads(row["metadata"])["error_information"] == { + "error_code": "499", + "error_message": "Client disconnected the request", + "error_class": "ClientDisconnected", + } + + @pytest.mark.usefixtures("local_model_cost_map") @pytest.mark.parametrize("stream", [1, "true"]) async def test_aspeech_stream_other_than_true_buffers_the_audio( From 7433b11d1941116a2e08037bcf6226ebac95cd27 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 14 Oct 2026 22:19:40 +0000 Subject: [PATCH 11/20] fix(tts): keep the caller's metadata and positional base_url intact, add speech stream integration cells --- litellm/main.py | 2 +- litellm/utils.py | 13 +- .../test_audio_speech_stream_chaos.py | 312 ++++++++ .../test_audio_speech_stream_wire.py | 678 ++++++++++++++++++ ..._provider_passthrough_stream_relay_wire.py | 144 ++++ .../sdk/test_speech_stream_sdk_wire.py | 224 ++++++ tests/unit/test_utils.py | 19 + 7 files changed, 1384 insertions(+), 8 deletions(-) create mode 100644 tests/integration/providers/test_audio_speech_stream_chaos.py create mode 100644 tests/integration/providers/test_audio_speech_stream_wire.py create mode 100644 tests/integration/providers/test_provider_passthrough_stream_relay_wire.py create mode 100644 tests/integration/sdk/test_speech_stream_sdk_wire.py diff --git a/litellm/main.py b/litellm/main.py index 16b759f6edc..eb27f889844 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8273,8 +8273,8 @@ def speech( headers: dict | None = None, custom_llm_provider: str | None = None, aspeech: bool | None = None, - stream: bool = False, base_url: str | None = None, + stream: bool = False, **kwargs, ) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]: user: Final = kwargs.get("user", None) diff --git a/litellm/utils.py b/litellm/utils.py index b3ffa5597d1..8a0abe34bb7 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1420,17 +1420,16 @@ def _mark_speech_stream_closed_early(logging_obj: LiteLLMLoggingObject) -> None: if not isinstance(litellm_params, dict): return request_metadata: Final = tuple( - metadata for metadata in map(litellm_params.get, ("metadata", "litellm_metadata")) if isinstance(metadata, dict) + (key, metadata) + for key in ("metadata", "litellm_metadata") + if isinstance(metadata := litellm_params.get(key), dict) ) - if not request_metadata: - litellm_params["metadata"] = { + for key, existing in request_metadata or (("metadata", {}),): + litellm_params[key] = { + **existing, "client_disconnected": True, "error_information": dict(_CLIENT_DISCONNECTED_ERROR_INFORMATION), } - return - for metadata in request_metadata: - metadata["client_disconnected"] = True - metadata["error_information"] = dict(_CLIENT_DISCONNECTED_ERROR_INFORMATION) def _async_speech_end_logger( diff --git a/tests/integration/providers/test_audio_speech_stream_chaos.py b/tests/integration/providers/test_audio_speech_stream_chaos.py new file mode 100644 index 00000000000..47ed30b3a96 --- /dev/null +++ b/tests/integration/providers/test_audio_speech_stream_chaos.py @@ -0,0 +1,312 @@ +import asyncio +import threading +import uuid +from collections import Counter +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + eventually, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.process import ( + OwnedProxy, + graceful_stop_seconds, + owned_proxy_process, +) +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(4 * graceful_stop_seconds() + 300) + +_MODEL: Final = "chaos-speech" +_SPEECH: Final = "/v1/audio/speech" +_FIRST: Final = b"audio-frame-0:" + bytes(range(64)) +_REST: Final = ( + b"audio-frame-1:" + bytes(range(64, 192)), + b"audio-frame-2:" + bytes(range(192, 256)), +) +_CLIP: Final = _FIRST + b"".join(_REST) +_BURST: Final = 24 +_HELD: Final = 8 +_SPEND_SQL: Final = ( + 'SELECT litellm_call_id, status, spend, metadata FROM "LiteLLM_SpendLogs" ' + "WHERE litellm_call_id = ANY(string_to_array(%s, ','))" +) + +Kind = Literal["complete", "cut", "hang_up", "buffered", "held", "held_buffered"] +_MIXED: Final[tuple[Kind, ...]] = ("complete", "cut", "hang_up", "buffered") +_GATED: Final[frozenset[str]] = frozenset({"hang_up", "held", "held_buffered"}) +_BUFFERED: Final[frozenset[str]] = frozenset({"buffered", "held_buffered"}) + + +@dataclass(frozen=True, slots=True) +class Clip: + kind: Kind + call_id: str + + @property + def text(self) -> str: + return f"{self.kind} {self.call_id}" + + +@dataclass(frozen=True, slots=True) +class Heard: + clip: Clip + status: int + body: bytes + error: str + + +def _clips(marker: str, count: int, kinds: Sequence[Kind]) -> tuple[Clip, ...]: + return tuple( + Clip(kinds[index % len(kinds)], f"{marker}-{index:02d}-{uuid.uuid4().hex[:8]}") for index in range(count) + ) + + +def _peer(release: threading.Event) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + kind: Final = string_value(object_value(JSON_OBJECT.validate_json(request.body))["input"]).split(" ", 1)[0] + return Reply( + content_type="audio/mpeg", + chunks=(_FIRST, *_REST), + gate_after_first=release if kind in _GATED else None, + abort_after=1 if kind == "cut" else None, + ) + + return respond + + +def _config(directory: Path, wire: Wire) -> Path: + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + settings: Final = object_value(config["litellm_settings"]) + merged: Final = { + **config, + "model_list": [ + { + "model_name": _MODEL, + "litellm_params": { + "model": "openai/gpt-4o-mini-tts", + "api_base": wire.url, + "api_key": "synthetic-key", + }, + "model_info": {"mode": "audio_speech"}, + } + ], + "litellm_settings": { + **settings, + "cache_params": { + **object_value(settings["cache_params"]), + "namespace": f"tts-{uuid.uuid4().hex}", + }, + }, + "router_settings": { + **object_value(config["router_settings"]), + "num_retries": 0, + }, + } + path: Final = directory / f"speech-chaos-{uuid.uuid4().hex}.yaml" + path.write_text(yaml.safe_dump(merged)) + return path + + +@contextmanager +def _owned( + gateway: Gateway, + directory: Path, + wire: Wire, + overrides: Mapping[str, str] = MappingProxyType({}), +) -> Iterator[OwnedProxy]: + with owned_proxy_process(gateway, directory, overrides, config=_config(directory, wire), workers=2) as proxy: + yield proxy + + +def _body(clip: Clip) -> dict[str, JsonValue]: + stream: Final[dict[str, JsonValue]] = {} if clip.kind in _BUFFERED else {"stream": True} + return {"model": _MODEL, "input": clip.text, "voice": "alloy", **stream} + + +async def _hear(client: httpx.AsyncClient, key: str, clip: Clip) -> Heard: + headers: Final = { + "Authorization": f"Bearer {key}", + "x-litellm-call-id": clip.call_id, + } + try: + async with client.stream("POST", _SPEECH, json=_body(clip), headers=headers) as response: + if clip.kind == "hang_up": + return Heard(clip, response.status_code, await anext(response.aiter_raw()), "") + return Heard( + clip, + response.status_code, + b"".join([chunk async for chunk in response.aiter_raw()]), + "", + ) + except httpx.TransportError as error: + return Heard(clip, -1, b"", type(error).__name__) + + +async def _burst(url: str, key: str, clips: Sequence[Clip]) -> tuple[Heard, ...]: + async with httpx.AsyncClient( + base_url=url, + timeout=90, + trust_env=False, + limits=httpx.Limits(max_connections=len(clips)), + ) as client: + return tuple(await asyncio.gather(*(_hear(client, key, clip) for clip in clips))) + + +def _rows(clips: Sequence[Clip], database_url: str | None) -> list[dict[str, JsonValue]]: + return read_rows( + _SPEND_SQL, + (",".join(clip.call_id for clip in clips),), + database_url=database_url, + ) + + +def _billed_once(clips: Sequence[Clip], database_url: str | None = None) -> dict[str, dict[str, JsonValue]]: + rows: Final = eventually( + lambda: _rows(clips, database_url), + lambda found: len(found) >= len(clips), + seconds=120, + return_last_on_timeout=True, + ) + counts: Final = Counter(string_value(row["litellm_call_id"]) for row in rows) + missing: Final = sorted(clip.call_id for clip in clips if clip.call_id not in counts) + assert missing == [] and set(counts.values()) == {1}, (missing, counts) + return {string_value(row["litellm_call_id"]): row for row in rows} + + +def _assert_heard(heard: Heard) -> None: + match heard.clip.kind: + case "cut": + assert (heard.status, heard.error) == (-1, "RemoteProtocolError"), heard + case "hang_up": + assert (heard.status, heard.body) == (200, _FIRST), heard + case _: + assert (heard.status, heard.body, heard.error) == (200, _CLIP, ""), heard + + +def _assert_billed(clip: Clip, row: dict[str, JsonValue]) -> None: + spend: Final = row["spend"] + assert isinstance(spend, float) and spend > 0, (clip, row) + match clip.kind: + case "cut": + assert row["status"] == "failure", (clip, row) + case "hang_up": + assert row["status"] == "success", (clip, row) + error: Final = object_value(object_value(row["metadata"])["error_information"]) + assert error["error_code"] == "499", (clip, error) + case _: + assert row["status"] == "success", (clip, row) + + +def _held_by(pid: int, wire: Wire) -> int: + port: Final = urlsplit(wire.url).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +def _marker() -> str: + return f"ttschaos{uuid.uuid4().hex[:8]}" + + +async def test_a_mixed_streamed_speech_burst_bills_every_call_once_and_leaks_no_key_slot( + gateway: Gateway, tmp_path: Path +) -> None: + release: Final = threading.Event() + marker: Final = _marker() + with wire_server(_peer(release)) as wire, _owned(gateway, tmp_path, wire) as proxy: + with proxy.gateway.scenario() as scenario: + key: Final = scenario.key(max_parallel_requests=_BURST) + url: Final = str(proxy.gateway.client.base_url) + first: Final = _clips(marker, _BURST, _MIXED) + heard: Final = await _burst(url, key, first) + release.set() + billed: Final = await asyncio.to_thread(_billed_once, first) + second: Final = _clips(f"{marker}again", _BURST, ("complete",)) + again: Final = await _burst(url, key, second) + await asyncio.to_thread(_billed_once, second) + for item in (*heard, *again): + _assert_heard(item) + for clip in first: + _assert_billed(clip, billed[clip.call_id]) + assert len(wire.drain()) == 2 * _BURST + + +async def test_a_worker_killed_mid_streamed_speech_leaves_the_sibling_serving_and_billing( + gateway: Gateway, tmp_path: Path +) -> None: + release: Final = threading.Event() + marker: Final = _marker() + held: Final = _clips(marker, _HELD, ("held", "held_buffered")) + after: Final = _clips(f"{marker}after", _HELD, ("complete", "buffered")) + with wire_server(_peer(release)) as wire, _owned(gateway, tmp_path, wire) as proxy: + workers: Final = eventually( + proxy.worker_pids, + lambda pids: len(pids) == 2, + seconds=graceful_stop_seconds(), + ) + url: Final = str(proxy.gateway.client.base_url) + async with asyncio.TaskGroup() as tasks: + burst: Final = tasks.create_task(_burst(url, proxy.gateway.key, held)) + try: + await asyncio.to_thread(eventually, wire.received.qsize, lambda size: size >= _HELD, 90) + held_by: Final = MappingProxyType({pid: _held_by(pid, wire) for pid in workers}) + victim: Final = max(workers, key=held_by.__getitem__) + with proxy.replacing_worker(victim) as kill: + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + finally: + release.set() + later: Final = await _burst(url, proxy.gateway.key, after) + survivors: Final = tuple(item for item in served if item.status != -1) + billed: Final = await asyncio.to_thread(_billed_once, (*(item.clip for item in survivors), *after)) + assert sum(held_by.values()) == _HELD and held_by[victim] > 0, held_by + assert len(survivors) == _HELD - held_by[victim], (held_by, served) + for item in (*survivors, *later): + _assert_heard(item) + _assert_billed(item.clip, billed[item.clip.call_id]) + + +async def test_a_proxy_restart_mid_streamed_speech_bills_every_served_clip_once_and_the_next_boot_serves( + gateway: Gateway, tmp_path: Path +) -> None: + release: Final = threading.Event() + marker: Final = _marker() + held: Final = _clips(marker, _HELD, ("held", "held_buffered")) + after: Final = _clips(f"{marker}after", _HELD, ("complete", "buffered")) + with wire_server(_peer(release)) as wire: + with _owned(gateway, tmp_path, wire) as first: + async with asyncio.TaskGroup() as tasks: + burst: Final = tasks.create_task(_burst(str(first.gateway.client.base_url), first.gateway.key, held)) + try: + await asyncio.to_thread(eventually, wire.received.qsize, lambda size: size >= _HELD, 90) + first.begin_shutdown() + finally: + release.set() + served: Final = await burst + with _owned(gateway, tmp_path, wire) as second: + later: Final = await _burst(str(second.gateway.client.base_url), second.gateway.key, after) + completed: Final = tuple(item for item in served if item.status == 200) + billed: Final = await asyncio.to_thread(_billed_once, (*(item.clip for item in completed), *after)) + assert len(completed) == _HELD, served + for item in (*completed, *later): + _assert_heard(item) + _assert_billed(item.clip, billed[item.clip.call_id]) diff --git a/tests/integration/providers/test_audio_speech_stream_wire.py b/tests/integration/providers/test_audio_speech_stream_wire.py new file mode 100644 index 00000000000..4750676dd45 --- /dev/null +++ b/tests/integration/providers/test_audio_speech_stream_wire.py @@ -0,0 +1,678 @@ +import asyncio +import base64 +import json +import threading +import uuid +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from dataclasses import dataclass +from itertools import accumulate +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import httpx +import openai +import psutil +import pytest +from integration._support.client import ( + JSON_OBJECT, + Gateway, + Scenario, + eventually, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +_PROVIDER_KEY: Final = "synthetic-tts-key" +_FIRST: Final = b"audio-frame-0:" + bytes(range(64)) +_REST: Final = ( + b"audio-frame-1:" + bytes(range(64, 192)), + b"audio-frame-2:" + bytes(range(192, 256)), +) +_CLIP: Final = _FIRST + b"".join(_REST) +_FIRST_BYTE_SECONDS: Final = 20 +_SPEECH: Final = "/v1/audio/speech" +_SPEND_SQL: Final = ( + 'SELECT status, call_type, model_group, spend, metadata FROM "LiteLLM_SpendLogs" WHERE litellm_call_id = %s' +) +_ELEVENLABS_VOICE: Final = "21m00Tcm4TlvDq8ikWAM" +_AZURE_DEPLOYMENT: Final = "gpt-4o-mini-tts" + + +@dataclass(frozen=True, slots=True) +class Provider: + model: str + parameters: Mapping[str, JsonValue] + streamed_path: str + buffered_path: str + + +PROVIDERS: Final = MappingProxyType( + { + "openai": Provider( + "openai/gpt-4o-mini-tts", + {"api_key": _PROVIDER_KEY}, + "/audio/speech", + "/audio/speech", + ), + "azure": Provider( + f"azure/{_AZURE_DEPLOYMENT}", + {"api_key": _PROVIDER_KEY, "api_version": "2025-04-01-preview"}, + f"/openai/deployments/{_AZURE_DEPLOYMENT}/audio/speech", + f"/openai/deployments/{_AZURE_DEPLOYMENT}/audio/speech", + ), + "elevenlabs": Provider( + "elevenlabs/eleven_v3", + {"api_key": _PROVIDER_KEY}, + f"/v1/text-to-speech/{_ELEVENLABS_VOICE}/stream", + f"/v1/text-to-speech/{_ELEVENLABS_VOICE}", + ), + "aws_polly": Provider( + "aws_polly/neural", + { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + }, + "/v1/speech", + "/v1/speech", + ), + "azure_speech": Provider( + "azure/speech/azure-tts", + {"api_key": _PROVIDER_KEY}, + "/cognitiveservices/v1", + "/cognitiveservices/v1", + ), + } +) + + +@dataclass(frozen=True, slots=True) +class Heard: + status: int + content_type: str + body: bytes + + +def _text() -> str: + return f"streamed speech {uuid.uuid4().hex}" + + +def _speech(model: str, text: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"model": model, "input": text, "voice": "alloy", **extra} + + +def _audio_peer(release: threading.Event | None = None, abort_after: int | None = None) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + return Reply( + content_type="audio/mpeg", + chunks=(_FIRST, *_REST), + gate_after_first=release, + abort_after=abort_after, + ) + + return respond + + +def _deployment(scenario: Scenario, wire: Wire, provider: Provider, **extra: JsonValue) -> str: + return scenario.model( + model=provider.model, + api_base=wire.url, + model_info={"mode": "audio_speech"}, + **provider.parameters, + **extra, + ) + + +def _headers(gateway: Gateway, call_id: str, key: str | None = None) -> dict[str, str]: + return { + "Authorization": f"Bearer {gateway.key if key is None else key}", + "x-litellm-call-id": call_id, + } + + +def _first_bytes(chunks: Iterator[bytes], size: int) -> bytes: + return next(heard for heard in accumulate(chunks) if len(heard) >= size) + + +async def _afirst_bytes(chunks: AsyncIterator[bytes], size: int, heard: bytes = b"") -> bytes: + if len(heard) >= size: + return heard + return await _afirst_bytes(chunks, size, heard + await anext(chunks)) + + +def _listen( + gateway: Gateway, + body: Mapping[str, JsonValue], + call_id: str, + release: threading.Event, +) -> Heard: + with gateway.client.stream( + "POST", + _SPEECH, + json=dict(body), + headers=_headers(gateway, call_id), + timeout=httpx.Timeout(60, read=_FIRST_BYTE_SECONDS), + ) as response: + chunks: Final = response.iter_raw() + first: Final = _first_bytes(chunks, len(_FIRST)) + release.set() + rest: Final = b"".join(chunks) + return Heard(response.status_code, response.headers.get("content-type", ""), first + rest) + + +def _spend_row(call_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows(_SPEND_SQL, (call_id,)), + lambda found: len(found) >= 1, + seconds=70, + ) + assert len(rows) == 1, rows + return rows[0] + + +def _buffered_spend(gateway: Gateway, model: str, text: str) -> float: + call_id: Final = str(uuid.uuid4()) + response: Final = gateway.request("POST", _SPEECH, _speech(model, text), headers={"x-litellm-call-id": call_id}) + assert response.status_code == 200, response.text[:400] + spend: Final = _spend_row(call_id)["spend"] + assert isinstance(spend, float) and spend > 0, spend + return spend + + +def _assert_billed(gateway: Gateway, model: str, call_id: str, text: str) -> dict[str, JsonValue]: + row: Final = _spend_row(call_id) + assert (row["status"], row["call_type"], row["model_group"]) == ("success", "aspeech", model), row + assert row["spend"] == pytest.approx(_buffered_spend(gateway, model, text)), row + return object_value(row["metadata"]) + + +def _failure_spend(model: str, call_id: str) -> JsonValue: + row: Final = _spend_row(call_id) + assert (row["status"], row["model_group"]) == ("failure", model), row + return row["spend"] + + +def _only_request(wire: Wire) -> Request: + requests: Final = wire.drain() + assert len(requests) == 1, [request.target for request in requests] + return requests[0] + + +def _assert_reached(wire: Wire, path: str, text: str) -> Request: + request: Final = _only_request(wire) + assert (request.method, urlsplit(request.target).path) == ("POST", path), request.target + assert text.encode() in request.body, request.body[:400] + return request + + +def _held_upstream_connections(wire: Wire) -> int: + port: Final = urlsplit(wire.url).port + return sum( + 1 + for connection in psutil.Process().net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.laddr and connection.laddr.port == port + ) + + +@pytest.mark.parametrize("provider", tuple(PROVIDERS)) +def test_streamed_speech_reaches_the_caller_before_the_clip_ends_and_bills_it_once( + gateway: Gateway, provider: str +) -> None: + release: Final = threading.Event() + with wire_server(_audio_peer(release)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS[provider]) + text: Final = _text() + call_id: Final = str(uuid.uuid4()) + heard: Final = _listen(gateway, _speech(model, text, stream=True), call_id, release) + assert (heard.status, heard.body) == (200, _CLIP), heard + assert heard.content_type.startswith("audio/"), heard.content_type + _assert_reached(wire, PROVIDERS[provider].streamed_path, text) + _assert_billed(gateway, model, call_id, text) + + +def test_openai_sdk_sync_client_hears_streamed_speech_before_the_clip_ends( + gateway: Gateway, +) -> None: + release: Final = threading.Event() + with wire_server(_audio_peer(release)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["openai"]) + text: Final = _text() + call_id: Final = str(uuid.uuid4()) + client: Final = openai.OpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.Client(trust_env=False, timeout=httpx.Timeout(60, read=_FIRST_BYTE_SECONDS)), + ) + with client.audio.speech.with_streaming_response.create( + model=model, + input=text, + voice="alloy", + extra_body={"stream": True}, + extra_headers={"x-litellm-call-id": call_id}, + ) as response: + chunks: Final = response.iter_bytes() + first: Final = _first_bytes(chunks, len(_FIRST)) + release.set() + assert first + b"".join(chunks) == _CLIP + client.close() + _assert_reached(wire, PROVIDERS["openai"].streamed_path, text) + _assert_billed(gateway, model, call_id, text) + + +async def _async_sdk_clip(gateway: Gateway, model: str, text: str, call_id: str, release: threading.Event) -> bytes: + client: Final = openai.AsyncOpenAI( + base_url=f"{gateway.client.base_url}/v1", + api_key=gateway.key, + max_retries=0, + http_client=httpx.AsyncClient(trust_env=False, timeout=httpx.Timeout(60, read=_FIRST_BYTE_SECONDS)), + ) + try: + async with client.audio.speech.with_streaming_response.create( + model=model, + input=text, + voice="alloy", + extra_body={"stream": True}, + extra_headers={"x-litellm-call-id": call_id}, + ) as response: + chunks: Final = response.iter_bytes() + first: Final = await _afirst_bytes(chunks, len(_FIRST)) + release.set() + return first + b"".join([chunk async for chunk in chunks]) + finally: + await client.close() + + +def test_openai_sdk_async_client_hears_streamed_speech_before_the_clip_ends( + gateway: Gateway, +) -> None: + release: Final = threading.Event() + with wire_server(_audio_peer(release)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["openai"]) + text: Final = _text() + call_id: Final = str(uuid.uuid4()) + assert asyncio.run(_async_sdk_clip(gateway, model, text, call_id, release)) == _CLIP + _assert_reached(wire, PROVIDERS["openai"].streamed_path, text) + _assert_billed(gateway, model, call_id, text) + + +@pytest.mark.parametrize("stream", [None, False], ids=["omitted", "false"]) +@pytest.mark.parametrize("provider", ["openai", "elevenlabs"]) +def test_buffered_speech_is_unchanged(gateway: Gateway, provider: str, stream: bool | None) -> None: + with wire_server(_audio_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS[provider]) + text: Final = _text() + call_id: Final = str(uuid.uuid4()) + extra: Final[dict[str, JsonValue]] = {} if stream is None else {"stream": stream} + response: Final = gateway.request( + "POST", + _SPEECH, + _speech(model, text, **extra), + headers={"x-litellm-call-id": call_id}, + ) + assert (response.status_code, response.content) == (200, _CLIP), response.text[:400] + _assert_reached(wire, PROVIDERS[provider].buffered_path, text) + _assert_billed(gateway, model, call_id, text) + + +def _mistral_peer(request: Request) -> Reply: + return Reply(body=json.dumps({"audio_data": base64.b64encode(_CLIP).decode()}).encode()) + + +def test_stream_true_on_a_provider_without_audio_streaming_answers_buffered_and_bills_it( + gateway: Gateway, +) -> None: + with wire_server(_mistral_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="mistral/voxtral-mini-tts-latest", + api_base=wire.url, + api_key=_PROVIDER_KEY, + model_info={"mode": "audio_speech"}, + ) + text: Final = _text() + call_id: Final = str(uuid.uuid4()) + response: Final = gateway.request( + "POST", + _SPEECH, + _speech(model, text, stream=True), + headers={"x-litellm-call-id": call_id}, + ) + assert (response.status_code, response.content) == (200, _CLIP), response.text[:400] + _assert_reached(wire, "/v1/audio/speech", text) + _assert_billed(gateway, model, call_id, text) + + +@pytest.mark.parametrize( + "stream", + ["true", 1, None, [], {}, "s" * 5120], + ids=["string-true", "int-1", "null", "empty-list", "empty-object", "5kb-string"], +) +def test_a_stream_value_other_than_literal_true_answers_buffered(gateway: Gateway, stream: JsonValue) -> None: + with wire_server(_audio_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["elevenlabs"]) + text: Final = _text() + call_id: Final = str(uuid.uuid4()) + response: Final = gateway.request( + "POST", + _SPEECH, + _speech(model, text, stream=stream), + headers={"x-litellm-call-id": call_id}, + ) + assert (response.status_code, response.content) == (200, _CLIP), response.text[:400] + request: Final = _assert_reached(wire, PROVIDERS["elevenlabs"].buffered_path, text) + assert "stream" not in object_value(json.loads(request.body)), request.body[:400] + _assert_billed(gateway, model, call_id, text) + + +def test_a_duplicated_stream_true_key_streams(gateway: Gateway) -> None: + with wire_server(_audio_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["elevenlabs"]) + text: Final = _text() + call_id: Final = str(uuid.uuid4()) + body: Final = json.dumps(_speech(model, text, stream=True))[:-1] + ', "stream": true}' + response: Final = gateway.client.post( + _SPEECH, + content=body.encode(), + headers={**_headers(gateway, call_id), "content-type": "application/json"}, + ) + assert (response.status_code, response.content) == (200, _CLIP), response.text[:400] + _assert_reached(wire, PROVIDERS["elevenlabs"].streamed_path, text) + _assert_billed(gateway, model, call_id, text) + + +def test_a_hang_up_mid_clip_closes_the_provider_stream_and_bills_the_clip_with_a_499_marker( + gateway: Gateway, +) -> None: + release: Final = threading.Event() + with wire_server(_audio_peer(release)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["openai"]) + text: Final = _text() + call_id: Final = str(uuid.uuid4()) + with gateway.client.stream( + "POST", + _SPEECH, + json=_speech(model, text, stream=True), + headers=_headers(gateway, call_id), + timeout=httpx.Timeout(60, read=_FIRST_BYTE_SECONDS), + ) as response: + assert response.status_code == 200 + assert _first_bytes(response.iter_raw(), len(_FIRST)) == _FIRST + eventually(lambda: _held_upstream_connections(wire), lambda held: held == 0) + release.set() + _assert_reached(wire, PROVIDERS["openai"].streamed_path, text) + metadata: Final = _assert_billed(gateway, model, call_id, text) + assert object_value(metadata["error_information"])["error_code"] == "499", metadata + + +def test_a_provider_cut_off_mid_clip_cuts_the_caller_off_and_bills_a_failure_at_full_cost( + gateway: Gateway, +) -> None: + with ( + wire_server(_audio_peer(abort_after=1)) as wire, + wire_server(_audio_peer()) as healthy, + gateway.scenario() as scenario, + ): + model: Final = _deployment(scenario, wire, PROVIDERS["openai"]) + twin: Final = _deployment(scenario, healthy, PROVIDERS["openai"]) + text: Final = _text() + call_id: Final = str(uuid.uuid4()) + with gateway.client.stream( + "POST", + _SPEECH, + json=_speech(model, text, stream=True), + headers=_headers(gateway, call_id), + ) as response: + assert response.status_code == 200 + chunks: Final = response.iter_raw() + assert _first_bytes(chunks, len(_FIRST)) == _FIRST + with pytest.raises(httpx.RemoteProtocolError): + b"".join(chunks) + _assert_reached(wire, PROVIDERS["openai"].streamed_path, text) + assert _failure_spend(model, call_id) == pytest.approx(_buffered_spend(gateway, twin, text)) + + +_REJECTIONS: Final = MappingProxyType( + { + "openai": json.dumps( + { + "error": { + "message": "Incorrect API key provided: synthetic", + "type": "invalid_request_error", + "code": "invalid_api_key", + } + } + ), + "elevenlabs": json.dumps( + { + "detail": { + "status": "invalid_api_key", + "message": "Invalid API key: synthetic", + } + } + ), + } +) + + +@pytest.mark.parametrize("provider", tuple(_REJECTIONS)) +def test_a_provider_rejection_before_any_audio_reaches_the_caller_with_its_status_and_message( + gateway: Gateway, provider: str +) -> None: + def respond(request: Request) -> Reply: + return Reply(status=401, body=_REJECTIONS[provider].encode()) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS[provider]) + text: Final = _text() + call_id: Final = str(uuid.uuid4()) + response: Final = gateway.request( + "POST", + _SPEECH, + _speech(model, text, stream=True), + headers={"x-litellm-call-id": call_id}, + ) + assert response.status_code == 401, response.text[:400] + assert "API key" in string_value(object_value(JSON_OBJECT.validate_json(response.content)["error"])["message"]) + _assert_reached(wire, PROVIDERS[provider].streamed_path, text) + assert _failure_spend(model, call_id) == 0.0 + + +def test_a_repeated_identical_streamed_request_reaches_the_provider_and_bills_each_call_once( + gateway: Gateway, +) -> None: + with wire_server(_audio_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["openai"]) + text: Final = _text() + call_ids: Final = (str(uuid.uuid4()), str(uuid.uuid4())) + responses: Final = tuple( + gateway.request( + "POST", + _SPEECH, + _speech(model, text, stream=True), + headers={"x-litellm-call-id": call_id}, + ) + for call_id in call_ids + ) + assert [(response.status_code, response.content) for response in responses] == [(200, _CLIP)] * 2 + received: Final = wire.drain() + assert [urlsplit(request.target).path for request in received] == ["/audio/speech"] * 2 + for call_id in call_ids: + _assert_billed(gateway, model, call_id, text) + + +def _admitted(gateway: Gateway, model: str, key: str) -> int: + return gateway.request("POST", _SPEECH, _speech(model, _text(), stream=True), key=key).status_code + + +def test_a_key_max_parallel_requests_slot_is_held_until_the_streamed_clip_ends( + gateway: Gateway, +) -> None: + release: Final = threading.Event() + with wire_server(_audio_peer(release)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["openai"]) + key: Final = scenario.key(max_parallel_requests=1) + with gateway.client.stream( + "POST", + _SPEECH, + json=_speech(model, _text(), stream=True), + headers=_headers(gateway, str(uuid.uuid4()), key), + timeout=httpx.Timeout(60, read=_FIRST_BYTE_SECONDS), + ) as held: + chunks: Final = held.iter_raw() + first: Final = _first_bytes(chunks, len(_FIRST)) + blocked: Final = gateway.request("POST", _SPEECH, _speech(model, _text(), stream=True), key=key) + release.set() + assert first + b"".join(chunks) == _CLIP + assert blocked.status_code == 429, blocked.text[:400] + eventually(lambda: _admitted(gateway, model, key), lambda status: status == 200) + assert len(wire.drain()) == 2 + + +def test_a_key_max_parallel_requests_slot_is_freed_after_each_completed_clip( + gateway: Gateway, +) -> None: + with wire_server(_audio_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["openai"]) + key: Final = scenario.key(max_parallel_requests=1) + for _ in range(3): + eventually(lambda: _admitted(gateway, model, key), lambda status: status == 200) + assert len(wire.drain()) == 3 + + +def test_a_key_max_parallel_requests_slot_is_freed_by_a_hang_up( + gateway: Gateway, +) -> None: + release: Final = threading.Event() + with wire_server(_audio_peer(release)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["openai"]) + key: Final = scenario.key(max_parallel_requests=1) + with gateway.client.stream( + "POST", + _SPEECH, + json=_speech(model, _text(), stream=True), + headers=_headers(gateway, str(uuid.uuid4()), key), + timeout=httpx.Timeout(60, read=_FIRST_BYTE_SECONDS), + ) as held: + assert _first_bytes(held.iter_raw(), len(_FIRST)) == _FIRST + eventually(lambda: _held_upstream_connections(wire), lambda count: count == 0) + release.set() + eventually(lambda: _admitted(gateway, model, key), lambda status: status == 200) + assert len(wire.drain()) == 2 + + +def test_a_deployment_max_parallel_requests_slot_is_held_until_the_streamed_clip_ends( + gateway: Gateway, +) -> None: + release: Final = threading.Event() + with wire_server(_audio_peer(release)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["openai"], max_parallel_requests=1) + with gateway.client.stream( + "POST", + _SPEECH, + json=_speech(model, _text(), stream=True), + headers=_headers(gateway, str(uuid.uuid4())), + timeout=httpx.Timeout(60, read=_FIRST_BYTE_SECONDS), + ) as held: + chunks: Final = held.iter_raw() + first: Final = _first_bytes(chunks, len(_FIRST)) + blocked: Final = gateway.request("POST", _SPEECH, _speech(model, _text(), stream=True)) + release.set() + assert first + b"".join(chunks) == _CLIP + assert blocked.status_code == 429, blocked.text[:400] + eventually(lambda: _admitted(gateway, model, gateway.key), lambda status: status == 200) + assert len(wire.drain()) == 2 + + +def test_a_deployment_max_parallel_requests_slot_is_freed_by_a_hang_up( + gateway: Gateway, +) -> None: + release: Final = threading.Event() + with wire_server(_audio_peer(release)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["openai"], max_parallel_requests=1) + with gateway.client.stream( + "POST", + _SPEECH, + json=_speech(model, _text(), stream=True), + headers=_headers(gateway, str(uuid.uuid4())), + timeout=httpx.Timeout(60, read=_FIRST_BYTE_SECONDS), + ) as held: + assert _first_bytes(held.iter_raw(), len(_FIRST)) == _FIRST + eventually(lambda: _held_upstream_connections(wire), lambda count: count == 0) + release.set() + eventually(lambda: _admitted(gateway, model, gateway.key), lambda status: status == 200) + assert len(wire.drain()) == 2 + + +_CHAT_ID: Final = "chatcmpl-nonboolstream" +_CHAT_USAGE: Final = {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6} + + +def _chat_chunk(delta: Mapping[str, JsonValue], finish_reason: str | None, **extra: JsonValue) -> bytes: + chunk: Final = { + "id": _CHAT_ID, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": dict(delta), "finish_reason": finish_reason}], + **extra, + } + return f"data: {json.dumps(chunk)}\n\n".encode() + + +def _chat_peer(request: Request) -> Reply: + if object_value(JSON_OBJECT.validate_json(request.body)).get("stream"): + return Reply( + content_type="text/event-stream", + chunks=( + _chat_chunk({"role": "assistant", "content": "hello"}, None), + _chat_chunk({}, "stop", usage=_CHAT_USAGE), + b"data: [DONE]\n\n", + ), + ) + completion: Final = { + "id": _CHAT_ID, + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hello"}, + "finish_reason": "stop", + } + ], + "usage": _CHAT_USAGE, + } + return Reply(body=json.dumps(completion).encode()) + + +@pytest.mark.parametrize("stream", ["true", 1], ids=["string-true", "int-1"]) +def test_a_chat_completion_with_a_non_boolean_stream_value_answers_and_lands_one_success_row( + gateway: Gateway, stream: JsonValue +) -> None: + with wire_server(_chat_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + call_id: Final = str(uuid.uuid4()) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": f"say hello {call_id}"}], + "stream": stream, + }, + headers={"x-litellm-call-id": call_id}, + ) + assert response.status_code == 200, response.text[:400] + assert "hello" in response.text, response.text[:400] + row: Final = _spend_row(call_id) + assert (row["status"], row["call_type"], row["model_group"]) == ( + "success", + "acompletion", + model, + ), row + assert isinstance(row["spend"], float) and row["spend"] > 0, row diff --git a/tests/integration/providers/test_provider_passthrough_stream_relay_wire.py b/tests/integration/providers/test_provider_passthrough_stream_relay_wire.py new file mode 100644 index 00000000000..10b05b114ee --- /dev/null +++ b/tests/integration/providers/test_provider_passthrough_stream_relay_wire.py @@ -0,0 +1,144 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.forward_proxy import tunnelling_forward_proxy +from integration._support.process import owned_proxy +from integration._support.tls import server_context, write_self_signed_cert +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +_KEY: Final = "synthetic-passthrough-key" +_GOOGLE_TOKEN: Final = "synthetic-google-access-token" +_VERTEX_HOST: Final = "us-central1-aiplatform.googleapis.com" +_VERTEX_MODEL: Final = "gemini-2.5-flash" +_FRAMES: Final = ( + b'data: {"id":"pt-1","choices":[{"index":0,"delta":{"content":"hel"}}]}\n\n', + b'data: {"id":"pt-1","choices":[{"index":0,"delta":{"content":"lo"}}]}\n\n', + b"data: [DONE]\n\n", +) +_GEMINI_FRAMES: Final = tuple( + f"data: {json.dumps(frame)}\r\n\r\n".encode() + for frame in ( + {"candidates": [{"content": {"role": "model", "parts": [{"text": "hel"}]}, "index": 0}]}, + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "lo"}]}, "index": 0, "finishReason": "STOP"} + ], + "usageMetadata": {"promptTokenCount": 4, "candidatesTokenCount": 2, "totalTokenCount": 6}, + "modelVersion": _VERTEX_MODEL, + }, + ) +) +_ROUTES: Final = MappingProxyType( + { + "openai": ("/openai/chat/completions", "/chat/completions"), + "mistral": ("/mistral/v1/chat/completions", "/v1/chat/completions"), + "cohere": ("/cohere/v2/chat", "/v2/chat"), + } +) + + +@dataclass(frozen=True, slots=True) +class _Relay: + wire: Wire + vertex: Wire + proxy: Gateway + + +def _sse(request: Request) -> Reply: + return Reply(content_type="text/event-stream", chunks=_FRAMES) + + +def _gemini_sse(request: Request) -> Reply: + return Reply(content_type="text/event-stream", chunks=_GEMINI_FRAMES) + + +@pytest.fixture(scope="module") +def relay(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Relay]: + directory: Final = tmp_path_factory.mktemp("passthrough-stream") + certificate, key = write_self_signed_cert(directory, (_VERTEX_HOST,)) + with ( + wire_server(_sse) as wire, + wire_server(_gemini_sse, tls=server_context(certificate, key)) as vertex, + tunnelling_forward_proxy(int(vertex.url.rsplit(":", 1)[1])) as tunnel, + gateway_from_environment() as rig, + ): + bases: Final = {f"{name}_API_BASE": wire.url for name in ("OPENAI", "MISTRAL", "COHERE")} + keys: Final = {f"{name}_API_KEY": _KEY for name in ("OPENAI", "MISTRAL", "COHERE")} + tunnelled: Final = { + "HTTPS_PROXY": tunnel.url, + "https_proxy": tunnel.url, + "NO_PROXY": "127.0.0.1,localhost", + "no_proxy": "127.0.0.1,localhost", + "SSL_CERT_FILE": str(certificate), + } + with owned_proxy(rig, directory, {**bases, **keys, **tunnelled}, workers=2) as proxy: + yield _Relay(wire, vertex, proxy) + + +def _relayed( + proxy: Gateway, path: str, body: dict[str, JsonValue], headers: Mapping[str, str] | None = None +) -> tuple[int, bytes]: + sent: Final = headers or {"Authorization": f"Bearer {proxy.key}"} + with proxy.client.stream("POST", path, json=body, headers=dict(sent)) as response: + return response.status_code, b"".join(response.iter_raw()) + + +def _only_post(wire: Wire) -> Request: + requests: Final = wire.drain() + assert [request.method for request in requests] == ["POST"], [request.target for request in requests] + return requests[0] + + +@pytest.mark.parametrize("route", tuple(_ROUTES)) +def test_a_provider_passthrough_stream_is_relayed_frame_for_frame(relay: _Relay, route: str) -> None: + proxy_path, upstream_path = _ROUTES[route] + marker: Final = f"passthrough-stream-{uuid.uuid4().hex}" + relayed: Final = _relayed( + relay.proxy, + proxy_path, + {"model": "relay-model", "stream": True, "messages": [{"role": "user", "content": marker}]}, + ) + assert relayed == (200, b"".join(_FRAMES)), relayed + forwarded: Final = _only_post(relay.wire) + assert urlsplit(forwarded.target).path == upstream_path, forwarded.target + assert marker.encode() in forwarded.body, forwarded.body[:400] + + +def test_a_vertex_passthrough_stream_is_relayed_frame_for_frame(relay: _Relay) -> None: + project: Final = f"relay-{uuid.uuid4().hex[:12]}" + model_path: Final = f"/v1/projects/{project}/locations/us-central1/publishers/google/models/{_VERTEX_MODEL}" + marker: Final = f"vertex-passthrough-stream-{uuid.uuid4().hex}" + relayed: Final = _relayed( + relay.proxy, + f"/vertex_ai{model_path}:streamGenerateContent", + {"contents": [{"role": "user", "parts": [{"text": marker}]}]}, + {"x-litellm-api-key": relay.proxy.key, "Authorization": f"Bearer {_GOOGLE_TOKEN}"}, + ) + assert relayed == (200, b"".join(_GEMINI_FRAMES)), relayed + forwarded: Final = _only_post(relay.vertex) + assert forwarded.target == f"{model_path}:streamGenerateContent?alt=sse", forwarded.target + assert forwarded.headers["authorization"] == f"Bearer {_GOOGLE_TOKEN}", forwarded.headers + assert marker.encode() in forwarded.body, forwarded.body[:400] + + +def test_an_azure_passthrough_stream_to_a_router_model_is_relayed_frame_for_frame(gateway: Gateway) -> None: + with wire_server(_sse) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="azure/relay-deployment", api_base=wire.url, api_version="2024-10-21") + marker: Final = f"azure-passthrough-stream-{uuid.uuid4().hex}" + relayed: Final = _relayed( + gateway, + f"/azure/openai/deployments/{model}/chat/completions?api-version=2024-10-21", + {"stream": True, "messages": [{"role": "user", "content": marker}]}, + ) + assert relayed == (200, b"".join(_FRAMES)), relayed + forwarded: Final = _only_post(wire) + assert urlsplit(forwarded.target).path.endswith("/chat/completions"), forwarded.target + assert marker.encode() in forwarded.body, forwarded.body[:400] diff --git a/tests/integration/sdk/test_speech_stream_sdk_wire.py b/tests/integration/sdk/test_speech_stream_sdk_wire.py new file mode 100644 index 00000000000..8f1c9728a55 --- /dev/null +++ b/tests/integration/sdk/test_speech_stream_sdk_wire.py @@ -0,0 +1,224 @@ +import asyncio +import threading +import uuid +from collections.abc import AsyncIterator, Callable, Iterator, Mapping +from dataclasses import dataclass +from itertools import accumulate +from typing import Final + +import litellm +import pytest +from integration._support.client import held +from integration._support.wire import Reply, Request, wire_server +from litellm.integrations.custom_logger import CustomLogger + +_BACKEND: Final = "openai/gpt-4o-mini-tts" +_PROVIDER_KEY: Final = "synthetic-tts-key" +_FIRST: Final = b"audio-frame-0:" + bytes(range(64)) +_REST: Final = ( + b"audio-frame-1:" + bytes(range(64, 192)), + b"audio-frame-2:" + bytes(range(192, 256)), +) +_CLIP: Final = _FIRST + b"".join(_REST) +_TIMEOUT_SECONDS: Final = 20 + + +@dataclass(frozen=True, slots=True) +class Logged: + call_id: object + error_code: object + exception: str + + +def _logged(kwargs: Mapping[str, object]) -> Logged: + payload: Final = kwargs.get("standard_logging_object") + error_information: Final = payload.get("error_information") if isinstance(payload, dict) else None + return Logged( + kwargs.get("litellm_call_id"), + error_information.get("error_code") if isinstance(error_information, dict) else None, + type(kwargs.get("exception")).__name__, + ) + + +class _Probe(CustomLogger): + def __init__(self) -> None: + super().__init__() + self._lock: Final = threading.Lock() + self._events: Final[list[tuple[str, Logged]]] = [] + + def _record(self, kind: str, kwargs: Mapping[str, object]) -> None: + with self._lock: + self._events.append((kind, _logged(kwargs))) + + async def async_log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: object, + end_time: object, + ) -> None: + self._record("async_success", kwargs) + + async def async_log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: object, + end_time: object, + ) -> None: + self._record("async_failure", kwargs) + + def log_success_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: object, + end_time: object, + ) -> None: + self._record("sync_success", kwargs) + + def log_failure_event( + self, + kwargs: dict[str, object], + response_obj: object, + start_time: object, + end_time: object, + ) -> None: + self._record("sync_failure", kwargs) + + def events(self) -> tuple[tuple[str, Logged], ...]: + with self._lock: + return tuple(self._events) + + +def _peer(release: threading.Event | None = None, abort_after: int | None = None) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert (request.method, request.target) == ("POST", "/audio/speech"), request.target + return Reply( + content_type="audio/mpeg", + chunks=(_FIRST, *_REST), + gate_after_first=release, + abort_after=abort_after, + ) + + return respond + + +def _first_bytes(chunks: Iterator[bytes], size: int) -> bytes: + return next(heard for heard in accumulate(chunks) if len(heard) >= size) + + +async def _afirst_bytes(chunks: AsyncIterator[bytes], size: int, heard: bytes = b"") -> bytes: + if len(heard) >= size: + return heard + return await _afirst_bytes(chunks, size, heard + await anext(chunks)) + + +def _probe() -> _Probe: + probe: Final = _Probe() + litellm.callbacks.append(probe) + return probe + + +def _logged_exactly(probe: _Probe, count: int) -> tuple[tuple[str, Logged], ...]: + return held(probe.events, lambda events: len(events) == count, holding=2) + + +async def _aspeech(api_base: str, call_id: str) -> litellm.HttpxBinaryResponseContent: + return await asyncio.wait_for( + litellm.aspeech( + model=_BACKEND, + input=f"streamed speech {call_id}", + voice="alloy", + api_base=api_base, + api_key=_PROVIDER_KEY, + stream=True, + litellm_call_id=call_id, + timeout=_TIMEOUT_SECONDS, + ), + _TIMEOUT_SECONDS, + ) + + +def _speech(api_base: str, call_id: str) -> litellm.HttpxBinaryResponseContent: + return litellm.speech( + model=_BACKEND, + input=f"streamed speech {call_id}", + voice="alloy", + api_base=api_base, + api_key=_PROVIDER_KEY, + stream=True, + litellm_call_id=call_id, + timeout=_TIMEOUT_SECONDS, + ) + + +async def test_aspeech_stream_returns_at_the_first_audio_and_logs_success_once_after_the_clip() -> None: + probe: Final = _probe() + release: Final = threading.Event() + call_id: Final = str(uuid.uuid4()) + with wire_server(_peer(release)) as wire: + response: Final = await _aspeech(wire.url, call_id) + chunks: Final = await response.aiter_bytes() + first: Final = await _afirst_bytes(chunks, len(_FIRST)) + assert probe.events() == () + release.set() + assert first + b"".join([chunk async for chunk in chunks]) == _CLIP + logged: Final = await asyncio.to_thread(_logged_exactly, probe, 1) + assert logged == (("async_success", Logged(call_id, "", "NoneType")),) + + +async def test_aspeech_stream_closed_early_logs_success_once_with_a_499_marker() -> None: + probe: Final = _probe() + release: Final = threading.Event() + call_id: Final = str(uuid.uuid4()) + with wire_server(_peer(release)) as wire: + response: Final = await _aspeech(wire.url, call_id) + assert await _afirst_bytes(await response.aiter_bytes(), len(_FIRST)) == _FIRST + await response.aclose() + logged: Final = await asyncio.to_thread(_logged_exactly, probe, 1) + release.set() + assert logged == (("async_success", Logged(call_id, "499", "NoneType")),) + + +async def test_aspeech_stream_cut_off_raises_a_mapped_error_and_logs_one_failure() -> None: + probe: Final = _probe() + call_id: Final = str(uuid.uuid4()) + with wire_server(_peer(abort_after=1)) as wire: + response: Final = await _aspeech(wire.url, call_id) + chunks: Final = await response.aiter_bytes() + assert await _afirst_bytes(chunks, len(_FIRST)) == _FIRST + with pytest.raises(litellm.APIConnectionError): + await anext(chunks) + logged: Final = await asyncio.to_thread(_logged_exactly, probe, 2) + assert sorted(logged) == [ + (kind, Logged(call_id, "500", "APIConnectionError")) for kind in ("async_failure", "sync_failure") + ] + + +def test_speech_stream_returns_at_the_first_audio_and_logs_success_once_after_the_clip() -> None: + probe: Final = _probe() + release: Final = threading.Event() + call_id: Final = str(uuid.uuid4()) + with wire_server(_peer(release)) as wire: + response: Final = _speech(wire.url, call_id) + chunks: Final = response.iter_bytes() + first: Final = _first_bytes(chunks, len(_FIRST)) + assert probe.events() == () + release.set() + assert first + b"".join(chunks) == _CLIP + logged: Final = _logged_exactly(probe, 1) + assert logged == (("sync_success", Logged(call_id, "", "NoneType")),) + + +def test_speech_stream_cut_off_raises_a_mapped_error_and_logs_one_failure() -> None: + probe: Final = _probe() + call_id: Final = str(uuid.uuid4()) + with wire_server(_peer(abort_after=1)) as wire: + response: Final = _speech(wire.url, call_id) + chunks: Final = response.iter_bytes() + assert _first_bytes(chunks, len(_FIRST)) == _FIRST + with pytest.raises(litellm.APIConnectionError): + next(chunks) + logged: Final = _logged_exactly(probe, 1) + assert logged == (("sync_failure", Logged(call_id, "500", "APIConnectionError")),) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index da1f51faf3f..94b4ff478d4 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -3876,6 +3876,25 @@ async def test_aspeech_stream_closed_mid_clip_marks_the_spend_row_as_a_client_di } +@pytest.mark.usefixtures("local_model_cost_map") +async def test_aspeech_stream_closed_mid_clip_leaves_the_callers_metadata_untouched( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + spend_rows: Final = _SpendRowLogger() + monkeypatch.setattr(litellm, "callbacks", [spend_rows]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + _mock_elevenlabs_audio(respx_mock, _PausedAsyncAudio(b"first", b"never sent", pause_seconds=60)) + caller_metadata: Final = {"user_api_key": "hashed-key"} + + response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=True, metadata=caller_metadata) + audio: Final = await response.aiter_bytes() + assert await anext(audio) == b"first" + await response.aclose() + await asyncio.wait_for(spend_rows.rows.get(), timeout=5) + + assert caller_metadata == {"user_api_key": "hashed-key"} + + @pytest.mark.usefixtures("local_model_cost_map") @pytest.mark.parametrize("stream", [1, "true"]) async def test_aspeech_stream_other_than_true_buffers_the_audio( From 24df2ae9e96985e3f87dc50a61b6fb073287827c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 00:49:29 +0000 Subject: [PATCH 12/20] fix(speech): keep stream keyword-only and buffer the speech health check Passing stream positionally skipped the logging wrapper, and a deployment with stream: true left the health check's provider stream open. Also gives each scripted chat reply its own id so the non-boolean stream cases do not collide on the spend log key, and adds the audit's round-2 cells --- .../health_check_helpers.py | 3 +- litellm/main.py | 1 + .../test_audio_speech_stream_wire.py | 74 +++++++++++++++++-- .../sdk/test_speech_stream_sdk_wire.py | 60 ++++++++++++++- 4 files changed, 128 insertions(+), 10 deletions(-) diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 03af9076c79..2709d474d85 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -254,8 +254,9 @@ class HealthCheckHelpers: ), "audio_speech": lambda: litellm.aspeech( **{ + "voice": "alloy", **_filter_model_params(model_params=model_params), - **({"voice": "alloy"} if "voice" not in _filter_model_params(model_params=model_params) else {}), + "stream": False, }, input=prompt or "test", ), diff --git a/litellm/main.py b/litellm/main.py index eb27f889844..43bcfc8c9e4 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8274,6 +8274,7 @@ def speech( custom_llm_provider: str | None = None, aspeech: bool | None = None, base_url: str | None = None, + *, stream: bool = False, **kwargs, ) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]: diff --git a/tests/integration/providers/test_audio_speech_stream_wire.py b/tests/integration/providers/test_audio_speech_stream_wire.py index 4750676dd45..623364ee797 100644 --- a/tests/integration/providers/test_audio_speech_stream_wire.py +++ b/tests/integration/providers/test_audio_speech_stream_wire.py @@ -319,6 +319,66 @@ def test_buffered_speech_is_unchanged(gateway: Gateway, provider: str, stream: b _assert_billed(gateway, model, call_id, text) +def test_a_deployment_level_stream_true_streams_speech_for_a_caller_that_omits_stream(gateway: Gateway) -> None: + release: Final = threading.Event() + with wire_server(_audio_peer(release)) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["elevenlabs"], stream=True) + text: Final = _text() + call_id: Final = str(uuid.uuid4()) + heard: Final = _listen(gateway, _speech(model, text), call_id, release) + assert (heard.status, heard.body) == (200, _CLIP), heard + _assert_reached(wire, PROVIDERS["elevenlabs"].streamed_path, text) + _assert_billed(gateway, model, call_id, text) + + +def test_a_caller_stream_false_overrides_a_deployment_level_stream_true(gateway: Gateway) -> None: + with wire_server(_audio_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS["elevenlabs"], stream=True) + buffered_twin: Final = _deployment(scenario, wire, PROVIDERS["elevenlabs"]) + text: Final = _text() + call_id: Final = str(uuid.uuid4()) + response: Final = gateway.request( + "POST", + _SPEECH, + _speech(model, text, stream=False), + headers={"x-litellm-call-id": call_id}, + ) + assert (response.status_code, response.content) == (200, _CLIP), response.text[:400] + _assert_reached(wire, PROVIDERS["elevenlabs"].buffered_path, text) + row: Final = _spend_row(call_id) + assert (row["status"], row["call_type"], row["model_group"]) == ("success", "aspeech", model), row + assert row["spend"] == pytest.approx(_buffered_spend(gateway, buffered_twin, text)), row + + +@pytest.mark.parametrize("provider", ["openai", "elevenlabs"]) +def test_a_health_check_on_a_deployment_level_streamed_speech_model_reports_it_healthy( + gateway: Gateway, provider: str +) -> None: + with wire_server(_audio_peer()) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS[provider], stream=True) + response: Final = gateway.request("GET", f"/health?model={model}", None) + assert response.status_code == 200, response.text[:400] + report: Final = object_value(response.json()) + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + assert [request.method for request in wire.drain()] == ["POST"] + + +def _long_audio(request: Request) -> Reply: + return Reply(content_type="audio/mpeg", chunks=(bytes(512 * 1024),) * 128) + + +@pytest.mark.parametrize("provider", ["openai", "elevenlabs"]) +def test_a_health_check_on_a_deployment_level_streamed_speech_model_leaves_no_provider_stream_open( + gateway: Gateway, provider: str +) -> None: + with wire_server(_long_audio) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire, PROVIDERS[provider], stream=True) + response: Final = gateway.request("GET", f"/health?model={model}", None) + assert response.status_code == 200, response.text[:400] + report: Final = object_value(response.json()) + assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report + + def _mistral_peer(request: Request) -> Reply: return Reply(body=json.dumps({"audio_data": base64.b64encode(_CLIP).decode()}).encode()) @@ -607,13 +667,14 @@ def test_a_deployment_max_parallel_requests_slot_is_freed_by_a_hang_up( assert len(wire.drain()) == 2 -_CHAT_ID: Final = "chatcmpl-nonboolstream" _CHAT_USAGE: Final = {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6} -def _chat_chunk(delta: Mapping[str, JsonValue], finish_reason: str | None, **extra: JsonValue) -> bytes: +def _chat_chunk( + chat_id: str, delta: Mapping[str, JsonValue], finish_reason: str | None, **extra: JsonValue +) -> bytes: chunk: Final = { - "id": _CHAT_ID, + "id": chat_id, "object": "chat.completion.chunk", "created": 1, "model": "gpt-4o-mini", @@ -624,17 +685,18 @@ def _chat_chunk(delta: Mapping[str, JsonValue], finish_reason: str | None, **ext def _chat_peer(request: Request) -> Reply: + chat_id: Final = f"chatcmpl-{uuid.uuid4().hex}" if object_value(JSON_OBJECT.validate_json(request.body)).get("stream"): return Reply( content_type="text/event-stream", chunks=( - _chat_chunk({"role": "assistant", "content": "hello"}, None), - _chat_chunk({}, "stop", usage=_CHAT_USAGE), + _chat_chunk(chat_id, {"role": "assistant", "content": "hello"}, None), + _chat_chunk(chat_id, {}, "stop", usage=_CHAT_USAGE), b"data: [DONE]\n\n", ), ) completion: Final = { - "id": _CHAT_ID, + "id": chat_id, "object": "chat.completion", "created": 1, "model": "gpt-4o-mini", diff --git a/tests/integration/sdk/test_speech_stream_sdk_wire.py b/tests/integration/sdk/test_speech_stream_sdk_wire.py index 8f1c9728a55..41d7c4c6dde 100644 --- a/tests/integration/sdk/test_speech_stream_sdk_wire.py +++ b/tests/integration/sdk/test_speech_stream_sdk_wire.py @@ -120,11 +120,17 @@ def _probe() -> _Probe: return probe +def _nothing_logged(events: tuple[tuple[str, Logged], ...]) -> bool: + return events == () + + def _logged_exactly(probe: _Probe, count: int) -> tuple[tuple[str, Logged], ...]: return held(probe.events, lambda events: len(events) == count, holding=2) -async def _aspeech(api_base: str, call_id: str) -> litellm.HttpxBinaryResponseContent: +async def _aspeech( + api_base: str, call_id: str, metadata: Mapping[str, str] | None = None +) -> litellm.HttpxBinaryResponseContent: return await asyncio.wait_for( litellm.aspeech( model=_BACKEND, @@ -133,6 +139,7 @@ async def _aspeech(api_base: str, call_id: str) -> litellm.HttpxBinaryResponseCo api_base=api_base, api_key=_PROVIDER_KEY, stream=True, + metadata=metadata, litellm_call_id=call_id, timeout=_TIMEOUT_SECONDS, ), @@ -161,7 +168,7 @@ async def test_aspeech_stream_returns_at_the_first_audio_and_logs_success_once_a response: Final = await _aspeech(wire.url, call_id) chunks: Final = await response.aiter_bytes() first: Final = await _afirst_bytes(chunks, len(_FIRST)) - assert probe.events() == () + await asyncio.to_thread(held, probe.events, _nothing_logged, holding=2) release.set() assert first + b"".join([chunk async for chunk in chunks]) == _CLIP logged: Final = await asyncio.to_thread(_logged_exactly, probe, 1) @@ -181,6 +188,53 @@ async def test_aspeech_stream_closed_early_logs_success_once_with_a_499_marker() assert logged == (("async_success", Logged(call_id, "499", "NoneType")),) +async def test_aspeech_stream_closed_early_leaves_the_callers_metadata_unmarked_for_its_next_call() -> None: + probe: Final = _probe() + release: Final = threading.Event() + caller_metadata: Final = {"tenant": "caller-owned"} + closed_early: Final = str(uuid.uuid4()) + completed: Final = str(uuid.uuid4()) + with wire_server(_peer(release)) as wire: + first: Final = await _aspeech(wire.url, closed_early, caller_metadata) + assert await _afirst_bytes(await first.aiter_bytes(), len(_FIRST)) == _FIRST + await first.aclose() + await asyncio.to_thread(_logged_exactly, probe, 1) + release.set() + assert caller_metadata.keys() & {"client_disconnected", "error_information"} == set(), caller_metadata + second: Final = await _aspeech(wire.url, completed, caller_metadata) + assert b"".join([chunk async for chunk in await second.aiter_bytes()]) == _CLIP + logged: Final = await asyncio.to_thread(_logged_exactly, probe, 2) + assert caller_metadata["tenant"] == "caller-owned" + assert logged == ( + ("async_success", Logged(closed_early, "499", "NoneType")), + ("async_success", Logged(completed, "", "NoneType")), + ) + + +def test_speech_takes_base_url_in_its_positional_slot() -> None: + probe: Final = _probe() + call_id: Final = str(uuid.uuid4()) + text: Final = f"positional base_url {call_id}" + with wire_server(_peer()) as wire: + positional_through_base_url: Final = ( + _BACKEND, + text, + "alloy", + _PROVIDER_KEY, + *(None,) * 6, + _TIMEOUT_SECONDS, + *(None,) * 7, + wire.url, + ) + response: Final = litellm.speech(*positional_through_base_url, litellm_call_id=call_id) + assert response.content == _CLIP + received: Final = wire.drain() + logged: Final = _logged_exactly(probe, 1) + assert [(request.method, request.target) for request in received] == [("POST", "/audio/speech")] + assert text.encode() in received[0].body + assert logged == (("sync_success", Logged(call_id, "", "NoneType")),) + + async def test_aspeech_stream_cut_off_raises_a_mapped_error_and_logs_one_failure() -> None: probe: Final = _probe() call_id: Final = str(uuid.uuid4()) @@ -204,7 +258,7 @@ def test_speech_stream_returns_at_the_first_audio_and_logs_success_once_after_th response: Final = _speech(wire.url, call_id) chunks: Final = response.iter_bytes() first: Final = _first_bytes(chunks, len(_FIRST)) - assert probe.events() == () + held(probe.events, _nothing_logged, holding=2) release.set() assert first + b"".join(chunks) == _CLIP logged: Final = _logged_exactly(probe, 1) From 25d31cdad17bef86e5bf947f9936e35312552215 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 15 Oct 2026 02:45:05 +0000 Subject: [PATCH 13/20] test(speech): check the health check calls the buffered path --- tests/integration/providers/test_audio_speech_stream_wire.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/integration/providers/test_audio_speech_stream_wire.py b/tests/integration/providers/test_audio_speech_stream_wire.py index 623364ee797..8cc340ee4ce 100644 --- a/tests/integration/providers/test_audio_speech_stream_wire.py +++ b/tests/integration/providers/test_audio_speech_stream_wire.py @@ -360,7 +360,9 @@ def test_a_health_check_on_a_deployment_level_streamed_speech_model_reports_it_h assert response.status_code == 200, response.text[:400] report: Final = object_value(response.json()) assert (report["healthy_count"], report["unhealthy_count"]) == (1, 0), report - assert [request.method for request in wire.drain()] == ["POST"] + assert [(request.method, urlsplit(request.target).path) for request in wire.drain()] == [ + ("POST", PROVIDERS[provider].buffered_path) + ] def _long_audio(request: Request) -> Reply: From 16fe486e604f66c860f2c33f627f4d41cde288ba Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 19 Oct 2026 10:19:11 -0700 Subject: [PATCH 14/20] test(integration): read the speech chaos spend rows from the worker database Main now gives each owned proxy a private database unless the test file opts into the worker's state. The chaos tests read billed rows from the worker database, so they opt in with the shares_worker_state marker like the other providers chaos files --- tests/integration/providers/test_audio_speech_stream_chaos.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/integration/providers/test_audio_speech_stream_chaos.py b/tests/integration/providers/test_audio_speech_stream_chaos.py index 47ed30b3a96..1686c10a1fb 100644 --- a/tests/integration/providers/test_audio_speech_stream_chaos.py +++ b/tests/integration/providers/test_audio_speech_stream_chaos.py @@ -30,7 +30,7 @@ from integration._support.process import ( from integration._support.wire import Reply, Request, Wire, wire_server from pydantic import JsonValue -pytestmark: Final = pytest.mark.timeout(4 * graceful_stop_seconds() + 300) +pytestmark: Final = [pytest.mark.timeout(4 * graceful_stop_seconds() + 300), pytest.mark.shares_worker_state] _MODEL: Final = "chaos-speech" _SPEECH: Final = "/v1/audio/speech" From 55fe2e0bd11b55b011f74cec00bd7d8993b3747e Mon Sep 17 00:00:00 2001 From: Gabriel Creti Date: Thu, 22 Oct 2026 12:22:45 +0200 Subject: [PATCH 15/20] fix(tts): release the router slot when a streamed read fails The deployment slot was only released through Response.aclose, which httpx does not call when the body read raises. A provider cutting a clip off mid-read left the slot held, so later Router.aspeech calls got 429. Release the slot on any stream end, and close the upstream stream when a read fails. Co-Authored-By: Claude --- .../litellm_core_utils/on_end_byte_stream.py | 19 +++++++++--- litellm/router.py | 23 +++++++------- .../test_router_max_parallel_requests.py | 31 +++++++++++++++++++ 3 files changed, 58 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/on_end_byte_stream.py b/litellm/litellm_core_utils/on_end_byte_stream.py index c8b2285cdd8..cccea630b4d 100644 --- a/litellm/litellm_core_utils/on_end_byte_stream.py +++ b/litellm/litellm_core_utils/on_end_byte_stream.py @@ -1,3 +1,4 @@ +import contextlib from collections.abc import AsyncIterator, Awaitable, Callable, Iterator from dataclasses import dataclass from typing import Final, TypeAlias @@ -41,13 +42,18 @@ class OnEndAsyncByteStream(httpx.AsyncByteStream): yield chunk except httpx.TransportError as error: mapped: Final = self._map_error(error) - await self._end(StreamFailed(mapped)) + await self._fail(mapped) raise mapped from error except Exception as error: - await self._end(StreamFailed(error)) + await self._fail(error) raise await self._end(StreamCompleted()) + async def _fail(self, error: Exception) -> None: + with contextlib.suppress(Exception): + await self._stream.aclose() + await self._end(StreamFailed(error)) + async def aclose(self) -> None: try: await self._stream.aclose() @@ -78,13 +84,18 @@ class OnEndSyncByteStream(httpx.SyncByteStream): yield from self._stream except httpx.TransportError as error: mapped: Final = self._map_error(error) - self._end(StreamFailed(mapped)) + self._fail(mapped) raise mapped from error except Exception as error: - self._end(StreamFailed(error)) + self._fail(error) raise self._end(StreamCompleted()) + def _fail(self, error: Exception) -> None: + with contextlib.suppress(Exception): + self._stream.close() + self._end(StreamFailed(error)) + def close(self) -> None: try: self._stream.close() diff --git a/litellm/router.py b/litellm/router.py index 2a2c4e7724a..c88d1f878b7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -109,6 +109,7 @@ from litellm.litellm_core_utils.hidden_params import ( ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.llm_cost_calc.utils import SERVICE_TIER_COST_KEY_SUFFIXES +from litellm.litellm_core_utils.on_end_byte_stream import OnEndAsyncByteStream, StreamEnd from litellm.litellm_core_utils.ptu_pricing import ( PTU_COST_ATTRIBUTION_ENV_VAR, declares_ptu, @@ -623,26 +624,26 @@ def _extend_trailing_text_block(blocks: Sequence[Mapping[str, object]], text: st async def _hold_slot_until_closed(response: object, slot: contextlib.AsyncExitStack) -> None: - """Keep a deployment slot held until a streamed binary response is closed. + """Keep a deployment slot held until a streamed binary response ends. A buffered response already has its body in hand, so its slot is released right - away. A streamed one is still pulling audio from the provider, so the slot has to - survive until whoever consumes it closes the response. + away. A streamed one is still pulling audio from the provider, so the slot is held + until the audio is read to the end, the read fails, or the response is closed. """ raw: Final = getattr(response, "response", None) - if not isinstance(raw, httpx.Response) or raw.is_stream_consumed: + if ( + not isinstance(raw, httpx.Response) + or raw.is_stream_consumed + or not isinstance(raw.stream, httpx.AsyncByteStream) + ): await slot.aclose() return - close: Final = raw.aclose - - async def _aclose_and_release() -> None: - try: - await close() - finally: + async def _release(_: StreamEnd) -> None: + with contextlib.suppress(Exception): # a release error must not mask the stream's own error await slot.aclose() - raw.aclose = _aclose_and_release # pyright: ignore[reportAttributeAccessIssue] # wrapping the bound method keeps every caller's close path releasing the slot + raw.stream = OnEndAsyncByteStream(raw.stream, map_error=lambda error: error, on_end=_release) def _is_fallback_hop(request_kwargs: Mapping[str, object]) -> bool: diff --git a/tests/unit/test_router/test_router_max_parallel_requests.py b/tests/unit/test_router/test_router_max_parallel_requests.py index 866a0843592..01e423ffb07 100644 --- a/tests/unit/test_router/test_router_max_parallel_requests.py +++ b/tests/unit/test_router/test_router_max_parallel_requests.py @@ -205,6 +205,37 @@ async def test_aspeech_holds_its_deployment_slot_until_the_audio_is_released( assert (while_audio_open, await _second_caller()) == ("refused" if stream else "admitted", "admitted") +class _AudioCutOffMidClip(httpx.AsyncByteStream): + async def __aiter__(self): + yield b"audio" + raise httpx.RemoteProtocolError("peer closed connection without sending complete message body") + + async def aclose(self) -> None: + return None + + +@pytest.mark.asyncio +async def test_aspeech_releases_its_deployment_slot_when_the_audio_read_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A provider cutting a streamed clip off mid-read frees the slot even when the caller never closes the response.""" + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _speech_router() + + def _client() -> litellm.AsyncHTTPHandler: + return litellm.AsyncHTTPHandler( + transport=httpx.MockTransport( + lambda _: httpx.Response(200, headers={"content-type": "audio/mpeg"}, stream=_AudioCutOffMidClip()) + ) + ) + + response: Final = await router.aspeech(model="tts", input="hello", voice="voice-id", client=_client(), stream=True) + with pytest.raises(litellm.APIConnectionError): + await response.aread() + + await router.aspeech(model="tts", input="hello", voice="voice-id", client=_client(), stream=True) + + @pytest.mark.asyncio @pytest.mark.parametrize( ("response", "releases_immediately"), From df58d9abf5fc511b445af97bef4f460a5c45b9cd Mon Sep 17 00:00:00 2001 From: Gabriel Creti Date: Thu, 22 Oct 2026 12:22:45 +0200 Subject: [PATCH 16/20] test(tts): drive speech stream tests with events instead of sleeps Co-Authored-By: Claude --- tests/unit/test_utils.py | 56 +++++++++++++++++++++++----------------- 1 file changed, 32 insertions(+), 24 deletions(-) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 9e20667e5d9..4079f24af02 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -3845,29 +3845,30 @@ class _SyncSpeechLogger(CustomLogger): self.logged_payloads.put(kwargs["standard_logging_object"]) -class _PausedAsyncAudio(httpx.AsyncByteStream): - def __init__(self, first: bytes, rest: bytes | Exception, pause_seconds: float) -> None: +class _GatedAsyncAudio(httpx.AsyncByteStream): + def __init__(self, first: bytes, rest: bytes | Exception, gate: asyncio.Event | None = None) -> None: self._first: Final = first self._rest: Final = rest - self._pause_seconds: Final = pause_seconds + self._gate: Final = gate async def __aiter__(self) -> AsyncIterator[bytes]: yield self._first - await asyncio.sleep(self._pause_seconds) + if self._gate is not None: + await self._gate.wait() if isinstance(self._rest, Exception): raise self._rest yield self._rest -class _PausedSyncAudio(httpx.SyncByteStream): - def __init__(self, first: bytes, rest: bytes, pause_seconds: float) -> None: +class _SyncAudio(httpx.SyncByteStream): + def __init__(self, first: bytes, rest: bytes) -> None: self._first: Final = first self._rest: Final = rest - self._pause_seconds: Final = pause_seconds + self.rest_sent_at: float | None = None def __iter__(self) -> Iterator[bytes]: yield self._first - time.sleep(self._pause_seconds) + self.rest_sent_at = time.time() yield self._rest @@ -3897,10 +3898,10 @@ def _audio_summary(num_bytes: int) -> dict[str, object]: return {"object": "binary", "content_type": "audio/mpeg", "num_bytes": num_bytes} -def _logged_seconds(payload: dict[str, object]) -> float: - start_time, end_time = payload["startTime"], payload["endTime"] - assert isinstance(start_time, float) and isinstance(end_time, float) - return end_time - start_time +def _logged_end(payload: dict[str, object]) -> float: + end_time: Final = payload["endTime"] + assert isinstance(end_time, float) + return end_time @pytest.fixture @@ -3936,14 +3937,19 @@ async def test_aspeech_logs_streamed_audio_like_buffered_audio( async def test_aspeech_stream_is_logged_when_the_audio_ends( respx_mock: respx.MockRouter, speech_logger: _SpeechLogger ) -> None: - _mock_elevenlabs_audio(respx_mock, _PausedAsyncAudio(b"first", b"rest", pause_seconds=0.3)) + rest_released: Final = asyncio.Event() + _mock_elevenlabs_audio(respx_mock, _GatedAsyncAudio(b"first", b"rest", rest_released)) response: Final = await _aspeech(stream=True) - assert await response.aread() == b"firstrest" + audio: Final = await response.aiter_bytes() + assert await anext(audio) == b"first" + audio_still_arriving_at: Final = time.time() + rest_released.set() + assert b"".join([chunk async for chunk in audio]) == b"rest" logged: Final = await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=5) assert logged["response"] == _audio_summary(len(b"firstrest")) - assert _logged_seconds(logged) >= 0.3 + assert _logged_end(logged) >= audio_still_arriving_at assert logged["error_information"] != { "error_code": "499", "error_message": "Client disconnected the request", @@ -3970,7 +3976,7 @@ async def test_aspeech_stream_broken_mid_clip_is_billed_as_a_failure( httpx.Response( status_code=200, headers={"content-type": "audio/mpeg"}, - stream=_PausedAsyncAudio(b"first", rest, pause_seconds=0), + stream=_GatedAsyncAudio(b"first", rest), ) for rest in (b"rest", upstream_error) ] @@ -3988,15 +3994,16 @@ async def test_aspeech_stream_broken_mid_clip_is_billed_as_a_failure( assert failed["status"] == "failure" assert isinstance(logged["response_cost"], float) and logged["response_cost"] > 0 assert failed["response_cost"] == logged["response_cost"] - with pytest.raises(TimeoutError): - await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=0.5) + await asyncio.sleep(0) + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + assert speech_logger.logged_payloads.empty() @pytest.mark.usefixtures("local_model_cost_map") async def test_aspeech_stream_closed_mid_clip_is_logged_with_the_audio_sent( respx_mock: respx.MockRouter, speech_logger: _SpeechLogger ) -> None: - _mock_elevenlabs_audio(respx_mock, _PausedAsyncAudio(b"first", b"never sent", pause_seconds=60)) + _mock_elevenlabs_audio(respx_mock, _GatedAsyncAudio(b"first", b"never sent", asyncio.Event())) response: Final = await _aspeech(stream=True) audio: Final = await response.aiter_bytes() @@ -4029,7 +4036,7 @@ async def test_aspeech_stream_closed_mid_clip_marks_the_spend_row_as_a_client_di spend_rows: Final = _SpendRowLogger() monkeypatch.setattr(litellm, "callbacks", [spend_rows]) monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - _mock_elevenlabs_audio(respx_mock, _PausedAsyncAudio(b"first", b"never sent", pause_seconds=60)) + _mock_elevenlabs_audio(respx_mock, _GatedAsyncAudio(b"first", b"never sent", asyncio.Event())) response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=True, metadata={"user_api_key": "hashed-key"}) audio: Final = await response.aiter_bytes() @@ -4052,7 +4059,7 @@ async def test_aspeech_stream_closed_mid_clip_leaves_the_callers_metadata_untouc spend_rows: Final = _SpendRowLogger() monkeypatch.setattr(litellm, "callbacks", [spend_rows]) monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) - _mock_elevenlabs_audio(respx_mock, _PausedAsyncAudio(b"first", b"never sent", pause_seconds=60)) + _mock_elevenlabs_audio(respx_mock, _GatedAsyncAudio(b"first", b"never sent", asyncio.Event())) caller_metadata: Final = {"user_api_key": "hashed-key"} response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=True, metadata=caller_metadata) @@ -4069,7 +4076,7 @@ async def test_aspeech_stream_closed_mid_clip_leaves_the_callers_metadata_untouc async def test_aspeech_stream_other_than_true_buffers_the_audio( respx_mock: respx.MockRouter, speech_logger: _SpeechLogger, stream: object ) -> None: - _mock_elevenlabs_audio(respx_mock, _PausedAsyncAudio(b"first", b"rest", pause_seconds=0)) + _mock_elevenlabs_audio(respx_mock, _GatedAsyncAudio(b"first", b"rest")) response: Final = await litellm.aspeech(**_SPEECH_ARGS, stream=stream) logged: Final = await asyncio.wait_for(speech_logger.logged_payloads.get(), timeout=5) @@ -4082,7 +4089,8 @@ async def test_aspeech_stream_other_than_true_buffers_the_audio( def test_speech_stream_is_logged_when_the_audio_ends( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch ) -> None: - _mock_elevenlabs_audio(respx_mock, _PausedSyncAudio(b"first", b"rest", pause_seconds=0.3)) + audio: Final = _SyncAudio(b"first", b"rest") + _mock_elevenlabs_audio(respx_mock, audio) speech_logger: Final = _SyncSpeechLogger() monkeypatch.setattr(litellm, "callbacks", [speech_logger]) @@ -4091,7 +4099,7 @@ def test_speech_stream_is_logged_when_the_audio_ends( logged: Final = speech_logger.logged_payloads.get(timeout=5) assert logged["response"] == _audio_summary(len(b"firstrest")) - assert _logged_seconds(logged) >= 0.3 + assert audio.rest_sent_at is not None and _logged_end(logged) >= audio.rest_sent_at class TestCallbackAsyncSyncSeparation: From 93ada24dfadd1ed09e0b9721eedaf968909c3abc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 3 Nov 2026 08:18:53 -0800 Subject: [PATCH 17/20] test(utils): capture streamed sync speech usage when the audio ends, not at the headers --- tests/unit/test_utils.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index cb9dc6bb16d..d2b8f909c86 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -4212,6 +4212,25 @@ def test_speech_stream_is_logged_when_the_audio_ends( assert audio.rest_sent_at is not None and _logged_end(logged) >= audio.rest_sent_at +@pytest.mark.usefixtures("local_model_cost_map") +def test_speech_stream_usage_is_captured_when_the_audio_ends( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + audio: Final = _SyncAudio(b"first", b"rest") + _mock_elevenlabs_audio(respx_mock, audio) + captured_at: Final[queue.Queue[float]] = queue.Queue() + monkeypatch.setattr( + litellm_logging.Logging, "usage_capture_callback", lambda *_: captured_at.put(time.time()) + ) + + response: Final = litellm.speech(**_SPEECH_ARGS, stream=True) + captured_before_read: Final = captured_at.qsize() + assert response.read() == b"firstrest" + + assert captured_before_read == 0 + assert audio.rest_sent_at is not None and captured_at.get(timeout=5) >= audio.rest_sent_at + + class TestCallbackAsyncSyncSeparation: """Test that LoggingCallbackManager auto-routes async callbacks to async lists.""" From 94fd515f9183ff43ef3255e8f98e26938ab26bf0 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 4 Nov 2026 13:08:18 -0800 Subject: [PATCH 18/20] 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 --- litellm/proxy/proxy_server.py | 110 +++++++++--------- .../proxy/proxy_server/test_routes_audio.py | 65 ++++++++++- 2 files changed, 119 insertions(+), 56 deletions(-) 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 From 4e0ee90bfea7c09d037ff1ee0dd7b1e993c2994d Mon Sep 17 00:00:00 2001 From: Gabriel Creti Date: Thu, 5 Nov 2026 11:28:35 +0100 Subject: [PATCH 19/20] fix(proxy): end streamed speech as failed when the route errors before responding Closing the unread audio on a pre-response error went through aclose, which always ends as StreamClosedEarly: the clip was success-logged with a 499 ClientDisconnected marker while the route returned 500. The route now ends the audio with the error, through every OnEndAsyncByteStream wrapper, so it is logged as a failure. A cancellation still closes it as a hang-up. Co-Authored-By: Claude --- .../litellm_core_utils/on_end_byte_stream.py | 11 ++++--- litellm/proxy/proxy_server.py | 12 +++++++ .../proxy/proxy_server/test_routes_audio.py | 31 ++++++++++++++++++- 3 files changed, 49 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/on_end_byte_stream.py b/litellm/litellm_core_utils/on_end_byte_stream.py index cccea630b4d..c8559329189 100644 --- a/litellm/litellm_core_utils/on_end_byte_stream.py +++ b/litellm/litellm_core_utils/on_end_byte_stream.py @@ -42,16 +42,19 @@ class OnEndAsyncByteStream(httpx.AsyncByteStream): yield chunk except httpx.TransportError as error: mapped: Final = self._map_error(error) - await self._fail(mapped) + await self.afail(mapped) raise mapped from error except Exception as error: - await self._fail(error) + await self.afail(error) raise await self._end(StreamCompleted()) - async def _fail(self, error: Exception) -> None: + async def afail(self, error: Exception) -> None: with contextlib.suppress(Exception): - await self._stream.aclose() + if isinstance(self._stream, OnEndAsyncByteStream): + await self._stream.afail(error) + else: + await self._stream.aclose() await self._end(StreamFailed(error)) async def aclose(self) -> None: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e6c6722e21b..28d7f8d0640 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -903,6 +903,7 @@ from starlette.background import BackgroundTask from starlette.exceptions import HTTPException as StarletteHTTPException from starlette.websockets import WebSocketState +from litellm.litellm_core_utils.on_end_byte_stream import OnEndAsyncByteStream from litellm.proxy import tracing_endpoints from litellm.proxy.middleware.admission_control_middleware import ( AdmissionControlMiddleware, @@ -13679,6 +13680,14 @@ async def _close_audio_speech_response(_response: HttpxBinaryResponseContent) -> await _response.aclose() +async def _fail_audio_speech_response(_response: HttpxBinaryResponseContent, error: Exception) -> None: + """Release the audio as failed, so a route error before any byte is sent is not logged as a caller hang-up.""" + with contextlib.suppress(Exception): + if isinstance(stream := _response.response.stream, OnEndAsyncByteStream): + await stream.afail(error) + await _close_audio_speech_response(_response) + + @router.post( "/v1/audio/speech", dependencies=[Depends(user_api_key_auth)], @@ -13799,6 +13808,9 @@ async def audio_speech( headers=custom_headers, background=BackgroundTask(_close_audio_speech_response, audio), ) + except Exception as error: + await _fail_audio_speech_response(response, error) + raise except BaseException: await _close_audio_speech_response(response) raise diff --git a/tests/unit/proxy/proxy_server/test_routes_audio.py b/tests/unit/proxy/proxy_server/test_routes_audio.py index 9e160399c78..8ed6179aa7a 100644 --- a/tests/unit/proxy/proxy_server/test_routes_audio.py +++ b/tests/unit/proxy/proxy_server/test_routes_audio.py @@ -18,6 +18,7 @@ import math import wave from collections.abc import AsyncIterator, Awaitable, Callable, Mapping from contextlib import AbstractContextManager +from functools import partial from typing import Final from unittest.mock import AsyncMock, MagicMock @@ -37,6 +38,7 @@ from litellm import Router from litellm.caching.caching import Cache from litellm.constants import AUDIO_SPEECH_CHUNK_SIZE from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.litellm_core_utils.on_end_byte_stream import OnEndAsyncByteStream, StreamEnd, StreamFailed from litellm.proxy import proxy_server from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler @@ -478,7 +480,7 @@ def _speech_request(payload: dict) -> Request: ) -def _audio_response(stream: _RecordingAudioStream) -> HttpxBinaryResponseContent: +def _audio_response(stream: httpx.AsyncByteStream) -> HttpxBinaryResponseContent: return HttpxBinaryResponseContent( httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}, stream=stream) ) @@ -645,6 +647,33 @@ async def test_audio_speech_route_frees_the_deployment_slot_when_a_response_hook assert first_upstream.closed is True +@pytest.mark.asyncio +async def test_audio_speech_route_ends_the_audio_as_failed_when_a_response_hook_fails( + monkeypatch: pytest.MonkeyPatch, +): + """A route error before any audio is sent is a failure, not a caller hang-up, through every + wrapper the router and the SDK put around the audio.""" + upstream: Final = _RecordingAudioStream([b"audio"], []) + ends: Final[list[tuple[str, StreamEnd]]] = [] + + async def _record(layer: str, end: StreamEnd) -> None: + ends.append((layer, end)) + + sdk_stream: Final = OnEndAsyncByteStream(upstream, map_error=lambda e: e, on_end=partial(_record, "sdk")) + router_stream: Final = OnEndAsyncByteStream(sdk_stream, map_error=lambda e: e, on_end=partial(_record, "router")) + _route_speech_to(monkeypatch, _audio_response(router_stream), response_headers_hook=_headers_from_the_audio_body) + + with pytest.raises(ProxyException): + await proxy_server.audio_speech( + request=_speech_request({"model": "tts-1", "input": "Hi", "voice": "alloy", "stream": True}), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert [(layer, type(end)) for layer, end in ends] == [("sdk", StreamFailed), ("router", StreamFailed)] + assert upstream.closed is True + + class _BreaksMidClipAudioStream(httpx.AsyncByteStream): async def __aiter__(self): yield b"a" * AUDIO_SPEECH_CHUNK_SIZE From a7e632cbc57a935ac84344ce03a5c9c1b460ba0f Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 5 Nov 2026 19:15:48 -0800 Subject: [PATCH 20/20] fix(proxy): close unread speech audio when a deployment guardrail rejects it A guardrail's post-call deployment hook runs before the audio stream's end logger is attached, so a rejection left the unread provider response open and held its connection. The failure path of the async wrapper now closes it, and a close that fails never replaces the rejection Also annotates the new test locals in test_routes_audio.py with Final --- litellm/utils.py | 9 +++++ .../sdk/test_speech_stream_sdk_wire.py | 24 ++++++++++++- .../proxy/proxy_server/test_routes_audio.py | 24 ++++++------- tests/unit/test_utils.py | 36 +++++++++++++++++++ 4 files changed, 80 insertions(+), 13 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index d4d2b61f5d3..4a72d4f05ff 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7,6 +7,7 @@ import ast import asyncio import base64 import binascii +import contextlib import contextvars import copy import datetime @@ -1674,6 +1675,13 @@ def _unread_speech_upstream(result: object, call_type: str, stream_requested: bo return result.unread_response() +async def _close_unread_speech(result: object, call_type: str, stream_requested: bool) -> None: + unread_speech: Final = _unread_speech_upstream(result, call_type, stream_requested) + if unread_speech is not None: + with contextlib.suppress(Exception): + await unread_speech.aclose() + + def _bill_failed_speech_stream(logging_obj: LiteLLMLoggingObject, result: object) -> None: logging_obj.record_partial_usage_for_failure(Usage(), logging_obj.response_cost_calculator(result=result) or 0.0) @@ -2770,6 +2778,7 @@ def client(original_function): return _replayed_when_run_unstreamed(result, transcription_runs_unstreamed) except Exception as e: traceback_exception: Final = traceback.format_exc() + await _close_unread_speech(result, call_type, "stream" in kwargs and kwargs["stream"] is True) # Reuse the timestamp taken right when the deployment call itself failed, before # the failure hook ran, so a slow callback doesn't inflate the reported duration. end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with diff --git a/tests/integration/sdk/test_speech_stream_sdk_wire.py b/tests/integration/sdk/test_speech_stream_sdk_wire.py index 41d7c4c6dde..3dcd29f1637 100644 --- a/tests/integration/sdk/test_speech_stream_sdk_wire.py +++ b/tests/integration/sdk/test_speech_stream_sdk_wire.py @@ -8,9 +8,11 @@ from typing import Final import litellm import pytest -from integration._support.client import held +from integration._support.client import eventually, held from integration._support.wire import Reply, Request, wire_server +from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import CallTypes, LLMResponseTypes _BACKEND: Final = "openai/gpt-4o-mini-tts" _PROVIDER_KEY: Final = "synthetic-tts-key" @@ -211,6 +213,26 @@ async def test_aspeech_stream_closed_early_leaves_the_callers_metadata_unmarked_ ) +class _SpeechRejected(Exception): + pass + + +class _RejectingGuardrail(CustomGuardrail): + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: LLMResponseTypes, call_type: CallTypes | None + ) -> None: + raise _SpeechRejected + + +async def test_aspeech_stream_rejected_by_a_deployment_guardrail_closes_the_provider_connection() -> None: + litellm.callbacks.append(_RejectingGuardrail()) + with wire_server(_peer(), keep_alive=True) as wire: + with pytest.raises(_SpeechRejected): + await _aspeech(wire.url, str(uuid.uuid4())) + await asyncio.to_thread(eventually, wire.open_connections, lambda still_open: still_open == 0) + assert wire.connections() == 1 + + def test_speech_takes_base_url_in_its_positional_slot() -> None: probe: Final = _probe() call_id: Final = str(uuid.uuid4()) diff --git a/tests/unit/proxy/proxy_server/test_routes_audio.py b/tests/unit/proxy/proxy_server/test_routes_audio.py index 8ed6179aa7a..0782de39da0 100644 --- a/tests/unit/proxy/proxy_server/test_routes_audio.py +++ b/tests/unit/proxy/proxy_server/test_routes_audio.py @@ -447,14 +447,14 @@ class _RecordingAudioStream(httpx.AsyncByteStream): async def test_audio_speech_chunk_generator_forwards_an_unread_upstream_stream(): """A streaming upstream response reaches the client piece by piece as it arrives, with no wait for a full AUDIO_SPEECH_CHUNK_SIZE before the first audio goes out.""" - produced: list[bytes] = [] + produced: Final[list[bytes]] = [] first, second = b"a" * (AUDIO_SPEECH_CHUNK_SIZE // 8), b"b" * AUDIO_SPEECH_CHUNK_SIZE - stream = _RecordingAudioStream([first, second], produced) - response = HttpxBinaryResponseContent( + stream: Final = _RecordingAudioStream([first, second], produced) + response: Final = HttpxBinaryResponseContent( httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}, stream=stream) ) - chunks = proxy_server._audio_speech_chunk_generator(response, on_upstream_error=AsyncMock()) + chunks: Final = proxy_server._audio_speech_chunk_generator(response, on_upstream_error=AsyncMock()) assert await anext(chunks) == first assert produced == [first] @@ -536,7 +536,7 @@ def _route_speech_to( @pytest.fixture def patched_streaming_speech(monkeypatch): """Route /v1/audio/speech to an upstream whose audio body is left unread.""" - stream = _RecordingAudioStream([b"a" * AUDIO_SPEECH_CHUNK_SIZE], []) + stream: Final = _RecordingAudioStream([b"a" * AUDIO_SPEECH_CHUNK_SIZE], []) _route_speech_to(monkeypatch, _audio_response(stream)) yield stream @@ -546,8 +546,8 @@ async def test_audio_speech_releases_its_upstream_when_the_client_walks_away(): """A client that drops mid-clip still has its upstream connection released. Starlette raises ClientDisconnect out of the response before any background task runs, so the body iterator's own cleanup is what has to cover this.""" - stream = _RecordingAudioStream([b"a" * AUDIO_SPEECH_CHUNK_SIZE] * 3, []) - response = _audio_response(stream) + stream: Final = _RecordingAudioStream([b"a" * AUDIO_SPEECH_CHUNK_SIZE] * 3, []) + response: Final = _audio_response(stream) await _serve_until_disconnect( StreamingResponse( @@ -567,19 +567,19 @@ async def test_audio_speech_releases_its_upstream_when_the_client_walks_away(): async def test_audio_speech_route_releases_an_upstream_whose_body_is_never_read(monkeypatch, patched_streaming_speech): """A response the server never iterates leaves the body iterator's cleanup unrun and out of reach of GC, so the route itself has to arm the release.""" - stream = patched_streaming_speech - captured: list[StreamingResponse] = [] + stream: Final = patched_streaming_speech + captured: Final[list[StreamingResponse]] = [] - original = proxy_server.StreamingResponse + original: Final = proxy_server.StreamingResponse def _capture(*args, **kwargs): - response = original(*args, **kwargs) + response: Final = original(*args, **kwargs) captured.append(response) return response monkeypatch.setattr(proxy_server, "StreamingResponse", _capture) - request = _speech_request({"model": "tts-1", "input": "Hi", "voice": "alloy", "stream": True}) + request: Final = _speech_request({"model": "tts-1", "input": "Hi", "voice": "alloy", "stream": True}) await proxy_server.audio_speech( request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test") ) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 35c53f0f030..481ee065471 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -4229,6 +4229,42 @@ async def test_aspeech_stream_other_than_true_buffers_the_audio( assert logged["response"] == _audio_summary(len(b"firstrest")) +class _SpeechRejected(Exception): + pass + + +class _RejectingGuardrail(CustomGuardrail): + async def async_post_call_success_deployment_hook(self, request_data, response, call_type) -> None: + raise _SpeechRejected + + +class _AudioThatFailsToClose(httpx.AsyncByteStream): + def __init__(self) -> None: + self.close_attempts: int = 0 + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"first" + + async def aclose(self) -> None: + self.close_attempts += 1 + raise httpx.ReadError("the provider connection broke while closing") + + +@pytest.mark.usefixtures("local_model_cost_map") +async def test_aspeech_stream_failing_to_close_keeps_the_guardrail_rejection( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + audio: Final = _AudioThatFailsToClose() + monkeypatch.setattr(litellm, "callbacks", [_RejectingGuardrail()]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + _mock_elevenlabs_audio(respx_mock, audio) + + with pytest.raises(_SpeechRejected): + await _aspeech(stream=True) + + assert audio.close_attempts == 1 + + @pytest.mark.usefixtures("local_model_cost_map") def test_speech_stream_is_logged_when_the_audio_ends( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch