mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-01 07:03:02 +00:00
fix(streaming): normalize stop after OpenAI-compatible tool calls
This commit is contained in:
parent
87d91d78b7
commit
ec136826c5
3 changed files with 95 additions and 15 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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'
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue