diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index f2da2bc4a1f..c9028118f46 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1954,7 +1954,16 @@ class CustomStreamWrapper: raise # Re-raise StopIteration else: self.sent_last_chunk = True - processed_chunk: Final = self.finish_reason_handler() + try: + processed_chunk: Final = self.finish_reason_handler() + except litellm.exceptions.IncompleteStreamError as strict_error: + self._record_partial_usage_for_failure() + if self.logging_obj is not None: + threading.Thread( + target=self.logging_obj.failure_handler, + args=(strict_error, traceback.format_exc()), + ).start() + raise # The logged response is built from self.chunks; keep a finish_reason the provider sent on its # last content chunk (stripped there), but never add the synthetic "stop" used when it sent none. if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: @@ -2222,7 +2231,17 @@ class CustomStreamWrapper: raise StopAsyncIteration # Re-raise StopIteration else: self.sent_last_chunk = True - processed_chunk: Final = self.finish_reason_handler() + try: + processed_chunk: Final = self.finish_reason_handler() + except litellm.exceptions.IncompleteStreamError as strict_error: + self._record_partial_usage_for_failure() + if self.logging_obj is not None: + asyncio.create_task( + self.logging_obj.dispatch_failure_handlers( + strict_error, traceback.format_exc(), prefer_async_handlers=True + ) + ) + raise # The logged response is built from self.chunks; keep a finish_reason the provider sent on its # last content chunk (stripped there), but never add the synthetic "stop" used when it sent none. if self.received_finish_reason is not None or self.intermittent_finish_reason is not None: diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 76ca0bd0f56..649bba238bf 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -5248,3 +5248,40 @@ def test_an_empty_finish_reason_still_marks_the_stream_finished(): assert wrapper.stream_reported_finished is True assert wrapper.finish_reason_handler().choices[0].finish_reason == "stop" + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_strict_incomplete_stream_records_partial_usage_before_raising( + sync_mode: bool, logging_obj: Logging +): + chunks = [ + ModelResponseStream( + id="chatcmpl-1", + created=1, + model=None, + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, index=0, delta=Delta(content="partial", role="assistant") + ) + ], + ) + ] + response = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=chunks), + model="gpt-4o", + custom_llm_provider="openai", + logging_obj=logging_obj, + ) + response.strict_stream_completion = True + recorded = MagicMock() + response._record_partial_usage_for_failure = recorded + + with pytest.raises(litellm.exceptions.IncompleteStreamError): + if sync_mode: + list(response) + else: + [c async for c in response] + + recorded.assert_called_once()