feat(vertex): add Lyria model support

This commit is contained in:
Emerson Gomes 2026-06-19 16:43:02 -05:00
parent 8c58b93572
commit 3ead9d1688
No known key found for this signature in database
GPG key ID: D3DF28AB5D1B5E17
6 changed files with 339 additions and 0 deletions

View file

@ -0,0 +1,3 @@
from .transformation import VertexAIInteractionsConfig
__all__ = ["VertexAIInteractionsConfig"]

View file

@ -45511,6 +45511,87 @@
"output_cost_per_token": 4e-07,
"supports_tool_choice": true
},
"vertex_ai/lyria-002": {
"litellm_provider": "vertex_ai",
"max_audio_length_hours": 0.009111111111111111,
"max_audio_per_prompt": 4,
"mode": "chat",
"output_cost_per_second": 0.002,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"audio"
],
"supports_audio_output": true
},
"vertex_ai/lyria-3-clip-preview": {
"input_cost_per_token": 0,
"litellm_provider": "vertex_ai",
"max_input_tokens": 131072,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"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_endpoints": [
"/v1beta/interactions"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"audio"
],
"supported_regions": [
"global"
],
"supports_audio_input": false,
"supports_audio_output": true,
"supports_function_calling": false,
"supports_image_input": true,
"supports_prompt_caching": false,
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": true,
"supports_web_search": false
},
"vertex_ai/lyria-3-pro-preview": {
"input_cost_per_token": 0,
"litellm_provider": "vertex_ai",
"max_input_tokens": 131072,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"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_endpoints": [
"/v1beta/interactions"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"audio"
],
"supported_regions": [
"global"
],
"supports_audio_input": false,
"supports_audio_output": true,
"supports_function_calling": false,
"supports_image_input": true,
"supports_prompt_caching": false,
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": true,
"supports_web_search": false
},
"vertex_ai/meta/llama-3.1-405b-instruct-maas": {
"input_cost_per_token": 5e-06,
"litellm_provider": "vertex_ai-llama_models",

View file

@ -46,6 +46,7 @@ else:
# Define EndpointType locally to avoid import issues
EndpointType = Any
_LYRIA_SECONDS_PER_AUDIO_PREDICTION = 30
class VertexPassthroughLoggingHandler:
@ -270,6 +271,16 @@ class VertexPassthroughLoggingHandler:
_json_response: Final[dict[str, object]] = httpx_response.json()
litellm_prediction_response: ModelResponse | EmbeddingResponse | ImageResponse = ModelResponse()
if VertexPassthroughLoggingHandler._is_lyria_predict_response(
model=model,
json_response=_json_response,
):
return VertexPassthroughLoggingHandler._handle_lyria_predict_response(
json_response=_json_response,
logging_obj=logging_obj,
model=model,
kwargs=kwargs,
)
if vertex_image_generation_class.is_image_generation_response(_json_response):
litellm_prediction_response = vertex_image_generation_class.process_image_generation_response(
_json_response,
@ -323,6 +334,64 @@ class VertexPassthroughLoggingHandler:
"kwargs": kwargs,
}
@staticmethod
def _handle_lyria_predict_response(
json_response: dict,
logging_obj: LiteLLMLoggingObj,
model: str,
kwargs: dict,
) -> PassThroughEndpointLoggingTypedDict:
prediction_count: Final = (
VertexPassthroughLoggingHandler._get_lyria_audio_prediction_count(
json_response=json_response
)
)
model_info: Final = litellm.model_cost.get(f"vertex_ai/{model}", {})
response_cost: Final = (
model_info.get("output_cost_per_second", 0.0)
* _LYRIA_SECONDS_PER_AUDIO_PREDICTION
* prediction_count
)
logging_obj.model = model
logging_obj.model_call_details["model"] = model
logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai"
logging_obj.custom_llm_provider = "vertex_ai"
logging_obj.model_call_details["response_cost"] = response_cost
kwargs["response_cost"] = response_cost
kwargs["model"] = model
kwargs["custom_llm_provider"] = "vertex_ai"
standard_pass_through_response_object: Final[StandardPassThroughResponseObject] = {
"response": json_response,
}
return {
"result": standard_pass_through_response_object,
"kwargs": kwargs,
}
@staticmethod
def _is_lyria_predict_response(model: str, json_response: dict) -> bool:
return (
model == "lyria-002"
and VertexPassthroughLoggingHandler._get_lyria_audio_prediction_count(
json_response=json_response
)
> 0
)
@staticmethod
def _get_lyria_audio_prediction_count(json_response: dict) -> int:
predictions: Final = json_response.get("predictions")
if not isinstance(predictions, list):
return 0
return sum(
1
for prediction in predictions
if isinstance(prediction, dict) and prediction.get("audioContent")
)
@staticmethod
def _extract_embed_content_input(request_body: dict | None, batch: bool) -> str:
"""Extract raw input text from an :embedContent or :batchEmbedContents request body for token counting."""

View file

@ -45511,6 +45511,87 @@
"output_cost_per_token": 4e-07,
"supports_tool_choice": true
},
"vertex_ai/lyria-002": {
"litellm_provider": "vertex_ai",
"max_audio_length_hours": 0.009111111111111111,
"max_audio_per_prompt": 4,
"mode": "chat",
"output_cost_per_second": 0.002,
"source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing#lyria",
"supported_modalities": [
"text"
],
"supported_output_modalities": [
"audio"
],
"supports_audio_output": true
},
"vertex_ai/lyria-3-clip-preview": {
"input_cost_per_token": 0,
"litellm_provider": "vertex_ai",
"max_input_tokens": 131072,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"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_endpoints": [
"/v1beta/interactions"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"audio"
],
"supported_regions": [
"global"
],
"supports_audio_input": false,
"supports_audio_output": true,
"supports_function_calling": false,
"supports_image_input": true,
"supports_prompt_caching": false,
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": true,
"supports_web_search": false
},
"vertex_ai/lyria-3-pro-preview": {
"input_cost_per_token": 0,
"litellm_provider": "vertex_ai",
"max_input_tokens": 131072,
"max_output_tokens": 8192,
"max_tokens": 8192,
"mode": "chat",
"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_endpoints": [
"/v1beta/interactions"
],
"supported_modalities": [
"text",
"image"
],
"supported_output_modalities": [
"audio"
],
"supported_regions": [
"global"
],
"supports_audio_input": false,
"supports_audio_output": true,
"supports_function_calling": false,
"supports_image_input": true,
"supports_prompt_caching": false,
"supports_response_schema": false,
"supports_system_messages": false,
"supports_vision": true,
"supports_web_search": false
},
"vertex_ai/meta/llama-3.1-405b-instruct-maas": {
"input_cost_per_token": 5e-06,
"litellm_provider": "vertex_ai-llama_models",

View file

@ -0,0 +1,68 @@
from datetime import datetime
from unittest.mock import MagicMock
import httpx
import litellm
import pytest
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import (
VertexPassthroughLoggingHandler,
)
def test_lyria_predict_response_preserves_audio_response_and_logs_cost(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setitem(
litellm.model_cost,
"vertex_ai/lyria-002",
{"output_cost_per_second": 0.002},
)
logging_obj = MagicMock()
logging_obj.model_call_details = {}
response = httpx.Response(
status_code=200,
json={
"predictions": [
{
"audioContent": "clip-1",
"mimeType": "audio/wav",
},
{
"audioContent": "clip-2",
"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 result["result"] == {
"response": {
"predictions": [
{
"audioContent": "clip-1",
"mimeType": "audio/wav",
},
{
"audioContent": "clip-2",
"mimeType": "audio/wav",
},
]
}
}
assert result["kwargs"]["model"] == "lyria-002"
assert result["kwargs"]["custom_llm_provider"] == "vertex_ai"
assert result["kwargs"]["response_cost"] == pytest.approx(0.12)
assert logging_obj.model == "lyria-002"
assert logging_obj.model_call_details["response_cost"] == pytest.approx(0.12)

View file

@ -1067,6 +1067,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid():
"/v1/images/variations",
"/v1/images/edits",
"/v1/batch",
"/v1beta/interactions",
"/v1/audio/transcriptions",
"/v1/audio/speech",
"/v1/ocr",
@ -2833,6 +2834,42 @@ def test_gemini_lyria_3_preview_models_in_cost_map():
assert clip["output_cost_per_image"] == 0.04
def test_vertex_ai_lyria_models_in_cost_map():
import json
from pathlib import Path
json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json"
with open(json_path) as f:
model_cost = json.load(f)
lyria_2 = model_cost.get("vertex_ai/lyria-002")
clip = model_cost.get("vertex_ai/lyria-3-clip-preview")
pro = model_cost.get("vertex_ai/lyria-3-pro-preview")
assert lyria_2 is not None
assert clip is not None
assert pro is not None
assert lyria_2["litellm_provider"] == "vertex_ai"
assert clip["litellm_provider"] == "vertex_ai"
assert pro["litellm_provider"] == "vertex_ai"
assert lyria_2["output_cost_per_second"] == 0.002
assert lyria_2["supported_modalities"] == ["text"]
assert lyria_2["supported_output_modalities"] == ["audio"]
assert lyria_2["supports_audio_output"] is True
assert clip["output_cost_per_image"] == 0.04
assert pro["output_cost_per_image"] == 0.08
assert clip["supported_endpoints"] == ["/v1beta/interactions"]
assert pro["supported_endpoints"] == ["/v1beta/interactions"]
assert clip["supported_modalities"] == ["text", "image"]
assert pro["supported_modalities"] == ["text", "image"]
assert clip["supported_regions"] == ["global"]
assert pro["supported_regions"] == ["global"]
assert clip["supports_audio_output"] is True
assert pro["supports_audio_output"] is True
assert clip["supports_image_input"] is True
assert pro["supports_image_input"] is True
def test_model_info_for_fireworks_short_form_models():
"""
Test that fireworks_ai short-form model entries (fireworks_ai/<model>)