scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

Merge 94fd515f91 into 1967b41368

Browse source
This commit is contained in:
Gabriel Creti 2026-11-04 21:23:52 +00:00 • committed by GitHub
commit e759a1c68d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
25 changed files with 2945 additions and 101 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/litellm_core_utils/health_check_helpers.py
View file

@ -323,8 +323,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",
),

12
litellm/litellm_core_utils/logging_utils.py
View file

@ -22,6 +22,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
@ -341,6 +342,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."""

109
litellm/litellm_core_utils/on_end_byte_stream.py Normal file
View file

@ -0,0 +1,109 @@
import contextlib
from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
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[[StreamEnd], Awaitable[None]],
) -> None:
self._stream: Final = stream
self._map_error: Final = map_error
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 httpx.TransportError as error:
mapped: Final = self._map_error(error)
await self._fail(mapped)
raise mapped from error
except Exception as 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()
finally:
await self._end(StreamClosedEarly())
async def _end(self, end: StreamEnd) -> None:
if self._ended:
return
self._ended = True
await self._on_end(end)
class OnEndSyncByteStream(httpx.SyncByteStream):
def __init__(
self,
stream: httpx.SyncByteStream,
map_error: Callable[[httpx.TransportError], Exception],
on_end: Callable[[StreamEnd], 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._fail(mapped)
raise mapped from error
except Exception as 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()
finally:
self._end(StreamClosedEarly())
def _end(self, end: StreamEnd) -> None:
if self._ended:
return
self._ended = True
self._on_end(end)

4
litellm/llms/aws_polly/text_to_speech/transformation.py
View file

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

34
litellm/llms/azure/azure.py
View file

@ -18,7 +18,11 @@ from pydantic import BaseModel
import litellm
from litellm.constants import 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.llms.custom_httpx.http_handler import (
AsyncHTTPHandler,
HTTPHandler,
@ -1220,6 +1224,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)
@ -1239,6 +1244,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(
@ -1260,13 +1266,13 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
},
)
response: Final = azure_client.audio.speech.create(
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
else azure_client.audio.speech.create(**params).response
)
return HttpxBinaryResponseContent(response=response.response)
return HttpxBinaryResponseContent(response=raw_response)
async def async_audio_speech(
self,
@ -1284,6 +1290,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,
@ -1304,14 +1311,13 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM):
},
)
azure_response: Final = await azure_client.audio.speech.create(
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
else (await azure_client.audio.speech.create(**params)).response
)
return HttpxBinaryResponseContent(response=azure_response.response)
return HttpxBinaryResponseContent(response=raw_response)
def get_headers(
self,

4
litellm/llms/azure/text_to_speech/transformation.py
View file

@ -34,6 +34,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"
@ -74,6 +76,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig):
aspeech: bool,
api_base: str | None,
api_key: str | None,
stream: bool = False,
**kwargs: object,
) -> Union[
"HttpxBinaryResponseContent",
@ -130,6 +133,7 @@ class AzureAVATextToSpeechConfig(BaseTextToSpeechConfig):
extra_headers=extra_headers,
client=None,
_is_async=aspeech,
stream=stream,
)
return response

4
litellm/llms/base_llm/text_to_speech/transformation.py
View file

@ -1,7 +1,7 @@
import types
from abc import ABC, abstractmethod
from collections.abc import Mapping
from typing import TYPE_CHECKING, Any, Final, TypedDict
from typing import TYPE_CHECKING, Any, ClassVar, Final, TypedDict
import httpx
@ -36,6 +36,8 @@ class TextToSpeechRequestData(TypedDict, total=False):
class BaseTextToSpeechConfig(ABC):
supports_streaming: ClassVar[bool] = False
def __init__(self):
pass

11
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -12594,6 +12594,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"],
@ -12615,6 +12616,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):
@ -12662,6 +12664,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:
@ -12670,6 +12674,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(
@ -12677,6 +12682,7 @@ class BaseLLMHTTPHandler:
headers=headers,
data=request_data["ssml_body"],
timeout=timeout,
stream=should_stream,
)
else:
raise ValueError(
@ -12710,6 +12716,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.
@ -12763,6 +12770,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:
@ -12771,6 +12780,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(
@ -12778,6 +12788,7 @@ class BaseLLMHTTPHandler:
headers=headers,
data=request_data["ssml_body"],
timeout=timeout,
stream=should_stream,
)
else:
raise ValueError(

7
litellm/llms/elevenlabs/text_to_speech/transformation.py
View file

@ -37,8 +37,11 @@ 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"
TTS_STREAM_PATH_SUFFIX = "/stream"
DEFAULT_OUTPUT_FORMAT = "pcm_44100"
VOICE_MAPPINGS = {
"alloy": "21m00Tcm4TlvDq8ikWAM", # Rachel
@ -63,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:
"""
@ -304,7 +308,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 = join_api_base(base_url, f"{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 = join_api_base(base_url, f"{self.TTS_ENDPOINT_PATH}/{encoded_voice_id}{stream_suffix}")
query_params: Final = litellm_params.get(self.ELEVENLABS_QUERY_PARAMS_KEY, {})
if query_params:

37
litellm/llms/openai/openai.py
View file

@ -36,7 +36,11 @@ from litellm.litellm_core_utils.chat_completion_agentic_loop import (
)
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, MockResponseIterator
from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException
from litellm.types.utils import (
@ -1666,6 +1670,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
aspeech: bool | None = None,
client=None,
shared_session: Optional["ClientSession"] = None,
stream: bool = False,
litellm_params: Mapping[str, object] | None = None,
) -> HttpxBinaryResponseContent:
if aspeech is not None and aspeech is True:
@ -1683,6 +1688,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
logging_obj=logging_obj,
client=client,
shared_session=shared_session,
stream=stream,
litellm_params=litellm_params,
)
@ -1707,14 +1713,14 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
},
)
response: Final = sync_client.audio.speech.create(
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
else sync_client.audio.speech.create(**params).response
)
speech_response: Final = HttpxBinaryResponseContent(response=response.response)
set_provider_response_headers_in_hidden_params(speech_response, response.response.headers)
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(
@ -1732,6 +1738,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
logging_obj: LiteLLMLoggingObj,
client=None,
shared_session: Optional["ClientSession"] = None,
stream: bool = False,
litellm_params: Mapping[str, object] | None = None,
) -> HttpxBinaryResponseContent:
openai_client: Final = cast(
@ -1757,14 +1764,14 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM):
},
)
response: Final = await openai_client.audio.speech.create(
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
else (await openai_client.audio.speech.create(**params)).response
)
speech_response: Final = HttpxBinaryResponseContent(response=response.response)
set_provider_response_headers_in_hidden_params(speech_response, response.response.headers)
speech_response: Final = HttpxBinaryResponseContent(response=raw_response)
set_provider_response_headers_in_hidden_params(speech_response, raw_response.headers)
return speech_response

10
litellm/main.py
View file

@ -8887,6 +8887,8 @@ 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]:
user: Final = kwargs.get("user", None)
@ -8900,6 +8902,7 @@ def speech(
model=model, custom_llm_provider=custom_llm_provider, api_base=caller_api_base
)
kwargs.pop("tags", [])
stream_audio: Final = stream is True
optional_params = {}
if response_format is not None:
@ -9024,6 +9027,7 @@ def speech(
client=client, # pass AsyncOpenAI, OpenAI client
aspeech=aspeech,
shared_session=shared_session,
stream=stream_audio,
litellm_params=litellm_params_dict,
)
elif custom_llm_provider in AZURE_OPENAI_AUDIO_PROVIDERS:
@ -9057,6 +9061,7 @@ def speech(
aspeech=aspeech or False,
api_base=api_base,
api_key=api_key,
stream=stream_audio,
**kwargs,
)
else:
@ -9104,6 +9109,7 @@ def speech(
client=client, # pass AsyncOpenAI, OpenAI client
aspeech=aspeech,
litellm_params=litellm_params_dict,
stream=stream_audio,
)
elif custom_llm_provider == "elevenlabs":
from litellm.llms.elevenlabs.text_to_speech.transformation import (
@ -9129,6 +9135,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_audio:
litellm_params_dict[ElevenLabsTextToSpeechConfig.ELEVENLABS_STREAM_KEY] = True
if api_base is not None:
litellm_params_dict["api_base"] = api_base
@ -9148,6 +9156,7 @@ def speech(
extra_headers=extra_headers,
client=client,
_is_async=aspeech or False,
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 (
@ -9349,6 +9358,7 @@ def speech(
aspeech=aspeech or False,
api_base=api_base,
api_key=api_key,
stream=stream_audio,
**kwargs,
)

155
litellm/proxy/proxy_server.py
View file

@ -899,6 +899,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
@ -13628,16 +13629,54 @@ 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]],
) -> 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)
async for chunk in _generator:
yield chunk
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 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)
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):
await _response.aclose()
@router.post(
@ -13662,7 +13701,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()
@ -13701,52 +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,
)
return StreamingResponse(
_audio_speech_chunk_generator(response),
media_type=media_type,
headers=custom_headers,
)
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),
)
except BaseException:
await _close_audio_speech_response(response)
raise
except Exception as e:
await proxy_logging_obj.post_call_failure_hook(

37
litellm/router_python_core.py
View file

@ -158,6 +158,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.pattern_match_deployments import PatternMatchRouter
from litellm.litellm_core_utils.ptu_pricing import (
PTU_COST_ATTRIBUTION_ENV_VAR,
@ -704,6 +705,29 @@ def _extend_trailing_text_block(blocks: Sequence[Mapping[str, object]], text: st
return [*blocks, {"type": "text", "text": text}]
async def _hold_slot_until_closed(response: object, slot: contextlib.AsyncExitStack) -> None:
"""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 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
or not isinstance(raw.stream, httpx.AsyncByteStream)
):
await slot.aclose()
return
async def _release(_: StreamEnd) -> None:
with contextlib.suppress(Exception): # a release error must not mask the stream's own error
await slot.aclose()
raw.stream = OnEndAsyncByteStream(raw.stream, map_error=lambda error: error, on_end=_release)
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
@ -6328,6 +6352,10 @@ class PythonCore(DeprecatedRouterMembers):
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:
```
@ -6399,7 +6427,10 @@ class PythonCore(DeprecatedRouterMembers):
)
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 self._host.invoke("aspeech")(
**{
**data,
@ -6409,6 +6440,10 @@ class PythonCore(DeprecatedRouterMembers):
**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)

3
litellm/types/llms/openai.py
View file

@ -169,6 +169,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)

190
litellm/utils.py
View file

@ -75,6 +75,7 @@ from litellm.constants import (
HUGGINGFACE_TOKENIZER_LOAD_RETRY_SECONDS,
INITIAL_RETRY_DELAY,
JITTER,
LITELLM_HTTP_STATUS_CLIENT_DISCONNECTED,
MAX_RETRY_DELAY,
MAX_TOKEN_TRIMMING_ATTEMPTS,
MINIMUM_PROMPT_CACHE_TOKEN_COUNT_OVERRIDE,
@ -101,6 +102,14 @@ from litellm.litellm_core_utils.get_model_cost_map import AliasedCostEntry, GetM
from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases
from litellm.litellm_core_utils.model_retirement import is_hidden_as_retired
from litellm.litellm_core_utils.offline_mode import is_offline_mode
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.litellm_core_utils.tool_fields import (
@ -239,6 +248,7 @@ from litellm.types.llms.openai import (
ChatCompletionNamedToolChoiceParam,
ChatCompletionToolParam,
ChatCompletionToolParamFunctionChunk,
HttpxBinaryResponseContent,
OpenAITextCompletionUserMessage,
OpenAIWebSearchOptions,
)
@ -275,6 +285,7 @@ from litellm.types.utils import (
SandboxProviders,
SearchProviders,
SelectTokenizerResponse,
StandardLoggingPayloadErrorInformation,
StreamingChoices,
TranscriptionResponse,
Usage,
@ -316,7 +327,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, Sized
from collections.abc import AsyncIterator, Awaitable, Callable, Collection, Iterable, Iterator, Mapping, Sequence, Sized
from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, runtime_checkable
from typing_extensions import assert_never
@ -1653,6 +1664,119 @@ def dispatch_success_logging(
)
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
_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[Mapping[str, object] | None] = logging_obj.model_call_details.get("litellm_params")
if not isinstance(litellm_params, dict):
return
request_metadata: Final = tuple(
(key, metadata)
for key in ("metadata", "litellm_metadata")
if isinstance(metadata := litellm_params.get(key), dict)
)
for key, existing in request_metadata or (("metadata", {}),):
litellm_params[key] = {
**existing,
"client_disconnected": True,
"error_information": dict(_CLIENT_DISCONNECTED_ERROR_INFORMATION),
}
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[[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
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
def _sync_speech_end_logger(
logging_obj: LiteLLMLoggingObject,
result: object,
start_time: datetime.datetime,
is_litellm_internal_call: bool,
executor: BoundedLoggingThreadPoolExecutor,
) -> Callable[[StreamEnd], None]:
ctx: Final = contextvars.copy_context()
def log_speech_end(end: StreamEnd) -> None:
end_time: Final = datetime.datetime.now() # noqa: DTZ005 # matches the naive start_time the wrapper took
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)
return log_speech_end
def _schedule_async_success_logging(
logging_obj: LiteLLMLoggingObject,
result: object,
@ -2274,11 +2398,26 @@ 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] = litellm_utils.executor
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,
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] = litellm_utils.executor
executor.submit(
ctx.run,
logging_obj.success_handler,
@ -2603,14 +2742,30 @@ 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, "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,
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 _replayed_when_run_unstreamed(result, transcription_runs_unstreamed)
except Exception as e:
@ -2767,7 +2922,16 @@ _STREAMING_CALL_TYPES: Final = frozenset(
}
)
_NEVER_STREAMING_CALL_TYPES: Final = frozenset(
_SPEECH_CALL_TYPES: Final = frozenset(
{
CallTypes.speech,
CallTypes.aspeech,
CallTypes.speech.value,
CallTypes.aspeech.value,
}
)
_NEVER_STREAMING_CALL_TYPES: Final = _SPEECH_CALL_TYPES | frozenset(
{
CallTypes.extraction,
CallTypes.aextraction,
@ -2788,6 +2952,10 @@ 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 and extraction are excluded. Extraction never streams, and "stream=True" on speech asks
for the audio bytes to be forwarded as they arrive while the return value is still a single
response object, so both are logged like any other non-streaming call.
"""
if call_type in _NEVER_STREAMING_CALL_TYPES or call_type in _RESULT_DECIDES_STREAMING_CALL_TYPES:
return False

312
tests/integration/providers/test_audio_speech_stream_chaos.py Normal file
View file

@ -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), pytest.mark.shares_worker_state]
_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])

742
tests/integration/providers/test_audio_speech_stream_wire.py Normal file
View file

@ -0,0 +1,742 @@
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 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, urlsplit(request.target).path) for request in wire.drain()] == [
("POST", PROVIDERS[provider].buffered_path)
]
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())
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_USAGE: Final = {"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6}
def _chat_chunk(
chat_id: str, 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:
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(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,
"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

144
tests/integration/providers/test_provider_passthrough_stream_relay_wire.py Normal file
View file

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

278
tests/integration/sdk/test_speech_stream_sdk_wire.py Normal file
View file

@ -0,0 +1,278 @@
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 _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, metadata: Mapping[str, str] | None = None
) -> 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,
metadata=metadata,
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))
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)
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_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())
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))
held(probe.events, _nothing_logged, holding=2)
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")),)

98
tests/unit/llms/custom_httpx/test_llm_http_handler.py
View file

@ -5895,6 +5895,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()

31
tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py
View file

@ -57,6 +57,37 @@ def test_speech_keeps_an_internal_prefixed_kwarg_out_of_the_elevenlabs_request(r
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"
@pytest.mark.parametrize(
("response_format", "expected_output_format"),
[("wav", "wav_24000"), ("pcm", "pcm_24000")],

105
tests/unit/llms/openai/test_openai_audio_speech.py Normal file
View file

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

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

@ -9,30 +9,37 @@ Pins (PR2):
from __future__ import annotations
import asyncio
import contextlib
import gc
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
import httpx
import orjson
import pytest
import respx
from fastapi import HTTPException
from fastapi import HTTPException, Request, Response
from fastapi.testclient import TestClient
from openai import AsyncOpenAI
from starlette.background import BackgroundTask
from starlette.datastructures import FormData, UploadFile
from starlette.responses import StreamingResponse
import litellm
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.proxy import proxy_server
from litellm.proxy._types import ProxyErrorTypes, ProxyException
from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth
from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler
from litellm.proxy.route_llm_request import route_request
from litellm.proxy.utils import _failure_fields_to_lift
from litellm.types.guardrails import GuardrailEventHooks, GuardrailStreamScope
@ -419,6 +426,256 @@ def test_audio_transcription_error(client, auth_as, patched_transcription_error,
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 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 // 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)
)
chunks = proxy_server._audio_speech_chunk_generator(response, on_upstream_error=AsyncMock())
assert await anext(chunks) == first
assert produced == [first]
assert b"".join([chunk async for chunk in chunks]) == second
assert stream.closed is True
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)
)
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)
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(
proxy_server,
"proxy_logging_obj",
MagicMock(
pre_call_hook=AsyncMock(side_effect=lambda **kw: kw["data"]),
post_call_failure_hook=failure_hook,
post_call_response_headers_hook=response_headers_hook or 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 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
@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, on_upstream_error=AsyncMock()),
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
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
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(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)
assert failure_hook.await_args.kwargs["request_data"]["response_cost"] == 0.25
@pytest.mark.parametrize("path", ["/v1/audio/transcriptions", "/audio/transcriptions"])
def test_audio_transcription_stream_returns_sse(client, auth_as, patched_transcription_stream, path):
files = {"file": ("audio.mp3", b"\x00\x01\x02", "audio/mpeg")}
@ -541,7 +798,11 @@ def _streams_when_asked(request: httpx.Request) -> httpx.Response:
def _sse_events(body: str) -> tuple[object, ...]:
assert body.endswith("data: [DONE]\n\n")
return tuple(json.loads(frame.removeprefix("data: ")) for frame in body.removesuffix("data: [DONE]\n\n").split("\n\n") if frame)
return tuple(
json.loads(frame.removeprefix("data: "))
for frame in body.removesuffix("data: [DONE]\n\n").split("\n\n")
if frame
)
@pytest.mark.usefixtures("patched_transcription_stream")

132
tests/unit/test_router/test_router_max_parallel_requests.py
View file

@ -1,10 +1,16 @@
from __future__ import annotations
import asyncio
import contextlib
from typing import Final
import litellm
import httpx
import pytest
import litellm
from litellm.router_python_core 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
@ -139,3 +145,127 @@ 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")
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"),
[
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"]

315
tests/unit/test_utils.py
View file

@ -13,11 +13,12 @@ import queue
import threading
import time
import uuid
from collections.abc import Callable, Iterator, Mapping
from collections.abc import AsyncIterator, Callable, Iterator, Mapping
from concurrent.futures import Future, ThreadPoolExecutor
from datetime import datetime, timedelta, timezone
from functools import partial
from pathlib import Path, PurePath
from types import MappingProxyType
from typing import Final, SupportsIndex, cast
from unittest import mock
from unittest.mock import AsyncMock, MagicMock, Mock, patch
@ -56,11 +57,13 @@ 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
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 (
CredentialItem,
@ -3954,10 +3957,318 @@ 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
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 _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._gate: Final = gate
async def __aiter__(self) -> AsyncIterator[bytes]:
yield self._first
if self._gate is not None:
await self._gate.wait()
if isinstance(self._rest, Exception):
raise self._rest
yield self._rest
class _SyncAudio(httpx.SyncByteStream):
def __init__(self, first: bytes, rest: bytes) -> None:
self._first: Final = first
self._rest: Final = rest
self.rest_sent_at: float | None = None
def __iter__(self) -> Iterator[bytes]:
yield self._first
self.rest_sent_at = time.time()
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",
}
)
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)
)
def _audio_summary(num_bytes: int) -> dict[str, object]:
return {"object": "binary", "content_type": "audio/mpeg", "num_bytes": num_bytes}
def _logged_end(payload: dict[str, object]) -> float:
end_time: Final = payload["endTime"]
assert isinstance(end_time, float)
return end_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, speech_logger: _SpeechLogger
) -> None:
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"})
)
async def logged_speech(stream: bool) -> dict[str, object]:
response: Final = await _aspeech(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"]
@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:
rest_released: Final = asyncio.Event()
_mock_elevenlabs_audio(respx_mock, _GatedAsyncAudio(b"first", b"rest", rest_released))
response: Final = await _aspeech(stream=True)
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_end(logged) >= audio_still_arriving_at
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(
("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:
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=_GatedAsyncAudio(b"first", rest),
)
for rest in (b"rest", upstream_error)
]
)
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"]
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, _GatedAsyncAudio(b"first", b"never sent", asyncio.Event()))
response: Final = await _aspeech(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"))
assert logged["error_information"] == {
"error_code": "499",
"error_message": "Client disconnected the request",
"error_class": "ClientDisconnected",
}
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, _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()
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")
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, _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)
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(
respx_mock: respx.MockRouter, speech_logger: _SpeechLogger, stream: object
) -> None:
_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)
assert response.content == b"firstrest"
assert logged["response"] == _audio_summary(len(b"firstrest"))
@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:
audio: Final = _SyncAudio(b"first", b"rest")
_mock_elevenlabs_audio(respx_mock, audio)
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 audio.rest_sent_at is not None and _logged_end(logged) >= audio.rest_sent_at
class _InlineExecutor:
def submit(self, fn: Callable[..., object], /, *args: object, **kwargs: object) -> Future[object]:
future: Final = Future[object]()
future.set_result(fn(*args, **kwargs))
return future
@pytest.mark.usefixtures("local_model_cost_map")
def test_speech_stream_success_handler_runs_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)
speech_logger: Final = _SyncSpeechLogger()
monkeypatch.setattr(litellm, "callbacks", [speech_logger])
monkeypatch.setattr("litellm.utils.executor", _InlineExecutor())
response: Final = litellm.speech(**_SPEECH_ARGS, stream=True)
logged_at_headers: Final = speech_logger.logged_payloads.qsize()
assert response.read() == b"firstrest"
assert logged_at_headers == 0
assert speech_logger.logged_payloads.get_nowait()["response"] == _audio_summary(len(b"firstrest"))
class TestCallbackAsyncSyncSeparation:
"""Test that LoggingCallbackManager auto-routes async callbacks to async lists."""