fix(google_genai): price gemini streams with the gemini rate card

Native generateContent streaming billed every Google stream at Vertex AI
rates. For gemini-3.1-flash-image and gemini-3.1-flash-image-preview, the
only models whose gemini/ and vertex_ai/ rate cards currently differ, that
is exactly 2x the correct amount, and any future model priced differently
across the two surfaces inherits the bug.

Two things combined. The google_genai streaming iterator hands the
pass-through logging handler a url_route with no hostname, so the handler
cannot sniff the provider off it and falls back to vertex_ai. Separately,
_create_vertex_response_logging_payload_for_generate_content accepted a
custom_llm_provider argument that all three call sites resolved correctly,
then hardcoded "vertex_ai" in its completion_cost call and applied the
argument only to the log label, so a row could read gemini while the amount
came from the Vertex rate card.

The helper now costs with the provider it is given and stamps it on kwargs
the way the Gemini handler already does. The provider is threaded from the
iterator through _route_streaming_logging_to_handler so callers that know
it skip the hostname sniffing, and pass-through callers that only have the
upstream URL keep the existing fallback.

Fixes #36042

litellm/google_genai/streaming_iterator.py

@@ -66,6 +66,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
         request_body: dict,
         model: str,
         hidden_params: dict[str, Any] | None = None,
+        custom_llm_provider: str | None = None,
     ):
         self.litellm_logging_obj = litellm_logging_obj
         self.request_body = request_body
@@ -73,11 +74,18 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
         self.collected_chunks: list[bytes] = []
         self.model = model
         self._hidden_params: dict[str, Any] = hidden_params or {}
+        self.custom_llm_provider = custom_llm_provider
 
     async def _handle_async_streaming_logging(
         self,
     ):
-        """Handle the logging after all chunks have been collected."""
+        """Handle the logging after all chunks have been collected.
+
+        ``url_route`` below carries no hostname, so the pass-through handler cannot
+        sniff the provider off it and would price every Google stream at Vertex rates.
+        ``custom_llm_provider`` is forwarded so costing uses the provider already
+        resolved for this request.
+        """
         from litellm.proxy.pass_through_endpoints.streaming_handler import (
             PassThroughStreamingHandler,
         )
@@ -94,6 +102,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
                 raw_bytes=self.collected_chunks,
                 end_time=end_time,
                 model=self.model,
+                custom_llm_provider=self.custom_llm_provider,
             )
         )
 
@@ -119,12 +128,12 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent
             request_body=request_body or {},
             model=model,
             hidden_params=hidden_params,
+            custom_llm_provider=custom_llm_provider,
         )
         self.response = response
         self.model = model
         self.generate_content_provider_config = generate_content_provider_config
         self.litellm_metadata = litellm_metadata
-        self.custom_llm_provider = custom_llm_provider
         # Gemini streamGenerateContent uses SSE line framing; iter_lines keeps
         # large inlineData payloads (e.g. image/jpeg) intact within one event.
         self.stream_iterator = response.iter_lines()
@@ -170,12 +179,12 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
             request_body=request_body or {},
             model=model,
             hidden_params=hidden_params,
+            custom_llm_provider=custom_llm_provider,
         )
         self.response = response
         self.model = model
         self.generate_content_provider_config = generate_content_provider_config
         self.litellm_metadata = litellm_metadata
-        self.custom_llm_provider = custom_llm_provider
         # Gemini streamGenerateContent uses SSE line framing; aiter_lines keeps
         # large inlineData payloads (e.g. image/jpeg) intact within one event.
         self.stream_iterator = response.aiter_lines()

litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py

@@ -396,6 +396,7 @@ class VertexPassthroughLoggingHandler:
         all_chunks: list[str],
         model: str | None,
         end_time: datetime,
+        custom_llm_provider: str | None = None,
     ) -> PassThroughEndpointLoggingTypedDict:
         """
         Takes raw chunks from Vertex passthrough endpoint and logs them in litellm callbacks
@@ -403,6 +404,11 @@ class VertexPassthroughLoggingHandler:
         - Builds complete response from chunks
         - Creates standard logging object
         - Logs in litellm callbacks
+
+        ``custom_llm_provider`` is the provider the caller already resolved. Callers that
+        only know the upstream URL leave it unset and the provider is sniffed from the
+        hostname instead. Native google_genai streams pass it explicitly because their
+        url_route carries no hostname to sniff.
         """
         kwargs: dict[str, Any] = {}
         model = model or VertexPassthroughLoggingHandler.extract_model_from_url(url_route)
@@ -429,7 +435,9 @@ class VertexPassthroughLoggingHandler:
             start_time=start_time,
             end_time=end_time,
             logging_obj=litellm_logging_obj,
-            custom_llm_provider=VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(url_route),
+            custom_llm_provider=(
+                custom_llm_provider or VertexPassthroughLoggingHandler._get_custom_llm_provider_from_url(url_route)
+            ),
         )
 
         return {
@@ -592,11 +600,12 @@ class VertexPassthroughLoggingHandler:
         response_cost: Final = litellm.completion_cost(
             completion_response=litellm_model_response,
             model=model,
-            custom_llm_provider="vertex_ai",
+            custom_llm_provider=custom_llm_provider,
         )
 
         kwargs["response_cost"] = response_cost
         kwargs["model"] = model
+        kwargs["custom_llm_provider"] = custom_llm_provider
 
         # pretty print standard logging object
         verbose_proxy_logger.debug("kwargs= %s", kwargs)

litellm/proxy/pass_through_endpoints/streaming_handler.py

@@ -129,6 +129,7 @@ class PassThroughStreamingHandler:
         raw_bytes: list[bytes],
         end_time: datetime,
         model: str | None = None,
+        custom_llm_provider: str | None = None,
     ):
         """
         Route the logging for the collected chunks to the appropriate handler
@@ -137,6 +138,9 @@ class PassThroughStreamingHandler:
         - Anthropic
         - Vertex AI
         - OpenAI
+
+        ``custom_llm_provider`` lets a caller that already knows the provider skip the
+        URL-hostname sniffing the pass-through handlers fall back to.
         """
         try:
             (
@@ -152,6 +156,7 @@ class PassThroughStreamingHandler:
                 raw_bytes=raw_bytes,
                 end_time=end_time,
                 model=model,
+                custom_llm_provider=custom_llm_provider,
             )
             # Always reached from an async context (anthropic_messages,
             # google_genai, and proxy pass-through stream tasks). prefer_async_handlers
@@ -179,6 +184,7 @@ class PassThroughStreamingHandler:
         raw_bytes: list[bytes],
         end_time: datetime,
         model: str | None,
+        custom_llm_provider: str | None = None,
     ) -> tuple[PassThroughEndpointLoggingResultValues, dict]:
         """
         Synchronous, CPU-bound reconstruction of the standard logging payload
@@ -217,6 +223,7 @@ class PassThroughStreamingHandler:
                     all_chunks=all_chunks,
                     end_time=end_time,
                     model=model,
+                    custom_llm_provider=custom_llm_provider,
                 )
             )
             standard_logging_response_object = vertex_passthrough_logging_handler_result["result"]

tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py

@@ -1,5 +1,6 @@
+import asyncio
 import json
-from unittest.mock import MagicMock
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
@@ -53,9 +54,7 @@ async def test_async_streaming_iterator_yields_complete_sse_events():
     assert chunk.startswith(b"data: ")
     assert chunk.endswith(b"\n\n")
     assert (
-        json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][0][
-            "inlineData"
-        ]["mimeType"]
+        json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][0]["inlineData"]["mimeType"]
         == "image/jpeg"
     )
 
@@ -76,9 +75,9 @@ def test_sync_streaming_iterator_yields_complete_sse_events():
     chunk = next(iterator)
     assert chunk.startswith(b"data: ")
     assert chunk.endswith(b"\n\n")
-    assert json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][
-        0
-    ]["inlineData"]["data"].startswith("A")
+    assert json.loads(chunk[len(b"data: ") : -2])["candidates"][0]["content"]["parts"][0]["inlineData"][
+        "data"
+    ].startswith("A")
 
 
 @pytest.mark.asyncio
@@ -126,3 +125,42 @@ async def test_async_streaming_iterator_forwards_sse_comment_events():
 
     chunk = await iterator.__anext__()
     assert chunk == b": keepalive\n\n"
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("custom_llm_provider", ["gemini", "vertex_ai"])
+async def test_async_streaming_iterator_forwards_custom_llm_provider_to_logging(
+    custom_llm_provider,
+):
+    """
+    The url_route handed to the pass-through logging handler carries no hostname, so
+    the handler cannot sniff the provider off it and used to price every Google stream
+    at Vertex rates. The iterator must forward the provider it already resolved.
+    """
+    mock_response = MagicMock()
+
+    async def _aiter_lines():
+        yield 'data: {"text":"hi"}'
+        yield ""
+
+    mock_response.aiter_lines = _aiter_lines
+
+    iterator = AsyncGoogleGenAIGenerateContentStreamingIterator(
+        response=mock_response,
+        model="gemini-3.1-flash-image",
+        logging_obj=MagicMock(spec=LiteLLMLoggingObj),
+        generate_content_provider_config=MagicMock(),
+        litellm_metadata={},
+        custom_llm_provider=custom_llm_provider,
+    )
+
+    with patch(
+        "litellm.proxy.pass_through_endpoints.streaming_handler.PassThroughStreamingHandler._route_streaming_logging_to_handler",
+        new=AsyncMock(),
+    ) as mock_route:
+        async for _ in iterator:
+            pass
+        await asyncio.sleep(0)
+
+    mock_route.assert_called_once()
+    assert mock_route.call_args.kwargs["custom_llm_provider"] == custom_llm_provider

tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py

@@ -0,0 +1,125 @@
+import json
+from datetime import datetime
+from unittest.mock import MagicMock
+
+import pytest
+
+import litellm
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import (
+    VertexPassthroughLoggingHandler,
+)
+from litellm.types.utils import ModelResponse, Usage
+
+# gemini/ and vertex_ai/ price this model differently (2x), so it is the only way to
+# observe which rate card a code path actually used.
+MODEL = "gemini-3.1-flash-image"
+PROMPT_TOKENS = 1000
+COMPLETION_TOKENS = 500
+
+
+def _expected_cost(prefix: str) -> float:
+    entry = litellm.model_cost[f"{prefix}/{MODEL}"]
+    return PROMPT_TOKENS * entry["input_cost_per_token"] + COMPLETION_TOKENS * entry["output_cost_per_token"]
+
+
+def _model_response() -> ModelResponse:
+    response = ModelResponse(model=MODEL)
+    response.usage = Usage(
+        prompt_tokens=PROMPT_TOKENS,
+        completion_tokens=COMPLETION_TOKENS,
+        total_tokens=PROMPT_TOKENS + COMPLETION_TOKENS,
+    )
+    return response
+
+
+def _logging_obj() -> LiteLLMLoggingObj:
+    logging_obj = MagicMock(spec=LiteLLMLoggingObj)
+    logging_obj.litellm_call_id = "call-id"
+    logging_obj.model_call_details = {}
+    logging_obj.optional_params = {}
+    return logging_obj
+
+
+def test_rate_cards_differ_so_the_assertions_below_are_meaningful():
+    assert _expected_cost("gemini") != _expected_cost("vertex_ai")
+
+
+@pytest.mark.parametrize("provider", ["gemini", "vertex_ai"])
+def test_generate_content_payload_prices_with_the_provider_it_is_given(provider):
+    """
+    The helper accepts custom_llm_provider and every call site resolves it, but it used
+    to hardcode "vertex_ai" in the completion_cost call and apply the argument only to
+    the log label. A gemini request was therefore labelled gemini and billed at Vertex
+    rates.
+    """
+    logging_obj = _logging_obj()
+
+    kwargs = VertexPassthroughLoggingHandler._create_vertex_response_logging_payload_for_generate_content(
+        litellm_model_response=_model_response(),
+        model=MODEL,
+        kwargs={},
+        start_time=datetime.now(),
+        end_time=datetime.now(),
+        logging_obj=logging_obj,
+        custom_llm_provider=provider,
+    )
+
+    assert kwargs["response_cost"] == pytest.approx(_expected_cost(provider))
+    assert kwargs["custom_llm_provider"] == provider
+    assert logging_obj.model_call_details["custom_llm_provider"] == provider
+
+
+def _sse_chunks() -> list[str]:
+    return [
+        "data: "
+        + json.dumps(
+            {
+                "candidates": [
+                    {
+                        "content": {"parts": [{"text": "hi"}], "role": "model"},
+                        "finishReason": "STOP",
+                        "index": 0,
+                    }
+                ],
+                "usageMetadata": {
+                    "promptTokenCount": PROMPT_TOKENS,
+                    "candidatesTokenCount": COMPLETION_TOKENS,
+                    "totalTokenCount": PROMPT_TOKENS + COMPLETION_TOKENS,
+                },
+                "modelVersion": MODEL,
+                "responseId": "r1",
+            }
+        )
+    ]
+
+
+def _collect(url_route: str, custom_llm_provider: str | None) -> dict:
+    return VertexPassthroughLoggingHandler._handle_logging_vertex_collected_chunks(
+        litellm_logging_obj=_logging_obj(),
+        passthrough_success_handler_obj=MagicMock(),
+        url_route=url_route,
+        request_body={},
+        endpoint_type=MagicMock(),
+        start_time=datetime.now(),
+        all_chunks=_sse_chunks(),
+        model=MODEL,
+        end_time=datetime.now(),
+        custom_llm_provider=custom_llm_provider,
+    )["kwargs"]
+
+
+def test_collected_chunks_use_an_explicitly_passed_provider():
+    """Native google_genai streams know their provider but send a hostname-less url_route."""
+    assert _collect("/v1/generateContent", "gemini")["response_cost"] == pytest.approx(_expected_cost("gemini"))
+
+
+def test_collected_chunks_fall_back_to_sniffing_the_url_when_no_provider_is_passed():
+    """Pass-through callers only know the upstream URL, so hostname sniffing must still work."""
+    gemini_url = "https://generativelanguage.googleapis.com/v1beta/models/x:streamGenerateContent"
+    assert _collect(gemini_url, None)["response_cost"] == pytest.approx(_expected_cost("gemini"))
+
+
+def test_collected_chunks_still_price_vertex_at_vertex_rates():
+    vertex_url = "https://us-central1-aiplatform.googleapis.com/v1/projects/p/locations/l/publishers/google/models/x:streamGenerateContent"
+    assert _collect(vertex_url, None)["response_cost"] == pytest.approx(_expected_cost("vertex_ai"))