mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-08 08:08:37 +00:00
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:
parent
8781b42c35
commit
09857e30c0
2 changed files with 52 additions and 2 deletions
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue