diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 6cb0331620f..6cf36e3e660 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -440,7 +440,7 @@ class BaseResponsesAPIStreamingIterator: and _response_obj is not None and _response_obj.usage is None ): - _response_obj.usage = _estimate_usage_from_text( + _response_obj.usage = _estimate_usage_safely( self.model or "", self.request_data.get("input"), self.request_data, @@ -1381,6 +1381,24 @@ def _estimate_usage_from_text( ) +def _estimate_usage_safely( + model: str, + request_input: object, + responses_api_request: Mapping[str, object], + generated_text: str, +) -> ResponseAPIUsage | None: + try: + return _estimate_usage_from_text( + model=model, + request_input=request_input, + responses_api_request=responses_api_request, + generated_text=generated_text, + ) + except Exception as e: + verbose_logger.debug("Could not estimate usage from stream text, billing $0: %s", e) + return None + + def _stamp_responses_usage_cost( response_obj: ResponsesAPIResponse | None, logging_obj: LiteLLMLoggingObj | None ) -> None: diff --git a/tests/test_litellm/responses/test_streaming_iterator.py b/tests/test_litellm/responses/test_streaming_iterator.py index c3721353d42..234e63ca3be 100644 --- a/tests/test_litellm/responses/test_streaming_iterator.py +++ b/tests/test_litellm/responses/test_streaming_iterator.py @@ -792,3 +792,27 @@ async def test_completed_event_without_usage_counts_multimodal_input_as_messages usage = iterator.completed_response.response.usage assert usage is not None assert usage.input_tokens < json_count / 2 + + +@pytest.mark.asyncio +async def test_completed_event_survives_a_failing_usage_estimate(): + """A raising token_counter must not break a stream that previously completed: + the estimate is best-effort and falls back to usage None.""" + response = _responses_api_response_without_usage() + iterator = _make_iterator( + sse_events=[ + _sse_event({"type": "response.output_text.delta", "delta": "hello world"}), + _sse_event({"type": "response.completed", "response": {}}), + ], + logging_obj=_logging_obj_stub(), + config=_mock_config_with_completed_response(response), + request_data={"input": "count these input tokens please"}, + ) + + with patch.object(litellm, "token_counter", side_effect=RuntimeError("tokenizer exploded")): + yielded: list = [] + async for chunk in iterator: + yielded.append(chunk) + + assert yielded + assert iterator.completed_response.response.usage is None