fix(vertex-ai): fall back to bundled Lyria costs

Prefer runtime model_cost when both numeric fields are present, then bundled Lyria metadata so stale maps still bill 0.002 * 30

litellm/llms/vertex_ai/common_utils.py

@@ -28,6 +28,8 @@ class VertexAILyriaModelInfo(TypedDict):
     vertex_ai_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"]]
     supported_audio_formats: ReadOnly[tuple[Literal["mp3", "wav"], ...]]
     output_cost_per_image: NotRequired[ReadOnly[float]]
+    output_cost_per_second: NotRequired[ReadOnly[float]]
+    audio_seconds_per_prediction: NotRequired[ReadOnly[float]]
 
 
 _VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER: Final = TypeAdapter(VertexAILyriaModelInfo)

litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py

@@ -1,5 +1,6 @@
 import asyncio
 import re
+from collections.abc import Mapping
 from datetime import datetime
 from typing import TYPE_CHECKING, Any, Final, cast
 from urllib.parse import urlparse
@@ -10,7 +11,10 @@ import litellm
 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_location_from_url
+from litellm.llms.vertex_ai.common_utils import (
+    get_vertex_ai_lyria_model_info,
+    get_vertex_location_from_url,
+)
 from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
     ModelResponseIterator as VertexModelResponseIterator,
 )
@@ -388,8 +392,18 @@ class VertexPassthroughLoggingHandler:
 
     @staticmethod
     def _get_audio_prediction_unit_cost(model: str) -> float | None:
-        model_info: Final = litellm.model_cost.get(f"vertex_ai/{model}")
-        if model_info is None:
+        runtime_unit_cost: Final = VertexPassthroughLoggingHandler._audio_prediction_unit_cost_from_model_info(
+            model_info=litellm.model_cost.get(f"vertex_ai/{model}")
+        )
+        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)
+        )
+
+    @staticmethod
+    def _audio_prediction_unit_cost_from_model_info(model_info: object) -> float | None:
+        if not isinstance(model_info, Mapping):
             return None
         output_cost_per_second: Final = model_info.get("output_cost_per_second")
         audio_seconds_per_prediction: Final = model_info.get("audio_seconds_per_prediction")

tests/test_litellm/llms/vertex_ai/test_vertex_passthrough_logging_handler.py

@@ -1,4 +1,5 @@
 from datetime import datetime
+from typing import Final
 from unittest.mock import MagicMock
 
 import httpx
@@ -143,3 +144,41 @@ def test_audio_predict_response_supports_bytes_base64_encoded(
 
     assert result["kwargs"]["response_cost"] == pytest.approx(0.06)
     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(
+    monkeypatch: pytest.MonkeyPatch,
+) -> 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)
+    logging_obj = MagicMock()
+    logging_obj.model_call_details = {}
+    response = httpx.Response(
+        status_code=200,
+        json={
+            "predictions": [
+                {
+                    "audioContent": "clip",
+                    "mimeType": "audio/wav",
+                }
+            ]
+        },
+    )
+
+    result = VertexPassthroughLoggingHandler.vertex_passthrough_handler(
+        httpx_response=response,
+        logging_obj=logging_obj,
+        url_route="/v1/projects/test/locations/us-central1/publishers/google/models/lyria-002:predict",
+        result=response.text,
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+        cache_hit=False,
+        request_body={"instances": [{"prompt": "ambient piano"}]},
+    )
+
+    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)