diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index e402023d240..90dd7e279d5 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -805,6 +805,10 @@ class CustomStreamWrapper: "function_call" in completion_obj and completion_obj["function_call"] is not None ) + or ( + "reasoning_content" in completion_obj + and completion_obj["reasoning_content"] is not None + ) or ( "tool_calls" in model_response.choices[0].delta and model_response.choices[0].delta["tool_calls"] is not None @@ -1134,7 +1138,10 @@ class CustomStreamWrapper: ): if self.received_finish_reason is not None: _chunk_has_content = isinstance(chunk, dict) and ( - bool(chunk.get("text", "")) or chunk.get("tool_use") is not None + _chunk_has_content = isinstance(chunk, dict) and ( + bool(chunk.get("text", "")) + or chunk.get("tool_use") is not None + or bool(chunk.get("reasoning_content", "")) ) if not _chunk_has_content and ( not isinstance(chunk, dict) @@ -1166,6 +1173,14 @@ class CustomStreamWrapper: ): completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]] + 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 ( "provider_specific_fields" in anthropic_response_obj and anthropic_response_obj["provider_specific_fields"] is not None @@ -2393,16 +2408,22 @@ def convert_generic_chunk_to_model_response_stream( ) -> ModelResponseStream: from litellm.types.utils import Delta + delta_kwargs: Dict[str, Any] = { + "content": chunk["text"], + "tool_calls": chunk.get("tool_use", None), + } + + reasoning_content = chunk.get("reasoning_content", None) + if reasoning_content is not None: + delta_kwargs["reasoning_content"] = reasoning_content + model_response_stream = ModelResponseStream( id=str(uuid.uuid4()), model="", choices=[ StreamingChoices( index=chunk.get("index", 0), - delta=Delta( - content=chunk["text"], - tool_calls=chunk.get("tool_use", None), - ), + delta=Delta(**delta_kwargs), ) ], finish_reason=chunk["finish_reason"] if chunk["is_finished"] else None, diff --git a/litellm/llms/databricks/streaming_utils.py b/litellm/llms/databricks/streaming_utils.py index eebe3182881..c6dedcc6827 100644 --- a/litellm/llms/databricks/streaming_utils.py +++ b/litellm/llms/databricks/streaming_utils.py @@ -24,10 +24,14 @@ class ModelResponseIterator: is_finished = False finish_reason = "" usage: Optional[ChatCompletionUsageBlock] = None + reasoning_content: Optional[str] = None if processed_chunk.choices[0].delta.content is not None: # type: ignore text = processed_chunk.choices[0].delta.content # type: ignore + if getattr(processed_chunk.choices[0].delta, "reasoning_content", None) is not None: # type: ignore + reasoning_content = processed_chunk.choices[0].delta.reasoning_content # type: ignore + if ( processed_chunk.choices[0].delta.tool_calls is not None # type: ignore and len(processed_chunk.choices[0].delta.tool_calls) > 0 # type: ignore @@ -68,6 +72,7 @@ class ModelResponseIterator: finish_reason=finish_reason, usage=usage, index=0, + reasoning_content=reasoning_content, ) 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 cd5806b3ab7..80c0e3284a0 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -273,6 +273,7 @@ class GenericStreamingChunk(TypedDict, total=False): finish_reason: Required[str] usage: Required[Optional[ChatCompletionUsageBlock]] index: int + reasoning_content: Optional[str] # use this dict if you want to return any provider specific fields in the response provider_specific_fields: Optional[Dict[str, Any]] diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index aad3de306c7..10aa4438565 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -192,6 +192,50 @@ def test_is_chunk_non_empty_with_annotations( ) +def test_is_chunk_non_empty_with_reasoning_content_in_completion_obj( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """Reasoning-only completion_obj chunks must not be dropped before Delta creation.""" + chunk = { + "id": "reasoning-only-chunk", + "object": "chat.completion.chunk", + "created": 1741037890, + "model": "deepseek-reasoner", + "choices": [ + { + "index": 0, + "delta": {"content": None}, + "logprobs": None, + "finish_reason": None, + } + ], + } + assert ( + initialized_custom_stream_wrapper.is_chunk_non_empty( + completion_obj={"content": "", "reasoning_content": "Thinking..."}, + model_response=ModelResponseStream(**chunk), + response_obj={}, + ) + is True + ) + + +def test_return_processed_chunk_logic_keeps_reasoning_only_completion_obj( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """Reasoning-only chunks should be returned once completion_obj carries reasoning_content.""" + completion_obj = {"content": "", "reasoning_content": "Thinking..."} + returned_chunk = initialized_custom_stream_wrapper.return_processed_chunk_logic( + completion_obj=completion_obj, + model_response=ModelResponseStream(), + response_obj={}, + ) + + assert returned_chunk is not None + assert returned_chunk.choices[0].delta.reasoning_content == "Thinking..." + assert returned_chunk.choices[0].delta.content == "" + + def test_optional_combine_thinking_block_in_choices( initialized_custom_stream_wrapper: CustomStreamWrapper, ): diff --git a/tests/test_litellm/llms/databricks/test_streaming_utils.py b/tests/test_litellm/llms/databricks/test_streaming_utils.py new file mode 100644 index 00000000000..d928a42ee50 --- /dev/null +++ b/tests/test_litellm/llms/databricks/test_streaming_utils.py @@ -0,0 +1,105 @@ +""" +Tests for the generic ModelResponseIterator in litellm/llms/databricks/streaming_utils.py. + +Verifies that reasoning_content is correctly extracted from streaming chunks, +fixing the issue where OpenAI-like providers (Watsonx, Cerebras, etc.) lost +reasoning_content during streaming. +""" + +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../../../..") +) # Adds the parent directory to the system path + +from unittest.mock import MagicMock + +from litellm.llms.databricks.streaming_utils import ModelResponseIterator + + +class TestModelResponseIteratorReasoningContent: + """Test reasoning_content extraction in the generic chunk_parser.""" + + def test_chunk_parser_extracts_reasoning_content(self): + """Verify reasoning_content is extracted from a streaming chunk delta.""" + handler = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True + ) + + chunk = { + "id": "chatcmpl-test-123", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-oss-120b", + "choices": [ + { + "index": 0, + "delta": { + "reasoning_content": "Let me think about this step by step.", + "content": None, + }, + "finish_reason": None, + } + ], + } + + result = handler.chunk_parser(chunk) + + assert result["reasoning_content"] == "Let me think about this step by step." + assert result["text"] == "" + + def test_chunk_parser_reasoning_content_none_when_absent(self): + """Verify reasoning_content is None when not present in the chunk.""" + handler = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True + ) + + chunk = { + "id": "chatcmpl-test-456", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-oss-120b", + "choices": [ + { + "index": 0, + "delta": {"content": "Hello world"}, + "finish_reason": None, + } + ], + } + + result = handler.chunk_parser(chunk) + + assert result["reasoning_content"] is None + assert result["text"] == "Hello world" + + def test_chunk_parser_both_content_and_reasoning(self): + """Verify both text and reasoning_content can be extracted simultaneously.""" + handler = ModelResponseIterator( + streaming_response=MagicMock(), sync_stream=True + ) + + chunk = { + "id": "chatcmpl-test-789", + "object": "chat.completion.chunk", + "created": 1700000000, + "model": "gpt-oss-120b", + "choices": [ + { + "index": 0, + "delta": { + "content": "The answer is 42.", + "reasoning_content": "Computing...", + }, + "finish_reason": None, + } + ], + } + + result = handler.chunk_parser(chunk) + + assert result["text"] == "The answer is 42." + assert result["reasoning_content"] == "Computing..."