fix(vertex-live): bill Gemini Live sessions end to end (internal copy of #37075)

litellm/llms/gemini/realtime/transformation.py

@@ -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 @@ def map_openai_params(self, optional_params: dict, non_default_params: dict) ->
                 )
             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 @@ def transform_response_done_event(
                     {**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()
 

litellm/proxy/pass_through_endpoints/pass_through_endpoints.py

@@ -2090,6 +2090,22 @@ def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Calla
     return json.dumps({**message, "setup": {**setup, "model": rewritten_model}})  # mutable-ok: one-shot json payload
 
 
+def _resolved_vertex_live_setup(
+    setup_data: Mapping[str, object], setup_model_rewriter: Callable[[str], str] | None
+) -> Mapping[str, object]:
+    """
+    Give the model extractor the same fully qualified path the upstream will receive.
+
+    Clients may name a bare gateway alias, which the rewriter turns into a ``projects/...`` path before
+    it reaches Vertex. The extractor only reads a path containing ``/models/``, so running it on the raw
+    frame logs the session as ``unknown`` at no cost, which is precisely the supported client form
+    """
+    setup_model: Final = setup_data.get("model")
+    if setup_model_rewriter is None or not isinstance(setup_model, str):
+        return setup_data
+    return {**setup_data, "model": setup_model_rewriter(setup_model)}
+
+
 def _truncated_close_reason(reason: str) -> str:
     """
     Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character
@@ -2314,7 +2330,9 @@ async def forward_client_to_upstream() -> None:
                                             setup_data,
                                         )
                                         if isinstance(setup_data, dict) and "model" in setup_data:
-                                            extracted_model = _extract_model_from_vertex_ai_setup(setup_data)
+                                            extracted_model = _extract_model_from_vertex_ai_setup(
+                                                _resolved_vertex_live_setup(setup_data, setup_model_rewriter)
+                                            )
                                             if extracted_model:
                                                 kwargs["model"] = extracted_model
                                                 kwargs["custom_llm_provider"] = "vertex_ai-language-models"

litellm/responses/litellm_completion_transformation/transformation.py

@@ -2831,6 +2831,22 @@ def _transform_chat_completion_usage_to_responses_usage(
                 if cache_write_tokens is not None
                 else MappingProxyType({})
             )
+            # The cost path reads the grounding counters off the input details, and a realtime
+            # session's usage is rebuilt from its own response.done, so dropping them here bills
+            # no per-query grounding fee at all.
+            grounding_request_counts: Final[Mapping[str, int]] = MappingProxyType(
+                {
+                    counter: count
+                    for counter, count in (
+                        ("web_search_requests", getattr(prompt_details, "web_search_requests", None)),
+                        (
+                            "google_maps_grounding_requests",
+                            getattr(prompt_details, "google_maps_grounding_requests", None),
+                        ),
+                    )
+                    if count is not None
+                }
+            )
             response_usage.input_tokens_details = InputTokensDetails(
                 cached_tokens=prompt_details.cached_tokens if prompt_details.cached_tokens is not None else 0,
                 text_tokens=prompt_details.text_tokens,
@@ -2839,6 +2855,7 @@ def _transform_chat_completion_usage_to_responses_usage(
                     cached_tokens_details if isinstance(cached_tokens_details, CachedTokensDetails) else None
                 ),
                 **cache_write_extra,
+                **grounding_request_counts,
             )
 
         # Translate completion_tokens_details to output_tokens_details

litellm/responses/utils.py

@@ -1183,6 +1183,10 @@ def _transform_response_api_usage_to_chat_usage(
                         response_api_usage.input_tokens_details, "cached_tokens_details", None
                     ),
                     cache_write_tokens=getattr(response_api_usage.input_tokens_details, "cache_write_tokens", None),
+                    web_search_requests=getattr(response_api_usage.input_tokens_details, "web_search_requests", None),
+                    google_maps_grounding_requests=getattr(
+                        response_api_usage.input_tokens_details, "google_maps_grounding_requests", None
+                    ),
                 )
         completion_tokens_details: CompletionTokensDetailsWrapper | None = None
         output_tokens_details: Final[OutputTokensDetails | None] = getattr(