scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-29 06:33:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 200

fix(tts): stream only on a literal true, mark hang-ups, and cut off without a fake DB outage

Browse source
This commit is contained in:
mateo-berri 2026-10-14 18:58:39 +00:00
parent 7f09947a65
commit ab48060051
6 changed files with 140 additions and 53 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
litellm/litellm_core_utils/on_end_byte_stream.py
View file

@ -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)

13
litellm/main.py
View file

@ -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,
)

23
litellm/proxy/proxy_server.py
View file

@ -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)

80
litellm/utils.py
View file

@ -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)

4
tests/unit/proxy/proxy_server/test_routes_audio.py
View file

@ -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)

24
tests/unit/test_utils.py
View file

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