From ec136826c551cdcfe05e0d7eff1fa9041f3abea8 Mon Sep 17 00:00:00 2001 From: Tin Date: Fri, 30 Oct 2026 19:51:19 -0700 Subject: [PATCH] fix(streaming): normalize stop after OpenAI-compatible tool calls --- .../litellm_core_utils/streaming_handler.py | 24 +++--- .../test_convert_to_streaming_response.py | 6 +- .../test_streaming_handler.py | 80 +++++++++++++++++++ 3 files changed, 95 insertions(+), 15 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 484edd869cb..857aa2b4fef 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -1302,8 +1302,8 @@ class CustomStreamWrapper: if _is_delta_empty: model_response.choices[0].delta = Delta(content=self._close_open_think_block()) # get any function call arguments - model_response.choices[0].finish_reason = map_finish_reason( - finish_reason=self.received_finish_reason + model_response.choices[0].finish_reason = self._normalize_tool_call_finish_reason( + map_finish_reason(finish_reason=self.received_finish_reason) ) # ensure consistent output to openai self.sent_last_chunk = True @@ -1737,6 +1737,8 @@ class CustomStreamWrapper: # if function returned but type set to None - mistral's api returns type: None tool["type"] = "function" model_response.choices[0].delta = Delta(**_json_delta) + if model_response.choices[0].delta.tool_calls: + self.tool_call = True except Exception as e: verbose_logger.exception( "litellm.CustomStreamWrapper.chunk_creator(): Exception occured - %s", e @@ -1966,6 +1968,11 @@ class CustomStreamWrapper: if self.logging_obj.is_sync_litellm_request(litellm_params): self.logging_obj.success_handler(processed_chunk, None, None, cache_hit) + def _normalize_tool_call_finish_reason(self, finish_reason: str) -> str: + if finish_reason == "stop" and self.tool_call: + return "tool_calls" + return finish_reason + def finish_reason_handler(self): model_response: Final = self.model_response_creator( hidden_params=None if self.received_model_id is None else {"received_model_id": self.received_model_id} @@ -1981,16 +1988,9 @@ class CustomStreamWrapper: if served_tier is not None: setattr(model_response, "service_tier", served_tier) # noqa: B010 # pydantic extra, not a declared field _finish_reason: Final = self.received_finish_reason or self.intermittent_finish_reason - if _finish_reason is not None: - model_response.choices[0].finish_reason = _finish_reason - else: - model_response.choices[0].finish_reason = "stop" - - ## if tool use - if ( - model_response.choices[0].finish_reason == "stop" and self.tool_call - ): # don't overwrite for other - potential error finish reasons - model_response.choices[0].finish_reason = "tool_calls" + model_response.choices[0].finish_reason = self._normalize_tool_call_finish_reason( + _finish_reason if _finish_reason is not None else "stop" + ) close_tag: Final = self._close_open_think_block() if close_tag is not None: model_response.choices[0].delta.content = close_tag diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py index 3713869d7ca..0121aa0e120 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py @@ -563,7 +563,7 @@ async def test_a_cached_answer_streams_its_tool_call_after_the_reasoning_and_the assert reasoning and max(reasoning) < min(text) assert called == [max(text) + 1] assert deltas[called[0]]["tool_calls"] == [tool_call] - assert deltas[-1]["finish"] == "stop" + assert [delta["finish"] for delta in deltas if delta["finish"]] == ["tool_calls"] @pytest.mark.asyncio @@ -586,7 +586,7 @@ async def test_a_cached_tool_turn_without_text_streams_the_reasoning_before_the_ assert "".join(str(deltas[i]["reasoning_content"]) for i in reasoning) == "I need the basket totals" assert reasoning and called and max(reasoning) < min(called) assert [deltas[i]["tool_calls"] for i in called] == [[tool_call]] - assert deltas[-1]["finish"] == "stop" + assert [delta["finish"] for delta in deltas if delta["finish"]] == ["tool_calls"] @pytest.mark.asyncio @@ -610,4 +610,4 @@ async def test_a_cached_tool_turn_keeps_its_whitespace_only_text_between_the_rea assert "".join(str(deltas[i]["content"]) for i in text) == whitespace assert reasoning and text and called and max(reasoning) < min(text) and max(text) < min(called) assert [deltas[i]["tool_calls"] for i in called] == [[tool_call]] - assert deltas[-1]["finish"] == "stop" + assert [delta["finish"] for delta in deltas if delta["finish"]] == ["tool_calls"] diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 4d1699c7fe4..365ed3452cf 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -7953,6 +7953,86 @@ def test_aamazing_unit_test_custom_stream_wrapper_n(): ), f"idx={idx} translated chunk = {chunk_dict} != openai chunk = {chunks[idx]}" +@pytest.mark.parametrize("async_stream", [False, True]) +@pytest.mark.parametrize("transport", ["sdk", "http_handler"]) +@pytest.mark.parametrize("separate_finish", [False, True]) +@pytest.mark.parametrize( + "delta_kind,finish_reason,expected", + [ + ("tool", "stop", "tool_calls"), + ("tool", "length", "length"), + ("tool", "content_filter", "content_filter"), + ("tool", "tool_calls", "tool_calls"), + ("empty_tools", "stop", "stop"), + ("text", "stop", "stop"), + ("function", "stop", "stop"), + ], +) +async def test_openai_compatible_tool_stream_finish_reason( + async_stream: bool, + transport: str, + separate_finish: bool, + delta_kind: str, + finish_reason: str, + expected: str, + logging_obj: Logging, +) -> None: + function: Final = {"name": "get_weather", "arguments": "{}"} + delta: Final = { + "tool": {"tool_calls": [{"index": 0, "id": "call_1", "type": "function", "function": function}]}, + "empty_tools": {"content": "Hello", "tool_calls": []}, + "text": {"content": "Hello"}, + "function": {"function_call": function}, + }[delta_kind] + choices: Final = ( + {"index": 0, "delta": {"role": "assistant", **delta}, "finish_reason": None if separate_finish else finish_reason}, + *(({"index": 0, "delta": {}, "finish_reason": finish_reason},) if separate_finish else ()), + ) + body: Final = "".join( + "data: " + json.dumps({ + "id": "chatcmpl-tool-stop", "object": "chat.completion.chunk", "created": 1, + "model": "m", "choices": [choice], + }) + "\n\n" + for choice in choices + ) + "data: [DONE]\n\n" + raw_response: Final = httpx.Response( + 200, content=body, headers={"content-type": "text/event-stream"}, + request=httpx.Request("POST", "https://provider.example/v1/chat/completions"), + ) + client: Final = AsyncOpenAI(api_key="unused") if async_stream else OpenAI(api_key="unused") + stream: Final = ( + ( + AsyncStream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + if isinstance(client, AsyncOpenAI) + else Stream(cast_to=ChatCompletionChunk, response=raw_response, client=client) + ) + if transport == "sdk" + else OpenAIChatCompletionStreamingHandler( + streaming_response=raw_response.aiter_lines() if async_stream else raw_response.iter_lines(), + sync_stream=not async_stream, + ) + ) + response: Final = CustomStreamWrapper( + completion_stream=stream, model="m", logging_obj=logging_obj, + custom_llm_provider="openai" if transport == "sdk" else "hosted_vllm", + ) + try: + received: Final = [chunk async for chunk in response] if async_stream else list(response) + assert [chunk.choices[0].finish_reason for chunk in received if chunk.choices[0].finish_reason] == [expected] + assembled: Final = litellm.stream_chunk_builder(received) + assert assembled.choices[0].finish_reason == expected + assert litellm.stream_chunk_builder(response.chunks).choices[0].finish_reason == expected + if delta_kind == "tool": + assert [tool.model_dump() for tool in assembled.choices[0].message.tool_calls] == [ + {"id": "call_1", "type": "function", "function": function} + ] + finally: + if isinstance(client, AsyncOpenAI): + await client.close() + else: + client.close() + + def test_unit_test_custom_stream_wrapper_function_call(): """ Test if model returns a tool call, the finish reason is correctly set to 'tool_calls'