diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 35d98b87591..9d33f86d841 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1930,6 +1930,10 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() + # 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: + self.chunks.append(processed_chunk) if self.stream_options is None: # add usage as hidden param usage = calculate_total_usage(chunks=self.chunks) processed_chunk._hidden_params["usage"] = usage @@ -2194,6 +2198,10 @@ class CustomStreamWrapper: else: self.sent_last_chunk = True processed_chunk: Final = self.finish_reason_handler() + # 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: + self.chunks.append(processed_chunk) if self.stream_options is None: usage: Final = calculate_total_usage(chunks=self.chunks) processed_chunk._hidden_params["usage"] = usage # pyright: ignore[reportPrivateUsage] # sync parity diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index d07e8822eb0..f88e082d577 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -5030,3 +5030,116 @@ async def test_openai_stream_relays_the_served_service_tier_on_every_chunk_inclu assert [chunk.get("service_tier") for chunk in relayed] == ["default"] * len(relayed), relayed assert relayed[-1]["usage"]["total_tokens"] == 11 + + +def _last_chunk_carries_finish_reason_wrapper( + logging_obj: Logging, finish_reason: str, sync_stream: bool +) -> CustomStreamWrapper: + """An OpenAI-compatible SSE body whose LAST chunk carries both a delta and the finish_reason, as vLLM emits + when speculative decoding finishes a reply in one engine step.""" + from litellm.llms.openai.chat.gpt_transformation import ( + OpenAIChatCompletionStreamingHandler, + ) + + def line(delta: dict, finish: Optional[str] = None) -> str: + chunk = { + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "m", + "choices": [{"index": 0, "delta": delta, "logprobs": None, "finish_reason": finish}], + } + return f"data: {json.dumps(chunk)}" + + if finish_reason == "tool_calls": + lines = [ + line({"role": "assistant", "content": ""}), + line({"tool_calls": [{"id": "call_1", "type": "function", "index": 0, "function": {"name": "bash", "arguments": ""}}]}), + line({"tool_calls": [{"index": 0, "function": {"arguments": '{"command": "ls'}}]}), + line({"tool_calls": [{"index": 0, "function": {"arguments": '"}'}}]}, "tool_calls"), + ] + else: + lines = [ + line({"role": "assistant", "content": "Hello, this reply is"}), + line({"content": " cut off"}, "length"), + ] + lines.append("data: [DONE]") + + if sync_stream: + streaming_response = iter(lines) + else: + + async def _stream(): + for item in lines: + yield item + + streaming_response = _stream() + return CustomStreamWrapper( + completion_stream=OpenAIChatCompletionStreamingHandler( + streaming_response=streaming_response, sync_stream=sync_stream + ), + model="m", + logging_obj=logging_obj, + custom_llm_provider="hosted_vllm", + ) + + +@pytest.mark.parametrize("finish_reason", ["tool_calls", "length"]) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_logged_response_keeps_finish_reason_from_last_content_chunk( + finish_reason: str, sync_mode: bool, logging_obj: Logging +): + """The client already got the right finish_reason here; the complete response built from ``chunks`` for + callbacks and SpendLogs used to say "stop" instead (tool calls and truncated replies both mislogged).""" + response = _last_chunk_carries_finish_reason_wrapper( + logging_obj, finish_reason, sync_stream=sync_mode + ) + if sync_mode: + received = list(response) + else: + received = [chunk async for chunk in response] + + assert [c.choices[0].finish_reason for c in received if c.choices and c.choices[0].finish_reason] == [ + finish_reason + ] + logged = litellm.stream_chunk_builder(chunks=response.chunks) + assert logged.choices[0].finish_reason == finish_reason + if finish_reason == "tool_calls": + tool_calls = logged.choices[0].message.tool_calls + assert len(tool_calls) == 1 + assert tool_calls[0].function.arguments == '{"command": "ls"}' + else: + assert logged.choices[0].message.content == "Hello, this reply is cut off" + + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_no_synthetic_finish_reason_logged_when_provider_sent_none(sync_mode: bool, logging_obj: Logging): + """A stream that ends before the provider sent any finish_reason (e.g. an Anthropic stream cut after + message_start) must not gain one in ``chunks``: the response builder relies on its absence to estimate usage + instead of taking the provider's placeholder.""" + chunks = [ + ModelResponseStream( + id="chatcmpl-1", + created=1, + model=None, + object="chat.completion.chunk", + choices=[StreamingChoices(finish_reason=None, index=0, delta=Delta(content=text, role="assistant"))], + ) + for text in ("partial", " reply") + ] + response = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=chunks), + model="bedrock/m", + custom_llm_provider="bedrock", + logging_obj=logging_obj, + ) + if sync_mode: + list(response) + else: + [c async for c in response] + + assert response.received_finish_reason is None + assert all(not (c.choices and c.choices[0].finish_reason) for c in response.chunks)