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