From ec9cdf783aa8f430912bd910dfaa62afb3ec83e2 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Sat, 4 Oct 2025 14:37:02 -0700 Subject: [PATCH] Revert "Add streamGenerateContent cost tracking in passthrough (#15199)" (#15202) This reverts commit 8095de506a80db01cb69cca63a2f59324a4bc491. --- .../gemini_passthrough_logging_handler.py | 91 +++---------------- .../pass_through_endpoints.py | 3 - .../streaming_handler.py | 22 ----- .../pass_through_endpoints.py | 1 - ...test_gemini_passthrough_logging_handler.py | 87 ------------------ 5 files changed, 15 insertions(+), 189 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py index ecd5b0eb09..8c96c2ab96 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py @@ -128,64 +128,6 @@ class GeminiPassthroughLoggingHandler: "kwargs": kwargs, } - @staticmethod - def _parse_gemini_streaming_json( - all_chunks: List[str], - gemini_iterator: GeminiModelResponseIterator, - ) -> List[Any]: - """ - Parse Gemini streaming chunks from fragmented JSON strings. - - Gemini's streaming format sends a single JSON array that may be split across - multiple lines. This method: - 1. Joins all fragmented string chunks into complete JSON - 2. Parses the JSON array/object - 3. Transforms each item using Gemini's chunk_parser - - Args: - all_chunks: Raw string chunks from the streaming response - gemini_iterator: GeminiModelResponseIterator instance for parsing - - Returns: - List of parsed chunks in OpenAI format, or empty list if parsing fails - """ - parsed_chunks = [] - - verbose_proxy_logger.debug(f"Gemini streaming: Processing {len(all_chunks)} raw chunks") - - # Gemini streaming response is a single JSON array that may be split across lines - # Join all chunks back together to reconstruct the complete JSON - combined_chunk = "".join(all_chunks) - - # Parse the combined JSON string - try: - dict_chunk = json.loads(combined_chunk) - verbose_proxy_logger.debug(f"Parsed JSON object: {type(dict_chunk)}") - - # Gemini returns an array of response objects - if isinstance(dict_chunk, list): - for item in dict_chunk: - try: - # Call chunk_parser directly with the dict, not _common_chunk_parsing_logic - parsed_chunk = gemini_iterator.chunk_parser(chunk=item) - if parsed_chunk is not None: - parsed_chunks.append(parsed_chunk) - except Exception as e: - verbose_proxy_logger.error(f"Error parsing Gemini chunk item: {e}", exc_info=True) - continue - else: - # Single object response - parsed_chunk = gemini_iterator.chunk_parser(chunk=dict_chunk) - if parsed_chunk is not None: - parsed_chunks.append(parsed_chunk) - - except json.JSONDecodeError as e: - verbose_proxy_logger.error(f"Failed to parse Gemini streaming response as JSON: {e}") - return [] - - verbose_proxy_logger.debug(f"Total parsed chunks: {len(parsed_chunks)}") - return parsed_chunks - @staticmethod def _build_complete_streaming_response( all_chunks: List[str], @@ -193,20 +135,20 @@ class GeminiPassthroughLoggingHandler: model: str, url_route: str, ) -> Optional[Union[ModelResponse, TextCompletionResponse]]: - if "generateContent" not in url_route and "streamGenerateContent" not in url_route: + parsed_chunks = [] + if "generateContent" in url_route or "streamGenerateContent" in url_route: + gemini_iterator: Any = GeminiModelResponseIterator( + streaming_response=None, + sync_stream=False, + logging_obj=litellm_logging_obj, + ) + chunk_parsing_logic: Any = gemini_iterator._common_chunk_parsing_logic + parsed_chunks = [chunk_parsing_logic(chunk) for chunk in all_chunks] + else: + return None + + if len(parsed_chunks) == 0: return None - - gemini_iterator: Any = GeminiModelResponseIterator( - streaming_response=None, - sync_stream=False, - logging_obj=litellm_logging_obj, - ) - - # Parse the streaming chunks - parsed_chunks = GeminiPassthroughLoggingHandler._parse_gemini_streaming_json( - all_chunks=all_chunks, - gemini_iterator=gemini_iterator, - ) all_openai_chunks = [] for parsed_chunk in parsed_chunks: @@ -214,10 +156,7 @@ class GeminiPassthroughLoggingHandler: continue all_openai_chunks.append(parsed_chunk) - complete_streaming_response = litellm.stream_chunk_builder( - chunks=all_openai_chunks, - logging_obj=litellm_logging_obj, - ) + complete_streaming_response = litellm.stream_chunk_builder(chunks=all_openai_chunks) return complete_streaming_response @@ -246,7 +185,7 @@ class GeminiPassthroughLoggingHandler: response_cost = litellm.completion_cost( completion_response=litellm_model_response, model=model, - custom_llm_provider=custom_llm_provider, + custom_llm_provider="gemini", ) kwargs["response_cost"] = response_cost diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index f41ee21db8..005b474424 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -301,9 +301,6 @@ class HttpPassThroughEndpointHelpers(BasePassthroughUtils): or ("rawPredict") in url or ("streamRawPredict") in url ): - # Check if it's Gemini (Google AI Studio) or Vertex AI - if parsed_url.hostname and parsed_url.hostname.endswith("generativelanguage.googleapis.com"): - return EndpointType.GEMINI return EndpointType.VERTEX_AI elif parsed_url.hostname == "api.anthropic.com": return EndpointType.ANTHROPIC diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 792b496a66..2d5b0a686c 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -105,28 +105,6 @@ class PassThroughStreamingHandler: anthropic_passthrough_logging_handler_result["result"] ) kwargs = anthropic_passthrough_logging_handler_result["kwargs"] - elif endpoint_type == EndpointType.GEMINI: - from litellm.proxy.pass_through_endpoints.llm_provider_handlers.gemini_passthrough_logging_handler import ( - GeminiPassthroughLoggingHandler, - ) - - gemini_passthrough_logging_handler_result = ( - 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.VERTEX_AI: vertex_passthrough_logging_handler_result = ( VertexPassthroughLoggingHandler._handle_logging_vertex_collected_chunks( diff --git a/litellm/types/passthrough_endpoints/pass_through_endpoints.py b/litellm/types/passthrough_endpoints/pass_through_endpoints.py index 85c8c0f335..c99775f2a6 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -6,7 +6,6 @@ from typing_extensions import TypedDict class EndpointType(str, Enum): VERTEX_AI = "vertex-ai" - GEMINI = "gemini" ANTHROPIC = "anthropic" OPENAI = "openai" GENERIC = "generic" diff --git a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py index 3854e3944c..6f87d8f6ab 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py @@ -285,90 +285,3 @@ class TestGeminiPassthroughLoggingHandler: assert call_kwargs["response_cost"] is not None assert call_kwargs["model"] == "gemini-1.5-flash" assert call_kwargs["custom_llm_provider"] == "gemini" - - @patch("litellm.completion_cost") - @patch("litellm.stream_chunk_builder") - def test_gemini_streaming_cost_calculation(self, mock_stream_chunk_builder, mock_completion_cost): - """Test that Gemini streaming passthrough correctly calculates cost with logging_obj""" - # Arrange - mock_completion_cost.return_value = 0.000025 - mock_logging_obj = self._create_mock_logging_obj() - - # Mock the stream_chunk_builder to return a response with usage - from litellm.utils import ModelResponse, Usage - mock_response = ModelResponse() - mock_usage = Usage(prompt_tokens=5, completion_tokens=10, total_tokens=15) - mock_response.usage = mock_usage - mock_stream_chunk_builder.return_value = mock_response - - # Mock fragmented JSON chunks (as they come from the streaming response) - fragmented_chunks = [ - '[{"candidates": [', - '{"content": {"parts": [{"text": "Hello"}], "role": "model"},', - '"finishReason": "STOP", "index": 0}],', - '"usageMetadata": {"promptTokenCount": 5, "candidatesTokenCount": 10, "totalTokenCount": 15}', - '}]' - ] - - # Act - result = GeminiPassthroughLoggingHandler._build_complete_streaming_response( - all_chunks=fragmented_chunks, - litellm_logging_obj=mock_logging_obj, - model="gemini-1.5-flash", - url_route="https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash:streamGenerateContent" - ) - - # Assert - assert result is not None - assert result == mock_response - - # Verify stream_chunk_builder was called with logging_obj for cost injection - mock_stream_chunk_builder.assert_called_once() - call_args = mock_stream_chunk_builder.call_args - - # Check that logging_obj was passed for cost injection - assert "logging_obj" in call_args.kwargs - assert call_args.kwargs["logging_obj"] == mock_logging_obj - - # Verify the chunks were properly reconstructed from fragmented JSON - # The first argument should be the chunks list - chunks_arg = call_args[0][0] if call_args[0] else call_args.kwargs.get("chunks", []) - assert len(chunks_arg) > 0 # Should have parsed chunks - - def test_gemini_streaming_json_parsing(self): - """Test that fragmented JSON chunks are correctly joined and parsed""" - # Arrange - # Mock fragmented JSON chunks that simulate how Gemini streaming response gets split - fragmented_chunks = [ - '[{"candidates": [', - '{"content": {"parts": [{"text": "Test response"}], "role": "model"},', - '"finishReason": "STOP", "index": 0}],', - '"usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 7, "totalTokenCount": 10}', - '}]' - ] - - # Mock the gemini iterator's chunk_parser method - mock_iterator = MagicMock() - mock_iterator.chunk_parser.return_value = {"candidates": [{"content": {"parts": [{"text": "Test response"}]}}]} - - # Act - result = GeminiPassthroughLoggingHandler._parse_gemini_streaming_json( - all_chunks=fragmented_chunks, - gemini_iterator=mock_iterator - ) - - # Assert - # The method should successfully join the fragmented JSON and return parsed chunks - assert isinstance(result, list) - assert len(result) == 1 # Should have one parsed chunk - - # Verify that the combined JSON is valid - combined_json = "".join(fragmented_chunks) - parsed_json = json.loads(combined_json) - assert isinstance(parsed_json, list) - assert len(parsed_json) == 1 - assert "candidates" in parsed_json[0] - assert "usageMetadata" in parsed_json[0] - - # Verify the iterator's chunk_parser was called - mock_iterator.chunk_parser.assert_called_once()