scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #36055 from BerriAI/devin_ai_fix_gemini_stream_billing_36042

Browse source 
fix(google_genai): price streamed generateContent with the provider that served it
This commit is contained in:
Mateo Wang 2026-08-26 18:05:50 -07:00 • committed by GitHub
commit 77765fd302
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 198 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

12
litellm/google_genai/streaming_iterator.py
View file

@ -2,6 +2,7 @@ import asyncio
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final
import litellm
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.proxy.pass_through_endpoints.success_handler import (
PassThroughEndpointLogging,
@ -65,6 +66,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
litellm_logging_obj: LiteLLMLoggingObj,
request_body: dict,
model: str,
custom_llm_provider: str,
hidden_params: dict[str, Any] | None = None,
):
self.litellm_logging_obj = litellm_logging_obj
@ -72,6 +74,10 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
self.start_time = datetime.now()
self.collected_chunks: list[bytes] = []
self.model = model
self.custom_llm_provider = custom_llm_provider
self.endpoint_type: Final = (
EndpointType.GEMINI if custom_llm_provider == litellm.LlmProviders.GEMINI.value else EndpointType.VERTEX_AI
)
self._hidden_params: dict[str, Any] = hidden_params or {}
async def _handle_async_streaming_logging(
@ -89,7 +95,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator:
passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ,
url_route="/v1/generateContent",
request_body=self.request_body or {},
endpoint_type=EndpointType.VERTEX_AI,
endpoint_type=self.endpoint_type,
start_time=self.start_time,
raw_bytes=self.collected_chunks,
end_time=end_time,
@ -118,13 +124,13 @@ class GoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateContent
litellm_logging_obj=logging_obj,
request_body=request_body or {},
model=model,
custom_llm_provider=custom_llm_provider,
hidden_params=hidden_params,
)
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()
@ -169,13 +175,13 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo
litellm_logging_obj=logging_obj,
request_body=request_body or {},
model=model,
custom_llm_provider=custom_llm_provider,
hidden_params=hidden_params,
)
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()

2
litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py
View file

@ -615,7 +615,7 @@ 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,
vertex_location=vertex_location,
)

23
litellm/proxy/pass_through_endpoints/streaming_handler.py
View file

@ -17,6 +17,9 @@ from litellm.types.utils import StandardPassThroughResponseObject
from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
AnthropicPassthroughLoggingHandler,
)
from .llm_provider_handlers.gemini_passthrough_logging_handler import (
GeminiPassthroughLoggingHandler,
)
from .llm_provider_handlers.openai_passthrough_logging_handler import (
OpenAIPassthroughLoggingHandler,
)
@ -243,6 +246,26 @@ class PassThroughStreamingHandler:
)
standard_logging_response_object = vertex_passthrough_logging_handler_result["result"]
kwargs = vertex_passthrough_logging_handler_result["kwargs"]
elif endpoint_type == EndpointType.GEMINI:
gemini_passthrough_logging_handler_result: Final = (
GeminiPassthroughLoggingHandler._handle_logging_gemini_collected_chunks( # pyright: ignore[reportPrivateUsage] # mirrors sibling handler dispatch
litellm_logging_obj=litellm_logging_obj,
passthrough_success_handler_obj=passthrough_success_handler_obj,
url_route=url_route,
request_body=request_body,
endpoint_type=endpoint_type,
start_time=start_time,
all_chunks=all_chunks,
end_time=end_time,
model=model,
)
)
standard_logging_response_object = ( # rebind-ok: branch bind in shared if/elif dispatch
gemini_passthrough_logging_handler_result["result"]
)
kwargs = ( # rebind-ok: branch bind in shared if/elif dispatch
gemini_passthrough_logging_handler_result["kwargs"]
)
elif endpoint_type == EndpointType.OPENAI:
openai_passthrough_logging_handler_result: Final = (
OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks(

1
litellm/types/passthrough_endpoints/pass_through_endpoints.py
View file

@ -22,6 +22,7 @@ LITELLM_PASS_THROUGH_ENDPOINT_MARKER: Final = "__litellm_pass_through_endpoint__
class EndpointType(str, Enum):
VERTEX_AI = "vertex-ai"
GEMINI = "gemini"
ANTHROPIC = "anthropic"
OPENAI = "openai"
GENERIC = "generic"

40
tests/test_litellm/google_genai/test_google_genai_streaming_iterator.py
View file

@ -8,6 +8,36 @@ from litellm.google_genai.streaming_iterator import (
GoogleGenAIGenerateContentStreamingIterator,
)
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
@pytest.mark.parametrize(
"custom_llm_provider, expected_endpoint_type",
[("gemini", EndpointType.GEMINI), ("vertex_ai", EndpointType.VERTEX_AI)],
)
@pytest.mark.parametrize(
"iterator_cls",
[
AsyncGoogleGenAIGenerateContentStreamingIterator,
GoogleGenAIGenerateContentStreamingIterator,
],
)
def test_streaming_logging_targets_the_provider_that_served_the_request(
iterator_cls: type,
custom_llm_provider: str,
expected_endpoint_type: EndpointType,
):
"""Routing every google stream through the vertex handler bills gemini/* at vertex_ai/ rates."""
iterator = iterator_cls(
response=MagicMock(),
model="gemini-3.1-flash-image",
logging_obj=MagicMock(spec=LiteLLMLoggingObj),
generate_content_provider_config=MagicMock(),
litellm_metadata={},
custom_llm_provider=custom_llm_provider,
)
assert iterator.endpoint_type is expected_endpoint_type
def _large_inline_data_event() -> str:
@ -53,9 +83,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 +104,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

130
tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py Normal file
View file

@ -0,0 +1,130 @@
import json
from collections.abc import Iterator
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.proxy.pass_through_endpoints.streaming_handler import (
PassThroughStreamingHandler,
)
from litellm.proxy.pass_through_endpoints.success_handler import (
PassThroughEndpointLogging,
)
from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType
MODEL = "gemini-stream-pricing-probe"
PROMPT_TOKENS = 1000
COMPLETION_TOKENS = 1000
GEMINI_INPUT_RATE = 1e-07
GEMINI_OUTPUT_RATE = 4e-07
VERTEX_INPUT_RATE = 1.5e-07
VERTEX_OUTPUT_RATE = 6e-07
GEMINI_COST = PROMPT_TOKENS * GEMINI_INPUT_RATE + COMPLETION_TOKENS * GEMINI_OUTPUT_RATE
VERTEX_COST = PROMPT_TOKENS * VERTEX_INPUT_RATE + COMPLETION_TOKENS * VERTEX_OUTPUT_RATE
@pytest.fixture(autouse=True)
def divergent_rate_cards(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]:
monkeypatch.setitem(
litellm.model_cost,
f"gemini/{MODEL}",
{
"input_cost_per_token": GEMINI_INPUT_RATE,
"output_cost_per_token": GEMINI_OUTPUT_RATE,
"litellm_provider": "gemini",
"mode": "chat",
},
)
monkeypatch.setitem(
litellm.model_cost,
f"vertex_ai/{MODEL}",
{
"input_cost_per_token": VERTEX_INPUT_RATE,
"output_cost_per_token": VERTEX_OUTPUT_RATE,
"litellm_provider": "vertex_ai",
"mode": "chat",
},
)
litellm.get_model_info.cache_clear()
yield
litellm.get_model_info.cache_clear()
def _chunks() -> list[str]:
payload = {
"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,
}
return [f"data: {json.dumps(payload)}"]
def _logging_obj() -> LiteLLMLoggingObj:
logging_obj = MagicMock(spec=LiteLLMLoggingObj)
logging_obj.model_call_details = {}
logging_obj.optional_params = {}
logging_obj.litellm_call_id = "test-call-id"
return logging_obj
@pytest.mark.parametrize(
"endpoint_type, expected_provider, expected_cost",
[
(EndpointType.GEMINI, "gemini", GEMINI_COST),
(EndpointType.VERTEX_AI, "vertex_ai", VERTEX_COST),
],
)
def test_streaming_generate_content_bills_against_the_requested_provider(
endpoint_type, expected_provider, expected_cost
):
logging_obj = _logging_obj()
_, kwargs = PassThroughStreamingHandler._build_passthrough_logging_result(
litellm_logging_obj=logging_obj,
passthrough_success_handler_obj=PassThroughEndpointLogging(),
url_route="/v1/generateContent",
request_body={},
endpoint_type=endpoint_type,
start_time=datetime.now(),
raw_bytes=[chunk.encode("utf-8") for chunk in _chunks()],
end_time=datetime.now(),
model=MODEL,
)
assert kwargs["response_cost"] == pytest.approx(expected_cost)
assert logging_obj.model_call_details["custom_llm_provider"] == expected_provider
def test_vertex_generate_content_payload_prices_gemini_urls_at_gemini_rates():
logging_obj = _logging_obj()
result = VertexPassthroughLoggingHandler._handle_logging_vertex_collected_chunks(
litellm_logging_obj=logging_obj,
passthrough_success_handler_obj=PassThroughEndpointLogging(),
url_route=f"https://generativelanguage.googleapis.com/v1beta/models/{MODEL}:streamGenerateContent",
request_body={},
endpoint_type=EndpointType.VERTEX_AI,
start_time=datetime.now(),
all_chunks=_chunks(),
model=MODEL,
end_time=datetime.now(),
)
assert result["kwargs"]["response_cost"] == pytest.approx(GEMINI_COST)
assert logging_obj.model_call_details["custom_llm_provider"] == "gemini"