diff --git a/litellm/litellm_core_utils/health_check_helpers.py b/litellm/litellm_core_utils/health_check_helpers.py index 2c7d00552b4..d25c7d7aea2 100644 --- a/litellm/litellm_core_utils/health_check_helpers.py +++ b/litellm/litellm_core_utils/health_check_helpers.py @@ -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", ), diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 88d84201beb..f4102036a37 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -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.""" diff --git a/litellm/litellm_core_utils/on_end_byte_stream.py b/litellm/litellm_core_utils/on_end_byte_stream.py new file mode 100644 index 00000000000..c8559329189 --- /dev/null +++ b/litellm/litellm_core_utils/on_end_byte_stream.py @@ -0,0 +1,112 @@ +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.afail(mapped) + raise mapped from error + except Exception as error: + await self.afail(error) + raise + await self._end(StreamCompleted()) + + async def afail(self, error: Exception) -> None: + with contextlib.suppress(Exception): + if isinstance(self._stream, OnEndAsyncByteStream): + await self._stream.afail(error) + else: + await self._stream.aclose() + await self._end(StreamFailed(error)) + + async def aclose(self) -> None: + 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) diff --git a/litellm/llms/aws_polly/text_to_speech/transformation.py b/litellm/llms/aws_polly/text_to_speech/transformation.py index 3b52a8df3f2..026d6b9d40e 100644 --- a/litellm/llms/aws_polly/text_to_speech/transformation.py +++ b/litellm/llms/aws_polly/text_to_speech/transformation.py @@ -35,6 +35,8 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): Reference: https://docs.aws.amazon.com/polly/latest/dg/API_SynthesizeSpeech.html """ + supports_streaming = True + def __init__(self): BaseTextToSpeechConfig.__init__(self) BaseAWSLLM.__init__(self) @@ -82,6 +84,7 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): aspeech: bool, api_base: str | None, api_key: str | None, + stream: bool = False, **kwargs: object, ) -> Union[ "HttpxBinaryResponseContent", @@ -128,6 +131,7 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): extra_headers=extra_headers, client=None, _is_async=aspeech, + stream=stream, ) return response diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 69ffac5d6f9..53a576cbedf 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -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, diff --git a/litellm/llms/azure/text_to_speech/transformation.py b/litellm/llms/azure/text_to_speech/transformation.py index acae650ed75..9fecc30833f 100644 --- a/litellm/llms/azure/text_to_speech/transformation.py +++ b/litellm/llms/azure/text_to_speech/transformation.py @@ -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 diff --git a/litellm/llms/base_llm/text_to_speech/transformation.py b/litellm/llms/base_llm/text_to_speech/transformation.py index 6df83966bf8..879d1622b0d 100644 --- a/litellm/llms/base_llm/text_to_speech/transformation.py +++ b/litellm/llms/base_llm/text_to_speech/transformation.py @@ -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 diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a10db31606..dba5f085cb0 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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( diff --git a/litellm/llms/elevenlabs/text_to_speech/transformation.py b/litellm/llms/elevenlabs/text_to_speech/transformation.py index 23ccf41c9da..afd7b55dd41 100644 --- a/litellm/llms/elevenlabs/text_to_speech/transformation.py +++ b/litellm/llms/elevenlabs/text_to_speech/transformation.py @@ -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: diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index f330d7fc6d4..f0d17cfa2af 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -37,7 +37,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 ( @@ -1673,6 +1677,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: @@ -1690,6 +1695,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): logging_obj=logging_obj, client=client, shared_session=shared_session, + stream=stream, litellm_params=litellm_params, ) @@ -1714,14 +1720,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( @@ -1739,6 +1745,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( @@ -1764,14 +1771,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 diff --git a/litellm/main.py b/litellm/main.py index 6b9964a450f..1aa7b06007e 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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, ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index c1af5d6f980..6f8c48b83e6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -904,9 +904,11 @@ 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 +from litellm.litellm_core_utils.on_end_byte_stream import OnEndAsyncByteStream from litellm.proxy import tracing_endpoints from litellm.proxy.middleware.admission_control_middleware import ( AdmissionControlMiddleware, @@ -13692,16 +13694,62 @@ 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() + + +async def _fail_audio_speech_response(_response: HttpxBinaryResponseContent, error: Exception) -> None: + """Release the audio as failed, so a route error before any byte is sent is not logged as a caller hang-up.""" + with contextlib.suppress(Exception): + if isinstance(stream := _response.response.stream, OnEndAsyncByteStream): + await stream.afail(error) + await _close_audio_speech_response(_response) @router.post( @@ -13726,7 +13774,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() @@ -13765,52 +13813,71 @@ 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 Exception as error: + await _fail_audio_speech_response(response, error) + raise + except BaseException: + await _close_audio_speech_response(response) + raise except Exception as e: await proxy_logging_obj.post_call_failure_hook( diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index c6ca4a1f399..632f3e509a2 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -2516,6 +2516,9 @@ class RustCore( return await self._drive(plan_for("atranscription"), request_from_call("atranscription", args, kwargs)) async def aspeech(self, *args: object, **kwargs: object) -> object: # kwargs-ok: endpoint payload + """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.""" request: Final = dict(request_from_call("aspeech", args, kwargs)) request.setdefault("voice", None) return await self._drive(plan_for("aspeech"), request) @@ -5163,7 +5166,7 @@ class _Driver: compacted: Final = await compacted_attempt(self._core._port, self._plan, deployment, call) # pyright: ignore[reportPrivateUsage] # see above except Exception as refused: # noqa: BLE001 # raised before the provider call, so no cooldown reads it return {**self._failed(step, deployment, call, refused), "recorded": True} - # The deployment's max_parallel_requests slot; an open chat stream takes it along. + # The deployment's max_parallel_requests slot; an open chat or speech stream takes it along. port = self._core._port # pyright: ignore[reportPrivateUsage] # the driver serves its own core # rebind-ok: the same port on every attempt with enter_deployment_slot(port, self._plan, deployment, call.kwargs) as slot: # A search tool is no deployment: nothing to check; Python's `_arerank` and diff --git a/litellm/router_rust/invoke/slots.py b/litellm/router_rust/invoke/slots.py index 08e88b5b2f1..1c63a5c0910 100644 --- a/litellm/router_rust/invoke/slots.py +++ b/litellm/router_rust/invoke/slots.py @@ -1,19 +1,22 @@ """ A deployment's ``max_parallel_requests`` slot on the Rust router core (``_deployment_slot`` and ``_get_client(client_type="max_parallel_requests")``): the attempt holds one of the deployment's -slots around its provider call, and an open chat stream keeps it until the stream is exhausted, -closed, or falls back. A caller arriving while every slot is taken gets a 429, which the core -records as that attempt's failure. +slots around its provider call, an open chat stream keeps it until the stream is exhausted, closed, +or falls back, and a streamed speech response keeps it until its audio ends. A caller arriving while +every slot is taken gets a 429, which the core records as that attempt's failure. The limit object is the one the Python router caches (``MaxParallelRequestsLimit``), kept in the core's routing cache under the same key, so its in-flight count is shared by every attempt. """ import contextlib -from collections.abc import Mapping +from collections.abc import Callable, Mapping from typing import TYPE_CHECKING, Final +import httpx + from litellm._internal_context import with_service_target +from litellm.litellm_core_utils.on_end_byte_stream import OnEndAsyncByteStream, StreamEnd from litellm.router_rust.invoke.shaping import ROUTER_CLIENTS_TARGET, as_dict if TYPE_CHECKING: @@ -104,4 +107,28 @@ def deployment_slot( return slot -__all__ = ["deployment_slot", "enter_deployment_slot", "max_parallel_requests_limit"] +def hold_slot_until_closed( + response: object, attempt_slot: contextlib.ExitStack, on_end: Callable[[StreamEnd], None] +) -> bool: + """Whether ``response`` took the attempt's slot. A body already in hand leaves the slot with the + attempt. A body still arriving from the provider takes it, or the deployment's + max_parallel_requests would only bound how many responses can start: it is held until the body + is read to the end, the read fails, or the response is closed, and ``on_end`` runs then.""" + raw: Final = getattr(response, "response", None) + if ( + not isinstance(raw, httpx.Response) + or raw.is_stream_consumed + or not isinstance(raw.stream, httpx.AsyncByteStream) + ): + return False + held: Final = attempt_slot.pop_all() + + async def release(end: StreamEnd) -> None: + held.close() + on_end(end) + + raw.stream = OnEndAsyncByteStream(raw.stream, map_error=lambda error: error, on_end=release) + return True + + +__all__ = ["deployment_slot", "enter_deployment_slot", "hold_slot_until_closed", "max_parallel_requests_limit"] diff --git a/litellm/router_rust/invoke/streaming.py b/litellm/router_rust/invoke/streaming.py index d8b4f87501a..b144d210955 100644 --- a/litellm/router_rust/invoke/streaming.py +++ b/litellm/router_rust/invoke/streaming.py @@ -15,7 +15,8 @@ Each wrapper also reports its stream's outcome to the core when the stream ends, logging reports it to the Python router's callbacks: the attempt that returned the stream counted only its request (the attempt answer is ``deferred``), so the success with the stream's final tokens, or a failure the stream raised after it opened, is recorded then (``StreamAttempt.report_success`` / -``report_failure``). A stream closed before it ends reports nothing. +``report_failure``). A stream closed before it ends reports nothing, except a speech stream, which +litellm logs as a success. The wrappers decide nothing about routing: what may take over, the retry budget and every deployment pick come from the core. They hold live Python iterators, frames and hidden params, @@ -64,7 +65,9 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import ( ) from litellm.litellm_core_utils.core_helpers import as_str_mapping, without from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params, set_hidden_params +from litellm.litellm_core_utils.on_end_byte_stream import StreamClosedEarly, StreamCompleted, StreamEnd, StreamFailed from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.router_rust.invoke.slots import hold_slot_until_closed from litellm.router_rust.surface.health import failure_already_logged, logged_call from litellm.types.utils import ModelResponseStream, Usage @@ -110,7 +113,7 @@ class StreamAttempt: committed_budget: int | None = None # The group a pre-routing hook selected for ``group``, which a fallback lookup keys on first. routed: str | None = None - # The deployment's max_parallel_requests slot the attempt holds; an open chat stream takes it. + # The deployment's max_parallel_requests slot the attempt holds; an open chat or speech stream takes it. slot: contextlib.ExitStack | None = None # Set by the wrapper that reports the stream's outcome when it ends; the attempt's answer is then # deferred and counts only the request and the tokens read so far (``counted_tokens``). @@ -1842,6 +1845,25 @@ def _reported_transcription_stream(attempt: StreamAttempt, response: object) -> return response +def _reported_speech_stream(attempt: StreamAttempt, response: object) -> object: + """A speech response still pulling audio from the provider keeps its deployment's slot until the + audio ends and reports its attempt's outcome then: a read that fails counts against the deployment + and cools it down by the rules, as the router's failure callback does, and a caller hanging up + counts nothing against it. A buffered response is done at the attempt, which counts it as any other.""" + + def report(end: StreamEnd) -> None: + match end: + case StreamFailed(error=error): + attempt.report_failure(error) + case StreamCompleted() | StreamClosedEarly(): + attempt.report_success(0, response) + + slot: Final = attempt.slot if attempt.slot is not None else contextlib.ExitStack() + if hold_slot_until_closed(response, slot, report): + attempt.take_outcome() + return response + + # --- the attempt hook --------------------------------------------------------------------------- @@ -1856,6 +1878,8 @@ async def open_attempt_stream(attempt: StreamAttempt, response: object) -> objec first_chunk: Final = await _aread_first_chunk(response) held_slot: Final = attempt.slot.pop_all() if attempt.slot is not None else contextlib.ExitStack() return wrap_chat_stream(attempt, response, first_chunk, held_slot) + if attempt.op == "aspeech": + return _reported_speech_stream(attempt, response) if not attempt.kwargs.get("stream"): return response if attempt.op == "atranscription": diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index c9739eee91d..27af0c4e177 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -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) diff --git a/litellm/utils.py b/litellm/utils.py index 53b02ebddac..4a72d4f05ff 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -7,6 +7,7 @@ import ast import asyncio import base64 import binascii +import contextlib import contextvars import copy import datetime @@ -75,6 +76,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 +103,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 +249,7 @@ from litellm.types.llms.openai import ( ChatCompletionNamedToolChoiceParam, ChatCompletionToolParam, ChatCompletionToolParamFunctionChunk, + HttpxBinaryResponseContent, OpenAITextCompletionUserMessage, OpenAIWebSearchOptions, ) @@ -275,6 +286,7 @@ from litellm.types.utils import ( SandboxProviders, SearchProviders, SelectTokenizerResponse, + StandardLoggingPayloadErrorInformation, StreamingChoices, TranscriptionResponse, Usage, @@ -316,7 +328,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 +1665,126 @@ 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() + + +async def _close_unread_speech(result: object, call_type: str, stream_requested: bool) -> None: + unread_speech: Final = _unread_speech_upstream(result, call_type, stream_requested) + if unread_speech is not None: + with contextlib.suppress(Exception): + await unread_speech.aclose() + + +def _bill_failed_speech_stream(logging_obj: LiteLLMLoggingObject, result: object) -> None: + logging_obj.record_partial_usage_for_failure(Usage(), logging_obj.response_cost_calculator(result=result) or 0.0) + + +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 +2406,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,18 +2750,35 @@ 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: traceback_exception: Final = traceback.format_exc() + await _close_unread_speech(result, call_type, "stream" in kwargs and kwargs["stream"] is True) # Reuse the timestamp taken right when the deployment call itself failed, before # the failure hook ran, so a slow callback doesn't inflate the reported duration. end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with @@ -2767,7 +2931,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 +2961,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 diff --git a/tests/integration/providers/test_audio_speech_stream_chaos.py b/tests/integration/providers/test_audio_speech_stream_chaos.py new file mode 100644 index 00000000000..1686c10a1fb --- /dev/null +++ b/tests/integration/providers/test_audio_speech_stream_chaos.py @@ -0,0 +1,312 @@ +import asyncio +import threading +import uuid +from collections import Counter +from collections.abc import Callable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final, Literal +from urllib.parse import urlsplit + +import httpx +import psutil +import pytest +import yaml +from integration._support.client import ( + JSON_OBJECT, + Gateway, + eventually, + object_value, + string_value, +) +from integration._support.database import read_rows +from integration._support.process import ( + OwnedProxy, + graceful_stop_seconds, + owned_proxy_process, +) +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = [pytest.mark.timeout(4 * graceful_stop_seconds() + 300), 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]) diff --git a/tests/integration/providers/test_audio_speech_stream_wire.py b/tests/integration/providers/test_audio_speech_stream_wire.py new file mode 100644 index 00000000000..8cc340ee4ce --- /dev/null +++ b/tests/integration/providers/test_audio_speech_stream_wire.py @@ -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 diff --git a/tests/integration/providers/test_provider_passthrough_stream_relay_wire.py b/tests/integration/providers/test_provider_passthrough_stream_relay_wire.py new file mode 100644 index 00000000000..10b05b114ee --- /dev/null +++ b/tests/integration/providers/test_provider_passthrough_stream_relay_wire.py @@ -0,0 +1,144 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Final +from urllib.parse import urlsplit + +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.forward_proxy import tunnelling_forward_proxy +from integration._support.process import owned_proxy +from integration._support.tls import server_context, write_self_signed_cert +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +_KEY: Final = "synthetic-passthrough-key" +_GOOGLE_TOKEN: Final = "synthetic-google-access-token" +_VERTEX_HOST: Final = "us-central1-aiplatform.googleapis.com" +_VERTEX_MODEL: Final = "gemini-2.5-flash" +_FRAMES: Final = ( + b'data: {"id":"pt-1","choices":[{"index":0,"delta":{"content":"hel"}}]}\n\n', + b'data: {"id":"pt-1","choices":[{"index":0,"delta":{"content":"lo"}}]}\n\n', + b"data: [DONE]\n\n", +) +_GEMINI_FRAMES: Final = tuple( + f"data: {json.dumps(frame)}\r\n\r\n".encode() + for frame in ( + {"candidates": [{"content": {"role": "model", "parts": [{"text": "hel"}]}, "index": 0}]}, + { + "candidates": [ + {"content": {"role": "model", "parts": [{"text": "lo"}]}, "index": 0, "finishReason": "STOP"} + ], + "usageMetadata": {"promptTokenCount": 4, "candidatesTokenCount": 2, "totalTokenCount": 6}, + "modelVersion": _VERTEX_MODEL, + }, + ) +) +_ROUTES: Final = MappingProxyType( + { + "openai": ("/openai/chat/completions", "/chat/completions"), + "mistral": ("/mistral/v1/chat/completions", "/v1/chat/completions"), + "cohere": ("/cohere/v2/chat", "/v2/chat"), + } +) + + +@dataclass(frozen=True, slots=True) +class _Relay: + wire: Wire + vertex: Wire + proxy: Gateway + + +def _sse(request: Request) -> Reply: + return Reply(content_type="text/event-stream", chunks=_FRAMES) + + +def _gemini_sse(request: Request) -> Reply: + return Reply(content_type="text/event-stream", chunks=_GEMINI_FRAMES) + + +@pytest.fixture(scope="module") +def relay(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Relay]: + directory: Final = tmp_path_factory.mktemp("passthrough-stream") + certificate, key = write_self_signed_cert(directory, (_VERTEX_HOST,)) + with ( + wire_server(_sse) as wire, + wire_server(_gemini_sse, tls=server_context(certificate, key)) as vertex, + tunnelling_forward_proxy(int(vertex.url.rsplit(":", 1)[1])) as tunnel, + gateway_from_environment() as rig, + ): + bases: Final = {f"{name}_API_BASE": wire.url for name in ("OPENAI", "MISTRAL", "COHERE")} + keys: Final = {f"{name}_API_KEY": _KEY for name in ("OPENAI", "MISTRAL", "COHERE")} + tunnelled: Final = { + "HTTPS_PROXY": tunnel.url, + "https_proxy": tunnel.url, + "NO_PROXY": "127.0.0.1,localhost", + "no_proxy": "127.0.0.1,localhost", + "SSL_CERT_FILE": str(certificate), + } + with owned_proxy(rig, directory, {**bases, **keys, **tunnelled}, workers=2) as proxy: + yield _Relay(wire, vertex, proxy) + + +def _relayed( + proxy: Gateway, path: str, body: dict[str, JsonValue], headers: Mapping[str, str] | None = None +) -> tuple[int, bytes]: + sent: Final = headers or {"Authorization": f"Bearer {proxy.key}"} + with proxy.client.stream("POST", path, json=body, headers=dict(sent)) as response: + return response.status_code, b"".join(response.iter_raw()) + + +def _only_post(wire: Wire) -> Request: + requests: Final = wire.drain() + assert [request.method for request in requests] == ["POST"], [request.target for request in requests] + return requests[0] + + +@pytest.mark.parametrize("route", tuple(_ROUTES)) +def test_a_provider_passthrough_stream_is_relayed_frame_for_frame(relay: _Relay, route: str) -> None: + proxy_path, upstream_path = _ROUTES[route] + marker: Final = f"passthrough-stream-{uuid.uuid4().hex}" + relayed: Final = _relayed( + relay.proxy, + proxy_path, + {"model": "relay-model", "stream": True, "messages": [{"role": "user", "content": marker}]}, + ) + assert relayed == (200, b"".join(_FRAMES)), relayed + forwarded: Final = _only_post(relay.wire) + assert urlsplit(forwarded.target).path == upstream_path, forwarded.target + assert marker.encode() in forwarded.body, forwarded.body[:400] + + +def test_a_vertex_passthrough_stream_is_relayed_frame_for_frame(relay: _Relay) -> None: + project: Final = f"relay-{uuid.uuid4().hex[:12]}" + model_path: Final = f"/v1/projects/{project}/locations/us-central1/publishers/google/models/{_VERTEX_MODEL}" + marker: Final = f"vertex-passthrough-stream-{uuid.uuid4().hex}" + relayed: Final = _relayed( + relay.proxy, + f"/vertex_ai{model_path}:streamGenerateContent", + {"contents": [{"role": "user", "parts": [{"text": marker}]}]}, + {"x-litellm-api-key": relay.proxy.key, "Authorization": f"Bearer {_GOOGLE_TOKEN}"}, + ) + assert relayed == (200, b"".join(_GEMINI_FRAMES)), relayed + forwarded: Final = _only_post(relay.vertex) + assert forwarded.target == f"{model_path}:streamGenerateContent?alt=sse", forwarded.target + assert forwarded.headers["authorization"] == f"Bearer {_GOOGLE_TOKEN}", forwarded.headers + assert marker.encode() in forwarded.body, forwarded.body[:400] + + +def test_an_azure_passthrough_stream_to_a_router_model_is_relayed_frame_for_frame(gateway: Gateway) -> None: + with wire_server(_sse) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="azure/relay-deployment", api_base=wire.url, api_version="2024-10-21") + marker: Final = f"azure-passthrough-stream-{uuid.uuid4().hex}" + relayed: Final = _relayed( + gateway, + f"/azure/openai/deployments/{model}/chat/completions?api-version=2024-10-21", + {"stream": True, "messages": [{"role": "user", "content": marker}]}, + ) + assert relayed == (200, b"".join(_FRAMES)), relayed + forwarded: Final = _only_post(wire) + assert urlsplit(forwarded.target).path.endswith("/chat/completions"), forwarded.target + assert marker.encode() in forwarded.body, forwarded.body[:400] diff --git a/tests/integration/sdk/test_speech_stream_sdk_wire.py b/tests/integration/sdk/test_speech_stream_sdk_wire.py new file mode 100644 index 00000000000..3dcd29f1637 --- /dev/null +++ b/tests/integration/sdk/test_speech_stream_sdk_wire.py @@ -0,0 +1,300 @@ +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 eventually, held +from integration._support.wire import Reply, Request, wire_server +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger +from litellm.types.utils import CallTypes, LLMResponseTypes + +_BACKEND: Final = "openai/gpt-4o-mini-tts" +_PROVIDER_KEY: Final = "synthetic-tts-key" +_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")), + ) + + +class _SpeechRejected(Exception): + pass + + +class _RejectingGuardrail(CustomGuardrail): + async def async_post_call_success_deployment_hook( + self, request_data: dict[str, object], response: LLMResponseTypes, call_type: CallTypes | None + ) -> None: + raise _SpeechRejected + + +async def test_aspeech_stream_rejected_by_a_deployment_guardrail_closes_the_provider_connection() -> None: + litellm.callbacks.append(_RejectingGuardrail()) + with wire_server(_peer(), keep_alive=True) as wire: + with pytest.raises(_SpeechRejected): + await _aspeech(wire.url, str(uuid.uuid4())) + await asyncio.to_thread(eventually, wire.open_connections, lambda still_open: still_open == 0) + assert wire.connections() == 1 + + +def test_speech_takes_base_url_in_its_positional_slot() -> None: + probe: Final = _probe() + call_id: Final = str(uuid.uuid4()) + 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")),) diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index dd0e5762845..39682242c01 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -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() diff --git a/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py b/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py index 88f0ae1d524..4a9b0a5c4be 100644 --- a/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py +++ b/tests/unit/llms/elevenlabs/test_elevenlabs_text_to_speech_transformation.py @@ -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")], diff --git a/tests/unit/llms/openai/test_openai_audio_speech.py b/tests/unit/llms/openai/test_openai_audio_speech.py new file mode 100644 index 00000000000..27ae0ffd543 --- /dev/null +++ b/tests/unit/llms/openai/test_openai_audio_speech.py @@ -0,0 +1,105 @@ +"""Streaming behaviour of the OpenAI and Azure OpenAI ``/audio/speech`` paths.""" + +from __future__ import annotations + +from typing import Final +from unittest.mock import Mock + +import httpx +import pytest +from openai import AsyncAzureOpenAI, AsyncOpenAI + +from litellm.llms.azure.azure import AzureChatCompletion +from litellm.llms.openai.openai import OpenAIChatCompletion + + +class _RecordingAudioStream(httpx.AsyncByteStream): + def __init__(self, chunks: list[bytes], produced: list[bytes]): + self._chunks = chunks + self._produced = produced + + async def __aiter__(self): + for chunk in self._chunks: + self._produced.append(chunk) + yield chunk + + +def _audio_client(produced: list[bytes]) -> httpx.AsyncClient: + def _respond(request: httpx.Request) -> httpx.Response: + return httpx.Response( + status_code=200, + headers={"content-type": "audio/mpeg"}, + stream=_RecordingAudioStream([b"one", b"two", b"three"], produced), + ) + + return httpx.AsyncClient(transport=httpx.MockTransport(_respond)) + + +async def _openai_speech(produced: list[bytes], *, stream: bool): + return await OpenAIChatCompletion().async_audio_speech( + model="tts-1", + input="hello", + voice="alloy", + optional_params={}, + api_key="fake-key", + api_base=None, + organization=None, + project=None, + max_retries=0, + timeout=10.0, + logging_obj=Mock(), + client=AsyncOpenAI(api_key="fake-key", http_client=_audio_client(produced)), + stream=stream, + ) + + +@pytest.mark.asyncio +async def test_openai_audio_speech_yields_chunks_before_upstream_finishes(): + produced: Final[list[bytes]] = [] + + response: Final = await _openai_speech(produced, stream=True) + audio: Final = await response.aiter_bytes() + + assert await anext(audio) == b"one" + assert produced == [b"one"] + + +@pytest.mark.asyncio +async def test_openai_audio_speech_buffers_without_stream(): + produced: Final[list[bytes]] = [] + + response: Final = await _openai_speech(produced, stream=False) + + assert produced == [b"one", b"two", b"three"] + assert response.content == b"onetwothree" + + +@pytest.mark.asyncio +async def test_azure_audio_speech_yields_chunks_before_upstream_finishes(): + produced: Final[list[bytes]] = [] + + response: Final = await AzureChatCompletion().async_audio_speech( + model="tts-1", + input="hello", + voice="alloy", + optional_params={}, + api_key="fake-key", + api_base="https://fake.openai.azure.com", + api_version="2024-02-01", + azure_ad_token=None, + azure_ad_token_provider=None, + max_retries=0, + timeout=10.0, + logging_obj=Mock(), + client=AsyncAzureOpenAI( + api_key="fake-key", + api_version="2024-02-01", + azure_endpoint="https://fake.openai.azure.com", + http_client=_audio_client(produced), + ), + stream=True, + ) + audio: Final = await response.aiter_bytes() + + assert await anext(audio) == b"one" + assert produced == [b"one"] diff --git a/tests/unit/proxy/proxy_server/test_routes_audio.py b/tests/unit/proxy/proxy_server/test_routes_audio.py index 9fa617c0dc3..0782de39da0 100644 --- a/tests/unit/proxy/proxy_server/test_routes_audio.py +++ b/tests/unit/proxy/proxy_server/test_routes_audio.py @@ -9,30 +9,39 @@ 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 functools import partial 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.litellm_core_utils.on_end_byte_stream import OnEndAsyncByteStream, StreamEnd, StreamFailed 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 +428,283 @@ 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: Final[list[bytes]] = [] + first, second = b"a" * (AUDIO_SPEECH_CHUNK_SIZE // 8), b"b" * AUDIO_SPEECH_CHUNK_SIZE + stream: Final = _RecordingAudioStream([first, second], produced) + response: Final = HttpxBinaryResponseContent( + httpx.Response(status_code=200, headers={"content-type": "audio/mpeg"}, stream=stream) + ) + + chunks: Final = 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: httpx.AsyncByteStream) -> 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: Final = _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: Final = _RecordingAudioStream([b"a" * AUDIO_SPEECH_CHUNK_SIZE] * 3, []) + response: Final = _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: Final = patched_streaming_speech + captured: Final[list[StreamingResponse]] = [] + + original: Final = proxy_server.StreamingResponse + + def _capture(*args, **kwargs): + response: Final = original(*args, **kwargs) + captured.append(response) + return response + + monkeypatch.setattr(proxy_server, "StreamingResponse", _capture) + + request: Final = _speech_request({"model": "tts-1", "input": "Hi", "voice": "alloy", "stream": True}) + await proxy_server.audio_speech( + request=request, fastapi_response=Response(), user_api_key_dict=UserAPIKeyAuth(api_key="sk-test") + ) + + 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 + + +@pytest.mark.asyncio +async def test_audio_speech_route_ends_the_audio_as_failed_when_a_response_hook_fails( + monkeypatch: pytest.MonkeyPatch, +): + """A route error before any audio is sent is a failure, not a caller hang-up, through every + wrapper the router and the SDK put around the audio.""" + upstream: Final = _RecordingAudioStream([b"audio"], []) + ends: Final[list[tuple[str, StreamEnd]]] = [] + + async def _record(layer: str, end: StreamEnd) -> None: + ends.append((layer, end)) + + sdk_stream: Final = OnEndAsyncByteStream(upstream, map_error=lambda e: e, on_end=partial(_record, "sdk")) + router_stream: Final = OnEndAsyncByteStream(sdk_stream, map_error=lambda e: e, on_end=partial(_record, "router")) + _route_speech_to(monkeypatch, _audio_response(router_stream), response_headers_hook=_headers_from_the_audio_body) + + with pytest.raises(ProxyException): + await proxy_server.audio_speech( + request=_speech_request({"model": "tts-1", "input": "Hi", "voice": "alloy", "stream": True}), + fastapi_response=Response(), + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test"), + ) + + assert [(layer, type(end)) for layer, end in ends] == [("sdk", StreamFailed), ("router", StreamFailed)] + assert upstream.closed is True + + +class _BreaksMidClipAudioStream(httpx.AsyncByteStream): + async def __aiter__(self): + yield b"a" * AUDIO_SPEECH_CHUNK_SIZE + 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 +827,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") diff --git a/tests/unit/test_router/test_router_max_parallel_requests.py b/tests/unit/test_router/test_router_max_parallel_requests.py index 1b350f31f23..15839866547 100644 --- a/tests/unit/test_router/test_router_max_parallel_requests.py +++ b/tests/unit/test_router/test_router_max_parallel_requests.py @@ -1,9 +1,21 @@ from __future__ import annotations +import asyncio +import contextlib +from collections.abc import Callable from typing import Final -import litellm +import httpx import pytest + +import litellm +from litellm.router_rust.invoke.slots import hold_slot_until_closed +from litellm.router_rust.invoke.streaming import _RECORDINGS +from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + get_deployment_failures_for_current_minute, + get_deployment_successes_for_current_minute, +) +from litellm.types.llms.openai import HttpxBinaryResponseContent from litellm.utils import calculate_max_parallel_requests @@ -91,3 +103,265 @@ 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 past the attempt, everything else frees it.""" + released: Final[list[str]] = [] + + with contextlib.ExitStack() as attempt_slot: + attempt_slot.callback(released.append, "released") + hold_slot_until_closed(response, attempt_slot, lambda _end: None) + + assert bool(released) is releases_immediately + if not releases_immediately: + await response.aclose() + assert released == ["released"] + + +class _AudioStalledMidClip(httpx.AsyncByteStream): + async def __aiter__(self): + yield b"audio" + raise httpx.ReadTimeout("the provider stopped sending audio") + + async def aclose(self) -> None: + return None + + +def _tts_deployment(model_name: str, api_base: str, **params: object) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": { + "model": "elevenlabs/eleven_multilingual_v2", + "api_key": "fake-key", + "api_base": api_base, + "max_parallel_requests": 1, + **params, + }, + "model_info": {"id": api_base}, + } + + +def _client_answering(answer: Callable[[httpx.Request], httpx.Response]) -> litellm.AsyncHTTPHandler: + return litellm.AsyncHTTPHandler(transport=httpx.MockTransport(answer)) + + +def _audio(stream: httpx.AsyncByteStream) -> httpx.Response: + return httpx.Response(200, headers={"content-type": "audio/mpeg"}, stream=stream) + + +@pytest.mark.asyncio +async def test_aspeech_falls_back_on_a_failure_before_the_headers_and_frees_the_failed_slot() -> None: + """A provider refusing a streamed clip before any audio fails the attempt, so the fallback serves it and the + failed deployment's only slot is free for the next caller.""" + router: Final = litellm.Router( + model_list=[ + _tts_deployment("tts", "http://refusing.test"), + _tts_deployment("tts-backup", "http://backup.test"), + ], + fallbacks=[{"tts": ["tts-backup"]}], + num_retries=0, + disable_cooldowns=True, + ) + hosts: Final[list[str]] = [] + + def _answer(request: httpx.Request) -> httpx.Response: + hosts.append(request.url.host) + if request.url.host == "refusing.test": + return httpx.Response(500, json={"detail": "synthesis failed"}) + return _audio(_OneChunkAudioStream()) + + async def _clip() -> bytes: + response: Final = await router.aspeech( + model="tts", input="hello", voice="voice-id", client=_client_answering(_answer), stream=True + ) + return await response.aread() + + assert (await _clip(), await _clip()) == (b"audio", b"audio") + assert hosts == ["refusing.test", "backup.test", "refusing.test", "backup.test"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("audio", "read_it", "counted"), + [ + pytest.param(_OneChunkAudioStream, True, (1, 0), id="read-to-the-end"), + pytest.param(_AudioCutOffMidClip, True, (0, 1), id="cut-off-mid-read"), + pytest.param(_AudioCutOffMidClip, False, (1, 0), id="caller-hung-up"), + ], +) +async def test_aspeech_counts_a_streamed_clip_against_its_deployment_when_the_audio_ends( + audio: type[httpx.AsyncByteStream], read_it: bool, counted: tuple[int, int] +) -> None: + """The attempt returns at the headers, so only the audio's end says how the clip went: a provider cutting it + off is a failure, and a caller walking away says nothing bad about the deployment.""" + router: Final = litellm.Router(model_list=[_tts_deployment("tts", "http://tts.test")], num_retries=0) + + def _outcome() -> tuple[int, int]: + return ( + get_deployment_successes_for_current_minute(router, "http://tts.test"), + get_deployment_failures_for_current_minute(router, "http://tts.test"), + ) + + clip: Final = await router.aspeech( + model="tts", + input="hello", + voice="voice-id", + client=_client_answering(lambda _: _audio(audio())), + stream=True, + ) + at_the_headers: Final = _outcome() + if read_it: + with contextlib.suppress(litellm.APIConnectionError): + await clip.aread() + else: + await clip.aclose() + await asyncio.gather(*_RECORDINGS) + + assert (at_the_headers, _outcome()) == ((0, 0), counted) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("read_it", "next_host"), + [ + pytest.param(True, "backup.test", id="provider-stalled-mid-clip"), + pytest.param(False, "first.test", id="caller-hung-up"), + ], +) +async def test_aspeech_cools_a_deployment_down_for_a_clip_that_stalled_after_the_headers( + read_it: bool, next_host: str +) -> None: + """A stalled clip is a 200 to the attempt, so the failure the audio raises is what takes the deployment out of + rotation for the next caller.""" + router: Final = litellm.Router( + model_list=[ + _tts_deployment("tts", "http://first.test", order=1), + _tts_deployment("tts", "http://backup.test", order=2), + ], + num_retries=0, + allowed_fails=0, + ) + hosts: Final[list[str]] = [] + + def _answer(request: httpx.Request) -> httpx.Response: + hosts.append(request.url.host) + return _audio(_AudioStalledMidClip()) + + async def _clip() -> HttpxBinaryResponseContent: + return await router.aspeech( + model="tts", input="hello", voice="voice-id", client=_client_answering(_answer), stream=True + ) + + first_clip: Final = await _clip() + if read_it: + with pytest.raises(litellm.Timeout): + await first_clip.aread() + else: + await first_clip.aclose() + await asyncio.gather(*_RECORDINGS) + await (await _clip()).aclose() + + assert hosts == ["first.test", next_host] diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index c41c4db2c02..481ee065471 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -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, @@ -3956,10 +3959,354 @@ 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")) + + +class _SpeechRejected(Exception): + pass + + +class _RejectingGuardrail(CustomGuardrail): + async def async_post_call_success_deployment_hook(self, request_data, response, call_type) -> None: + raise _SpeechRejected + + +class _AudioThatFailsToClose(httpx.AsyncByteStream): + def __init__(self) -> None: + self.close_attempts: int = 0 + + async def __aiter__(self) -> AsyncIterator[bytes]: + yield b"first" + + async def aclose(self) -> None: + self.close_attempts += 1 + raise httpx.ReadError("the provider connection broke while closing") + + +@pytest.mark.usefixtures("local_model_cost_map") +async def test_aspeech_stream_failing_to_close_keeps_the_guardrail_rejection( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + audio: Final = _AudioThatFailsToClose() + monkeypatch.setattr(litellm, "callbacks", [_RejectingGuardrail()]) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + _mock_elevenlabs_audio(respx_mock, audio) + + with pytest.raises(_SpeechRejected): + await _aspeech(stream=True) + + assert audio.close_attempts == 1 + + +@pytest.mark.usefixtures("local_model_cost_map") +def test_speech_stream_is_logged_when_the_audio_ends( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> 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."""