diff --git a/litellm/utils.py b/litellm/utils.py index 27e67f4a81..6fdd2d88bc 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -2081,6 +2081,14 @@ def _is_async_request( return False +_STREAMING_CALL_TYPES = frozenset({ + CallTypes.generate_content_stream, + CallTypes.agenerate_content_stream, + CallTypes.generate_content_stream.value, + CallTypes.agenerate_content_stream.value, +}) + + def _is_streaming_request( kwargs: Dict[str, Any], call_type: Union[CallTypes, str], @@ -2093,23 +2101,7 @@ def _is_streaming_request( """ if "stream" in kwargs and kwargs["stream"] is True: return True - - ######################################################### - # Check if it's a google genai streaming request - if isinstance(call_type, str): - # check if it can be casted to CallTypes - try: - call_type = CallTypes(call_type) - except ValueError: - return False - - if ( - call_type == CallTypes.generate_content_stream - or call_type == CallTypes.agenerate_content_stream - ): - return True - ######################################################### - return False + return call_type in _STREAMING_CALL_TYPES def _select_tokenizer( diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 71e043eb22..352125d16c 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -18,10 +18,12 @@ from litellm.types.utils import ( ModelResponseStream, StreamingChoices, ) +from litellm.types.utils import CallTypes from litellm.utils import ( ProviderConfigManager, TextCompletionStreamWrapper, _check_provider_match, + _is_streaming_request, get_llm_provider, get_optional_params_image_gen, is_cached_message, @@ -3270,3 +3272,35 @@ class TestDropParamsWithPromptCacheKey: assert "prompt_cache_key" not in result # temperature should remain (it's supported by Bedrock) assert result.get("temperature") == 0.7 + + +class TestIsStreamingRequest: + def test_stream_true_in_kwargs(self): + assert _is_streaming_request(kwargs={"stream": True}, call_type="acompletion") is True + + def test_stream_false_in_kwargs(self): + assert _is_streaming_request(kwargs={"stream": False}, call_type="acompletion") is False + + def test_no_stream_in_kwargs(self): + assert _is_streaming_request(kwargs={}, call_type="acompletion") is False + + def test_generate_content_stream_string(self): + assert _is_streaming_request(kwargs={}, call_type=CallTypes.generate_content_stream.value) is True + + def test_agenerate_content_stream_string(self): + assert _is_streaming_request(kwargs={}, call_type=CallTypes.agenerate_content_stream.value) is True + + def test_generate_content_stream_enum(self): + assert _is_streaming_request(kwargs={}, call_type=CallTypes.generate_content_stream) is True + + def test_agenerate_content_stream_enum(self): + assert _is_streaming_request(kwargs={}, call_type=CallTypes.agenerate_content_stream) is True + + def test_non_streaming_call_type_string(self): + assert _is_streaming_request(kwargs={}, call_type="acompletion") is False + + def test_non_streaming_call_type_enum(self): + assert _is_streaming_request(kwargs={}, call_type=CallTypes.acompletion) is False + + def test_stream_true_overrides_non_streaming_call_type(self): + assert _is_streaming_request(kwargs={"stream": True}, call_type=CallTypes.acompletion) is True