fix(speech): send the deployment's extra_headers and chat's caching rule through the TTS bridge

Backport of #50803 to lts/1.108.x.
Cherry-picked from 33584bcc0d (litellm_speech_bridge_api_base).

litellm/endpoints/speech/speech_to_completion_bridge/handler.py

@@ -2,6 +2,7 @@
 Handler for transforming /chat/completions api requests to litellm.responses requests
 """
 
+from collections.abc import Mapping
 from typing import TYPE_CHECKING, Final
 
 from typing_extensions import TypedDict
@@ -89,6 +90,7 @@ class SpeechToCompletionBridgeHandler:
         logging_obj: "LiteLLMLoggingObj",
         custom_llm_provider: str,
         api_base: str | None,
+        extra_headers: Mapping[str, str] | None,
         timeout: "float | httpx.Timeout | None",
         cache_controls: "CacheControlsHint",
     ) -> "HttpxBinaryResponseContent":
@@ -110,7 +112,13 @@ class SpeechToCompletionBridgeHandler:
             model=model,
             input=input,
             optional_params=optional_params,
-            litellm_params={**litellm_params, "api_base": api_base, "timeout": timeout, **cache_controls},
+            litellm_params={
+                **litellm_params,
+                "api_base": api_base,
+                "extra_headers": extra_headers,
+                "timeout": timeout,
+                **cache_controls,
+            },
             headers=headers,
             litellm_logging_obj=logging_obj,
             custom_llm_provider=custom_llm_provider,

litellm/main.py

@@ -149,6 +149,7 @@ from litellm.types.utils import (
     BaseUrlRoutingHint,
     CacheControlsHint,
     CustomPricingLiteLLMParams,
+    DeploymentHeadersHint,
     ModelResponseStream,
     ProviderRoutingHints,
     RawRequestTypedDict,
@@ -4139,6 +4140,7 @@ _MODEL_INFO_ADAPTER: Final = TypeAdapter(Mapping[str, object])
 _PROVIDER_ROUTING_HINTS: Final = TypeAdapter(ProviderRoutingHints)
 _BASE_URL_HINT: Final = TypeAdapter(BaseUrlRoutingHint)
 _CACHE_CONTROLS_HINT: Final = TypeAdapter(CacheControlsHint)
+_DEPLOYMENT_HEADERS_HINT: Final = TypeAdapter(DeploymentHeadersHint)
 
 
 _AWS_BEDROCK_CLIENT_CREDENTIAL_KEYS: Final = (
@@ -8748,6 +8750,7 @@ def speech(
                 logging_obj=logging_obj,
                 custom_llm_provider=custom_llm_provider,
                 api_base=api_base,
+                extra_headers=_DEPLOYMENT_HEADERS_HINT.validate_python(kwargs).get("extra_headers"),
                 timeout=timeout,
                 cache_controls=_CACHE_CONTROLS_HINT.validate_python(kwargs),
             )
@@ -8802,6 +8805,7 @@ def speech(
             logging_obj=logging_obj,
             custom_llm_provider=custom_llm_provider,
             api_base=api_base,
+            extra_headers=_DEPLOYMENT_HEADERS_HINT.validate_python(kwargs).get("extra_headers"),
             timeout=timeout,
             cache_controls=_CACHE_CONTROLS_HINT.validate_python(kwargs),
         )

litellm/types/utils.py

@@ -43,7 +43,6 @@ from pydantic import (
     OnErrorOmit,
     PrivateAttr,
     SkipValidation,
-    StrictBool,
     field_serializer,
     field_validator,
 )
@@ -4874,4 +4873,8 @@ class BaseUrlRoutingHint(TypedDict, total=False):
 
 class CacheControlsHint(TypedDict, total=False):
     cache: ReadOnly[OnErrorOmit[Mapping[str, object] | None]]
-    caching: ReadOnly[OnErrorOmit[StrictBool | None]]
+    caching: ReadOnly[OnErrorOmit[JsonValue]]
+
+
+class DeploymentHeadersHint(TypedDict, total=False):
+    extra_headers: ReadOnly[OnErrorOmit[Mapping[str, str] | None]]

tests/unit/llms/vertex_ai/text_to_speech/test_transformation.py

@@ -1140,10 +1140,12 @@ async def test_aspeech_vertex_ai_gemini_tts_routes_to_configured_api_base(
         vertex_credentials=_AUTHORIZED_USER,
         vertex_location="global",
         api_base="https://vertex.gateway.internal",
+        extra_headers={"x-gateway-key": "gateway-secret"},
     )
 
     assert gateway_route.call_count == 1
     assert gateway_route.calls.last.request.headers["authorization"] == "Bearer minted-google-token"
+    assert gateway_route.calls.last.request.headers["x-gateway-key"] == "gateway-secret"
     assert _PCM_AUDIO in response.content
 
 

tests/unit/test_main.py

@@ -4798,6 +4798,7 @@ async def test_speech_gemini_tts_routes_to_configured_api_base(
         "input": "hello from litellm",
         "voice": "Kore",
         "api_key": "fake-gemini-key",
+        "extra_headers": {"x-gateway-key": "gateway-secret"},
         endpoint_param: GEMINI_GATEWAY_BASE,
     }
 
@@ -4805,6 +4806,7 @@ async def test_speech_gemini_tts_routes_to_configured_api_base(
 
     assert gateway_route.call_count == 1
     assert gateway_route.calls.last.request.headers["x-goog-api-key"] == "fake-gemini-key"
+    assert gateway_route.calls.last.request.headers["x-gateway-key"] == "gateway-secret"
     assert b"pcm-audio-bytes" in response.content
 
 
@@ -4838,6 +4840,34 @@ def test_speech_gemini_tts_without_a_timeout_waits_as_long_as_the_same_chat_comp
     assert speech_call.request.extensions["timeout"] == chat_call.request.extensions["timeout"]
 
 
+async def test_speech_gemini_tts_behind_the_router_waits_as_long_as_the_same_chat_completion(
+    respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
+):
+    monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
+    route: Final = respx_mock.post(url__regex=GEMINI_DEFAULT_GENERATE_CONTENT_URL).mock(
+        return_value=httpx.Response(200, json=_gemini_tts_generate_content_response())
+    )
+    router: Final = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gemini-tts",
+                "litellm_params": {"model": f"gemini/{GEMINI_TTS_MODEL}", "api_key": "fake-gemini-key"},
+            }
+        ]
+    )
+
+    await router.aspeech(model="gemini-tts", input="hello", voice="Kore")
+    await router.acompletion(
+        model="gemini-tts",
+        messages=[{"role": "user", "content": "hello"}],
+        modalities=["audio"],
+        audio={"voice": "Kore", "format": "pcm16"},
+    )
+
+    speech_call, chat_call = route.calls
+    assert speech_call.request.extensions["timeout"] == chat_call.request.extensions["timeout"]
+
+
 @pytest.mark.parametrize(
     ("first_controls", "second_controls", "expected_upstream_calls"),
     [
@@ -4845,6 +4875,8 @@ def test_speech_gemini_tts_without_a_timeout_waits_as_long_as_the_same_chat_comp
         ({}, {"cache": {"no-cache": True}}, 2),
         ({"cache": {"no-store": True}}, {}, 2),
         ({}, {"caching": False}, 2),
+        ({}, {"caching": "yes"}, 2),
+        ({}, {"caching": 0}, 2),
     ],
 )
 def test_speech_gemini_tts_honors_the_request_cache_controls(
@@ -4871,8 +4903,11 @@ def test_speech_gemini_tts_honors_the_request_cache_controls(
     assert route.call_count == expected_upstream_calls
 
 
-@pytest.mark.parametrize("malformed_controls", [{"caching": "yes"}, {"caching": 1}, {"cache": {1: True}}])
-def test_speech_gemini_tts_without_a_cache_ignores_malformed_cache_controls(
+@pytest.mark.parametrize(
+    "malformed_controls",
+    [{"caching": "yes"}, {"caching": 1}, {"cache": {1: True}}, {"extra_headers": {"x-gateway-key": 5}}],
+)
+def test_speech_gemini_tts_without_a_cache_ignores_malformed_controls(
     respx_mock: respx.MockRouter,
     monkeypatch: pytest.MonkeyPatch,
     malformed_controls: Mapping[str, object],