scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-09 22:31:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 86

fix(vertex): make Lyria routing and billing data-driven

Browse source
This commit is contained in:
Emerson Gomes 2026-07-15 20:28:00 -05:00
parent b96844dd0c
commit f18cb0cdb4
No known key found for this signature in database
GPG key ID: D3DF28AB5D1B5E17
8 changed files with 172 additions and 20 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

35
litellm/llms/vertex_ai/common_utils.py
View file

@ -1,9 +1,12 @@
import re
from copy import deepcopy
from enum import Enum
from functools import lru_cache
from typing import Any, Final, Literal, cast, get_type_hints
import httpx
from pydantic import TypeAdapter, ValidationError
from typing_extensions import NotRequired, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -21,6 +24,38 @@ from litellm.types.utils import TokenCountResponse
from litellm.utils import supports_response_schema, supports_system_messages
class VertexAILyriaModelInfo(TypedDict):
vertex_ai_audio_api: Literal["lyria_predict", "lyria_interactions"]
supported_audio_formats: tuple[Literal["mp3", "wav"], ...]
output_cost_per_image: NotRequired[float]
_VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER = TypeAdapter(VertexAILyriaModelInfo)
def _validate_vertex_ai_lyria_model_info(raw_model_info: object) -> VertexAILyriaModelInfo | None:
if raw_model_info is None:
return None
try:
return _VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER.validate_python(raw_model_info)
except ValidationError:
return None
@lru_cache(maxsize=32)
def _get_bundled_vertex_ai_lyria_model_info(model_key: str) -> VertexAILyriaModelInfo | None:
from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap
bundled_model_info = GetModelCostMap.load_local_model_cost_map().get(model_key)
return _validate_vertex_ai_lyria_model_info(bundled_model_info)
def get_vertex_ai_lyria_model_info(model: str) -> VertexAILyriaModelInfo | None:
model_key = model if model.startswith("vertex_ai/") else f"vertex_ai/{model}"
runtime_model_info = _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)
class VertexAIError(BaseLLMException):
def __init__(
self,

36
litellm/llms/vertex_ai/text_to_speech/transformation.py
View file

@ -21,6 +21,10 @@ from litellm.llms.base_llm.text_to_speech.transformation import (
BaseTextToSpeechConfig,
TextToSpeechRequestData,
)
from litellm.llms.vertex_ai.common_utils import (
VertexAILyriaModelInfo,
get_vertex_ai_lyria_model_info,
)
from litellm.llms.vertex_ai.vertex_llm_base import VertexBase
from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
from litellm.types.llms.vertex_ai_text_to_speech import (
@ -476,15 +480,16 @@ class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
LYRIA_MODELS = {
"lyria-002",
"lyria-3-clip-preview",
"lyria-3-pro-preview",
}
@classmethod
def is_lyria_model(cls, model: str) -> bool:
return model.removeprefix("vertex_ai/") in cls.LYRIA_MODELS
return get_vertex_ai_lyria_model_info(model=model) is not None
@staticmethod
def _get_model_info(model: str) -> VertexAILyriaModelInfo:
model_info = get_vertex_ai_lyria_model_info(model=model)
if model_info is None:
raise ValueError(f"Vertex AI model {model!r} does not declare a Lyria audio API")
return model_info
def get_supported_openai_params(self, model: str) -> list:
return ["response_format"]
@ -499,6 +504,7 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
) -> tuple[str | None, dict]:
mapped_params = dict(optional_params)
base_model = model.removeprefix("vertex_ai/")
model_info = self._get_model_info(model=model)
unsupported_params = [param for param in ("speed", "instructions") if mapped_params.get(param) is not None]
if unsupported_params:
if drop_params or litellm.drop_params:
@ -514,9 +520,7 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
),
)
response_format = mapped_params.get("response_format")
supported_formats = (
{"wav"} if base_model == "lyria-002" else {"mp3", "wav"} if base_model == "lyria-3-pro-preview" else {"mp3"}
)
supported_formats = frozenset(model_info["supported_audio_formats"])
if response_format is not None and response_format not in supported_formats:
if drop_params or litellm.drop_params:
mapped_params.pop("response_format", None)
@ -538,6 +542,7 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
litellm_params: dict,
) -> str:
base_model = model.removeprefix("vertex_ai/")
model_info = self._get_model_info(model=model)
project = self.safe_get_vertex_ai_project(litellm_params)
if project is None:
_, project = self._ensure_access_token(
@ -545,7 +550,7 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
project_id=None,
custom_llm_provider="vertex_ai",
)
if base_model.startswith("lyria-3-"):
if model_info["vertex_ai_audio_api"] == "lyria_interactions":
from litellm.llms.vertex_ai.interactions.transformation import (
VertexAIInteractionsConfig,
)
@ -581,7 +586,8 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
}
)
base_model = model.removeprefix("vertex_ai/")
if base_model == "lyria-002":
model_info = self._get_model_info(model=model)
if model_info["vertex_ai_audio_api"] == "lyria_predict":
request_body = {
"instances": [{"prompt": input}],
"parameters": {"sample_count": 1},
@ -605,9 +611,10 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
response_json = raw_response.json()
base_model = model.removeprefix("vertex_ai/")
model_info = self._get_model_info(model=model)
audio_data: str | None = None
mime_type: str | None = None
if base_model == "lyria-002":
if model_info["vertex_ai_audio_api"] == "lyria_predict":
predictions = response_json.get("predictions") or []
if predictions:
audio_data = predictions[0].get("audioContent") or predictions[0].get("bytesBase64Encoded")
@ -621,7 +628,8 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
mime_type = content.get("mime_type")
if audio_data is None:
raise ValueError(f"No generated audio found in Vertex AI {base_model} response")
mime_type = mime_type or ("audio/wav" if base_model == "lyria-002" else "audio/mpeg")
default_format = model_info["supported_audio_formats"][0]
mime_type = mime_type or {"mp3": "audio/mpeg", "wav": "audio/wav"}[default_format]
response = HttpxBinaryResponseContent(
httpx.Response(
status_code=raw_response.status_code,

19
litellm/model_prices_and_context_window_backup.json
View file

@ -45519,6 +45519,9 @@
"mode": "audio_speech",
"output_cost_per_second": 0.002,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria",
"supported_audio_formats": [
"wav"
],
"supported_endpoints": [
"/v1/audio/speech"
],
@ -45528,7 +45531,8 @@
"supported_output_modalities": [
"audio"
],
"supports_audio_output": true
"supports_audio_output": true,
"vertex_ai_audio_api": "lyria_predict"
},
"vertex_ai/lyria-3-clip-preview": {
"input_cost_per_token": 0,
@ -45540,6 +45544,9 @@
"output_cost_per_image": 0.04,
"output_cost_per_token": 0,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria",
"supported_audio_formats": [
"mp3"
],
"supported_endpoints": [
"/v1beta/interactions",
"/v1/audio/speech"
@ -45562,7 +45569,8 @@
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"vertex_ai_audio_api": "lyria_interactions"
},
"vertex_ai/lyria-3-pro-preview": {
"input_cost_per_token": 0,
@ -45574,6 +45582,10 @@
"output_cost_per_image": 0.08,
"output_cost_per_token": 0,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria",
"supported_audio_formats": [
"mp3",
"wav"
],
"supported_endpoints": [
"/v1beta/interactions",
"/v1/audio/speech"
@ -45596,7 +45608,8 @@
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"vertex_ai_audio_api": "lyria_interactions"
},
"vertex_ai/meta/llama-3.1-405b-instruct-maas": {
"input_cost_per_token": 5e-06,

2
litellm/types/utils.py
View file

@ -168,6 +168,8 @@ class ProviderSpecificModelInfo(TypedDict, total=False):
default_reasoning_effort: ReadOnly[Literal["none", "minimal", "low", "medium", "high", "xhigh"] | None]
supports_output_config: bool | None
supports_image_size: bool | None
supported_audio_formats: list[Literal["mp3", "wav"]] | None
vertex_ai_audio_api: Literal["lyria_predict", "lyria_interactions"] | None
bedrock_output_config_effort_ceiling: Literal["low", "medium", "high", "max", "xhigh"] | None
bedrock_converse_supports_strict_tools: bool | None

2
litellm/utils.py
View file

@ -5940,6 +5940,8 @@ def _get_model_info_helper(
provider_specific_entry=_model_info.get("provider_specific_entry", None),
uses_embed_content=_model_info.get("uses_embed_content", None),
supports_image_size=_model_info.get("supports_image_size", None),
supported_audio_formats=_model_info.get("supported_audio_formats", None),
vertex_ai_audio_api=_model_info.get("vertex_ai_audio_api", None),
)
for cost_key, cost_value in _model_info.items():
if cost_key not in returned_model_info and _ABOVE_THRESHOLD_COST_KEY.search(cost_key) is not None:

19
model_prices_and_context_window.json
View file

@ -45519,6 +45519,9 @@
"mode": "audio_speech",
"output_cost_per_second": 0.002,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria",
"supported_audio_formats": [
"wav"
],
"supported_endpoints": [
"/v1/audio/speech"
],
@ -45528,7 +45531,8 @@
"supported_output_modalities": [
"audio"
],
"supports_audio_output": true
"supports_audio_output": true,
"vertex_ai_audio_api": "lyria_predict"
},
"vertex_ai/lyria-3-clip-preview": {
"input_cost_per_token": 0,
@ -45540,6 +45544,9 @@
"output_cost_per_image": 0.04,
"output_cost_per_token": 0,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria",
"supported_audio_formats": [
"mp3"
],
"supported_endpoints": [
"/v1beta/interactions",
"/v1/audio/speech"
@ -45562,7 +45569,8 @@
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"vertex_ai_audio_api": "lyria_interactions"
},
"vertex_ai/lyria-3-pro-preview": {
"input_cost_per_token": 0,
@ -45574,6 +45582,10 @@
"output_cost_per_image": 0.08,
"output_cost_per_token": 0,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria",
"supported_audio_formats": [
"mp3",
"wav"
],
"supported_endpoints": [
"/v1beta/interactions",
"/v1/audio/speech"
@ -45596,7 +45608,8 @@
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": true,
"supports_web_search": false
"supports_web_search": false,
"vertex_ai_audio_api": "lyria_interactions"
},
"vertex_ai/meta/llama-3.1-405b-instruct-maas": {
"input_cost_per_token": 5e-06,

62
tests/test_litellm/llms/vertex_ai/text_to_speech/test_transformation.py
View file

@ -180,6 +180,68 @@ class TestVertexAILyriaTextToSpeechConfig:
assert isinstance(config, VertexAILyriaTextToSpeechConfig)
@pytest.mark.parametrize(
("model", "vertex_ai_audio_api", "supported_audio_formats", "expected_url"),
[
(
"future-lyria-predict",
"lyria_predict",
["wav"],
"https://us-central1-aiplatform.googleapis.com/v1/projects/music-project/locations/"
"us-central1/publishers/google/models/future-lyria-predict:predict",
),
(
"future-music-interactions",
"lyria_interactions",
["mp3", "wav"],
"https://aiplatform.googleapis.com/v1beta1/projects/music-project/locations/global/interactions",
),
],
)
def test_dispatches_from_model_metadata(
self,
monkeypatch,
model,
vertex_ai_audio_api,
supported_audio_formats,
expected_url,
):
monkeypatch.setitem(
litellm.model_cost,
f"vertex_ai/{model}",
{
"vertex_ai_audio_api": vertex_ai_audio_api,
"supported_audio_formats": supported_audio_formats,
},
)
config = ProviderConfigManager.get_provider_text_to_speech_config(
model=model,
provider=LlmProviders.VERTEX_AI,
)
assert isinstance(config, VertexAILyriaTextToSpeechConfig)
assert (
config.get_complete_url(
model=model,
api_base=None,
litellm_params={
"vertex_project": "music-project",
"vertex_location": "us-central1",
},
)
== expected_url
)
def test_vertex_chirp_does_not_select_lyria_config(self):
config = ProviderConfigManager.get_provider_text_to_speech_config(
model="chirp",
provider=LlmProviders.VERTEX_AI,
)
assert isinstance(config, VertexAITextToSpeechConfig)
assert not isinstance(config, VertexAILyriaTextToSpeechConfig)
def test_get_complete_url_for_lyria_2(self):
config = VertexAILyriaTextToSpeechConfig()

17
tests/test_litellm/test_utils.py
View file

@ -1048,6 +1048,17 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"supports_sampling_params": {"type": "boolean"},
"supports_output_config": {"type": "boolean"},
"supports_speed": {"type": "boolean"},
"supported_audio_formats": {
"type": "array",
"items": {
"type": "string",
"enum": ["mp3", "wav"],
},
},
"vertex_ai_audio_api": {
"type": "string",
"enum": ["lyria_predict", "lyria_interactions"],
},
"bedrock_output_config_effort_ceiling": {
"type": "string",
"enum": ["low", "medium", "high", "max", "xhigh"],
@ -2863,9 +2874,15 @@ def test_vertex_ai_lyria_models_in_cost_map():
assert lyria_2["supported_modalities"] == ["text"]
assert lyria_2["supported_output_modalities"] == ["audio"]
assert lyria_2["supports_audio_output"] is True
assert lyria_2["supported_audio_formats"] == ["wav"]
assert lyria_2["vertex_ai_audio_api"] == "lyria_predict"
assert lyria_2["supported_endpoints"] == ["/v1/audio/speech"]
assert clip["output_cost_per_image"] == 0.04
assert pro["output_cost_per_image"] == 0.08
assert clip["supported_audio_formats"] == ["mp3"]
assert pro["supported_audio_formats"] == ["mp3", "wav"]
assert clip["vertex_ai_audio_api"] == "lyria_interactions"
assert pro["vertex_ai_audio_api"] == "lyria_interactions"
assert clip["supported_endpoints"] == [
"/v1beta/interactions",
"/v1/audio/speech",