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 8c96c2ab96a..ecd5b0eb094 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,6 +128,64 @@ 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], @@ -135,20 +193,20 @@ class GeminiPassthroughLoggingHandler: model: str, url_route: str, ) -> Optional[Union[ModelResponse, TextCompletionResponse]]: - 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: + if "generateContent" not in url_route and "streamGenerateContent" not in url_route: 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: @@ -156,7 +214,10 @@ class GeminiPassthroughLoggingHandler: continue all_openai_chunks.append(parsed_chunk) - complete_streaming_response = litellm.stream_chunk_builder(chunks=all_openai_chunks) + complete_streaming_response = litellm.stream_chunk_builder( + chunks=all_openai_chunks, + logging_obj=litellm_logging_obj, + ) return complete_streaming_response @@ -185,7 +246,7 @@ class GeminiPassthroughLoggingHandler: response_cost = litellm.completion_cost( completion_response=litellm_model_response, model=model, - custom_llm_provider="gemini", + custom_llm_provider=custom_llm_provider, ) 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 005b4744248..f41ee21db80 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -301,6 +301,9 @@ 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 2d5b0a686ce..792b496a664 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -105,6 +105,28 @@ 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 c99775f2a6c..85c8c0f3351 100644 --- a/litellm/types/passthrough_endpoints/pass_through_endpoints.py +++ b/litellm/types/passthrough_endpoints/pass_through_endpoints.py @@ -6,6 +6,7 @@ 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 6f87d8f6ab5..3854e3944ce 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,3 +285,90 @@ 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()