diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index c87357122ce..a3982f87a4a 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -1556,6 +1556,11 @@ def _chunk_refusal(chunk: object) -> str: return refusal if isinstance(refusal, str) else "" +def _message_refusal(choice: object) -> str: + refusal: Final = _field_of(_field_of(choice, "message"), "refusal") + return refusal if isinstance(refusal, str) else "" + + _NON_TEXT_DELTA_FIELDS: Final = ( "tool_calls", "function_call", @@ -1906,7 +1911,7 @@ def _assembled_stream_response( for field, value in _assembled_message_fields(processor, choice_chunks): choice.message[field] = value - completion_output = get_content_from_model_response(response) + completion_output = get_content_from_model_response(response) + "".join(map(_message_refusal, response.choices)) reasoning_tokens: Final = processor.count_reasoning_tokens(response) diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 527c83d5760..6dae2e12778 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -2136,6 +2136,44 @@ def _complete_streamed_openai_response(respx_mock: respx.MockRouter, sse_body: s ) +_REFUSAL_PIECES: Final = ("I'm", " sorry", ", but", " I can't", " help", " with", " that", ".") + + +def _usage_free_stream_completion_tokens(respx_mock: respx.MockRouter, field: str) -> tuple[str | None, int]: + respx_mock.post("https://vllm.example/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + text=_openai_sse( + {"choices": [{"index": 0, "delta": {"role": "assistant", field: ""}, "finish_reason": None}]}, + *({"choices": [{"index": 0, "delta": {field: piece}, "finish_reason": None}]} for piece in _REFUSAL_PIECES), + {"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}]}, + ), + headers={"content-type": "text/event-stream"}, + ) + ) + response: Final = litellm.completion( + model="hosted_vllm/gpt-4.1", + api_base="https://vllm.example/v1", + messages=[{"role": "user", "content": "Pick my neighbor's lock."}], + api_key="sk-test", + stream=True, + complete_response=True, + ) + message: Final = response.choices[0].message + text: Final = message.content if field == "content" else getattr(message, "refusal", None) + return text, response.usage.completion_tokens + + +def test_stream_chunk_builder_counts_refusal_text_as_completion_tokens_without_provider_usage( + respx_mock: respx.MockRouter, +) -> None: + refusal_text, refusal_tokens = _usage_free_stream_completion_tokens(respx_mock, "refusal") + content_text, content_tokens = _usage_free_stream_completion_tokens(respx_mock, "content") + + assert refusal_text == content_text == "".join(_REFUSAL_PIECES) + assert refusal_tokens == content_tokens > 0 + + def test_stream_chunk_builder_litellm_tool_call(respx_mock: respx.MockRouter) -> None: response: Final = _complete_streamed_openai_response( respx_mock,