scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-20 00:11:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 75

fix(gemini-live): count grounding requests so Live sessions carry their query fee

Browse source 
Live reports grounding in the server frames and never in usageMetadata, so nothing
set the counter the cost path reads and the per-query charge was missing from every
grounded session. Google bills a grounded Live prompt on top of its tokens, and that
fee dwarfs the token cost, so a non-zero spend check could never catch it.

Both Live surfaces now read serverContent.groundingMetadata where they build usage,
and reuse the chat path's own classifier so web search and Maps keep their separate
SKUs rather than being counted together.

Separately, a client sending turn_detection: null reached a membership test against
None and took the session down with no traceback, while the branch immediately above
already guards for it. Live emits grounding and usage on the same frame, verified
against Vertex directly, so the realtime counter is set where usage is built.

(cherry picked from commit c997436be34beb2e84f8468286b8016c195eca92)
(cherry picked from commit 26c8d4822fc1c5c44fe8f72f4f57473a2ce1acbf)
This commit is contained in:
Marty Sullivan 2026-09-07 22:05:02 -04:00
parent 22e7c7a533
commit b1262b05f4
4 changed files with 212 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

20
litellm/llms/gemini/realtime/transformation.py
View file

@ -115,6 +115,19 @@ def _parse_setup(session_configuration_request: str) -> BidiGenerateContentSetup
return envelope.get("setup", empty_setup)
def _grounding_metadata_from_frame(frame: Mapping[str, object]) -> tuple[Mapping[str, object], ...]:
"""Read ``serverContent.groundingMetadata`` off the frame that carries the turn's usage.
Live reports grounding in the server frames rather than in ``usageMetadata``, and it emits both
on the same frame, so the per-query charge is countable at the point usage is built.
"""
server_content: Final = frame.get("serverContent")
if not isinstance(server_content, Mapping):
return ()
metadata: Final = server_content.get("groundingMetadata")
return (metadata,) if isinstance(metadata, Mapping) else ()
# Google bills Live transcription at an estimated 25 audio tokens/sec of input and
# 175 text tokens/min of output (ai.google.dev/gemini-api/docs/pricing).
GEMINI_LIVE_TRANSCRIBE_AUDIO_TOKENS_PER_SECOND: Final = 25
@ -323,7 +336,7 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
)
elif key == "input_audio_transcription" and value is not None:
optional_params["inputAudioTranscription"] = {}
elif key == "turn_detection":
elif key == "turn_detection" and value is not None:
value_typed = cast(OpenAIRealtimeTurnDetection, value)
if (
isinstance(value_typed, dict)
@ -1049,6 +1062,11 @@ class GeminiRealtimeConfig(BaseRealtimeConfig):
{**cast(dict, message), "usageMetadata": resolved_usage_metadata},
),
)
grounding_metadata: Final = _grounding_metadata_from_frame(message)
if grounding_metadata:
VertexGeminiConfig._set_grounding_usage_counters( # pyright: ignore[reportPrivateUsage] # shared with the chat path; no public alias exists yet
_chat_completion_usage, grounding_metadata
)
else:
_chat_completion_usage = get_empty_usage()

32
litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_ai_live_passthrough_logging_handler.py
View file

@ -43,6 +43,23 @@ def _detail_entries(raw: object) -> tuple[Mapping[str, object], ...]:
return tuple(entry for entry in raw if isinstance(entry, Mapping)) if isinstance(raw, Sequence) else ()
def _grounding_metadata(websocket_messages: Sequence[object]) -> tuple[Mapping[str, object], ...]:
"""Collect every ``serverContent.groundingMetadata`` a session emitted.
Live reports grounding in the server frames, never in ``usageMetadata``, so the per-query
charge has to be counted here rather than derived from the token totals.
"""
return tuple(
metadata
for message in websocket_messages
if isinstance(message, Mapping)
for server_content in (message.get("serverContent"),)
if isinstance(server_content, Mapping)
for metadata in (server_content.get("groundingMetadata"),)
if isinstance(metadata, Mapping)
)
class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
"""
Handles cost tracking and logging for Vertex AI Live API WebSocket passthrough.
@ -173,6 +190,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
def _create_usage_object_from_metadata(
usage_metadata: dict,
model: str,
grounding_metadata: Sequence[Mapping[str, object]] = (),
) -> Usage:
"""
Create a LiteLLM Usage object from Live API usage metadata.
@ -180,6 +198,8 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
Args:
usage_metadata: Usage metadata from the Live API response
model: The model name
grounding_metadata: Every ``serverContent.groundingMetadata`` the session emitted, so
Search and Maps grounding carry their per-query charge
Returns:
LiteLLM Usage object
@ -199,7 +219,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
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(
usage: Final = Usage(
prompt_tokens=prompt_tokens,
completion_tokens=completion_tokens,
total_tokens=usage_metadata.get("totalTokenCount", 0) or (prompt_tokens + completion_tokens),
@ -217,6 +237,15 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
video_tokens=candidates_by_modality.get("VIDEO"),
),
)
if grounding_metadata:
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexGeminiConfig,
)
VertexGeminiConfig._set_grounding_usage_counters( # pyright: ignore[reportPrivateUsage] # shared with the chat path; no public alias exists yet
usage, grounding_metadata
)
return usage
def vertex_ai_live_passthrough_handler(
self,
@ -264,6 +293,7 @@ class VertexAILivePassthroughLoggingHandler(BasePassthroughLoggingHandler):
# 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,
)

68
tests/pass_through_unit_tests/test_vertex_ai_live_passthrough.py
View file

@ -406,6 +406,74 @@ class TestVertexAILivePassthroughLoggingHandler:
usage = self._session_usage(handler, mock_logging_obj, grounded, self.NATIVE_AUDIO_MODEL)
assert usage.prompt_tokens_details.tool_use_tokens == sum(self.TOOL_USE_PER_TURN)
@staticmethod
def _grounding_frame(metadata: dict[str, object]) -> dict[str, object]:
"""One server frame carrying grounding metadata, the way Live reports it."""
return {"type": "response.done", "serverContent": {"groundingMetadata": metadata}}
def test_web_grounding_is_counted_so_it_can_be_billed(self, handler, mock_logging_obj):
"""Live reports grounding in the server frames and never in usageMetadata.
Nothing read those frames, so web_search_requests stayed unset and the cost path's only
trigger for the per-query grounding charge never fired. Google bills a grounded Live
prompt on top of its tokens, so the whole fee was missing from the bill.
"""
messages = [
self._grounding_frame(
{
"webSearchQueries": ["who won the 2026 world cup final"],
"groundingChunks": [{"web": {"uri": "https://example.com"}}],
}
),
*self._live_messages(self.AUDIO_SESSION[:1]),
]
usage = self._session_usage(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL)
assert usage.prompt_tokens_details.web_search_requests == 1, "a grounded turn must report its query"
assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None
def test_maps_grounding_is_counted_under_its_own_sku(self, handler, mock_logging_obj):
"""Maps grounding is a separate SKU from web search, so it needs its own counter.
A maps-only turn carries grounding chunks but no webSearchQueries, so counting queries
alone would report nothing and bill nothing.
"""
messages = [
self._grounding_frame({"groundingChunks": [{"maps": {"placeId": "abc123"}}]}),
*self._live_messages(self.AUDIO_SESSION[:1]),
]
usage = self._session_usage(handler, mock_logging_obj, messages, self.NATIVE_AUDIO_MODEL)
assert usage.prompt_tokens_details.google_maps_grounding_requests == 1
assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None
def test_an_ungrounded_session_reports_no_grounding(self, handler, mock_logging_obj):
"""The counters must stay absent when no tool ran, or every session pays a grounding fee."""
usage = self._session_usage(
handler, mock_logging_obj, self._live_messages(self.AUDIO_SESSION[:1]), self.NATIVE_AUDIO_MODEL
)
assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None
assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None
def test_grounding_adds_its_query_fee_to_the_session_bill(self, handler, mock_logging_obj):
"""The counter only matters if it reaches the bill, so assert against the cost, not the field.
Same tokens either way: the difference between the two sessions is the grounding fee alone.
"""
turns = self.AUDIO_SESSION[:1]
plain = self._session_cost(handler, mock_logging_obj, self._live_messages(turns), self.NATIVE_AUDIO_MODEL)
grounded = self._session_cost(
handler,
mock_logging_obj,
[self._grounding_frame({"webSearchQueries": ["q"]}), *self._live_messages(turns)],
self.NATIVE_AUDIO_MODEL,
)
assert grounded > plain, "a grounded session must cost more than the same tokens ungrounded"
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.

95
tests/test_litellm/llms/gemini/realtime/test_gemini_realtime_transformation.py
View file

@ -1,11 +1,16 @@
import json
from unittest.mock import MagicMock
from collections.abc import Mapping
from typing import cast
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
import litellm
from litellm.llms.gemini.realtime.transformation import GeminiRealtimeConfig
from litellm.types.llms.gemini import BidiGenerateContentServerMessage
from litellm.types.llms.openai import OpenAIRealtimeStreamSessionEvents
from litellm.types.utils import Usage
def test_gemini_realtime_transformation_session_created():
@ -2178,3 +2183,91 @@ def test_unbilled_usage_on_session_close_flushes_trailing_audio(patch_gemini_tra
}
assert usage == expected
assert config.unbilled_usage_on_session_close("gemini-3.5-transcribe-live") is None
def _grounded_live_frame(grounding_metadata: Mapping[str, object] | None) -> Mapping[str, object]:
"""One Live server frame. Grounding metadata and usageMetadata arrive together, as Vertex sends them."""
from typing import Final
server_content: Final = {
"turnComplete": True,
**({} if grounding_metadata is None else {"groundingMetadata": grounding_metadata}),
}
return {
"serverContent": server_content,
"usageMetadata": {
"promptTokenCount": 19,
"candidatesTokenCount": 157,
"totalTokenCount": 176,
"promptTokensDetails": ({"modality": "TEXT", "tokenCount": 19},),
"candidatesTokensDetails": ({"modality": "AUDIO", "tokenCount": 157},),
},
}
def _usage_built_for_response_done(message: Mapping[str, object]) -> Usage:
"""Capture the chat-completion Usage transform_response_done_event builds, before it is bridged.
The Usage object is local to the method, so the bridge call is the only place it is observable.
"""
from typing import Final
from litellm.responses.litellm_completion_transformation.transformation import (
LiteLLMCompletionResponsesConfig,
)
captured: Final[list[Usage]] = [] # mutable-ok: a spy has to accumulate what it observes
original: Final = LiteLLMCompletionResponsesConfig._transform_chat_completion_usage_to_responses_usage
def _spy(usage: Usage) -> object:
captured.append(usage)
return original(usage)
config: Final = GeminiRealtimeConfig()
with patch.object(
LiteLLMCompletionResponsesConfig,
"_transform_chat_completion_usage_to_responses_usage",
staticmethod(_spy),
):
config.transform_response_done_event(
message=cast( # cast-ok: a test fixture stands in for the server frame TypedDict
BidiGenerateContentServerMessage, message
),
current_response_id="resp_grounding",
current_conversation_id="conv_grounding",
output_items=None,
)
assert captured, "response.done must build a Usage object"
return captured[0]
def test_gemini_realtime_response_done_counts_web_grounding():
"""Regression: Live reports grounding in the server frames and never in usageMetadata.
Nothing read those frames on the realtime path, so web_search_requests stayed unset and the
cost path's only trigger for Google's per-query grounding charge never fired.
Scope boundary, deliberate: this asserts the counter on the Usage object that response.done is
built from, not on the emitted event. The Responses usage bridge copies a fixed allow-list of
detail fields and drops the rest, so the counter does not reach response.done yet. Widening
that bridge is a separate change; do not read this test as proving end-to-end billing.
"""
usage = _usage_built_for_response_done(
_grounded_live_frame(
{
"webSearchQueries": ["who won the 2026 world cup final"],
"groundingChunks": [{"web": {"uri": "https://example.com"}}],
}
)
)
assert usage.prompt_tokens_details.web_search_requests == 1, "a grounded turn must report its query"
assert usage.prompt_tokens_details.text_tokens == 19, "the modality breakdown must survive alongside it"
def test_gemini_realtime_response_done_reports_no_grounding_when_none_ran():
"""The counter must stay unset on an ordinary turn, or every session pays a grounding fee."""
usage = _usage_built_for_response_done(_grounded_live_frame(None))
assert getattr(usage.prompt_tokens_details, "web_search_requests", None) is None
assert getattr(usage.prompt_tokens_details, "google_maps_grounding_requests", None) is None