From 82edb9e901c7b87a99a7b5c318b3db333ca6165b Mon Sep 17 00:00:00 2001 From: Emerson Gomes Date: Sun, 6 Sep 2026 00:09:27 -0500 Subject: [PATCH] fix(vertex): preserve Lyria pricing fallback and audio MIME --- litellm/cost_calculator.py | 14 +++---- litellm/llms/vertex_ai/common_utils.py | 22 ++++++++++- .../text_to_speech/transformation.py | 4 +- .../vertex_passthrough_logging_handler.py | 6 +-- ...test_vertex_passthrough_logging_handler.py | 27 +++++++++++--- .../text_to_speech/test_transformation.py | 17 +++++++++ tests/test_litellm/test_cost_calculator.py | 37 +++++++++++++++++-- 7 files changed, 103 insertions(+), 24 deletions(-) diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 5dcbb1d5f37..fa8e1a5264c 100644 --- a/litellm/cost_calculator.py +++ b/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: diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index b649d8dac0e..3895ea25434 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/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): diff --git a/litellm/llms/vertex_ai/text_to_speech/transformation.py b/litellm/llms/vertex_ai/text_to_speech/transformation.py index 2047be2ea37..5a93aca3ac9 100644 --- a/litellm/llms/vertex_ai/text_to_speech/transformation.py +++ b/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 }, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index f7625d71168..afdbaf3b06c 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/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: diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py b/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py index f3d28d58c22..179301a9c81 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py +++ b/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) diff --git a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py b/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py index 457c3f76dbb..4f4df69780a 100644 --- a/tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py +++ b/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"], diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index ffc71c76d03..aa630c4916b 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/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):