mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
fix(vertex-live): price each grounded turn's query fee on the /vertex_ai/live passthrough
This commit is contained in:
mateo-berri
parent
228d87db63
commit
83594427fc
2 changed files with 174 additions and 19 deletions
|
|
@ -7,11 +7,12 @@ Supports different modalities: text, audio, video, and web search.
|
|||
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
from itertools import chain
|
||||
from itertools import chain, pairwise
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from typing import Final, Literal, TypeAlias
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.proxy.pass_through_endpoints.llm_provider_handlers.base_passthrough_logging_handler import (
|
||||
BasePassthroughLoggingHandler,
|
||||
)
|
||||
|
|
@ -20,6 +21,7 @@ from litellm.proxy.pass_through_endpoints.llm_provider_handlers.openai_passthrou
|
|||
)
|
||||
from litellm.types.utils import (
|
||||
CompletionTokensDetailsWrapper,
|
||||
CostBreakdown,
|
||||
LlmProviders,
|
||||
ModelResponse,
|
||||
PromptTokensDetailsWrapper,
|
||||
|
|
@ -60,6 +62,35 @@ def _grounding_metadata(websocket_messages: Sequence[object]) -> tuple[Mapping[s
|
|||
)
|
||||
|
||||
|
||||
def _turns(websocket_messages: Sequence[object]) -> tuple[tuple[object, ...], ...]:
|
||||
"""Split a session at every ``usageMetadata`` frame; frames after the last one never got their usage."""
|
||||
closes: Final = tuple(
|
||||
index + 1
|
||||
for index, message in enumerate(websocket_messages)
|
||||
if isinstance(message, Mapping) and isinstance(message.get("usageMetadata"), dict)
|
||||
)
|
||||
return tuple(tuple(websocket_messages[start:end]) for start, end in pairwise((0, *closes)))
|
||||
|
||||
|
||||
_SummedField: TypeAlias = Literal[
|
||||
"input_cost",
|
||||
"output_cost",
|
||||
"tool_usage_cost",
|
||||
"cache_read_cost",
|
||||
"cache_creation_cost",
|
||||
"reasoning_cost",
|
||||
"original_cost",
|
||||
"discount_amount",
|
||||
"margin_fixed_amount",
|
||||
"margin_total_amount",
|
||||
]
|
||||
|
||||
|
||||
def _summed(breakdowns: Sequence[CostBreakdown], field: _SummedField) -> float | None:
|
||||
values: Final = tuple(value for breakdown in breakdowns if (value := breakdown.get(field)) is not None)
|
||||
return sum(values) if values else None
|
||||
|
||||
|
||||
class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
||||
"""
|
||||
Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough.
|
||||
|
|
@ -141,7 +172,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
|
||||
@staticmethod
|
||||
def _extract_usage_metadata_from_websocket_messages(
|
||||
websocket_messages: list[dict],
|
||||
websocket_messages: Sequence[object],
|
||||
) -> dict | None:
|
||||
"""
|
||||
Extract and aggregate usage metadata from a list of WebSocket messages.
|
||||
|
|
@ -158,9 +189,11 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
Dictionary containing aggregated usage metadata, or None if not found
|
||||
"""
|
||||
snapshots: Final = tuple(
|
||||
message["usageMetadata"]
|
||||
metadata
|
||||
for message in websocket_messages
|
||||
if isinstance(message, dict) and isinstance(message.get("usageMetadata"), dict)
|
||||
if isinstance(message, Mapping)
|
||||
for metadata in (message.get("usageMetadata"),)
|
||||
if isinstance(metadata, dict)
|
||||
)
|
||||
|
||||
if not snapshots:
|
||||
|
|
@ -247,10 +280,72 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
)
|
||||
return usage
|
||||
|
||||
def _session_usage(self, websocket_messages: Sequence[object], model: str) -> Usage | None:
|
||||
usage_metadata: Final = self._extract_usage_metadata_from_websocket_messages(websocket_messages)
|
||||
if usage_metadata is None:
|
||||
return None
|
||||
return self._create_usage_object_from_metadata(
|
||||
usage_metadata=usage_metadata,
|
||||
grounding_metadata=_grounding_metadata(websocket_messages),
|
||||
model=model,
|
||||
)
|
||||
|
||||
def _turn_cost(
|
||||
self,
|
||||
turn: Sequence[object],
|
||||
model: str,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> tuple[float, CostBreakdown] | None:
|
||||
usage: Final = self._session_usage(turn, model)
|
||||
if usage is None:
|
||||
return None
|
||||
cost: Final = logging_obj._response_cost_calculator( # pyright: ignore[reportPrivateUsage] # the call's own calculator keeps custom pricing and the deployment's region in step with the spend row
|
||||
result=ModelResponse(model=model, usage=usage),
|
||||
litellm_model_name=model,
|
||||
)
|
||||
if cost is None:
|
||||
return None
|
||||
breakdown: Final = logging_obj.cost_breakdown
|
||||
return None if breakdown is None else (cost, breakdown)
|
||||
|
||||
def _session_cost(
|
||||
self,
|
||||
websocket_messages: Sequence[object],
|
||||
model: str,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
) -> float | None:
|
||||
"""Price each turn on its own tokens and grounding, so two grounded turns pay the query fee twice."""
|
||||
turn_costs: Final = tuple(self._turn_cost(turn, model, logging_obj) for turn in _turns(websocket_messages))
|
||||
priced: Final = tuple(turn_cost for turn_cost in turn_costs if turn_cost is not None)
|
||||
if not priced or len(priced) != len(turn_costs):
|
||||
return None
|
||||
breakdowns: Final = tuple(breakdown for _, breakdown in priced)
|
||||
first: Final = breakdowns[0]
|
||||
total_cost: Final = sum(cost for cost, _ in priced)
|
||||
logging_obj.set_cost_breakdown(
|
||||
input_cost=_summed(breakdowns, "input_cost") or 0.0,
|
||||
output_cost=_summed(breakdowns, "output_cost") or 0.0,
|
||||
total_cost=total_cost,
|
||||
cost_for_built_in_tools_cost_usd_dollar=_summed(breakdowns, "tool_usage_cost") or 0.0,
|
||||
original_cost=_summed(breakdowns, "original_cost"),
|
||||
discount_percent=first.get("discount_percent"),
|
||||
discount_amount=_summed(breakdowns, "discount_amount"),
|
||||
margin_percent=first.get("margin_percent"),
|
||||
margin_fixed_amount=_summed(breakdowns, "margin_fixed_amount"),
|
||||
margin_total_amount=_summed(breakdowns, "margin_total_amount"),
|
||||
cache_read_cost=_summed(breakdowns, "cache_read_cost"),
|
||||
cache_creation_cost=_summed(breakdowns, "cache_creation_cost"),
|
||||
reasoning_cost=_summed(breakdowns, "reasoning_cost"),
|
||||
service_tier=first.get("service_tier"),
|
||||
data_residency=first.get("data_residency"),
|
||||
vertex_location=first.get("vertex_location"),
|
||||
)
|
||||
return total_cost
|
||||
|
||||
def vertex_ai_live_passthrough_handler(
|
||||
self,
|
||||
websocket_messages: list[dict],
|
||||
logging_obj,
|
||||
websocket_messages: Sequence[object],
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
url_route: str,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
|
|
@ -274,28 +369,25 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
"""
|
||||
try:
|
||||
# Extract model from request body or kwargs
|
||||
model: Final = kwargs.get("model", "gemini-2.0-flash-live-preview-04-09")
|
||||
requested_model: Final = kwargs.get("model")
|
||||
model: Final = (
|
||||
requested_model if isinstance(requested_model, str) else "gemini-2.0-flash-live-preview-04-09"
|
||||
)
|
||||
custom_llm_provider: Final = kwargs.get("custom_llm_provider", "vertex_ai")
|
||||
verbose_proxy_logger.debug(
|
||||
"Vertex AI Live API model: %s, custom_llm_provider: %s", model, custom_llm_provider
|
||||
)
|
||||
|
||||
# Extract usage metadata from WebSocket messages
|
||||
usage_metadata: Final = self._extract_usage_metadata_from_websocket_messages(websocket_messages)
|
||||
usage: Final = self._session_usage(websocket_messages, model)
|
||||
|
||||
if not usage_metadata:
|
||||
if usage is None:
|
||||
verbose_proxy_logger.warning("No usage metadata found in Vertex AI Live API WebSocket messages")
|
||||
return {
|
||||
"result": None,
|
||||
"kwargs": kwargs,
|
||||
}
|
||||
|
||||
# Create Usage object for standard LiteLLM logging
|
||||
usage: Final = self._create_usage_object_from_metadata(
|
||||
usage_metadata=usage_metadata,
|
||||
grounding_metadata=_grounding_metadata(websocket_messages),
|
||||
model=model,
|
||||
)
|
||||
response_cost: Final = self._session_cost(websocket_messages, model, logging_obj)
|
||||
|
||||
# Create a mock ModelResponse for standard logging
|
||||
litellm_model_response: Final = ModelResponse(
|
||||
|
|
@ -306,6 +398,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
usage=usage,
|
||||
choices=[],
|
||||
)
|
||||
if response_cost is not None:
|
||||
litellm_model_response._hidden_params["response_cost"] = response_cost # pyright: ignore[reportPrivateUsage] # the logger reads the cost off the response's hidden params; the constructor's hidden_params kwarg is reset by pydantic
|
||||
|
||||
kwargs["model"] = model
|
||||
kwargs["custom_llm_provider"] = custom_llm_provider
|
||||
|
|
@ -314,7 +408,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
|
|||
import re
|
||||
|
||||
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]"
|
||||
safe_model: Final = model if allowed_pattern.match(model) else "[REDACTED]"
|
||||
verbose_proxy_logger.debug(
|
||||
"Vertex AI Live API passthrough cost tracking - Model: %s, "
|
||||
"Prompt tokens: %s %s, Completion tokens: %s %s",
|
||||
|
|
|
|||
|
|
@ -24,7 +24,7 @@ from litellm.proxy.pass_through_endpoints.success_handler import (
|
|||
PassThroughEndpointLogging,
|
||||
)
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.types.utils import LlmProviders, Usage
|
||||
from litellm.types.utils import CostBreakdown, LlmProviders, Usage
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
|
||||
|
||||
|
|
@ -47,6 +47,7 @@ class TestVertexAILivePassthroughLoggingHandler:
|
|||
"""Create a mock logging object"""
|
||||
mock = MagicMock(spec=LiteLLMLoggingObj)
|
||||
mock.model_call_details = {}
|
||||
mock._response_cost_calculator.return_value = None
|
||||
return mock
|
||||
|
||||
@pytest.fixture
|
||||
|
|
@ -474,6 +475,64 @@ class TestVertexAILivePassthroughLoggingHandler:
|
|||
|
||||
assert grounded > plain, "a grounded session must cost more than the same tokens ungrounded"
|
||||
|
||||
def _priced_logging_obj(self) -> LiteLLMLoggingObj:
|
||||
"""A real logging object, since the session's price is handed to it turn by turn."""
|
||||
logging_obj = LiteLLMLoggingObj(
|
||||
model=self.NATIVE_AUDIO_MODEL,
|
||||
messages=[],
|
||||
stream=True,
|
||||
call_type="pass_through_endpoint",
|
||||
start_time=datetime.now(),
|
||||
litellm_call_id="live-session",
|
||||
function_id="live",
|
||||
)
|
||||
logging_obj.update_environment_variables(
|
||||
model=self.NATIVE_AUDIO_MODEL,
|
||||
user="u",
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
call_type="pass_through_endpoint",
|
||||
)
|
||||
logging_obj.model_call_details["custom_llm_provider"] = "vertex_ai"
|
||||
return logging_obj
|
||||
|
||||
def _billed_session(
|
||||
self, handler: VertexAILivePassthroughLoggingHandler, messages: list[dict[str, object]]
|
||||
) -> tuple[float, CostBreakdown]:
|
||||
logging_obj = self._priced_logging_obj()
|
||||
result = handler.vertex_ai_live_passthrough_handler(
|
||||
websocket_messages=messages,
|
||||
logging_obj=logging_obj,
|
||||
url_route="/vertex_ai/live",
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
request_body={},
|
||||
model=self.NATIVE_AUDIO_MODEL,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
assert result["result"] is not None, "the handler must produce a usage-bearing response to bill"
|
||||
assert logging_obj.cost_breakdown is not None, "the session's price must reach the logging object"
|
||||
return result["result"]._hidden_params["response_cost"], logging_obj.cost_breakdown
|
||||
|
||||
def test_each_grounded_turn_pays_its_own_query_fee(self, handler):
|
||||
"""Google charges the grounding fee per grounded prompt, not per session.
|
||||
|
||||
Summing the session into one usage collapsed two grounded turns into one query, so the
|
||||
second question was answered for free. The bill now grows by one fee per grounded turn.
|
||||
"""
|
||||
head, turn = self._live_messages(self.AUDIO_SESSION[:1])
|
||||
grounding = self._grounding_frame({"webSearchQueries": ["q"]})
|
||||
|
||||
plain_cost, _ = self._billed_session(handler, [head, turn, turn])
|
||||
one_cost, one_breakdown = self._billed_session(handler, [head, grounding, turn, turn])
|
||||
two_cost, two_breakdown = self._billed_session(handler, [head, grounding, turn, grounding, turn])
|
||||
|
||||
fee = one_cost - plain_cost
|
||||
assert fee > 0, "a grounded turn must cost more than the same tokens ungrounded"
|
||||
assert two_cost - plain_cost == pytest.approx(2 * fee), "two grounded turns must pay the fee twice"
|
||||
assert two_breakdown["total_cost"] == pytest.approx(two_cost)
|
||||
assert two_breakdown["tool_usage_cost"] == pytest.approx(2 * one_breakdown["tool_usage_cost"])
|
||||
|
||||
def test_reporting_tool_use_tokens_does_not_move_the_bill(self, handler, mock_logging_obj):
|
||||
"""Deliberate boundary: these tokens are reported here, and priced nowhere.
|
||||
|
||||
|
|
@ -676,6 +735,7 @@ class TestVertexAILivePassthroughIntegration:
|
|||
"""Create a mock logging object"""
|
||||
mock = MagicMock(spec=LiteLLMLoggingObj)
|
||||
mock.model_call_details = {}
|
||||
mock._response_cost_calculator.return_value = None
|
||||
return mock
|
||||
|
||||
@patch(
|
||||
|
|
@ -809,6 +869,7 @@ class TestVertexAILivePassthroughErrorHandling:
|
|||
"""Create a mock logging object"""
|
||||
mock = MagicMock(spec=LiteLLMLoggingObj)
|
||||
mock.model_call_details = {}
|
||||
mock._response_cost_calculator.return_value = None
|
||||
return mock
|
||||
|
||||
def test_invalid_websocket_messages_format(self):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue