tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-07 05:07:46 +00:00
Code Issues Packages Projects Releases Wiki Activity

Revert "Add streamGenerateContent cost tracking in passthrough (#15199)" (#15202)

Browse Source 
This reverts commit 8095de506a.
This commit is contained in:
Ishaan Jaff
2025-10-04 14:37:02 -07:00
committed by GitHub
parent 82fd182df2
commit ec9cdf783a
5 changed files with 15 additions and 189 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py
+15 -76
View File
@@ -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
litellm/proxy/pass_through_endpoints/pass_through_endpoints.py
-3
View File
@@ -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
litellm/proxy/pass_through_endpoints/streaming_handler.py
-22
View File
@@ -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(
litellm/types/passthrough_endpoints/pass_through_endpoints.py
-1
View File
@@ -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"
tests/test_litellm/proxy/pass_through_endpoints/llm_provider_handlers/test_gemini_passthrough_logging_handler.py
-87
View File
@@ -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()