diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 0911c4bbdf4..70bad6395ef 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -307,6 +307,7 @@ class CustomStreamWrapper: self._provider_response_model: str | None = None self.received_finish_reason: str | None = None self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream + self.received_any_finish_reason = False # a later chunk can reset intermittent_finish_reason self.special_tokens = [ "<|assistant|>", "<|system|>", @@ -1534,6 +1535,7 @@ class CustomStreamWrapper: return _ProviderChunkEarlyReturn(None) completion_obj["content"] = response_obj["text"] self.intermittent_finish_reason = response_obj.get("finish_reason", None) + self.received_any_finish_reason |= self.intermittent_finish_reason is not None if response_obj["is_finished"]: if response_obj["finish_reason"] == "error": raise Exception( @@ -1909,6 +1911,7 @@ class CustomStreamWrapper: and (self.custom_llm_provider in _OPENAI_AZURE_CHAT_PROVIDERS or self.uses_openai_chat_stream) and self.received_finish_reason is None and self.intermittent_finish_reason is None + and not self.received_any_finish_reason ) @staticmethod diff --git a/litellm/llms/langgraph/chat/sse_iterator.py b/litellm/llms/langgraph/chat/sse_iterator.py index 499cbab3b7b..6b2f0c2a3a8 100644 --- a/litellm/llms/langgraph/chat/sse_iterator.py +++ b/litellm/llms/langgraph/chat/sse_iterator.py @@ -37,8 +37,9 @@ class LangGraphSSEStreamIterator: return self def __aiter__(self): - """Initialize async iteration.""" - self.async_line_iterator = self.response.aiter_lines() + """Initialize async iteration once; CustomStreamWrapper re-enters `async for` on every chunk.""" + if self.async_line_iterator is None: + self.async_line_iterator = self.response.aiter_lines() return self def _parse_sse_line(self, line: str) -> ModelResponseStream | None: diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 7aa0f7b01ad..2d1c2a581cf 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -29,6 +29,9 @@ from litellm.litellm_core_utils.streaming_handler import ( ) from litellm.llms.anthropic.chat.handler import ModelResponseIterator as AnthropicModelResponseIterator from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator as GeminiModelResponseIterator, +) from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response from litellm.types.utils import ( CompletionTokensDetailsWrapper, @@ -1759,6 +1762,48 @@ async def test_generic_chunk_stream_without_a_finish_reason_keeps_the_default_st assert "".join(c.choices[0].delta.content or "" for c in received if c.choices) == "partial reply" +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_gemini_stream_with_a_usage_chunk_after_its_finish_reason_still_completes( + sync_mode: bool, logging_obj: Logging +): + """Gemini marks every chunk unfinished, so its finishReason only shows up mid-stream before the usage chunk.""" + lines: Final = [ + f"data: {json.dumps(event)}" + for event in ( + {"candidates": [{"content": {"parts": [{"text": "Hello"}], "role": "model"}, "index": 0}]}, + { + "candidates": [ + {"content": {"parts": [{"text": " there"}], "role": "model"}, "finishReason": "STOP", "index": 0} + ] + }, + {"usageMetadata": {"promptTokenCount": 3, "candidatesTokenCount": 2, "totalTokenCount": 5}}, + ) + ] + + async def async_lines() -> AsyncIterator[str]: + for line in lines: + yield line + + logging_obj.update_environment_variables( + model="gemini-2.5-flash", user=None, optional_params={}, litellm_params={"litellm_call_id": "gemini-tail"} + ) + iterator: Final = GeminiModelResponseIterator( + streaming_response=iter(lines) if sync_mode else async_lines(), sync_stream=sync_mode, logging_obj=logging_obj + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=iter(iterator) if sync_mode else iterator, + model="gemini-2.5-flash", + custom_llm_provider="gemini", + logging_obj=logging_obj, + ) + + received: Final = await _collect(wrapper, sync_mode) + + assert "".join(c.choices[0].delta.content or "" for c in received if c.choices) == "Hello there" + assert [c.choices[0].finish_reason for c in received if c.choices and c.choices[0].finish_reason] == ["stop"] + + def test_streaming_handler_with_created_time_propagation( initialized_custom_stream_wrapper: CustomStreamWrapper, logging_obj: Logging ): diff --git a/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py b/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py index 238f7dd2f26..6612c71f436 100644 --- a/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py +++ b/tests/unit/llms/langgraph/chat/test_langgraph_transformation.py @@ -268,3 +268,28 @@ async def test_langgraph_streaming_acompletion_without_a_client_posts_through_th "assistant_id": "agent", "input": {"messages": [{"role": "human", "content": "What is 25 * 4?"}]}, } + + +@pytest.mark.asyncio +async def test_langgraph_async_stream_delivers_every_message_event_before_its_stop(monkeypatch): + async def events(): + yield b'data: ["messages", [[{"type": "ai", "content": "25 * 4"}, {}]]]\n\n' + yield b'data: ["messages", [[{"type": "ai", "content": " is 100"}, {}]]]\n\n' + yield b'data: ["metadata", {"run_id": "run-1"}]\n\n' + + cached_client = get_async_httpx_client(llm_provider="langgraph", params={}) + monkeypatch.setattr( + cached_client, + "client", + httpx.AsyncClient(transport=httpx.MockTransport(lambda request: httpx.Response(200, content=events()))), + ) + + response = await litellm.acompletion( + model="langgraph/agent", + messages=[{"role": "user", "content": "What is 25 * 4?"}], + api_base="http://langgraph.test", + stream=True, + ) + streamed = [(chunk.choices[0].delta.content, chunk.choices[0].finish_reason) async for chunk in response] + + assert streamed == [("25 * 4", None), (" is 100", None), (None, "stop")]