diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py index 0e2eb60704d..0ac04654e30 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py @@ -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", diff --git a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py index 1815ff134aa..dbb6c4cd702 100644 --- a/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py +++ b/tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py @@ -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):