diff --git a/litellm/google_genai/streaming_iterator.py b/litellm/google_genai/streaming_iterator.py index e03f7ee745f..79240172a62 100644 --- a/litellm/google_genai/streaming_iterator.py +++ b/litellm/google_genai/streaming_iterator.py @@ -2,10 +2,8 @@ import asyncio from datetime import datetime from typing import TYPE_CHECKING, Any, Final +from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj -from litellm.proxy.pass_through_endpoints.success_handler import ( - PassThroughEndpointLogging, -) from litellm.types.passthrough_endpoints.pass_through_endpoints import EndpointType if TYPE_CHECKING: @@ -15,8 +13,6 @@ if TYPE_CHECKING: else: BaseGoogleGenAIGenerateContentConfig = Any -GLOBAL_PASS_THROUGH_SUCCESS_HANDLER_OBJ: Final = PassThroughEndpointLogging() - def _encode_google_genai_sse_event(event_lines: list[str]) -> bytes: return ("\n".join(event_lines) + "\n\n").encode("utf-8") @@ -78,15 +74,25 @@ class BaseGoogleGenAIGenerateContentStreamingIterator: self, ): """Handle the logging after all chunks have been collected.""" - from litellm.proxy.pass_through_endpoints.streaming_handler import ( - PassThroughStreamingHandler, - ) + try: + from litellm.proxy.pass_through_endpoints.streaming_handler import ( + PassThroughStreamingHandler, + ) + from litellm.proxy.pass_through_endpoints.success_handler import ( + PassThroughEndpointLogging, + ) + except ImportError as e: + verbose_logger.debug( + "Skipping google_genai streaming cost logging; proxy extras not installed: %s", + e, + ) + return end_time: Final = datetime.now() 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, + passthrough_success_handler_obj=PassThroughEndpointLogging(), url_route="/v1/generateContent", request_body=self.request_body or {}, endpoint_type=EndpointType.VERTEX_AI, @@ -190,4 +196,4 @@ class AsyncGoogleGenAIGenerateContentStreamingIterator(BaseGoogleGenAIGenerateCo return chunk except StopAsyncIteration: await self._handle_async_streaming_logging() - raise StopAsyncIteration + raise StopAsyncIteration \ No newline at end of file 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..a706e77d931 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 json -from unittest.mock import MagicMock +import sys +from unittest.mock import MagicMock, patch import pytest @@ -126,3 +127,51 @@ async def test_async_streaming_iterator_forwards_sse_comment_events(): chunk = await iterator.__anext__() assert chunk == b": keepalive\n\n" + + +def test_streaming_iterator_import_does_not_pull_in_proxy(): + """SDK streaming must not import litellm.proxy at module scope. + + A plain `pip install litellm` has no proxy extras (fastapi, backoff), so a + module-level proxy import here crashes google_genai streaming (issue #36043). + """ + import importlib + + import litellm.google_genai.streaming_iterator as mod + + src = importlib.util.find_spec(mod.__name__).origin + with open(src) as f: + top_level = "".join( + line for line in f if not line.startswith((" ", "\t")) + ) + assert "litellm.proxy" not in top_level + + +@pytest.mark.asyncio +async def test_async_stream_completes_when_proxy_logging_unavailable(): + """End-of-stream cost logging must degrade, not kill the stream (issue #36043).""" + 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-test", + logging_obj=MagicMock(spec=LiteLLMLoggingObj), + generate_content_provider_config=MagicMock(), + litellm_metadata={}, + custom_llm_provider="gemini", + ) + + blocked = { + "litellm.proxy.pass_through_endpoints.streaming_handler": None, + "litellm.proxy.pass_through_endpoints.success_handler": None, + } + with patch.dict(sys.modules, blocked): + chunks = [chunk async for chunk in iterator] + + assert chunks == [b'data: {"text":"hi"}\n\n'] \ No newline at end of file