Merge 5573af3189 into 82ca10cc2d

litellm/litellm_core_utils/streaming_handler.py

@@ -914,6 +914,7 @@ class CustomStreamWrapper:
                 and model_response.choices[0].delta.role is not None
             )
             or (getattr(model_response.choices[0].delta, "reasoning_items", None) is not None)
+            or (getattr(model_response.choices[0].delta, "refusal", None) is not None)
         ):
             return True
         else:

tests/unit/litellm_core_utils/test_streaming_handler.py

@@ -2,6 +2,7 @@ import asyncio
 import json
 import threading
 import time
+from collections.abc import AsyncIterator
 from typing import Final, NoReturn, Optional
 from unittest.mock import AsyncMock, MagicMock, Mock, patch
 
@@ -6444,6 +6445,68 @@ def _last_chunk_carries_finish_reason_wrapper(
     )
 
 
+@pytest.mark.parametrize("provider", ["openai", "hosted_vllm"])
+@pytest.mark.parametrize("sync_mode", [True, False])
+@pytest.mark.parametrize("finish_with_refusal", [True, False])
+@pytest.mark.asyncio
+async def test_stream_preserves_refusal_only_deltas(
+    provider: str, sync_mode: bool, finish_with_refusal: bool, logging_obj: Logging
+) -> None:
+    from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
+
+    refusal_parts: Final = ("I can't", " help", " with", " that.")
+
+    def sse_line(delta: dict[str, str], finish_reason: str | None = None) -> str:
+        payload: Final = {
+            "id": "chatcmpl-refusal",
+            "object": "chat.completion.chunk",
+            "created": 1,
+            "model": "stream-test-model",
+            "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}],
+        }
+        return f"data: {json.dumps(payload)}"
+
+    refusal_lines: Final = tuple(
+        sse_line({"refusal": part}, "stop" if finish_with_refusal and index == len(refusal_parts) - 1 else None)
+        for index, part in enumerate(refusal_parts)
+    )
+    terminal_lines: Final = () if finish_with_refusal else (sse_line({}, "stop"),)
+    lines: Final = (
+        sse_line({"role": "assistant", "content": ""}),
+        *refusal_lines,
+        *terminal_lines,
+        "data: [DONE]",
+    )
+
+    async def async_lines() -> AsyncIterator[str]:
+        for line in lines:
+            yield line
+
+    response: Final = CustomStreamWrapper(
+        completion_stream=OpenAIChatCompletionStreamingHandler(
+            streaming_response=iter(lines) if sync_mode else async_lines(), sync_stream=sync_mode
+        ),
+        model="stream-test-model",
+        logging_obj=logging_obj,
+        custom_llm_provider=provider,
+    )
+    received: Final = tuple(response) if sync_mode else tuple([chunk async for chunk in response])
+    refusals: Final = tuple(
+        getattr(chunk.choices[0].delta, "refusal", None)
+        for chunk in received
+        if chunk.choices and getattr(chunk.choices[0].delta, "refusal", None) is not None
+    )
+    finishes: Final = tuple(
+        chunk.choices[0].finish_reason
+        for chunk in received
+        if chunk.choices and chunk.choices[0].finish_reason is not None
+    )
+
+    assert refusals == refusal_parts
+    assert finishes == ("stop",)
+    assert received[0].choices[0].delta.role == "assistant"
+
+
 @pytest.mark.parametrize("finish_reason", ["tool_calls", "length"])
 @pytest.mark.parametrize("sync_mode", [True, False])
 @pytest.mark.asyncio