From 4ca7a2597c2d074021fe5379cb2e141ae782ccfa Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 10 Nov 2026 12:11:33 -0800 Subject: [PATCH] fix(speech): await the TTS bridge under aspeech Under aspeech the Gemini and Vertex TTS bridge awaits acompletion instead of running sync completion in an executor thread, so a slow upstream no longer holds a worker thread. This completes 176808e268 now that c5fbd1c3c3 adds the internal sub-call scope. The inner chat call shares the speech call's logging object, so the bridge keeps the speech call's acompletion flag, and sync success callbacks still run once with the speech response --- .../speech_to_completion_bridge/handler.py | 41 ++++++++---- litellm/litellm_core_utils/litellm_logging.py | 14 ++++- litellm/main.py | 2 + tests/unit/test_main.py | 63 +++++++++++++++++++ 4 files changed, 107 insertions(+), 13 deletions(-) diff --git a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py index 3e9cd7740ec..5f88208d567 100644 --- a/litellm/endpoints/speech/speech_to_completion_bridge/handler.py +++ b/litellm/endpoints/speech/speech_to_completion_bridge/handler.py @@ -3,7 +3,7 @@ Handler for transforming /chat/completions api requests to litellm.responses req """ import json -from collections.abc import Mapping +from collections.abc import Coroutine, Mapping from typing import TYPE_CHECKING, Final from typing_extensions import TypedDict @@ -32,7 +32,7 @@ class SpeechToCompletionBridgeHandlerInputKwargs(TypedDict): model: str input: str voice: str | dict | None - optional_params: dict + optional_params: dict[str, object] litellm_params: dict logging_obj: "LiteLLMLoggingObj" headers: dict @@ -106,10 +106,10 @@ class SpeechToCompletionBridgeHandler: extra_headers: Mapping[str, object] | None, timeout: "float | httpx.Timeout | None", cache_controls: "CacheControlsHint", - ) -> "HttpxBinaryResponseContent": + aspeech: bool, + ) -> "HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]": received_args: Final = locals() - from litellm import completion - from litellm.types.utils import ModelResponse + from litellm import acompletion, completion validated_kwargs: Final = self.validate_input_kwargs(received_args) model = validated_kwargs["model"] @@ -138,19 +138,36 @@ class SpeechToCompletionBridgeHandler: voice=voice, ) + requested_response_format: Final = validated_kwargs["optional_params"].get("response_format") + response_format: Final = requested_response_format if isinstance(requested_response_format, str) else None + if aspeech: + return self._aspeech(acompletion(**request_data), response_format, logging_obj) + with internal_sub_call(): result: Final = completion( **request_data, ) + return self._to_speech_response(result, response_format) - requested_response_format: Final = optional_params.get("response_format") - if isinstance(result, ModelResponse): - return self.transformation_handler.transform_response( - model_response=result, - response_format=requested_response_format if isinstance(requested_response_format, str) else None, - ) - else: + async def _aspeech( + self, + pending_completion: Coroutine[object, object, object], + response_format: str | None, + logging_obj: "LiteLLMLoggingObj", + ) -> "HttpxBinaryResponseContent": + with logging_obj.preserve_entrypoint_acompletion(), internal_sub_call(): + result: Final = await pending_completion + return self._to_speech_response(result, response_format) + + def _to_speech_response(self, result: object, response_format: str | None) -> "HttpxBinaryResponseContent": + from litellm.types.utils import ModelResponse + + if not isinstance(result, ModelResponse): raise Exception(f"Unmapped response type. Got type: {type(result)}") + return self.transformation_handler.transform_response( + model_response=result, + response_format=response_format, + ) speech_to_completion_bridge_handler = SpeechToCompletionBridgeHandler() diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8cd8e55c292..d5028472ada 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -12,7 +12,8 @@ import subprocess import sys import time import traceback -from collections.abc import Awaitable, Callable, Iterable, Iterator, Mapping, Sequence +from collections.abc import Awaitable, Callable, Generator, Iterable, Iterator, Mapping, Sequence +from contextlib import contextmanager from datetime import datetime as dt_object from functools import lru_cache from types import MappingProxyType, TracebackType @@ -73,6 +74,7 @@ from litellm.litellm_core_utils.classifier_logging import ( is_classifier_call, ) from litellm.litellm_core_utils.core_helpers import ( + as_str_mapping, get_provider_response_headers_from_hidden_params, is_expected_client_error, proxy_stamped_used_client_oauth_token, @@ -1047,6 +1049,16 @@ class Logging(LiteLLMLoggingBaseClass): if "custom_llm_provider" in self.model_call_details: self.custom_llm_provider = self.model_call_details["custom_llm_provider"] + @contextmanager + def preserve_entrypoint_acompletion(self) -> Generator[None]: + acompletion: Final = (as_str_mapping(self._model_call_details_view.get("litellm_params")) or {}).get( + CallTypes.acompletion.value + ) + try: + yield + finally: + self.litellm_params[CallTypes.acompletion.value] = acompletion + def update_from_kwargs( self, kwargs: dict, diff --git a/litellm/main.py b/litellm/main.py index ae6dc13256f..a1ac2bec45b 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -8753,6 +8753,7 @@ def speech( extra_headers=_DEPLOYMENT_HEADERS_HINT.validate_python(kwargs).get("extra_headers"), timeout=timeout, cache_controls=_CACHE_CONTROLS_HINT.validate_python(kwargs), + aspeech=aspeech or False, ) # Vertex AI Text-to-Speech (Google Cloud TTS) @@ -8808,6 +8809,7 @@ def speech( extra_headers=_DEPLOYMENT_HEADERS_HINT.validate_python(kwargs).get("extra_headers"), timeout=timeout, cache_controls=_CACHE_CONTROLS_HINT.validate_python(kwargs), + aspeech=aspeech or False, ) elif custom_llm_provider == "runwayml": from litellm.llms.runwayml.text_to_speech.transformation import ( diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index 181845aabee..367a67dcfa1 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -4547,6 +4547,39 @@ async def test_aspeech_gemini_bridge_logs_only_the_speech_call_under_its_call_id _assert_logged_once_as_speech(events) + +@pytest.mark.asyncio +async def test_aspeech_gemini_bridge_runs_sync_function_callbacks_once_with_the_speech_response( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("GEMINI_API_KEY", raising=False) + monkeypatch.delenv("GOOGLE_API_KEY", raising=False) + logged: Final[list[tuple[type, object]]] = [] # mutable-ok: records what the sync callback receives + + def sync_callback(kwargs: dict[str, object], response_obj: object, start_time: object, end_time: object) -> None: + logged.append((type(response_obj), kwargs.get("litellm_call_id"))) + + monkeypatch.setattr(litellm, "success_callback", [sync_callback]) + respx_mock.post( + url__regex=r"https://generativelanguage\.googleapis\.com/v1beta/models/gemini-2\.5-flash-preview-tts:generateContent.*" + ).mock(return_value=httpx.Response(200, json=_gemini_tts_generate_content_response())) + + await litellm.aspeech( + model="gemini/gemini-2.5-flash-preview-tts", + input="sync callback check", + voice="Kore", + api_key="fake-gemini-key", + litellm_call_id="speech-call-id", + ) + for _ in range(100): + if logged: + break + await asyncio.sleep(0.05) + await asyncio.sleep(0.3) + + assert logged == [(HttpxBinaryResponseContent, "speech-call-id")] + @pytest.mark.asyncio async def test_aspeech_gemini_bridge_logs_a_cached_reply_only_as_the_speech_call( respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch @@ -5063,6 +5096,36 @@ def test_speech_gemini_tts_keeps_cached_audio_only_for_the_request_ttl( assert route.call_count == 2 +async def test_aspeech_gemini_tts_waits_for_a_stalled_upstream_without_holding_a_worker_thread( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream_reached: Final = asyncio.Event() + release_upstream: Final = asyncio.Event() + + async def stalled_upstream(request: httpx.Request) -> httpx.Response: + upstream_reached.set() + await release_upstream.wait() + return httpx.Response(200, json=_gemini_tts_generate_content_response()) + + respx_mock.post(url__regex=GEMINI_DEFAULT_GENERATE_CONTENT_URL).mock(side_effect=stalled_upstream) + loop: Final = asyncio.get_running_loop() + with ThreadPoolExecutor(max_workers=1) as single_worker: + loop.set_default_executor(single_worker) + speech_task: Final = asyncio.create_task( + litellm.aspeech(model=f"gemini/{GEMINI_TTS_MODEL}", input="hello", voice="Kore", api_key="fake-gemini-key") + ) + await asyncio.wait_for(upstream_reached.wait(), timeout=10) + + worker_answer: Final = await asyncio.wait_for(loop.run_in_executor(None, lambda: "free"), timeout=10) + + release_upstream.set() + response: Final = await asyncio.wait_for(speech_task, timeout=10) + + assert worker_answer == "free" + assert b"pcm-audio-bytes" in response.content + + def test_speech_gemini_tts_sends_number_and_boolean_deployment_headers_as_their_config_text( respx_mock: respx.MockRouter, ):