diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index e03f7ee745f..e1de1bea412 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/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() diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 2e3f7bb9aa6..4810b9d726a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/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) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index ff1c12d08d7..0851e1d0089 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/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"] diff --git a/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py b/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py index d74a05ec59c..e9365982837 100644 --- a/tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py +++ b/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 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_vertex_passthrough_logging_handler.py new file mode 100644 index 00000000000..a5cbca25bd0 --- /dev/null +++ b/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"))