diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index c63f28ccf6e..f3d13ce5ea5 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2047,6 +2047,7 @@ class CustomStreamWrapper: return processed_chunk except Exception as e: traceback_exception: Final = traceback.format_exc() + self._record_partial_usage_for_failure() # LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated threading.Thread(target=self.logging_obj.failure_handler, args=(e, traceback_exception)).start() self.close() diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index d0d7a7e1a47..7699aa46b7d 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -3855,6 +3855,63 @@ def test_chunk_creator_tool_calls_not_dropped_on_finish( assert initialized_custom_stream_wrapper.received_finish_reason == "tool_calls" +@pytest.mark.parametrize("sync_mode", [True, False]) +async def test_stream_aborted_after_content_bills_the_partial_output_on_its_failure(sync_mode: bool) -> None: + model: Final = "gpt-4o" + messages: Final = [{"role": "user", "content": "Write a long sentence about a fox. " * 30}] + text: Final = "The quick brown fox jumps over the lazy dog. " * 20 + logging_obj: Final = Logging( + model=model, + messages=messages, + stream=True, + call_type="completion", + start_time=time.time(), + litellm_call_id="aborted-call", + function_id="aborted-fn", + ) + logging_obj.update_environment_variables( + model=model, user=None, optional_params={}, litellm_params={"litellm_call_id": "aborted-call"} + ) + sse_events: Final = ( + { + "id": "chatcmpl-aborted", + "object": "chat.completion.chunk", + "created": 1, + "model": model, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + for delta, finish_reason in (({"role": "assistant", "content": text}, None), ({}, "abort")) + ) + raw_response: Final = httpx.Response( + 200, + content="".join(f"data: {json.dumps(event)}\n\n" for event in sse_events), + headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", "https://provider.example/v1/chat/completions"), + ) + wrapper: Final = CustomStreamWrapper( + completion_stream=OpenAIChatCompletionStreamingHandler( + streaming_response=raw_response.iter_lines() if sync_mode else raw_response.aiter_lines(), + sync_stream=sync_mode, + ), + model=model, + custom_llm_provider="openai", + logging_obj=logging_obj, + ) + + with pytest.raises(MidStreamFallbackError, match="abort"): + _chunks: Final = list(wrapper) if sync_mode else [chunk async for chunk in wrapper] + + expected: Final = ( + litellm.token_counter(model=model, messages=messages), + litellm.token_counter(model=model, text=text, count_response_tokens=True), + ) + usage: Final = logging_obj.model_call_details["combined_usage_object"] + assert (usage.prompt_tokens, usage.completion_tokens) == expected + rates: Final = litellm.model_cost[model] + spend: Final = expected[0] * rates["input_cost_per_token"] + expected[1] * rates["output_cost_per_token"] + assert logging_obj.model_call_details["response_cost"] == pytest.approx(spend) + + def test_record_partial_usage_for_failure_stashes_usage_and_cost(): """A stream that breaks mid-flight must surface the usage assembled from the chunks already delivered, plus its cost, on the logging object so the