diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 2df470edc59..12ba041aff8 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -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, ), ) ], diff --git a/litellm/llms/databricks/streaming_utils.py b/litellm/llms/databricks/streaming_utils.py index eca87f9dd5f..a95578a41d7 100644 --- a/litellm/llms/databricks/streaming_utils.py +++ b/litellm/llms/databricks/streaming_utils.py @@ -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}") diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cc7ef851111..1bdb2919ba7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -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] diff --git a/tests/unit/llms/databricks/test_databricks_streaming_utils.py b/tests/unit/llms/databricks/test_databricks_streaming_utils.py index 5c5527f7be8..0f52afbab3a 100644 --- a/tests/unit/llms/databricks/test_databricks_streaming_utils.py +++ b/tests/unit/llms/databricks/test_databricks_streaming_utils.py @@ -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" + diff --git a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py index 75b60b2dd6c..defa02daad9 100644 --- a/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py +++ b/tests/unit/llms/openai_like/chat/test_openai_like_chat_handler.py @@ -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 +