fix(vertex): preserve Lyria pricing fallback and audio MIME

litellm/cost_calculator.py

@@ -495,17 +495,15 @@ def cost_per_token(
 
     # see this https://learn.microsoft.com/en-us/azure/ai-services/openai/concepts/models
     if call_type == "speech" or call_type == "aspeech":
+        if custom_llm_provider in ("vertex_ai", "vertex_ai_beta"):
+            from litellm.llms.vertex_ai.common_utils import get_vertex_ai_lyria_generation_cost
+
+            lyria_generation_cost: Final = get_vertex_ai_lyria_generation_cost(model_without_prefix)
+            if lyria_generation_cost is not None:
+                return 0.0, lyria_generation_cost
         speech_model_info = litellm.get_model_info(model=model_without_prefix, custom_llm_provider=custom_llm_provider)
         prompt_cost: float = 0.0
         completion_cost: float = 0.0
-        if not speech_model_info.get("input_cost_per_character") and not speech_model_info.get("input_cost_per_token"):
-            output_cost_per_generation: Final = speech_model_info.get("output_cost_per_image")
-            speech_output_cost_per_second: Final = speech_model_info.get("output_cost_per_second")
-            audio_seconds_per_prediction: Final = speech_model_info.get("audio_seconds_per_prediction")
-            if output_cost_per_generation is not None:
-                return prompt_cost, float(output_cost_per_generation)
-            if speech_output_cost_per_second is not None and audio_seconds_per_prediction is not None:
-                return prompt_cost, float(speech_output_cost_per_second) * float(audio_seconds_per_prediction)
         cost_metric: Final = select_cost_metric_for_model(speech_model_info)
         if cost_metric == "cost_per_character":
             if prompt_characters is None:

litellm/llms/vertex_ai/common_utils.py

@@ -2,6 +2,7 @@ import re
 from copy import deepcopy
 from enum import Enum
 from functools import lru_cache
+from types import MappingProxyType
 from typing import Any, Final, Literal, cast, get_type_hints
 
 import httpx
@@ -55,7 +56,26 @@ def _get_bundled_vertex_ai_lyria_model_info(model_key: str) -> VertexAILyriaMode
 def get_vertex_ai_lyria_model_info(model: str) -> VertexAILyriaModelInfo | None:
     model_key: Final = model if model.startswith("vertex_ai/") else f"vertex_ai/{model}"
     runtime_model_info: Final = _validate_vertex_ai_lyria_model_info(litellm.model_cost.get(model_key))
-    return runtime_model_info or _get_bundled_vertex_ai_lyria_model_info(model_key)
+    bundled_model_info: Final = _get_bundled_vertex_ai_lyria_model_info(model_key)
+    if runtime_model_info is None:
+        return bundled_model_info
+    if bundled_model_info is None:
+        return runtime_model_info
+    return _validate_vertex_ai_lyria_model_info(MappingProxyType({**bundled_model_info, **runtime_model_info}))
+
+
+def get_vertex_ai_lyria_generation_cost(model: str) -> float | None:
+    model_info: Final = get_vertex_ai_lyria_model_info(model)
+    if model_info is None:
+        return None
+    generation_cost: Final = model_info.get("output_cost_per_image")
+    if generation_cost is not None:
+        return generation_cost
+    cost_per_second: Final = model_info.get("output_cost_per_second")
+    seconds_per_prediction: Final = model_info.get("audio_seconds_per_prediction")
+    if cost_per_second is None or seconds_per_prediction is None:
+        return None
+    return cost_per_second * seconds_per_prediction
 
 
 class VertexAIError(BaseLLMException):

litellm/llms/vertex_ai/text_to_speech/transformation.py

@@ -681,9 +681,11 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
                         )  # rebind-ok: interactions response supplies its audio MIME type
         if audio_data is None:
             raise ValueError(f"No generated audio found in Vertex AI {base_model} response")
+        decoded_audio: Final = base64.b64decode(audio_data)
         default_format: Final = model_info["supported_audio_formats"][0]
         mime_type = (
             mime_type
+            or speech_media_type_from_audio_bytes(decoded_audio)
             or {  # mutable-ok: short-lived lookup selects the default response MIME type; rebind-ok: absent provider MIME type falls back to model metadata
                 "mp3": "audio/mpeg",
                 "wav": "audio/wav",
@@ -692,7 +694,7 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
         response: Final = HttpxBinaryResponseContent(
             httpx.Response(
                 status_code=raw_response.status_code,
-                content=base64.b64decode(audio_data),
+                content=decoded_audio,
                 headers={  # mutable-ok: httpx requires a concrete response header dictionary
                     "content-type": mime_type
                 },

litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py

@@ -12,7 +12,7 @@ from litellm._logging import verbose_proxy_logger
 from litellm.constants import VERTEX_BATCH_PREDICTION_JOBS_ROUTE
 from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
 from litellm.llms.vertex_ai.common_utils import (
-    get_vertex_ai_lyria_model_info,
+    get_vertex_ai_lyria_generation_cost,
     get_vertex_location_from_url,
 )
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
@@ -397,9 +397,7 @@ class VertexPassthroughLoggingHandler:
         )
         if runtime_unit_cost is not None:
             return runtime_unit_cost
-        return VertexPassthroughLoggingHandler._audio_prediction_unit_cost_from_model_info(
-            model_info=get_vertex_ai_lyria_model_info(model=model)
-        )
+        return get_vertex_ai_lyria_generation_cost(model=model)
 
     @staticmethod
     def _audio_prediction_unit_cost_from_model_info(model_info: object) -> float | None:

tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py

@@ -146,13 +146,27 @@ def test_audio_predict_response_supports_bytes_base64_encoded(
     assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.06)
 
 
-def test_lyria_predict_cost_falls_back_to_bundled_map_when_runtime_map_omits_model(
+@pytest.mark.parametrize(
+    "missing_fields",
+    (
+        None,
+        ("output_cost_per_second",),
+        ("audio_seconds_per_prediction",),
+        ("output_cost_per_second", "audio_seconds_per_prediction"),
+    ),
+)
+def test_lyria_predict_cost_falls_back_to_bundled_map_when_runtime_metadata_is_incomplete(
     monkeypatch: pytest.MonkeyPatch,
+    missing_fields: tuple[str, ...] | None,
 ) -> None:
-    stale_runtime_model_cost: Final = {
-        key: value for key, value in litellm.model_cost.items() if key != "vertex_ai/lyria-002"
-    }
-    monkeypatch.setattr(litellm, "model_cost", stale_runtime_model_cost)
+    if missing_fields is None:
+        monkeypatch.delitem(litellm.model_cost, "vertex_ai/lyria-002")
+    else:
+        monkeypatch.setitem(
+            litellm.model_cost,
+            "vertex_ai/lyria-002",
+            {key: value for key, value in litellm.model_cost["vertex_ai/lyria-002"].items() if key not in missing_fields},
+        )
     logging_obj = MagicMock()
     logging_obj.model_call_details = {}
     response = httpx.Response(
@@ -178,7 +192,8 @@ def test_lyria_predict_cost_falls_back_to_bundled_map_when_runtime_map_omits_mod
         request_body={"instances": [{"prompt": "ambient piano"}]},
     )
 
-    assert "vertex_ai/lyria-002" not in litellm.model_cost
+    if missing_fields is None:
+        assert "vertex_ai/lyria-002" not in litellm.model_cost
     assert result["kwargs"]["model"] == "lyria-002"
     assert result["kwargs"]["response_cost"] == pytest.approx(0.06)
     assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.06)

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

@@ -169,6 +169,23 @@ def test_transform_text_to_speech_response_leaves_unknown_bytes_unlabeled():
 
 
 class TestVertexAILyriaTextToSpeechConfig:
+    def test_response_without_mime_type_uses_audio_container(self) -> None:
+        audio: Final = b"RIFF\x24\x00\x00\x00WAVEfmt \x10\x00\x00\x00"
+        raw_response: Final = httpx.Response(
+            200,
+            json={"outputs": [{"type": "audio", "data": base64.b64encode(audio).decode()}]},
+        )
+
+        response: Final = VertexAILyriaTextToSpeechConfig().transform_text_to_speech_response(
+            model="lyria-3-pro-preview",
+            raw_response=raw_response,
+            logging_obj=MagicMock(),
+        )
+
+        assert response.content == audio
+        assert response.response.headers["content-type"] == "audio/wav"
+        assert response._hidden_params["audio_mime_type"] == "audio/wav"
+
     @pytest.mark.parametrize(
         "model",
         ["lyria-002", "vertex_ai/lyria-3-clip-preview", "lyria-3-pro-preview"],

tests/test_litellm/test_cost_calculator.py

@@ -1,6 +1,7 @@
 
 import json
 from pathlib import Path
+from typing import Final
 
 import pytest
 
@@ -154,14 +155,42 @@ def test_cost_calculator_with_response_cost_in_additional_headers():
         ("vertex_ai/lyria-3-pro-preview", 0.08),
     ],
 )
-def test_vertex_lyria_speech_cost(model, expected_cost, _local_model_cost_map):
-    cost = completion_cost(
+@pytest.mark.parametrize("runtime_state", ("complete", "missing", "routing_only", "custom_zero", "custom_price"))
+@pytest.mark.parametrize("call_type", ("speech", "aspeech"))
+def test_vertex_lyria_speech_cost(
+    model: str,
+    expected_cost: float,
+    _local_model_cost_map: None,
+    monkeypatch: pytest.MonkeyPatch,
+    runtime_state: str,
+    call_type: str,
+) -> None:
+    model_info: Final = litellm.model_cost[model]
+    if runtime_state == "missing":
+        monkeypatch.delitem(litellm.model_cost, model)
+    elif runtime_state == "routing_only":
+        monkeypatch.setitem(
+            litellm.model_cost,
+            model,
+            {
+                key: value
+                for key, value in model_info.items()
+                if key not in ("output_cost_per_image", "output_cost_per_second", "audio_seconds_per_prediction")
+            },
+        )
+    elif runtime_state in ("custom_zero", "custom_price"):
+        cost_key: Final = "output_cost_per_image" if "output_cost_per_image" in model_info else "output_cost_per_second"
+        multiplier: Final = 0 if runtime_state == "custom_zero" else 2
+        monkeypatch.setitem(litellm.model_cost, model, {**model_info, cost_key: model_info[cost_key] * multiplier})
+
+    cost: Final = completion_cost(
         model=model,
         prompt="A bright synth track",
-        call_type="speech",
+        call_type=call_type,
     )
 
-    assert cost == pytest.approx(expected_cost)
+    expected: Final = 0 if runtime_state == "custom_zero" else expected_cost * (2 if runtime_state == "custom_price" else 1)
+    assert cost == pytest.approx(expected)
 
 
 def test_baseten_model_api_pricing_entries(_local_model_cost_map):