This commit is contained in:
Marty Sullivan 2026-08-27 19:56:22 -05:00 committed by GitHub
commit 0b62ebd38f
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 161 additions and 223 deletions

View file

@ -5,7 +5,9 @@ Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough e
Supports different modalities: text, audio, video, and web search.
"""
from collections.abc import Mapping, Sequence
from datetime import datetime
from types import MappingProxyType
from typing import Any, Final
from litellm._logging import verbose_proxy_logger
@ -15,8 +17,13 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrough_logging_handler import (
PassThroughEndpointLoggingTypedDict,
)
from litellm.types.utils import LlmProviders, ModelResponse, Usage
from litellm.utils import get_model_info
from litellm.types.utils import (
CompletionTokensDetailsWrapper,
LlmProviders,
ModelResponse,
PromptTokensDetailsWrapper,
Usage,
)
class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
@ -128,101 +135,14 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
return aggregated
@staticmethod
def _calculate_live_api_cost(
model: str,
usage_metadata: dict,
custom_llm_provider: str = "vertex_ai",
) -> float:
"""
Calculate cost for Vertex AI Live API based on usage metadata.
Args:
model: The model name (e.g., "gemini-2.0-flash-live-preview-04-09")
usage_metadata: Usage metadata from the Live API response
custom_llm_provider: The LLM provider (default: "vertex_ai")
Returns:
Total cost in USD
"""
try:
# Get model pricing information
model_info: Final = get_model_info(model=model, custom_llm_provider=custom_llm_provider)
verbose_proxy_logger.debug("Vertex AI Live API model info for '%s': %s", model, model_info)
# Check if pricing info is available
if not model_info or not model_info.get("input_cost_per_token"):
verbose_proxy_logger.error("No pricing info found for %s in local model pricing database", model)
return 0.0
total_cost = 0.0
# Extract token counts from usage metadata
prompt_token_count: Final = usage_metadata.get("promptTokenCount", 0)
candidates_token_count: Final = usage_metadata.get("candidatesTokenCount", 0)
# Calculate base text token costs
input_cost_per_token: Final = model_info.get("input_cost_per_token", 0.0)
output_cost_per_token: Final = model_info.get("output_cost_per_token", 0.0)
total_cost += prompt_token_count * input_cost_per_token
total_cost += candidates_token_count * output_cost_per_token
# Handle modality-specific costs if present
prompt_tokens_details: Final = usage_metadata.get("promptTokensDetails", [])
candidates_tokens_details: Final = usage_metadata.get("candidatesTokensDetails", [])
# Process prompt tokens by modality
for detail in prompt_tokens_details:
modality = detail.get("modality", "TEXT")
token_count = detail.get("tokenCount", 0)
if modality == "AUDIO":
audio_cost_per_token = model_info.get("input_cost_per_audio_token", 0.0)
total_cost += token_count * audio_cost_per_token
elif modality == "VIDEO":
# Video tokens are typically per second, but we'll treat as per token for now
video_cost_per_token = model_info.get("input_cost_per_video_per_second", 0.0)
total_cost += token_count * video_cost_per_token
# TEXT tokens are already handled above
# Process candidate tokens by modality
for detail in candidates_tokens_details:
modality = detail.get("modality", "TEXT")
token_count = detail.get("tokenCount", 0)
if modality == "AUDIO":
audio_cost_per_token = model_info.get("output_cost_per_audio_token", 0.0)
total_cost += token_count * audio_cost_per_token
elif modality == "VIDEO":
# Video tokens are typically per second, but we'll treat as per token for now
video_cost_per_token = model_info.get("output_cost_per_video_per_second", 0.0)
total_cost += token_count * video_cost_per_token
# TEXT tokens are already handled above
# Handle web search costs if present
tool_use_prompt_token_count: Final = usage_metadata.get("toolUsePromptTokenCount", 0)
if tool_use_prompt_token_count > 0:
# Web search typically has a fixed cost per request
web_search_cost: Final = model_info.get("web_search_cost_per_request", 0.0)
if isinstance(web_search_cost, (int, float)) and web_search_cost > 0:
total_cost += web_search_cost
else:
# Fallback to token-based pricing for tool use
total_cost += tool_use_prompt_token_count * input_cost_per_token
verbose_proxy_logger.debug(
f"Vertex AI Live API cost calculation - Model: {model}, "
f"Prompt tokens: {prompt_token_count}, "
f"Candidate tokens: {candidates_token_count}, "
f"Total cost: ${total_cost:.6f}"
)
return total_cost
except Exception as e:
verbose_proxy_logger.error("Error calculating Vertex AI Live API cost: %s", e)
return 0.0
def _tokens_by_modality(details: Sequence[Mapping[str, Any]]) -> Mapping[str, int]:
"""Sum a Live API ``*TokensDetails`` list into ``{modality: tokenCount}``."""
return MappingProxyType(
{
modality: sum(d.get("tokenCount", 0) for d in details if d.get("modality", "TEXT") == modality)
for modality in {d.get("modality", "TEXT") for d in details}
}
)
@staticmethod
def _create_usage_object_from_metadata(
@ -239,38 +159,35 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
Returns:
LiteLLM Usage object
"""
prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0)
completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0)
total_tokens: Final = usage_metadata.get("totalTokenCount", 0)
_ = model
# Create modality-specific token details if available
prompt_tokens_details: Final = usage_metadata.get("promptTokensDetails", [])
candidates_tokens_details: Final = usage_metadata.get("candidatesTokensDetails", [])
prompt_by_modality: Final = VertexAILivePassthroughLoggingHandler._tokens_by_modality(
usage_metadata.get("promptTokensDetails") or []
)
candidates_by_modality: Final = VertexAILivePassthroughLoggingHandler._tokens_by_modality(
usage_metadata.get("candidatesTokensDetails") or []
)
# Extract text tokens from details
text_prompt_tokens = 0
text_completion_tokens = 0
for detail in prompt_tokens_details:
if detail.get("modality") == "TEXT":
text_prompt_tokens = detail.get("tokenCount", 0)
break
for detail in candidates_tokens_details:
if detail.get("modality") == "TEXT":
text_completion_tokens = detail.get("tokenCount", 0)
break
# If no text tokens found in details, use total counts
if text_prompt_tokens == 0:
text_prompt_tokens = prompt_tokens
if text_completion_tokens == 0:
text_completion_tokens = completion_tokens
prompt_tokens: Final = usage_metadata.get("promptTokenCount", 0) or sum(prompt_by_modality.values())
completion_tokens: Final = usage_metadata.get("candidatesTokenCount", 0) or sum(candidates_by_modality.values())
return Usage(
prompt_tokens=text_prompt_tokens,
completion_tokens=text_completion_tokens,
total_tokens=total_tokens,
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=usage_metadata.get("totalTokenCount", 0) or (prompt_tokens + completion_tokens),
prompt_tokens_details=PromptTokensDetailsWrapper(
text_tokens=prompt_by_modality.get("TEXT"),
audio_tokens=prompt_by_modality.get("AUDIO"),
image_tokens=prompt_by_modality.get("IMAGE"),
video_tokens=prompt_by_modality.get("VIDEO"),
cached_tokens=usage_metadata.get("cachedContentTokenCount", 0) or 0,
),
completion_tokens_details=CompletionTokensDetailsWrapper(
text_tokens=candidates_by_modality.get("TEXT"),
audio_tokens=candidates_by_modality.get("AUDIO"),
image_tokens=candidates_by_modality.get("IMAGE"),
video_tokens=candidates_by_modality.get("VIDEO"),
),
)
def vertex_ai_live_passthrough_handler(
@ -316,14 +233,6 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
"kwargs": kwargs,
}
# Calculate cost using Live API specific pricing
response_cost: Final = self._calculate_live_api_cost(
model=model,
usage_metadata=usage_metadata,
custom_llm_provider=custom_llm_provider,
)
# Create Usage object for standard LiteLLM logging
usage: Final = self._create_usage_object_from_metadata(
usage_metadata=usage_metadata,
model=model,
@ -339,8 +248,6 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
choices=[],
)
# Update kwargs with cost information
kwargs["response_cost"] = response_cost
kwargs["model"] = model
kwargs["custom_llm_provider"] = custom_llm_provider
@ -350,10 +257,13 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
allowed_pattern: Final = re.compile(r"^[A-Za-z0-9._\-:]+$")
safe_model: Final = model if isinstance(model, str) and allowed_pattern.match(model) else "[REDACTED]"
verbose_proxy_logger.debug(
f"Vertex AI Live API passthrough cost tracking - "
f"Model: {safe_model}, Cost: ${response_cost:.6f}, "
f"Prompt tokens: {usage.prompt_tokens}, "
f"Completion tokens: {usage.completion_tokens}"
"Vertex AI Live API passthrough cost tracking - Model: %s, "
"Prompt tokens: %s %s, Completion tokens: %s %s",
safe_model,
usage.prompt_tokens,
usage.prompt_tokens_details,
usage.completion_tokens,
usage.completion_tokens_details,
)
return {

View file

@ -201,88 +201,117 @@ class TestVertexAILivePassthroughLoggingHandler:
assert text_prompt["tokenCount"] == 10
assert audio_prompt["tokenCount"] == 10
@patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info"
)
def test_calculate_cost_basic(self, mock_get_model_info, handler):
"""Test basic cost calculation"""
mock_get_model_info.return_value = {
"input_cost_per_token": 0.000001,
"output_cost_per_token": 0.000002,
}
def test_usage_carries_every_modality(self, handler):
"""Regression: the Usage object reported only TEXT, so audio/image billed as nothing.
prompt_tokens must be the full count and the details must name each modality,
because the cost calculator prices audio and image from *_tokens_details.
"""
usage_metadata = {
"promptTokenCount": 100,
"candidatesTokenCount": 50,
"totalTokenCount": 150,
}
cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata)
# The cost calculation may include additional factors, so we check it's reasonable
expected_min_cost = (100 * 0.000001) + (50 * 0.000002)
assert cost >= expected_min_cost
assert cost > 0
@patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info"
)
def test_calculate_cost_with_audio(self, mock_get_model_info, handler):
"""Test cost calculation with audio tokens"""
mock_get_model_info.return_value = {
"input_cost_per_token": 0.000001,
"output_cost_per_token": 0.000002,
"input_cost_per_audio_token": 0.0001,
"output_cost_per_audio_token": 0.0002,
}
usage_metadata = {
"promptTokenCount": 100,
"candidatesTokenCount": 50,
"totalTokenCount": 150,
"promptTokenCount": 1300,
"candidatesTokenCount": 124,
"totalTokenCount": 1424,
"promptTokensDetails": [
{"modality": "TEXT", "tokenCount": 80},
{"modality": "AUDIO", "tokenCount": 20},
{"modality": "TEXT", "tokenCount": 13},
{"modality": "AUDIO", "tokenCount": 127},
{"modality": "IMAGE", "tokenCount": 1160},
],
"candidatesTokensDetails": [
{"modality": "TEXT", "tokenCount": 30},
{"modality": "AUDIO", "tokenCount": 20},
{"modality": "TEXT", "tokenCount": 29},
{"modality": "AUDIO", "tokenCount": 95},
],
}
cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata)
usage = handler._create_usage_object_from_metadata(
usage_metadata=usage_metadata, model="gemini-live-2.5-flash"
)
# Should include both text and audio costs
assert cost > 0
assert cost > (100 * 0.000001) + (
50 * 0.000002
) # Should be higher due to audio
assert usage.prompt_tokens == 1300, "the full prompt count must survive, not just its text share"
assert usage.completion_tokens == 124
assert usage.prompt_tokens_details.text_tokens == 13
assert usage.prompt_tokens_details.audio_tokens == 127
assert usage.prompt_tokens_details.image_tokens == 1160
assert usage.completion_tokens_details.text_tokens == 29
assert usage.completion_tokens_details.audio_tokens == 95
@patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info"
def test_usage_sums_repeated_modality_entries(self, handler):
"""A modality may appear more than once across aggregated turns; sum, don't overwrite."""
usage = handler._create_usage_object_from_metadata(
usage_metadata={
"promptTokenCount": 40,
"candidatesTokenCount": 0,
"promptTokensDetails": [
{"modality": "IMAGE", "tokenCount": 10},
{"modality": "IMAGE", "tokenCount": 25},
{"modality": "TEXT", "tokenCount": 5},
],
},
model="gemini-live-2.5-flash",
)
assert usage.prompt_tokens_details.image_tokens == 35
assert usage.prompt_tokens_details.text_tokens == 5
@pytest.mark.parametrize(
"label,prompt_details,candidate_details",
[
("text only", [("TEXT", 6)], [("TEXT", 2)]),
("audio in", [("TEXT", 13), ("AUDIO", 127)], [("TEXT", 18)]),
("image in", [("TEXT", 10), ("IMAGE", 258)], [("TEXT", 24)]),
("frames in", [("TEXT", 11), ("IMAGE", 1032)], [("TEXT", 26)]),
("audio both ways", [("TEXT", 13), ("AUDIO", 127)], [("TEXT", 29), ("AUDIO", 95)]),
],
)
def test_calculate_cost_with_web_search(self, mock_get_model_info, handler):
"""Test cost calculation with web search (tool use)"""
mock_get_model_info.return_value = {
"input_cost_per_token": 0.000001,
"output_cost_per_token": 0.000002,
"web_search_cost_per_request": 0.01,
}
def test_live_session_bills_each_modality_at_its_own_rate(self, handler, label, prompt_details, candidate_details):
"""Every payload here is a real Vertex Live session's usageMetadata.
usage_metadata = {
"promptTokenCount": 100,
"candidatesTokenCount": 50,
"totalTokenCount": 150,
"toolUsePromptTokenCount": 10,
}
Before the fix these billed the text share only, from 1x (text) to 55x under.
The expected amount is derived from the entry's own rates rather than hardcoded,
so this stays correct as prices move, and it is asserted exactly, so dropping a
modality and double-charging one both fail.
"""
from litellm.cost_calculator import completion_cost
from litellm.types.utils import ModelResponse
from litellm.utils import get_model_info
cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata)
model = "gemini-live-2.5-flash-preview-native-audio-09-2025"
info = get_model_info(model=model, custom_llm_provider="vertex_ai")
# Should include web search cost
expected_base_cost = (100 * 0.000001) + (50 * 0.000002)
# The web search cost might be handled differently, so just check it's reasonable
assert cost >= expected_base_cost
assert cost > 0
text_in = info["input_cost_per_token"]
audio_in = info.get("input_cost_per_audio_token") or text_in
image_in = info.get("input_cost_per_image_token") or text_in
text_out = info["output_cost_per_token"]
audio_out = info.get("output_cost_per_audio_token") or text_out
rate_in = {"TEXT": text_in, "AUDIO": audio_in, "IMAGE": image_in}
rate_out = {"TEXT": text_out, "AUDIO": audio_out}
expected = sum(c * rate_in[m] for m, c in prompt_details) + sum(c * rate_out[m] for m, c in candidate_details)
usage = handler._create_usage_object_from_metadata(
usage_metadata={
"promptTokenCount": sum(c for _, c in prompt_details),
"candidatesTokenCount": sum(c for _, c in candidate_details),
"promptTokensDetails": [{"modality": m, "tokenCount": c} for m, c in prompt_details],
"candidatesTokensDetails": [{"modality": m, "tokenCount": c} for m, c in candidate_details],
},
model=model,
)
cost = completion_cost(
completion_response=ModelResponse(
id="x", object="chat.completion", created=0, model=model, usage=usage, choices=[]
),
model=f"vertex_ai/{model}",
custom_llm_provider="vertex_ai",
call_type="acompletion",
)
assert cost == pytest.approx(expected, rel=1e-9), label
text_only = sum(c for m, c in prompt_details if m == "TEXT") * text_in + sum(
c for m, c in candidate_details if m == "TEXT"
) * text_out
if any(m != "TEXT" for m, _ in prompt_details + candidate_details) and audio_in != text_in:
assert cost > text_only, f"{label}: non-text modalities must add cost"
def test_vertex_ai_live_passthrough_handler_integration(
self, handler, mock_logging_obj, sample_websocket_messages
@ -540,25 +569,24 @@ class TestVertexAILivePassthroughErrorHandling:
result = handler._extract_usage_metadata_from_websocket_messages(messages)
assert result is None
@patch(
"litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_ai_live_passthrough_logging_handler.get_model_info"
)
def test_cost_calculation_with_missing_model_info(self, mock_get_model_info):
"""Test cost calculation when model info is missing"""
def test_usage_without_modality_details(self):
"""Older payloads carry only the totals; fall back to them rather than reporting zero."""
handler = VertexAILivePassthroughLoggingHandler()
# Mock missing model info
mock_get_model_info.return_value = {}
usage = handler._create_usage_object_from_metadata(
usage_metadata={
"promptTokenCount": 100,
"candidatesTokenCount": 50,
"totalTokenCount": 150,
},
model="unknown-model",
)
usage_metadata = {
"promptTokenCount": 100,
"candidatesTokenCount": 50,
"totalTokenCount": 150,
}
# Should not raise an exception, should return 0 or handle gracefully
cost = handler._calculate_live_api_cost("unknown-model", usage_metadata)
assert cost == 0.0
assert usage.prompt_tokens == 100
assert usage.completion_tokens == 50
assert usage.total_tokens == 150
assert usage.prompt_tokens_details.audio_tokens is None
assert usage.prompt_tokens_details.image_tokens is None
def test_handler_with_none_websocket_messages(self, mock_logging_obj):
"""Test handler with None websocket messages"""