mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
fix(vertex-live): bill every modality on the /vertex_ai/live passthrough
The Live passthrough builds Usage from the TEXT-modality counts alone, so audio, image and video tokens never reach the cost calculator and bill as nothing. A one-turn audio session reported 13 text and 127 audio input tokens and billed the 13; a camera session reported 1043 prompt tokens and billed 11. Reporting the full per-modality breakdown fixes it, because the shared Gemini input and output cost path already prices audio, image and video from prompt_tokens_details and completion_tokens_details. On the native-audio entry that is a 6x difference per token in both directions, which is the whole gap. Aggregation across turns is unchanged. Google charges per turn for every token in the Live session context window, current turn plus all accumulated tokens from previous turns, so the existing summing is what Vertex bills and it stays as it is. That is worth stating because the cumulative promptTokensDetails looks like a restatement of one running total, and treating it that way would under-bill a multi-turn session. See the LiveAPI context-window note on https://cloud.google.com/vertex-ai/generative-ai/pricing. Live can also name the modality carrying the rest of a turn and omit its tokenCount. Reading that absent key as zero left the tokens inside candidatesTokenCount but outside the breakdown, so real speech was charged at the text output rate. A lone unpriced entry now takes whatever the turn's declared count leaves over. Two or more cannot be told apart, so they are still left to the calculator's text remainder. Server-side toolUsePromptTokenCount is now reported in prompt_tokens_details. It is deliberately kept out of prompt_tokens: no Gemini route prices tool-use tokens, and adding them there instead suppresses the cache-overlap correction and raises the bill for no extra work. Removes _calculate_live_api_cost, whose result never reached the bill. It set kwargs["response_cost"], which the standard logging path recomputes from the ModelResponse, and on a measured audio session it returned $0.000487 against a $0.0000425 row. Now that the modality counts reach the standard calculator, keeping a second hand-rolled pricing path would only ever double-charge. The rewrite of the aggregator is arithmetically identical to what it replaced. It sums the same three counts and the same per-modality details, still takes the remaining fields from the first turn, and drops nine LIT010, one C901 and 42 basedpyright findings in the process.
This commit is contained in:
Marty Sullivan
parent
ee7c7e14f3
commit
e2b41286d3
2 changed files with 375 additions and 286 deletions
|
|
@ -5,7 +5,10 @@ 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 itertools import chain
|
||||
from types import MappingProxyType
|
||||
from typing import Any, Final
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -15,8 +18,23 @@ 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,
|
||||
)
|
||||
|
||||
_AGGREGATED_FIELDS: Final = frozenset(
|
||||
{
|
||||
"promptTokenCount",
|
||||
"candidatesTokenCount",
|
||||
"totalTokenCount",
|
||||
"promptTokensDetails",
|
||||
"candidatesTokensDetails",
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
||||
|
|
@ -48,6 +66,56 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
"""Return the LLM provider name."""
|
||||
return LlmProviders.VERTEX_AI
|
||||
|
||||
@staticmethod
|
||||
def _resolve_detail_counts(
|
||||
details: Sequence[Mapping[str, Any]],
|
||||
declared_total: object,
|
||||
) -> tuple[tuple[str, int], ...]:
|
||||
"""
|
||||
Pair each of one turn's ``*TokensDetails`` entries with its token count.
|
||||
|
||||
Live sometimes names the modality that carries the rest of a turn without a
|
||||
``tokenCount``, and reading the absent key as zero drops those tokens from the
|
||||
breakdown, so real audio ends up priced as text. A lone unpriced entry therefore takes
|
||||
whatever the turn's declared count leaves over. Two or more cannot be told apart, so
|
||||
they are left out and the cost calculator charges the remainder as text.
|
||||
"""
|
||||
priced: Final = tuple(
|
||||
(str(detail.get("modality", "TEXT")), count)
|
||||
for detail in details
|
||||
if isinstance(count := detail.get("tokenCount"), int)
|
||||
)
|
||||
unpriced: Final = tuple(
|
||||
str(detail.get("modality", "TEXT")) for detail in details if not isinstance(detail.get("tokenCount"), int)
|
||||
)
|
||||
if len(unpriced) != 1 or not isinstance(declared_total, int):
|
||||
return priced
|
||||
residual: Final = declared_total - sum(count for _, count in priced)
|
||||
return priced if residual <= 0 else (*priced, (unpriced[0], residual))
|
||||
|
||||
@staticmethod
|
||||
def _sum_by_modality(counts: Sequence[tuple[str, int]]) -> Mapping[str, int]:
|
||||
"""Total the (modality, tokenCount) pairs of one or more turns per modality."""
|
||||
return MappingProxyType({modality: sum(c for m, c in counts if m == modality) for modality, _ in counts})
|
||||
|
||||
@staticmethod
|
||||
def _merged_modality_totals(
|
||||
snapshots: Sequence[Mapping[str, Any]],
|
||||
count_key: str,
|
||||
details_key: str,
|
||||
) -> Mapping[str, int]:
|
||||
"""Total every turn's per-modality counts, so the breakdown adds up the way the totals do."""
|
||||
return VertexAILivePassthroughLoggingHandler._sum_by_modality(
|
||||
tuple(
|
||||
chain.from_iterable(
|
||||
VertexAILivePassthroughLoggingHandler._resolve_detail_counts(
|
||||
snapshot.get(details_key) or [], snapshot.get(count_key)
|
||||
)
|
||||
for snapshot in snapshots
|
||||
)
|
||||
)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _extract_usage_metadata_from_websocket_messages(
|
||||
websocket_messages: list[dict],
|
||||
|
|
@ -55,175 +123,45 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
"""
|
||||
Extract and aggregate usage metadata from a list of WebSocket messages.
|
||||
|
||||
Live emits one ``usageMetadata`` per turn and Google charges per turn for every token in
|
||||
the session context window, which is the current turn's tokens plus all accumulated
|
||||
tokens from previous turns, so the turns add up rather than restating each other. See
|
||||
the Live API note under https://cloud.google.com/vertex-ai/generative-ai/pricing.
|
||||
|
||||
Args:
|
||||
websocket_messages: List of WebSocket messages from the Live API
|
||||
|
||||
Returns:
|
||||
Dictionary containing aggregated usage metadata, or None if not found
|
||||
"""
|
||||
all_usage_metadata: Final = []
|
||||
snapshots: Final = tuple(
|
||||
message["usageMetadata"]
|
||||
for message in websocket_messages
|
||||
if isinstance(message, dict) and isinstance(message.get("usageMetadata"), dict)
|
||||
)
|
||||
|
||||
# Collect all usage metadata messages
|
||||
for message in websocket_messages:
|
||||
if isinstance(message, dict) and "usageMetadata" in message:
|
||||
all_usage_metadata.append(message["usageMetadata"])
|
||||
|
||||
if not all_usage_metadata:
|
||||
if not snapshots:
|
||||
return None
|
||||
|
||||
# If only one usage metadata, return it as-is
|
||||
if len(all_usage_metadata) == 1:
|
||||
return all_usage_metadata[0]
|
||||
|
||||
# Aggregate multiple usage metadata messages
|
||||
aggregated: Final[dict[str, Any]] = {
|
||||
"promptTokenCount": 0,
|
||||
"candidatesTokenCount": 0,
|
||||
"totalTokenCount": 0,
|
||||
"promptTokensDetails": [],
|
||||
"candidatesTokensDetails": [],
|
||||
prompt_totals: Final = VertexAILivePassthroughLoggingHandler._merged_modality_totals(
|
||||
snapshots, "promptTokenCount", "promptTokensDetails"
|
||||
)
|
||||
candidate_totals: Final = VertexAILivePassthroughLoggingHandler._merged_modality_totals(
|
||||
snapshots, "candidatesTokenCount", "candidatesTokensDetails"
|
||||
)
|
||||
return {
|
||||
**{key: value for key, value in snapshots[0].items() if key not in _AGGREGATED_FIELDS},
|
||||
"promptTokenCount": sum(snapshot.get("promptTokenCount", 0) for snapshot in snapshots),
|
||||
"candidatesTokenCount": sum(snapshot.get("candidatesTokenCount", 0) for snapshot in snapshots),
|
||||
"totalTokenCount": sum(snapshot.get("totalTokenCount", 0) for snapshot in snapshots),
|
||||
"promptTokensDetails": [
|
||||
{"modality": modality, "tokenCount": count} for modality, count in prompt_totals.items() if count > 0
|
||||
],
|
||||
"candidatesTokensDetails": [
|
||||
{"modality": modality, "tokenCount": count} for modality, count in candidate_totals.items() if count > 0
|
||||
],
|
||||
}
|
||||
|
||||
# Aggregate token counts
|
||||
for usage in all_usage_metadata:
|
||||
aggregated["promptTokenCount"] += usage.get("promptTokenCount", 0)
|
||||
aggregated["candidatesTokenCount"] += usage.get("candidatesTokenCount", 0)
|
||||
aggregated["totalTokenCount"] += usage.get("totalTokenCount", 0)
|
||||
|
||||
# Aggregate token details by modality
|
||||
modality_totals: Final = {}
|
||||
|
||||
for usage in all_usage_metadata:
|
||||
# Process prompt tokens details
|
||||
for detail in usage.get("promptTokensDetails", []):
|
||||
modality = detail.get("modality", "TEXT")
|
||||
token_count = detail.get("tokenCount", 0)
|
||||
|
||||
if modality not in modality_totals:
|
||||
modality_totals[modality] = {"prompt": 0, "candidate": 0}
|
||||
modality_totals[modality]["prompt"] += token_count
|
||||
|
||||
# Process candidate tokens details
|
||||
for detail in usage.get("candidatesTokensDetails", []):
|
||||
modality = detail.get("modality", "TEXT")
|
||||
token_count = detail.get("tokenCount", 0)
|
||||
|
||||
if modality not in modality_totals:
|
||||
modality_totals[modality] = {"prompt": 0, "candidate": 0}
|
||||
modality_totals[modality]["candidate"] += token_count
|
||||
|
||||
# Convert aggregated modality totals back to details format
|
||||
for modality, totals in modality_totals.items():
|
||||
if totals["prompt"] > 0:
|
||||
aggregated["promptTokensDetails"].append({"modality": modality, "tokenCount": totals["prompt"]})
|
||||
if totals["candidate"] > 0:
|
||||
aggregated["candidatesTokensDetails"].append({"modality": modality, "tokenCount": totals["candidate"]})
|
||||
|
||||
# Add any additional fields from the first usage metadata
|
||||
first_usage: Final = all_usage_metadata[0]
|
||||
for key, value in first_usage.items():
|
||||
if key not in aggregated:
|
||||
aggregated[key] = value
|
||||
|
||||
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
|
||||
|
||||
@staticmethod
|
||||
def _create_usage_object_from_metadata(
|
||||
usage_metadata: dict,
|
||||
|
|
@ -239,38 +177,37 @@ 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)
|
||||
prompt_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality(
|
||||
VertexAILivePassthroughLoggingHandler._resolve_detail_counts(
|
||||
usage_metadata.get("promptTokensDetails") or [], usage_metadata.get("promptTokenCount")
|
||||
)
|
||||
)
|
||||
candidates_by_modality: Final = VertexAILivePassthroughLoggingHandler._sum_by_modality(
|
||||
VertexAILivePassthroughLoggingHandler._resolve_detail_counts(
|
||||
usage_metadata.get("candidatesTokensDetails") or [], usage_metadata.get("candidatesTokenCount")
|
||||
)
|
||||
)
|
||||
|
||||
# Create modality-specific token details if available
|
||||
prompt_tokens_details: Final = usage_metadata.get("promptTokensDetails", [])
|
||||
candidates_tokens_details: Final = usage_metadata.get("candidatesTokensDetails", [])
|
||||
|
||||
# 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"),
|
||||
tool_use_tokens=usage_metadata.get("toolUsePromptTokenCount") or None,
|
||||
),
|
||||
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,13 +253,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,
|
||||
|
|
@ -339,8 +269,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 +278,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 {
|
||||
|
|
|
|||
|
|
@ -201,88 +201,247 @@ 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 and 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 can 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
|
||||
|
||||
NATIVE_AUDIO_MODEL = "gemini-live-2.5-flash-preview-native-audio-09-2025"
|
||||
|
||||
# A four-turn native-audio session. Google charges per turn for the whole session context
|
||||
# window, so the prompt side repeats the accumulated audio while the candidates side reports
|
||||
# only that turn's own response. The last turn names AUDIO and omits its tokenCount, which is
|
||||
# the shape Live really emits at the end of a spoken answer.
|
||||
AUDIO_SESSION = (
|
||||
{"prompt": (14, 122), "candidates": (8, 20)},
|
||||
{"prompt": (21, 182), "candidates": (5, 50)},
|
||||
{"prompt": (24, 203), "candidates": (13, 27)},
|
||||
{"prompt": (24, 203), "candidates": (0, 3), "candidate_audio_token_count_missing": True},
|
||||
)
|
||||
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,
|
||||
}
|
||||
|
||||
usage_metadata = {
|
||||
"promptTokenCount": 100,
|
||||
"candidatesTokenCount": 50,
|
||||
"totalTokenCount": 150,
|
||||
"toolUsePromptTokenCount": 10,
|
||||
}
|
||||
@staticmethod
|
||||
def _live_messages(turns):
|
||||
"""Wrap (text, audio) prompt/candidate pairs as the server messages a Live session emits."""
|
||||
return [{"type": "session.created", "session": {"id": "s"}}] + [
|
||||
{
|
||||
"type": "response.done",
|
||||
"usageMetadata": {
|
||||
"promptTokenCount": sum(turn["prompt"]),
|
||||
"candidatesTokenCount": sum(turn["candidates"]),
|
||||
"totalTokenCount": sum(turn["prompt"]) + sum(turn["candidates"]),
|
||||
"promptTokensDetails": [
|
||||
{"modality": "TEXT", "tokenCount": turn["prompt"][0]},
|
||||
{"modality": "AUDIO", "tokenCount": turn["prompt"][1]},
|
||||
],
|
||||
"candidatesTokensDetails": (
|
||||
[{"modality": "AUDIO"}]
|
||||
if turn.get("candidate_audio_token_count_missing")
|
||||
else [
|
||||
{"modality": "TEXT", "tokenCount": turn["candidates"][0]},
|
||||
{"modality": "AUDIO", "tokenCount": turn["candidates"][1]},
|
||||
]
|
||||
),
|
||||
},
|
||||
}
|
||||
for turn in turns
|
||||
]
|
||||
|
||||
cost = handler._calculate_live_api_cost("gemini-1.5-pro", usage_metadata)
|
||||
@staticmethod
|
||||
def _session_usage(handler, mock_logging_obj, messages, model):
|
||||
result = handler.vertex_ai_live_passthrough_handler(
|
||||
websocket_messages=messages,
|
||||
logging_obj=mock_logging_obj,
|
||||
url_route="/vertex_ai/live",
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
request_body={},
|
||||
model=model,
|
||||
)
|
||||
assert result["result"] is not None, "the handler must produce a usage-bearing response to bill"
|
||||
return result["result"].usage
|
||||
|
||||
# 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
|
||||
@classmethod
|
||||
def _session_cost(cls, handler, mock_logging_obj, messages, model):
|
||||
from litellm.cost_calculator import completion_cost
|
||||
from litellm.types.utils import ModelResponse
|
||||
|
||||
usage = cls._session_usage(handler, mock_logging_obj, messages, model)
|
||||
return 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",
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _expected_session_cost(cls, turns):
|
||||
from litellm.utils import get_model_info
|
||||
|
||||
info = get_model_info(model=cls.NATIVE_AUDIO_MODEL, custom_llm_provider="vertex_ai")
|
||||
return (
|
||||
sum(turn["prompt"][0] for turn in turns) * info["input_cost_per_token"]
|
||||
+ sum(turn["prompt"][1] for turn in turns) * info["input_cost_per_audio_token"]
|
||||
+ sum(turn["candidates"][0] for turn in turns) * info["output_cost_per_token"]
|
||||
+ sum(turn["candidates"][1] for turn in turns) * info["output_cost_per_audio_token"]
|
||||
)
|
||||
|
||||
def test_every_turn_of_a_session_is_billed(self, handler, mock_logging_obj):
|
||||
"""Google charges per turn for the whole context window, so every turn adds to the bill.
|
||||
|
||||
Billing one snapshot instead gives away all the other turns: on this session the
|
||||
largest single turn is well under the session total, and its share of the audio is
|
||||
priced 6x the text rate, so the gap is money rather than rounding.
|
||||
"""
|
||||
turns = self.AUDIO_SESSION[:3]
|
||||
cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL)
|
||||
|
||||
assert cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9)
|
||||
widest_single_turn = max(self._expected_session_cost([turn]) for turn in turns)
|
||||
assert cost > widest_single_turn, "billing one snapshot drops every other turn of the session"
|
||||
|
||||
def test_audio_named_without_a_token_count_bills_at_the_audio_rate(self, handler, mock_logging_obj):
|
||||
"""Live can name the modality carrying the rest of a turn and omit its tokenCount.
|
||||
|
||||
Reading the absent key as zero left those tokens inside candidatesTokenCount but outside
|
||||
the breakdown, so the calculator charged real speech at the text output rate. At this
|
||||
entry's rates the last turn's 3 audio tokens are $0.0000360 rather than $0.0000060.
|
||||
"""
|
||||
turns = self.AUDIO_SESSION
|
||||
usage = self._session_usage(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL)
|
||||
|
||||
assert usage.completion_tokens_details.audio_tokens == 100, "the unpriced entry takes the turn's residual"
|
||||
assert usage.completion_tokens_details.text_tokens == 26
|
||||
assert usage.completion_tokens == 126
|
||||
|
||||
cost = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL)
|
||||
assert cost == pytest.approx(self._expected_session_cost(turns), rel=1e-9)
|
||||
|
||||
def test_server_side_tool_use_prompt_tokens_are_reported(self, handler, mock_logging_obj):
|
||||
"""toolUsePromptTokenCount was dropped, so a grounded session logged fewer tokens than it used.
|
||||
|
||||
It is reported, not billed. Nothing in the shared Gemini input-cost path prices
|
||||
tool-use tokens, and folding them into prompt_tokens here would suppress that
|
||||
path's cache-overlap correction and raise the bill instead.
|
||||
"""
|
||||
messages = self._live_messages(self.AUDIO_SESSION[:1])
|
||||
grounded = [dict(message) for message in messages]
|
||||
grounded[-1]["usageMetadata"] = {**grounded[-1]["usageMetadata"], "toolUsePromptTokenCount": 500}
|
||||
|
||||
usage = self._session_usage(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL)
|
||||
assert usage.prompt_tokens_details.tool_use_tokens == 500
|
||||
|
||||
plain_cost = self._session_cost(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL)
|
||||
grounded_cost = self._session_cost(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL)
|
||||
assert grounded_cost == pytest.approx(plain_cost, rel=1e-9), "reporting tool use must not move the bill"
|
||||
|
||||
@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_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.
|
||||
|
||||
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
|
||||
|
||||
model = self.NATIVE_AUDIO_MODEL
|
||||
info = get_model_info(model=model, custom_llm_provider="vertex_ai")
|
||||
|
||||
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 +699,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"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue