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] 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")