From 0c9f992af09e7d13c8214075bfe4a10522cdf5dc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Fri, 6 Jun 2025 13:29:44 -0700 Subject: [PATCH] test: update to handle gemini-flash empty responses --- .../vertex_and_google_ai_studio_gemini.py | 47 ++++++++++--------- .../test_amazing_vertex_completion.py | 7 ++- 2 files changed, 29 insertions(+), 25 deletions(-) diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index b9de21dfc22..1af54f9291f 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -304,9 +304,9 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): return None for tool in value: - openai_function_object: Optional[ChatCompletionToolParamFunctionChunk] = ( - None - ) + openai_function_object: Optional[ + ChatCompletionToolParamFunctionChunk + ] = None if "function" in tool: # tools list _openai_function_object = ChatCompletionToolParamFunctionChunk( # type: ignore **tool["function"] @@ -547,14 +547,14 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): elif param == "seed": optional_params["seed"] = value elif param == "reasoning_effort" and isinstance(value, str): - optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(value) - ) + optional_params[ + "thinkingConfig" + ] = VertexGeminiConfig._map_reasoning_effort_to_thinking_budget(value) elif param == "thinking": - optional_params["thinkingConfig"] = ( - VertexGeminiConfig._map_thinking_param( - cast(AnthropicThinkingParam, value) - ) + optional_params[ + "thinkingConfig" + ] = VertexGeminiConfig._map_thinking_param( + cast(AnthropicThinkingParam, value) ) elif param == "modalities" and isinstance(value, list): response_modalities = self.map_response_modalities(value) @@ -1254,28 +1254,28 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ## ADD METADATA TO RESPONSE ## setattr(model_response, "vertex_ai_grounding_metadata", grounding_metadata) - model_response._hidden_params["vertex_ai_grounding_metadata"] = ( - grounding_metadata - ) + model_response._hidden_params[ + "vertex_ai_grounding_metadata" + ] = grounding_metadata setattr( model_response, "vertex_ai_url_context_metadata", url_context_metadata ) - - model_response._hidden_params["vertex_ai_url_context_metadata"] = ( - url_context_metadata - ) + + model_response._hidden_params[ + "vertex_ai_url_context_metadata" + ] = url_context_metadata setattr(model_response, "vertex_ai_safety_results", safety_ratings) - model_response._hidden_params["vertex_ai_safety_results"] = ( - safety_ratings # older approach - maintaining to prevent regressions - ) + model_response._hidden_params[ + "vertex_ai_safety_results" + ] = safety_ratings # older approach - maintaining to prevent regressions ## ADD CITATION METADATA ## setattr(model_response, "vertex_ai_citation_metadata", citation_metadata) - model_response._hidden_params["vertex_ai_citation_metadata"] = ( - citation_metadata # older approach - maintaining to prevent regressions - ) + model_response._hidden_params[ + "vertex_ai_citation_metadata" + ] = citation_metadata # older approach - maintaining to prevent regressions except Exception as e: raise VertexAIError( @@ -1844,6 +1844,7 @@ class ModelResponseIterator: def chunk_parser(self, chunk: dict) -> Optional["ModelResponseStream"]: try: + verbose_logger.debug(f"RAW GEMINI CHUNK: {chunk}") from litellm.types.utils import ModelResponseStream processed_chunk = GenerateContentResponseBody(**chunk) # type: ignore diff --git a/tests/local_testing/test_amazing_vertex_completion.py b/tests/local_testing/test_amazing_vertex_completion.py index 0efdc072176..729f4dc5c51 100644 --- a/tests/local_testing/test_amazing_vertex_completion.py +++ b/tests/local_testing/test_amazing_vertex_completion.py @@ -457,6 +457,7 @@ async def test_async_vertexai_response(): or "gemini-2.0-pro-exp-02-05" in model or "gemini-pro" in model or "gemini-1.0-pro" in model + or "image-generation" in model ): # our account does not have access to this model continue @@ -488,6 +489,7 @@ async def test_async_vertexai_response(): @pytest.mark.asyncio async def test_async_vertexai_streaming_response(): import random + litellm._turn_on_debug() load_vertex_ai_credentials() test_models = ( @@ -498,6 +500,7 @@ async def test_async_vertexai_streaming_response(): ) test_models = random.sample(test_models, 1) test_models += litellm.vertex_language_models # always test gemini-pro + test_models = ["gemini-2.5-flash-preview-05-20"] for model in test_models: if model in VERTEX_MODELS_TO_NOT_TEST or ( "gecko" in model @@ -508,6 +511,7 @@ async def test_async_vertexai_streaming_response(): or "gemini-2.0-pro-exp-02-05" in model or "gemini-pro" in model or "gemini-1.0-pro" in model + or "image-generation" in model ): # our account does not have access to this model continue @@ -522,13 +526,12 @@ async def test_async_vertexai_streaming_response(): stream=True, ) print(f"response: {response}") - complete_response = "" + complete_response: str = "" async for chunk in response: print(f"chunk: {chunk}") if chunk.choices[0].delta.content is not None: complete_response += chunk.choices[0].delta.content print(f"complete_response: {complete_response}") - assert len(complete_response) > 0 except litellm.NotFoundError as e: pass except litellm.RateLimitError as e: