ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(proxy): omit OpenAI [DONE] on google-genai streamGenerateContent (#29426)

Browse Source 
* fix(proxy): omit OpenAI [DONE] on google-genai streamGenerateContent

google-genai SDK uses ?alt=sse and cannot parse the proxy's trailing
data: [DONE] chunk. Skip that terminator for agenerate_content_stream.

Co-authored-by: Cursor <cursoragent@cursor.com>

* fix(proxy): address Greptile review on google-genai stream fix

Always yield stream error_message; only gate data: [DONE] on the skip flag.
Set _litellm_skip_openai_stream_done in google_endpoints instead of common_request_processing.

Co-authored-by: Cursor <cursoragent@cursor.com>

---------

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
Sameer Kankute 2026-06-02 03:08:19 +05:30 committed by GitHub
parent 29270a36a5
commit c908505e6a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 110 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
litellm/proxy/google_endpoints/endpoints.py
View File

@ -105,6 +105,8 @@ async def google_stream_generate_content(
if "model" not in data:
data["model"] = model_name
data["stream"] = True
# google-genai SDK (?alt=sse) must not receive OpenAI's data: [DONE] terminator.
data["_litellm_skip_openai_stream_done"] = True
processor = ProxyBaseLLMRequestProcessing(data=data)
try:

7
litellm/proxy/proxy_server.py
View File

@ -7072,11 +7072,12 @@ async def async_data_generator( # noqa: PLR0915
# still flush their post-stream logging.
ProxyLogging._fire_deferred_stream_logging(request_data)
# Streaming is done, yield the [DONE] chunk
if error_message is not None:
yield error_message
done_message = "[DONE]"
yield f"data: {done_message}\n\n"
# OpenAI-compatible streams terminate with data: [DONE]; Google GenAI (?alt=sse) does not.
if not request_data.get("_litellm_skip_openai_stream_done"):
done_message = "[DONE]"
yield f"data: {done_message}\n\n"
except Exception as e:
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.async_data_generator(): Exception occured - {}".format(

104
tests/test_litellm/proxy/test_proxy_server.py
View File

@ -5165,6 +5165,110 @@ async def test_async_data_generator_passes_through_google_native_sse_bytes():
assert yielded_text[-1] == "data: [DONE]\n\n"
@pytest.mark.asyncio
async def test_async_data_generator_google_genai_stream_omits_openai_done():
"""
google-genai SDK streamGenerateContent?alt=sse must not receive data: [DONE].
"""
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.proxy_server import async_data_generator
from litellm.proxy.utils import ProxyLogging
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_request_data = {
"model": "gemini-2.0-flash",
"_litellm_skip_openai_stream_done": True,
}
gemini_event = (
b'data: {"candidates": [{"content": {"parts": [{"text": "Hi"}]}}]}\n\n'
)
class MockStream:
def __aiter__(self):
return self._stream()
async def _stream(self):
yield gemini_event
async def aclose(self):
pass
mock_response = MockStream()
mock_response.aclose = AsyncMock()
mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
mock_proxy_logging_obj.has_streaming_callbacks.return_value = False
mock_proxy_logging_obj.needs_iterator_wrap.return_value = False
mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False
mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = MagicMock()
mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock()
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock()
with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj):
with patch.object(ProxyLogging, "_fire_deferred_stream_logging"):
yielded_data = []
async for data in async_data_generator(
mock_response, mock_user_api_key_dict, mock_request_data
):
yielded_data.append(data)
yielded_text = [
chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk
for chunk in yielded_data
]
assert yielded_text == [gemini_event.decode("utf-8")]
assert "[DONE]" not in "".join(yielded_text)
@pytest.mark.asyncio
async def test_async_data_generator_google_genai_stream_forwards_error_without_done():
"""Stream errors must still reach the client when OpenAI [DONE] is skipped."""
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.proxy_server import async_data_generator
from litellm.proxy.utils import ProxyLogging
error_sse = 'data: {"error": {"message": "stream failed"}}\n\n'
mock_user_api_key_dict = MagicMock(spec=UserAPIKeyAuth)
mock_request_data = {
"model": "gemini-2.0-flash",
"_litellm_skip_openai_stream_done": True,
}
class MockStream:
def __aiter__(self):
return self._stream()
async def _stream(self):
yield error_sse
async def aclose(self):
pass
mock_response = MockStream()
mock_response.aclose = AsyncMock()
mock_proxy_logging_obj = MagicMock(spec=ProxyLogging)
mock_proxy_logging_obj.has_streaming_callbacks.return_value = False
mock_proxy_logging_obj.needs_iterator_wrap.return_value = False
mock_proxy_logging_obj.needs_per_chunk_streaming_hook.return_value = False
mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = MagicMock()
mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock()
mock_proxy_logging_obj.post_call_failure_hook = AsyncMock()
with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj):
with patch.object(ProxyLogging, "_fire_deferred_stream_logging"):
yielded_data = []
async for data in async_data_generator(
mock_response, mock_user_api_key_dict, mock_request_data
):
yielded_data.append(data)
yielded_text = [
chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk
for chunk in yielded_data
]
assert yielded_text == [error_sse]
assert "[DONE]" not in "".join(yielded_text)
@pytest.mark.asyncio
async def test_async_data_generator_cleanup_on_normal_completion():
"""