style(vertex-ai): satisfy Lyria quality gates

This commit is contained in:
Emerson Gomes 2026-08-12 13:33:33 -05:00
parent bd5123564c
commit 6ec53f2846
No known key found for this signature in database
GPG key ID: D3DF28AB5D1B5E17
7 changed files with 153 additions and 89 deletions

View file

@ -498,16 +498,14 @@ def cost_per_token(
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"
):
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")
output_cost_per_second: Final = speech_model_info.get("output_cost_per_second")
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 output_cost_per_second is not None and audio_seconds_per_prediction is not None:
return prompt_cost, float(output_cost_per_second) * float(audio_seconds_per_prediction)
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:

View file

@ -6,7 +6,7 @@ from typing import Any, Final, Literal, cast, get_type_hints
import httpx
from pydantic import TypeAdapter, ValidationError
from typing_extensions import NotRequired, TypedDict
from typing_extensions import NotRequired, ReadOnly, TypedDict
import litellm
from litellm._logging import verbose_logger
@ -25,12 +25,12 @@ 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_audio_api: ReadOnly[Literal["lyria_predict", "lyria_interactions"]]
supported_audio_formats: ReadOnly[tuple[Literal["mp3", "wav"], ...]]
output_cost_per_image: NotRequired[ReadOnly[float]]
_VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER = TypeAdapter(VertexAILyriaModelInfo)
_VERTEX_AI_LYRIA_MODEL_INFO_ADAPTER: Final = TypeAdapter(VertexAILyriaModelInfo)
def _validate_vertex_ai_lyria_model_info(raw_model_info: object) -> VertexAILyriaModelInfo | None:
@ -46,13 +46,13 @@ def _validate_vertex_ai_lyria_model_info(raw_model_info: object) -> VertexAILyri
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)
bundled_model_info: Final = 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))
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)

View file

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

View file

@ -8,7 +8,7 @@ Reference: https://cloud.google.com/text-to-speech/docs/reference/rest/v1/text/s
import base64
from collections.abc import Coroutine
from types import MappingProxyType
from typing import TYPE_CHECKING, Any, Final, Union
from typing import TYPE_CHECKING, Any, Final, TypeAlias, Union
import httpx
@ -40,6 +40,10 @@ else:
LiteLLMLoggingObj = Any
HttpxBinaryResponseContent = Any
_LyriaVoice: TypeAlias = (
str | dict | None
) # mutable-ok: inherited interface supports structured provider voice dictionaries
class VertexAITextToSpeechConfig(BaseTextToSpeechConfig, VertexBase):
"""
@ -486,26 +490,34 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
@staticmethod
def _get_model_info(model: str) -> VertexAILyriaModelInfo:
model_info = get_vertex_ai_lyria_model_info(model=model)
model_info: Final = 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"]
def get_supported_openai_params(
self, model: str
) -> list: # mutable-ok: inherited provider interface returns a concrete parameter list
return [ # mutable-ok: inherited provider interface requires a concrete parameter list
"response_format"
]
def map_openai_params(
self,
model: str,
optional_params: dict,
voice: str | dict | None = None,
optional_params: dict, # mutable-ok: inherited provider interface accepts a concrete parameter dictionary
voice: _LyriaVoice = None,
drop_params: bool = False,
kwargs: dict | None = None,
) -> 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]
kwargs: dict | None = None, # mutable-ok: inherited provider interface accepts a concrete keyword dictionary
) -> tuple[str | None, dict]: # mutable-ok: inherited provider interface returns concrete mapped parameters
mapped_params: Final = dict( # mutable-ok: mapping drops unsupported parameters before provider dispatch
optional_params
)
base_model: Final = model.removeprefix("vertex_ai/")
model_info: Final = self._get_model_info(model=model)
unsupported_params: Final = tuple(
param for param in ("speed", "instructions") if mapped_params.get(param) is not None
)
if unsupported_params:
if drop_params or litellm.drop_params:
for param in unsupported_params:
@ -519,8 +531,8 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
"from the call, set `litellm.drop_params = True`"
),
)
response_format = mapped_params.get("response_format")
supported_formats = frozenset(model_info["supported_audio_formats"])
response_format: Final = mapped_params.get("response_format")
supported_formats: Final = 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)
@ -539,17 +551,20 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
self,
model: str,
api_base: str | None,
litellm_params: dict,
litellm_params: dict, # mutable-ok: inherited provider interface accepts concrete LiteLLM parameters
) -> 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(
base_model: Final = model.removeprefix("vertex_ai/")
model_info: Final = self._get_model_info(model=model)
configured_project: Final = self.safe_get_vertex_ai_project(litellm_params)
project: Final = (
self._ensure_access_token(
credentials=self.safe_get_vertex_ai_credentials(litellm_params),
project_id=None,
custom_llm_provider="vertex_ai",
)
)[1]
if configured_project is None
else configured_project
)
if model_info["vertex_ai_audio_api"] == "lyria_interactions":
from litellm.llms.vertex_ai.interactions.transformation import (
VertexAIInteractionsConfig,
@ -558,10 +573,13 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
return VertexAIInteractionsConfig().get_complete_url(
api_base=api_base,
model=base_model,
litellm_params={**litellm_params, "vertex_project": project},
litellm_params={ # mutable-ok: interactions dispatch expects a concrete parameter dictionary
**litellm_params,
"vertex_project": project,
},
)
location = self.safe_get_vertex_ai_location(litellm_params) or self.get_default_vertex_location()
base_url = self.get_api_base(api_base=api_base, vertex_location=location).rstrip("/")
location: Final = self.safe_get_vertex_ai_location(litellm_params) or self.get_default_vertex_location()
base_url: Final = self.get_api_base(api_base=api_base, vertex_location=location).rstrip("/")
return f"{base_url}/v1/projects/{project}/locations/{location}/publishers/google/models/{base_model}:predict"
def transform_text_to_speech_request(
@ -569,9 +587,9 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
model: str,
input: str,
voice: str | None,
optional_params: dict,
litellm_params: dict,
headers: dict,
optional_params: dict, # mutable-ok: inherited provider interface accepts concrete mapped parameters
litellm_params: dict, # mutable-ok: inherited provider interface accepts concrete LiteLLM parameters
headers: dict, # mutable-ok: inherited provider interface accepts and updates concrete HTTP headers
) -> TextToSpeechRequestData:
access_token, project = self._ensure_access_token(
credentials=self.safe_get_vertex_ai_credentials(litellm_params),
@ -579,23 +597,32 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
custom_llm_provider="vertex_ai",
)
headers.update(
{
{ # mutable-ok: HTTP dispatch requires a concrete header dictionary
"Authorization": f"Bearer {access_token}",
"x-goog-user-project": project,
"Content-Type": "application/json",
}
)
base_model = model.removeprefix("vertex_ai/")
model_info = self._get_model_info(model=model)
base_model: Final = model.removeprefix("vertex_ai/")
model_info: Final = self._get_model_info(model=model)
if model_info["vertex_ai_audio_api"] == "lyria_predict":
request_body = {
"instances": [{"prompt": input}],
"parameters": {"sample_count": 1},
request_body = { # mutable-ok: predict dispatch requires a concrete provider request dictionary; rebind-ok: exactly one provider API shape initializes the request
"instances": [ # mutable-ok: predict dispatch requires a concrete instances list
{"prompt": input} # mutable-ok: predict dispatch requires a concrete instance dictionary
],
"parameters": { # mutable-ok: predict dispatch requires a concrete parameters dictionary
"sample_count": 1
},
}
else:
request_body = {"model": base_model, "input": input}
request_body = { # mutable-ok: interactions dispatch requires a concrete provider request dictionary; rebind-ok: exactly one provider API shape initializes the request
"model": base_model,
"input": input,
}
if optional_params.get("response_format") == "wav":
request_body["response_format"] = {
request_body[
"response_format"
] = { # mutable-ok: interactions dispatch requires a nested response-format dictionary
"type": "audio",
"mime_type": "audio/wav",
}
@ -609,33 +636,49 @@ class VertexAILyriaTextToSpeechConfig(VertexAITextToSpeechConfig):
) -> "HttpxBinaryResponseContent":
from litellm.types.llms.openai import HttpxBinaryResponseContent
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
response_json: Final = raw_response.json()
base_model: Final = model.removeprefix("vertex_ai/")
model_info: Final = self._get_model_info(model=model)
audio_data: str | None = None # rebind-ok: response parsing discovers audio data in provider-specific shapes
mime_type: str | None = None # rebind-ok: response parsing discovers the MIME type beside the audio payload
if model_info["vertex_ai_audio_api"] == "lyria_predict":
predictions = response_json.get("predictions") or []
predictions: Final = response_json.get("predictions") or ()
if predictions:
audio_data = predictions[0].get("audioContent") or predictions[0].get("bytesBase64Encoded")
mime_type = predictions[0].get("mimeType")
audio_data = predictions[0].get("audioContent") or predictions[0].get(
"bytesBase64Encoded"
) # rebind-ok: predict response supplies the generated audio value
mime_type = predictions[0].get("mimeType") # rebind-ok: predict response supplies its audio MIME type
else:
for step in response_json.get("steps") or response_json.get("outputs") or []:
content_items = step.get("content") or [] if step.get("type") == "model_output" else [step]
for step in response_json.get("steps") or response_json.get("outputs") or ():
content_items = step.get("content") or () if step.get("type") == "model_output" else (step,)
for content in content_items:
if content.get("type") == "audio" and content.get("data"):
audio_data = content["data"]
mime_type = content.get("mime_type")
audio_data = content[
"data"
] # rebind-ok: interactions response supplies the generated audio value
mime_type = content.get(
"mime_type"
) # 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")
default_format = model_info["supported_audio_formats"][0]
mime_type = mime_type or {"mp3": "audio/mpeg", "wav": "audio/wav"}[default_format]
response = HttpxBinaryResponseContent(
default_format: Final = model_info["supported_audio_formats"][0]
mime_type = (
mime_type
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",
}[default_format]
)
response: Final = HttpxBinaryResponseContent(
httpx.Response(
status_code=raw_response.status_code,
content=base64.b64decode(audio_data),
headers={"content-type": mime_type},
headers={ # mutable-ok: httpx requires a concrete response header dictionary
"content-type": mime_type
},
)
)
response._hidden_params = {"audio_mime_type": mime_type}
response._hidden_params = { # mutable-ok: response metadata is a concrete dictionary
"audio_mime_type": mime_type
}
return response

View file

@ -8261,9 +8261,13 @@ def speech(
# Vertex AI Text-to-Speech (Google Cloud TTS)
if text_to_speech_provider_config is None:
if VertexAILyriaTextToSpeechConfig.is_lyria_model(model):
text_to_speech_provider_config = VertexAILyriaTextToSpeechConfig()
text_to_speech_provider_config = (
VertexAILyriaTextToSpeechConfig()
) # rebind-ok: model metadata selects the Lyria provider implementation
else:
text_to_speech_provider_config = VertexAITextToSpeechConfig()
text_to_speech_provider_config = (
VertexAITextToSpeechConfig()
) # rebind-ok: non-Lyria Vertex models use the standard TTS implementation
# Cast to specific Vertex AI config type to access dispatch method
vertex_config: Final = cast(VertexAITextToSpeechConfig, text_to_speech_provider_config)

View file

@ -334,10 +334,10 @@ class VertexPassthroughLoggingHandler:
@staticmethod
def _handle_audio_predict_response(
json_response: dict,
json_response: dict, # mutable-ok: passthrough logging receives the decoded provider response dictionary
logging_obj: LiteLLMLoggingObj,
model: str,
kwargs: dict,
kwargs: dict, # mutable-ok: passthrough logging enriches the shared callback metadata dictionary
) -> PassThroughEndpointLoggingTypedDict:
prediction_count: Final = VertexPassthroughLoggingHandler._get_audio_prediction_count(
json_response=json_response
@ -346,26 +346,41 @@ class VertexPassthroughLoggingHandler:
VertexPassthroughLoggingHandler._get_audio_prediction_unit_cost(model=model) or 0.0
) * 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
logging_obj.model = model # rebind-ok: passthrough attribution records the resolved Vertex model
logging_obj.model_call_details[ # rebind-ok: passthrough attribution enriches callback metadata
"model"
] = model
logging_obj.model_call_details[ # rebind-ok: passthrough attribution enriches callback metadata
"custom_llm_provider"
] = "vertex_ai"
logging_obj.custom_llm_provider = ( # rebind-ok: attribution records the resolved provider
"vertex_ai"
)
logging_obj.model_call_details[ # rebind-ok: passthrough attribution enriches callback metadata
"response_cost"
] = response_cost
kwargs["response_cost"] = response_cost
kwargs["model"] = model
kwargs["custom_llm_provider"] = "vertex_ai"
kwargs[ # rebind-ok: callback metadata is enriched for downstream hooks
"response_cost"
] = response_cost
kwargs["model"] = model # rebind-ok: callback metadata records the resolved model
kwargs["custom_llm_provider"] = "vertex_ai" # rebind-ok: callback metadata records the resolved provider
standard_pass_through_response_object: Final[StandardPassThroughResponseObject] = {
standard_pass_through_response_object: Final[
StandardPassThroughResponseObject
] = { # mutable-ok: callback contract requires a concrete response dictionary
"response": json_response,
}
return {
return { # mutable-ok: passthrough logging contract requires a concrete result dictionary
"result": standard_pass_through_response_object,
"kwargs": kwargs,
}
@staticmethod
def _is_audio_predict_response(model: str, json_response: dict) -> bool:
def _is_audio_predict_response(
model: str,
json_response: dict, # mutable-ok: predicate inspects the decoded provider response dictionary without mutation
) -> bool:
return (
VertexPassthroughLoggingHandler._get_audio_prediction_count(json_response=json_response) > 0
and VertexPassthroughLoggingHandler._get_audio_prediction_unit_cost(model=model) is not None
@ -373,7 +388,9 @@ class VertexPassthroughLoggingHandler:
@staticmethod
def _get_audio_prediction_unit_cost(model: str) -> float | None:
model_info: Final = litellm.model_cost.get(f"vertex_ai/{model}", {})
model_info: Final = litellm.model_cost.get(f"vertex_ai/{model}")
if model_info is None:
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")
if not isinstance(output_cost_per_second, (int, float)) or not isinstance(
@ -383,7 +400,9 @@ class VertexPassthroughLoggingHandler:
return float(output_cost_per_second * audio_seconds_per_prediction)
@staticmethod
def _get_audio_prediction_count(json_response: dict) -> int:
def _get_audio_prediction_count(
json_response: dict, # mutable-ok: counter inspects the decoded provider response dictionary without mutation
) -> int:
predictions: Final = json_response.get("predictions")
if not isinstance(predictions, list):
return 0

View file

@ -168,8 +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
supported_audio_formats: ReadOnly[Sequence[Literal["mp3", "wav"]] | None]
vertex_ai_audio_api: ReadOnly[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
@ -312,9 +312,9 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False):
output_cost_per_video_per_second: float | None # only for vertex ai models
output_cost_per_audio_per_second: float | None # only for vertex ai models
output_cost_per_second: float | None # for OpenAI Speech models
audio_seconds_per_prediction: float | None
max_audio_length_hours: float | None
max_audio_per_prompt: int | None
audio_seconds_per_prediction: ReadOnly[float | None]
max_audio_length_hours: ReadOnly[float | None]
max_audio_per_prompt: ReadOnly[int | None]
output_cost_per_second_1080p: (
float | None
) # video_generation tier: key output_cost_per_second_<resolution> (e.g. 1080p, 720p)