From d972e112b0e5eb7c9ce8973b6d4b86ba8b7753a6 Mon Sep 17 00:00:00 2001 From: mohi-devhub Date: Fri, 10 Apr 2026 21:43:07 +0530 Subject: [PATCH] fix: replace bare raise Exception with descriptive ValueError in AdapterCompletionStreamWrapper - Replace bare 'raise Exception' with ValueError containing descriptive message - Replace print() with proper logging using logging.getLogger(__name__) - Re-raise ValueError explicitly to avoid catching it in generic Exception handler - Add test coverage for the streaming wrapper behavior --- litellm/types/utils.py | 17 ++- tests/test_litellm/types/test_types_utils.py | 107 +++++++++++++------ 2 files changed, 87 insertions(+), 37 deletions(-) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index cd5806b3ab7..21f1571243f 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1,4 +1,5 @@ import json +import logging import time from enum import Enum from typing import ( @@ -2481,23 +2482,33 @@ class AdapterCompletionStreamWrapper: try: for chunk in self.completion_stream: if chunk == "None" or chunk is None: - raise Exception + raise ValueError( + "Received None chunk in stream. This usually indicates " + "an issue with the adapter response transformation." + ) return chunk raise StopIteration except StopIteration: raise StopIteration + except ValueError: + raise except Exception as e: - print(f"AdapterCompletionStreamWrapper - {e}") # noqa + logging.getLogger(__name__).warning(f"AdapterCompletionStreamWrapper - {e}") async def __anext__(self): try: async for chunk in self.completion_stream: if chunk == "None" or chunk is None: - raise Exception + raise ValueError( + "Received None chunk in stream. This usually indicates " + "an issue with the adapter response transformation." + ) return chunk raise StopIteration except StopIteration: raise StopAsyncIteration + except ValueError: + raise class StandardLoggingUserAPIKeyMetadata(TypedDict): diff --git a/tests/test_litellm/types/test_types_utils.py b/tests/test_litellm/types/test_types_utils.py index adfa681dbd5..7068fb6ec5a 100644 --- a/tests/test_litellm/types/test_types_utils.py +++ b/tests/test_litellm/types/test_types_utils.py @@ -80,51 +80,51 @@ def test_usage_completion_tokens_details_text_tokens(): # Test data from the reported issue usage_data = { - 'completion_tokens': 77, - 'prompt_tokens': 11937, - 'total_tokens': 12014, - 'completion_tokens_details': { - 'accepted_prediction_tokens': None, - 'audio_tokens': None, - 'reasoning_tokens': 65, - 'rejected_prediction_tokens': None, - 'text_tokens': 12 + "completion_tokens": 77, + "prompt_tokens": 11937, + "total_tokens": 12014, + "completion_tokens_details": { + "accepted_prediction_tokens": None, + "audio_tokens": None, + "reasoning_tokens": 65, + "rejected_prediction_tokens": None, + "text_tokens": 12, + }, + "prompt_tokens_details": { + "audio_tokens": None, + "cached_tokens": None, + "text_tokens": 11937, + "image_tokens": None, }, - 'prompt_tokens_details': { - 'audio_tokens': None, - 'cached_tokens': None, - 'text_tokens': 11937, - 'image_tokens': None - } } # Create Usage object u = Usage(**usage_data) - + # Verify the object has the text_tokens field - assert hasattr(u.completion_tokens_details, 'text_tokens') + assert hasattr(u.completion_tokens_details, "text_tokens") assert u.completion_tokens_details.text_tokens == 12 - + # Get model_dump output dump_result = u.model_dump() - + # Verify text_tokens is present in the model_dump output - assert 'completion_tokens_details' in dump_result - assert 'text_tokens' in dump_result['completion_tokens_details'] - assert dump_result['completion_tokens_details']['text_tokens'] == 12 - + assert "completion_tokens_details" in dump_result + assert "text_tokens" in dump_result["completion_tokens_details"] + assert dump_result["completion_tokens_details"]["text_tokens"] == 12 + # Verify the full completion_tokens_details structure expected_completion_details = { - 'accepted_prediction_tokens': None, - 'audio_tokens': None, - 'reasoning_tokens': 65, - 'rejected_prediction_tokens': None, - 'text_tokens': 12, - 'image_tokens': None, - 'video_tokens': None + "accepted_prediction_tokens": None, + "audio_tokens": None, + "reasoning_tokens": 65, + "rejected_prediction_tokens": None, + "text_tokens": 12, + "image_tokens": None, + "video_tokens": None, } - assert dump_result['completion_tokens_details'] == expected_completion_details - + assert dump_result["completion_tokens_details"] == expected_completion_details + # Verify round-trip serialization works new_usage = Usage(**dump_result) assert new_usage.completion_tokens_details.text_tokens == 12 @@ -257,7 +257,9 @@ class TestNativeFinishReason: ) assert choice.finish_reason == "length" assert choice.provider_specific_fields["native_finish_reason"] == "max_tokens" - assert choice.provider_specific_fields["citations"] == [{"url": "http://example.com"}] + assert choice.provider_specific_fields["citations"] == [ + {"url": "http://example.com"} + ] def test_gemini_safety_reason_exposed(self): from litellm.types.utils import Choices @@ -279,6 +281,8 @@ class TestNativeFinishReason: choice = Choices(finish_reason="MAX_TOKENS") assert choice.finish_reason == "length" assert choice.provider_specific_fields["native_finish_reason"] == "MAX_TOKENS" + + def test_delta_maps_reasoning_to_reasoning_content(): """ Test that Delta maps 'reasoning' field to 'reasoning_content'. @@ -291,7 +295,9 @@ def test_delta_maps_reasoning_to_reasoning_content(): # When provider sends 'reasoning' (e.g., Cerebras gpt-oss streaming) delta = Delta(content=None, role="assistant", reasoning="thinking step by step") assert delta.reasoning_content == "thinking step by step" - assert not hasattr(delta, "reasoning"), "reasoning should not leak as an extra attribute" + assert not hasattr( + delta, "reasoning" + ), "reasoning should not leak as an extra attribute" # When provider sends 'reasoning_content' directly (e.g., NIM), it still works delta2 = Delta(content="hello", reasoning_content="direct reasoning") @@ -304,3 +310,36 @@ def test_delta_maps_reasoning_to_reasoning_content(): # When neither is present, reasoning_content is not set (OpenAI spec) delta4 = Delta(content="hello") assert not hasattr(delta4, "reasoning_content") + + +def test_adapter_completion_stream_wrapper(): + """Test AdapterCompletionStreamWrapper properly handles None chunks and logging.""" + from litellm.types.utils import AdapterCompletionStreamWrapper + + # Test 1: Normal chunk iteration works + normal_chunks = ["chunk1", "chunk2", "chunk3"] + wrapper = AdapterCompletionStreamWrapper(iter(normal_chunks)) + result = list(wrapper) + assert result == ["chunk1", "chunk2", "chunk3"] + + # Test 2: None chunk raises ValueError with descriptive message + none_chunks = ["valid", None, "after"] + wrapper = AdapterCompletionStreamWrapper(iter(none_chunks)) + with pytest.raises(ValueError) as exc_info: + for chunk in wrapper: + pass + assert "None chunk" in str(exc_info.value) + + # Test 3: "None" string raises ValueError + string_none_chunks = ["valid", "None", "after"] + wrapper = AdapterCompletionStreamWrapper(iter(string_none_chunks)) + with pytest.raises(ValueError) as exc_info: + for chunk in wrapper: + pass + assert "None chunk" in str(exc_info.value) + + # Test 4: Empty iterator raises StopIteration + empty_chunks = [] + wrapper = AdapterCompletionStreamWrapper(iter(empty_chunks)) + with pytest.raises(StopIteration): + next(wrapper)