From 0c5583b83f03d642f6ee4f42619b94023ce3d7e8 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Thu, 6 Aug 2026 05:18:49 +0000 Subject: [PATCH 1/3] fix(google_genai): price streamed generateContent with the provider that served it Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/google_genai/streaming_iterator.py | 14 ++- .../vertex_passthrough_logging_handler.py | 2 +- .../streaming_handler.py | 19 ++++ .../pass_through_endpoints.py | 1 + .../test_google_genai_streaming_iterator.py | 40 +++++++- .../test_streaming_handler.py | 99 +++++++++++++++++++ 6 files changed, 170 insertions(+), 5 deletions(-) create mode 100644 tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index e03f7ee745f..e2fac6a615b 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -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,7 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: self.start_time = datetime.now() self.collected_chunks: list[bytes] = [] self.model = model + self.custom_llm_provider = custom_llm_provider self._hidden_params: dict[str, Any] = hidden_params or {} async def _handle_async_streaming_logging( @@ -83,13 +86,18 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: ) end_time: Final = datetime.now() + endpoint_type: Final = ( + EndpointType.GEMINI + if self.custom_llm_provider == litellm.LlmProviders.GEMINI.value + else EndpointType.VERTEX_AI + ) asyncio.create_task( PassThroughStreamingHandler._route_streaming_logging_to_handler( litellm_logging_obj=self.litellm_logging_obj, 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=endpoint_type, start_time=self.start_time, raw_bytes=self.collected_chunks, end_time=end_time, @@ -118,13 +126,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 +177,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() 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 afd8684dd92..36455611c95 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 @@ -592,7 +592,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, ) kwargs["response_cost"] = response_cost diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index ff1c12d08d7..907c59d28cc 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -15,6 +15,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, ) @@ -221,6 +224,22 @@ 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( + 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 = gemini_passthrough_logging_handler_result["result"] + kwargs = gemini_passthrough_logging_handler_result["kwargs"] elif endpoint_type == EndpointType.OPENAI: openai_passthrough_logging_handler_result: Final = ( OpenAIPassthroughLoggingHandler._handle_logging_openai_collected_chunks( diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index f59ca0d9041..548702e4139 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -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" 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..91058767730 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 @@ -8,6 +9,43 @@ 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.asyncio +@pytest.mark.parametrize( + "custom_llm_provider, expected_endpoint_type", + [("gemini", EndpointType.GEMINI), ("vertex_ai", EndpointType.VERTEX_AI)], +) +async def test_streaming_logging_routes_to_the_provider_that_served_the_request( + custom_llm_provider, expected_endpoint_type +): + """Routing every google stream through the vertex handler bills gemini/* at vertex_ai/ rates.""" + mock_response = MagicMock() + + async def _aiter_lines(): + yield 'data: {"candidates": []}' + + 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) + assert mock_route.call_args.kwargs["endpoint_type"] == expected_endpoint_type def _large_inline_data_event() -> str: diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py new file mode 100644 index 00000000000..d0c28fd60a9 --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py @@ -0,0 +1,99 @@ +import json +from datetime import datetime +from unittest.mock import MagicMock + +import pytest + +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-3.1-flash-image" + +# gemini/ rate card: 2.5e-07 in, 1.5e-06 out. vertex_ai/ rate card is exactly 2x that. +GEMINI_COST = 1000 * 2.5e-07 + 1000 * 1.5e-06 +VERTEX_COST = 2 * GEMINI_COST + + +def _chunks() -> list[str]: + payload = { + "candidates": [ + { + "content": {"parts": [{"text": "hi"}], "role": "model"}, + "finishReason": "STOP", + "index": 0, + } + ], + "usageMetadata": { + "promptTokenCount": 1000, + "candidatesTokenCount": 1000, + "totalTokenCount": 2000, + }, + "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 +): + """A streamed gemini/* request must not be priced off the vertex_ai/ rate card.""" + 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(): + """The AI Studio host resolves to `gemini`, so the cost must follow it, not the vertex_ai default.""" + 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" From 057781a1879cd0676c5ed53f4b8872576239d7c4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 16:54:05 -0700 Subject: [PATCH 2/3] test(pass_through): pin stream pricing tests to injected divergent rate cards --- .../test_streaming_handler.py | 49 +++++++++++++++---- 1 file changed, 40 insertions(+), 9 deletions(-) diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py index d0c28fd60a9..dd9fbd9161f 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py @@ -1,9 +1,11 @@ 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, @@ -16,11 +18,42 @@ from litellm.proxy.pass_through_endpoints.success_handler import ( ) from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType -MODEL = "gemini-3.1-flash-image" +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 -# gemini/ rate card: 2.5e-07 in, 1.5e-06 out. vertex_ai/ rate card is exactly 2x that. -GEMINI_COST = 1000 * 2.5e-07 + 1000 * 1.5e-06 -VERTEX_COST = 2 * GEMINI_COST + +@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]: @@ -33,9 +66,9 @@ def _chunks() -> list[str]: } ], "usageMetadata": { - "promptTokenCount": 1000, - "candidatesTokenCount": 1000, - "totalTokenCount": 2000, + "promptTokenCount": PROMPT_TOKENS, + "candidatesTokenCount": COMPLETION_TOKENS, + "totalTokenCount": PROMPT_TOKENS + COMPLETION_TOKENS, }, "modelVersion": MODEL, } @@ -60,7 +93,6 @@ def _logging_obj() -> LiteLLMLoggingObj: def test_streaming_generate_content_bills_against_the_requested_provider( endpoint_type, expected_provider, expected_cost ): - """A streamed gemini/* request must not be priced off the vertex_ai/ rate card.""" logging_obj = _logging_obj() _, kwargs = PassThroughStreamingHandler._build_passthrough_logging_result( @@ -80,7 +112,6 @@ def test_streaming_generate_content_bills_against_the_requested_provider( def test_vertex_generate_content_payload_prices_gemini_urls_at_gemini_rates(): - """The AI Studio host resolves to `gemini`, so the cost must follow it, not the vertex_ai default.""" logging_obj = _logging_obj() result = VertexPassthroughLoggingHandler._handle_logging_vertex_collected_chunks( From e8ec34c4c83edbcf445aff83636b68dc62589059 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 26 Aug 2026 17:18:44 -0700 Subject: [PATCH 3/3] refactor(google_genai): pick the stream logging endpoint type at construction --- litellm/google_genai/streaming_iterator.py | 10 ++-- .../streaming_handler.py | 10 ++-- .../test_google_genai_streaming_iterator.py | 48 ++++++++----------- 3 files changed, 30 insertions(+), 38 deletions(-) diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index e2fac6a615b..a49e43e7bdc 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -75,6 +75,9 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: 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( @@ -86,18 +89,13 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: ) end_time: Final = datetime.now() - endpoint_type: Final = ( - EndpointType.GEMINI - if self.custom_llm_provider == litellm.LlmProviders.GEMINI.value - else EndpointType.VERTEX_AI - ) asyncio.create_task( PassThroughStreamingHandler._route_streaming_logging_to_handler( litellm_logging_obj=self.litellm_logging_obj, passthrough_success_handler_obj=GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ, url_route="/v1/generateContent", request_body=self.request_body or {}, - endpoint_type=endpoint_type, + endpoint_type=self.endpoint_type, start_time=self.start_time, raw_bytes=self.collected_chunks, end_time=end_time, diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index c4dcd086629..5ad41b00890 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -248,7 +248,7 @@ class PassThroughStreamingHandler: kwargs = vertex_passthrough_logging_handler_result["kwargs"] elif endpoint_type == EndpointType.GEMINI: gemini_passthrough_logging_handler_result: Final = ( - GeminiPassthroughLoggingHandler._handle_logging_gemini_collected_chunks( + 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, @@ -260,8 +260,12 @@ class PassThroughStreamingHandler: model=model, ) ) - standard_logging_response_object = gemini_passthrough_logging_handler_result["result"] - kwargs = gemini_passthrough_logging_handler_result["kwargs"] + 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( 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 91058767730..e8ec2848233 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,6 +1,5 @@ -import asyncio import json -from unittest.mock import AsyncMock, MagicMock, patch +from unittest.mock import MagicMock import pytest @@ -12,24 +11,25 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType -@pytest.mark.asyncio @pytest.mark.parametrize( "custom_llm_provider, expected_endpoint_type", [("gemini", EndpointType.GEMINI), ("vertex_ai", EndpointType.VERTEX_AI)], ) -async def test_streaming_logging_routes_to_the_provider_that_served_the_request( - custom_llm_provider, expected_endpoint_type +@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.""" - mock_response = MagicMock() - - async def _aiter_lines(): - yield 'data: {"candidates": []}' - - mock_response.aiter_lines = _aiter_lines - - iterator = AsyncGoogleGenAIGenerateContentStreamingIterator( - response=mock_response, + iterator = iterator_cls( + response=MagicMock(), model="gemini-3.1-flash-image", logging_obj=MagicMock(spec=LiteLLMLoggingObj), generate_content_provider_config=MagicMock(), @@ -37,15 +37,7 @@ async def test_streaming_logging_routes_to_the_provider_that_served_the_request( 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) - assert mock_route.call_args.kwargs["endpoint_type"] == expected_endpoint_type + assert iterator.endpoint_type is expected_endpoint_type def _large_inline_data_event() -> str: @@ -91,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" ) @@ -114,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