fix(streaming): let a finished generic chunk end every streamed choice

An OpenAI-compatible provider that streams n=2 without per-choice finish_reason
ends with `[DONE]`, which arrives as one finished generic chunk on choice 0. A
drop after it errored because choice 1 never finished, while SDK 2 ended the
stream cleanly. A finished generic chunk now ends every choice streamed so far
This commit is contained in:
mateo-berri 2026-10-30 01:35:30 -07:00
parent 8781b42c35
commit 09857e30c0
2 changed files with 52 additions and 2 deletions

View file

@ -108,6 +108,14 @@ def _choice_progress(chunk: object) -> tuple[tuple[int, bool], ...]:
return ()
def _ends_every_choice(chunk: object) -> bool:
match chunk:
case {"is_finished": True}:
return True
case _:
return False
def _stamp_usage_cost(
usage_chunk: ModelResponse | TextCompletionResponse | ModelResponseStream,
complete_response: ModelResponse | TextCompletionResponse | ModelResponseStream,
@ -2277,8 +2285,10 @@ class CustomStreamWrapper:
def _record_choice_progress(self, chunk: object) -> None:
progress: Final = _choice_progress(chunk)
self._streamed_choice_indices = self._streamed_choice_indices.union(index for index, _ in progress)
self._finished_choice_indices = self._finished_choice_indices.union(
index for index, finished in progress if finished
self._finished_choice_indices = (
self._streamed_choice_indices
if _ends_every_choice(chunk)
else self._finished_choice_indices.union(index for index, finished in progress if finished)
)
def _every_choice_finished(self) -> bool:

View file

@ -32,6 +32,7 @@ from litellm.litellm_core_utils.streaming_handler import (
_ProviderChunkParsed,
)
from litellm.llms.custom_httpx.http_handler import HTTPHandler
from litellm.llms.fireworks_ai.chat.transformation import FireworksAIChatCompletionStreamingHandler
from litellm.llms.ollama.completion.transformation import OllamaTextCompletionResponseIterator
from litellm.llms.openai.chat.gpt_transformation import OpenAIChatCompletionStreamingHandler
from litellm.responses.mcp.chat_completions_handler import _add_mcp_metadata_to_response
@ -5452,6 +5453,45 @@ async def test_ollama_generate_stream_dropped_after_its_done_line_ends_cleanly(l
assert [chunk.choices[0].finish_reason for chunk in collected if chunk.choices[0].finish_reason] == ["stop"]
def _fireworks_sse_line_for_choice(index: int, content: str) -> str:
chunk: Final = {
"id": "chatcmpl-two-choices",
"object": "chat.completion.chunk",
"created": 1,
"model": "accounts/fireworks/models/deepseek-v4p1-flash",
"choices": [{"index": index, "delta": {"role": "assistant", "content": content}, "finish_reason": None}],
}
return f"data: {json.dumps(chunk)}"
def _fireworks_n2_lines_without_finish_reasons_through_done_then(failure: Exception) -> Iterator[str]:
yield _fireworks_sse_line_for_choice(0, "Hi")
yield _fireworks_sse_line_for_choice(1, "Hel")
yield "data: [DONE]"
raise failure
@pytest.mark.asyncio
@_SYNC_AND_ASYNC
async def test_n2_stream_dropped_after_done_without_choice_finish_reasons_ends_cleanly(
logging_obj: Logging, sync_mode: bool
):
logging_obj.update_environment_variables(litellm_params={}, optional_params={"n": 2})
lines: Final = _fireworks_n2_lines_without_finish_reasons_through_done_then(_httpx2_read_error())
wrapper: Final = CustomStreamWrapper(
completion_stream=FireworksAIChatCompletionStreamingHandler(
streaming_response=lines if sync_mode else _as_async(lines), sync_stream=sync_mode
),
model="accounts/fireworks/models/deepseek-v4p1-flash",
custom_llm_provider="fireworks_ai",
logging_obj=logging_obj,
)
collected: Final = await _collect(wrapper, sync_mode)
assert "".join(chunk.choices[0].delta.content or "" for chunk in collected if chunk.choices) == "HiHel"
def _generic_chunks_with_a_second_content_block_then(failure: Exception) -> Iterator[GenericStreamingChunk]:
yield GenericStreamingChunk(text="Hi", is_finished=False, finish_reason="", usage=None, index=0)
yield GenericStreamingChunk(text="!", is_finished=False, finish_reason="", usage=None, index=1)