mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(openai_like): preserve reasoning_content, multi-tool calls, and usage in streaming
Fixes #48238 When streaming responses from OpenAI-like endpoints (including Vertex Model Garden dedicated endpoints 'vertex_ai/openai/...'), reasoning chunks and secondary tool call deltas were previously dropped because: 1. ModelResponseIterator.chunk_parser only checked delta.content and tool_calls[0]. 2. GenericStreamingChunk lacked reasoning_content and tool_calls fields. 3. CustomStreamWrapper considered reasoning-only chunks (with content=='') empty, dropping them before reaching the client. This patch: - Extends GenericStreamingChunk to support reasoning_content and tool_calls. - Updates ModelResponseIterator to extract reasoning_content and preserve all indexed tool calls. - Updates CustomStreamWrapper to treat reasoning_content chunks as non-empty and emit them with content=None. - Adds comprehensive unit and streaming regression tests.
This commit is contained in:
Rajasimhareddybolla
parent
dddbaa803d
commit
2856c56d20
5 changed files with 286 additions and 17 deletions
|
|
@ -971,6 +971,11 @@ class CustomStreamWrapper:
|
|||
if (
|
||||
"content" in completion_obj
|
||||
and (isinstance(completion_obj["content"], str) and len(completion_obj["content"]) > 0)
|
||||
or (
|
||||
"reasoning_content" in completion_obj
|
||||
and completion_obj["reasoning_content"] is not None
|
||||
and len(completion_obj["reasoning_content"]) > 0
|
||||
)
|
||||
or (
|
||||
"tool_calls" in completion_obj
|
||||
and completion_obj["tool_calls"] is not None
|
||||
|
|
@ -1214,7 +1219,13 @@ class CustomStreamWrapper:
|
|||
)
|
||||
else:
|
||||
## else
|
||||
completion_obj["content"] = model_response_str
|
||||
if model_response_str or (
|
||||
completion_obj.get("reasoning_content") is None
|
||||
and completion_obj.get("tool_calls") is None
|
||||
):
|
||||
completion_obj["content"] = model_response_str
|
||||
else:
|
||||
completion_obj["content"] = None
|
||||
if self.sent_first_chunk is False:
|
||||
completion_obj["role"] = "assistant"
|
||||
self.sent_first_chunk = True
|
||||
|
|
@ -1345,7 +1356,9 @@ class CustomStreamWrapper:
|
|||
if self.received_finish_reason is not None:
|
||||
_chunk_has_content: Final = isinstance(chunk, dict) and (
|
||||
bool(chunk.get("text", ""))
|
||||
or bool(chunk.get("reasoning_content", ""))
|
||||
or chunk.get("tool_use") is not None
|
||||
or bool(chunk.get("tool_calls"))
|
||||
# Usage-only final chunks are valid and needed to surface
|
||||
# finish_reason/usage to downstream translators.
|
||||
or chunk.get("usage") is not None
|
||||
|
|
@ -1354,6 +1367,8 @@ class CustomStreamWrapper:
|
|||
raise StopIteration
|
||||
anthropic_response_obj: Final[GChunk] = cast(GChunk, chunk)
|
||||
completion_obj["content"] = anthropic_response_obj["text"]
|
||||
if "reasoning_content" in anthropic_response_obj and anthropic_response_obj["reasoning_content"] is not None:
|
||||
completion_obj["reasoning_content"] = anthropic_response_obj["reasoning_content"]
|
||||
if anthropic_response_obj["is_finished"]:
|
||||
self.received_finish_reason = anthropic_response_obj["finish_reason"]
|
||||
|
||||
|
|
@ -1367,7 +1382,9 @@ class CustomStreamWrapper:
|
|||
litellm.Usage(**anthropic_response_obj["usage"]),
|
||||
)
|
||||
|
||||
if "tool_use" in anthropic_response_obj and anthropic_response_obj["tool_use"] is not None:
|
||||
if "tool_calls" in anthropic_response_obj and anthropic_response_obj["tool_calls"] is not None:
|
||||
completion_obj["tool_calls"] = anthropic_response_obj["tool_calls"]
|
||||
elif "tool_use" in anthropic_response_obj and anthropic_response_obj["tool_use"] is not None:
|
||||
completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]]
|
||||
|
||||
if (
|
||||
|
|
@ -2776,6 +2793,14 @@ def convert_generic_chunk_to_model_response_stream(
|
|||
) -> ModelResponseStream:
|
||||
from litellm.types.utils import Delta
|
||||
|
||||
tool_calls = chunk.get("tool_calls", None)
|
||||
if tool_calls is None and chunk.get("tool_use") is not None:
|
||||
tool_calls = [chunk["tool_use"]]
|
||||
|
||||
content = chunk["text"]
|
||||
if not content and (chunk.get("reasoning_content") is not None or tool_calls is not None):
|
||||
content = None
|
||||
|
||||
model_response_stream: Final = ModelResponseStream(
|
||||
id=str(uuid.uuid4()),
|
||||
model="",
|
||||
|
|
@ -2783,8 +2808,9 @@ def convert_generic_chunk_to_model_response_stream(
|
|||
StreamingChoices(
|
||||
index=chunk.get("index", 0),
|
||||
delta=Delta(
|
||||
content=chunk["text"],
|
||||
tool_calls=chunk.get("tool_use", None),
|
||||
content=content,
|
||||
reasoning_content=chunk.get("reasoning_content", None),
|
||||
tool_calls=tool_calls,
|
||||
),
|
||||
)
|
||||
],
|
||||
|
|
|
|||
|
|
@ -33,6 +33,8 @@ class ModelResponseIterator:
|
|||
|
||||
text = ""
|
||||
tool_use: ChatCompletionToolCallChunk | None = None
|
||||
tool_calls: list[ChatCompletionToolCallChunk] | None = None
|
||||
reasoning_content: str | None = None
|
||||
is_finished = False
|
||||
finish_reason = ""
|
||||
usage: ChatCompletionUsageBlock | None = None
|
||||
|
|
@ -54,21 +56,34 @@ class ModelResponseIterator:
|
|||
if processed_chunk.choices[0].delta.content is not None:
|
||||
text = processed_chunk.choices[0].delta.content
|
||||
|
||||
delta_reasoning: Final[str | None] = getattr(processed_chunk.choices[0].delta, "reasoning_content", None)
|
||||
if delta_reasoning is not None:
|
||||
reasoning_content = delta_reasoning
|
||||
|
||||
if (
|
||||
processed_chunk.choices[0].delta.tool_calls is not None
|
||||
and len(processed_chunk.choices[0].delta.tool_calls) > 0
|
||||
and processed_chunk.choices[0].delta.tool_calls[0].function is not None
|
||||
and processed_chunk.choices[0].delta.tool_calls[0].function.arguments is not None
|
||||
):
|
||||
tool_use = ChatCompletionToolCallChunk(
|
||||
id=processed_chunk.choices[0].delta.tool_calls[0].id,
|
||||
type="function",
|
||||
function=ChatCompletionToolCallFunctionChunk(
|
||||
name=processed_chunk.choices[0].delta.tool_calls[0].function.name,
|
||||
arguments=processed_chunk.choices[0].delta.tool_calls[0].function.arguments,
|
||||
),
|
||||
index=processed_chunk.choices[0].delta.tool_calls[0].index,
|
||||
)
|
||||
tool_calls = []
|
||||
for tc in processed_chunk.choices[0].delta.tool_calls:
|
||||
tc_fn = getattr(tc, "function", None)
|
||||
fn_chunk = (
|
||||
ChatCompletionToolCallFunctionChunk(
|
||||
name=getattr(tc_fn, "name", None),
|
||||
arguments=getattr(tc_fn, "arguments", "") or "",
|
||||
)
|
||||
if tc_fn is not None
|
||||
else ChatCompletionToolCallFunctionChunk(name=None, arguments="")
|
||||
)
|
||||
call_chunk = ChatCompletionToolCallChunk(
|
||||
id=getattr(tc, "id", None),
|
||||
type="function",
|
||||
function=fn_chunk,
|
||||
index=getattr(tc, "index", 0) if getattr(tc, "index", None) is not None else 0,
|
||||
)
|
||||
tool_calls.append(call_chunk)
|
||||
if tool_calls:
|
||||
tool_use = tool_calls[0]
|
||||
|
||||
if processed_chunk.choices[0].finish_reason is not None:
|
||||
is_finished = True
|
||||
|
|
@ -78,7 +93,7 @@ class ModelResponseIterator:
|
|||
if usage_chunk is not None:
|
||||
usage = _usage_block(usage_chunk)
|
||||
|
||||
return GenericStreamingChunk(
|
||||
generic_chunk: Final[GenericStreamingChunk] = GenericStreamingChunk(
|
||||
text=text,
|
||||
tool_use=tool_use,
|
||||
is_finished=is_finished,
|
||||
|
|
@ -86,6 +101,12 @@ class ModelResponseIterator:
|
|||
usage=usage,
|
||||
index=0,
|
||||
)
|
||||
if reasoning_content is not None:
|
||||
generic_chunk["reasoning_content"] = reasoning_content
|
||||
if tool_calls is not None:
|
||||
generic_chunk["tool_calls"] = tool_calls
|
||||
|
||||
return generic_chunk
|
||||
except json.JSONDecodeError:
|
||||
raise ValueError(f"Failed to decode JSON from chunk: {chunk}")
|
||||
|
||||
|
|
|
|||
|
|
@ -498,6 +498,8 @@ class ModelInfo(ModelInfoBase, total=False):
|
|||
class GenericStreamingChunk(TypedDict, total=False):
|
||||
text: Required[str]
|
||||
tool_use: ChatCompletionToolCallChunk | None
|
||||
tool_calls: list[ChatCompletionToolCallChunk] | None
|
||||
reasoning_content: str | None
|
||||
is_finished: Required[bool]
|
||||
finish_reason: Required[str]
|
||||
usage: Required[ChatCompletionUsageBlock | None]
|
||||
|
|
|
|||
|
|
@ -222,3 +222,73 @@ async def test_async_stream_without_done_marker_ends_after_the_last_chunk():
|
|||
chunks: Final = [chunk async for chunk in ModelResponseIterator(streaming_response=lines(), sync_stream=False)]
|
||||
|
||||
assert [(chunk["text"], chunk["finish_reason"]) for chunk in chunks] == [("144", ""), ("", "stop")]
|
||||
|
||||
|
||||
def test_chunk_parser_extracts_reasoning_content():
|
||||
"""Reasoning-only delta chunks must preserve reasoning_content."""
|
||||
iterator = ModelResponseIterator(streaming_response=None, sync_stream=True)
|
||||
chunk = {
|
||||
"id": "chatcmpl-reasoning",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "glm-4",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"role": "assistant", "reasoning_content": "Thinking about prime numbers..."},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
result = iterator.chunk_parser(chunk=chunk)
|
||||
|
||||
assert result["text"] == ""
|
||||
assert result.get("reasoning_content") == "Thinking about prime numbers..."
|
||||
|
||||
|
||||
def test_chunk_parser_extracts_multiple_tool_calls_with_indices():
|
||||
"""Multiple tool call deltas in a chunk must preserve all tool calls and their indices."""
|
||||
iterator = ModelResponseIterator(streaming_response=None, sync_stream=True)
|
||||
chunk = {
|
||||
"id": "chatcmpl-tools",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "glm-4",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
|
||||
},
|
||||
{
|
||||
"index": 1,
|
||||
"id": "call_2",
|
||||
"type": "function",
|
||||
"function": {"name": "get_time", "arguments": '{"city": "Paris"}'},
|
||||
},
|
||||
]
|
||||
},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
result = iterator.chunk_parser(chunk=chunk)
|
||||
|
||||
assert result.get("tool_calls") is not None
|
||||
assert len(result["tool_calls"]) == 2
|
||||
assert result["tool_calls"][0]["id"] == "call_1"
|
||||
assert result["tool_calls"][0]["function"]["name"] == "get_weather"
|
||||
assert result["tool_calls"][0]["index"] == 0
|
||||
assert result["tool_calls"][1]["id"] == "call_2"
|
||||
assert result["tool_calls"][1]["function"]["name"] == "get_time"
|
||||
assert result["tool_calls"][1]["index"] == 1
|
||||
# Backward compatibility with tool_use
|
||||
assert result["tool_use"]["id"] == "call_1"
|
||||
|
||||
|
|
|
|||
|
|
@ -1,12 +1,12 @@
|
|||
import json
|
||||
from typing import Final
|
||||
|
||||
from unittest.mock import Mock
|
||||
|
||||
import httpx
|
||||
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
from litellm.llms.openai_like.chat.handler import make_sync_call
|
||||
from litellm.utils import CustomStreamWrapper
|
||||
|
||||
_FAKE_STREAMED_REPLY: Final = {
|
||||
"id": "chatcmpl-fake-stream",
|
||||
|
|
@ -73,3 +73,153 @@ def test_fake_stream_keeps_reasoning_every_tool_call_and_token_detail():
|
|||
assert [call.function.name for call in delta.tool_calls] == ["get_weather", "get_time"]
|
||||
assert chunks[0].usage.prompt_tokens_details.cached_tokens == 64
|
||||
assert chunks[0].usage.completion_tokens_details.reasoning_tokens == 1920
|
||||
|
||||
|
||||
def test_real_stream_keeps_reasoning_multi_tool_calls_and_terminal_usage():
|
||||
events = [
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "c1",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "glm-4",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"role": "assistant", "reasoning_content": "Thinking step 1"},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
),
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "c2",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "glm-4",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"tool_calls": [
|
||||
{
|
||||
"index": 0,
|
||||
"id": "call_1",
|
||||
"type": "function",
|
||||
"function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
|
||||
},
|
||||
{
|
||||
"index": 1,
|
||||
"id": "call_2",
|
||||
"type": "function",
|
||||
"function": {"name": "get_time", "arguments": '{"city": "Paris"}'},
|
||||
},
|
||||
]
|
||||
},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
),
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "c3",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "glm-4",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {"content": "Here is the result"},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
),
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "c4",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "glm-4",
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||
}
|
||||
),
|
||||
"data: "
|
||||
+ json.dumps(
|
||||
{
|
||||
"id": "c5",
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 1,
|
||||
"model": "glm-4",
|
||||
"choices": [],
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 64,
|
||||
"total_tokens": 74,
|
||||
"completion_tokens_details": {"reasoning_tokens": 30},
|
||||
},
|
||||
}
|
||||
),
|
||||
"data: [DONE]",
|
||||
]
|
||||
sse_bytes = ("\n\n".join(events) + "\n\n").encode("utf-8")
|
||||
|
||||
def respond(request: httpx.Request) -> httpx.Response:
|
||||
return httpx.Response(200, headers={"content-type": "text/event-stream"}, content=sse_bytes)
|
||||
|
||||
mock_logging = Mock()
|
||||
mock_logging.model_call_details = {"litellm_params": {}}
|
||||
mock_logging.stream_options = {"include_usage": True}
|
||||
|
||||
raw_stream = make_sync_call(
|
||||
client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))),
|
||||
api_base="https://openai-like.example/v1/chat/completions",
|
||||
headers={},
|
||||
data=json.dumps({"model": "glm-4", "messages": []}),
|
||||
model="glm-4",
|
||||
messages=[],
|
||||
logging_obj=mock_logging,
|
||||
fake_stream=False,
|
||||
)
|
||||
wrapper = CustomStreamWrapper(
|
||||
completion_stream=raw_stream,
|
||||
model="glm-4",
|
||||
custom_llm_provider="vertex_ai",
|
||||
logging_obj=mock_logging,
|
||||
stream_options={"include_usage": True},
|
||||
)
|
||||
|
||||
chunks = list(wrapper)
|
||||
assert len(chunks) == 5
|
||||
|
||||
# 1. Reasoning-only delta must reach the client with reasoning_content and content=None
|
||||
reasoning_chunk = chunks[0]
|
||||
assert reasoning_chunk.choices[0].delta.reasoning_content == "Thinking step 1"
|
||||
assert reasoning_chunk.choices[0].delta.content is None
|
||||
|
||||
# 2. Multi tool call delta must preserve both tool calls and their indices
|
||||
tool_chunk = chunks[1]
|
||||
assert len(tool_chunk.choices[0].delta.tool_calls) == 2
|
||||
assert tool_chunk.choices[0].delta.tool_calls[0].function.name == "get_weather"
|
||||
assert tool_chunk.choices[0].delta.tool_calls[0].index == 0
|
||||
assert tool_chunk.choices[0].delta.tool_calls[1].function.name == "get_time"
|
||||
assert tool_chunk.choices[0].delta.tool_calls[1].index == 1
|
||||
|
||||
# 3. Content delta
|
||||
assert chunks[2].choices[0].delta.content == "Here is the result"
|
||||
|
||||
# 4. Finish reason
|
||||
assert chunks[3].choices[0].finish_reason == "stop"
|
||||
|
||||
# 5. Terminal usage chunk with provider usage details preserved
|
||||
assert chunks[4].usage is not None
|
||||
assert chunks[4].usage.completion_tokens == 64
|
||||
assert chunks[4].usage.prompt_tokens == 10
|
||||
assert chunks[4].usage.completion_tokens_details.reasoning_tokens == 30
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue