diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index bf8c64cb022..371291c862c 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -2516,6 +2516,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 e1b4e061218..5b82a486ede 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -4113,6 +4113,64 @@ async def test_stream_cut_before_a_finish_chunk_bills_the_token_estimate_on_its_ assert logging_obj.model_call_details["response_cost"] == pytest.approx(spend) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +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=datetime.datetime(2025, 3, 3, 21, 38, 10), + 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"): + await _collect(wrapper, sync_mode) + + 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) + + @pytest.mark.parametrize("sync_mode", [True, False]) @pytest.mark.asyncio async def test_stream_with_usage_after_the_finish_chunk_bills_only_the_reported_usage(sync_mode: bool) -> None: