From 3a8c4dabd91adc4d497c8772c19f1869c55a5edf Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 9 Jun 2025 18:20:19 -0700 Subject: [PATCH] fix(streaming_handler.py): respect is_finished = False in hidden params internal logic for preventing ending stream early --- litellm/litellm_core_utils/streaming_handler.py | 8 ++++---- tests/local_testing/test_amazing_vertex_completion.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 0b019cf5300..bbbfcfa53e9 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -442,9 +442,10 @@ class CustomStreamWrapper: is_finished = ( True # check if str_line._hidden_params["is_finished"] is True ) - if hasattr( - str_line, "_hidden_params" - ) and str_line._hidden_params.get("is_finished"): + if ( + hasattr(str_line, "_hidden_params") + and str_line._hidden_params.get("is_finished") is not None + ): is_finished = str_line._hidden_params.get("is_finished") finish_reason = str_line.choices[0].finish_reason @@ -1194,7 +1195,6 @@ class CustomStreamWrapper: if response_obj is None: return completion_obj["content"] = response_obj["text"] - print_verbose(f"completion obj content: {completion_obj['content']}") if response_obj["is_finished"]: if response_obj["finish_reason"] == "error": raise Exception( diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 95bf7e01a95..6247b9816be 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -3946,7 +3946,7 @@ def test_vertex_ai_streaming_response_id(): def test_vertex_ai_gemini_2_5_pro_streaming(): load_vertex_ai_credentials() - litellm._turn_on_debug() + # litellm._turn_on_debug() response = completion( model="vertex_ai/gemini-2.5-pro-preview-06-05", messages=[{"role": "user", "content": "Hi!"}],