From e60f1b8e08a6f026534f443b11922143954ca957 Mon Sep 17 00:00:00 2001 From: waani Date: Tue, 26 May 2026 18:10:33 +0800 Subject: [PATCH] test(gemini): add unit tests for function call preservation fixes - Empty promptFeedback.blockReason should not trigger content_filter (streaming and non-streaming) - SAFETY finishReason with functionCall parts produces tool_calls - functionCall thought_signatures not extracted to message level --- .../test_gemini_reasoning_content.py | 47 +++++ ...test_vertex_and_google_ai_studio_gemini.py | 176 ++++++++++++++++++ 2 files changed, 223 insertions(+) diff --git a/tests/local_testing/test_gemini_reasoning_content.py b/tests/local_testing/test_gemini_reasoning_content.py index d95a4577888..89b1b9c150d 100644 --- a/tests/local_testing/test_gemini_reasoning_content.py +++ b/tests/local_testing/test_gemini_reasoning_content.py @@ -141,3 +141,50 @@ def test_round_trip_without_thought_signature_still_works(): text_part = model_message["parts"][0] assert text_part["text"] == "Hi there" assert "thoughtSignature" not in text_part + + +def test_function_call_thought_signature_not_extracted_to_message_level(): + """ + Regression test: functionCall parts carry thoughtSignature for multi-turn + context preservation, but it should NOT appear in the message-level + provider_specific_fields.thought_signatures list — it is handled via + tool_call.provider_specific_fields.thought_signature instead. + + Non-functionCall parts with thoughtSignature should still be extracted. + """ + config = VertexGeminiConfig() + + # Parts: functionCall with thoughtSignature + text with thoughtSignature + parts = [ + {"functionCall": {"name": "get_weather", "args": {}}, "thoughtSignature": "sig-func-1"}, + {"text": "I'll check the weather", "thoughtSignature": "sig-text-1"}, + ] + + signatures = config._extract_thought_signatures_from_parts(parts) + + # Only the text part's signature should be extracted + assert signatures is not None + assert len(signatures) == 1, ( + f"Expected 1 signature (text only), got {len(signatures)}: {signatures}" + ) + assert signatures[0] == "sig-text-1" + + +def test_all_function_call_parts_skipped_in_thought_signatures(): + """ + Regression test: when ALL parts are functionCall, _extract_thought_signatures_from_parts + should return None (no signatures at message level). + """ + config = VertexGeminiConfig() + + parts = [ + {"functionCall": {"name": "fn1", "args": {}}, "thoughtSignature": "sig-1"}, + {"functionCall": {"name": "fn2", "args": {}}, "thoughtSignature": "sig-2"}, + ] + + signatures = config._extract_thought_signatures_from_parts(parts) + + # All parts are functionCall — should return None + assert signatures is None, ( + f"Expected None when all parts are functionCall, got {signatures}" + ) diff --git a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index 45b9f4293fa..bc9d31e2f56 100644 --- a/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -3811,6 +3811,56 @@ def test_chunk_parser_handles_prompt_feedback_block_with_usage(): ), f"total_tokens should be 8175, got {result.usage.total_tokens}" +def test_chunk_parser_empty_blockReason_does_not_trigger_content_filter(): + """ + Regression test: Gemini returns promptFeedback.blockReason as an empty string + when the prompt is NOT blocked. The old code used `"blockReason" in dict` which + checked key existence, falsely triggering content_filter for empty string. + """ + from unittest.mock import Mock + + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + # promptFeedback.blockReason is empty string (NOT blocked) + chunk = { + "promptFeedback": { + "blockReason": "", + "blockReasonMessage": "", + }, + "candidates": [ + { + "content": {"parts": [{"text": "Hello"}], "role": "model"}, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + }, + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=[], sync_stream=True, logging_obj=logging_obj + ) + + # Act + result = streaming_obj.chunk_parser(chunk) + + # Assert — should NOT be content_filter since blockReason is empty + assert result is not None + assert len(result.choices) == 1 + assert ( + result.choices[0].finish_reason != "content_filter" + ), "Empty blockReason should NOT trigger content_filter" + assert result.choices[0].delta.content == "Hello" + + def test_vertex_ai_traffic_type_preserved_in_hidden_params_streaming(): """Test trafficType is preserved in _hidden_params for streaming.""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( @@ -3837,6 +3887,132 @@ def test_vertex_ai_traffic_type_preserved_in_hidden_params_streaming(): ) +def test_non_streaming_empty_blockReason_does_not_trigger_blocked_response(): + """ + Regression test: non-streaming path must not short-circuit to + _handle_blocked_response when promptFeedback.blockReason is empty string. + """ + from unittest.mock import MagicMock + + from litellm import ModelResponse + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + # completion has promptFeedback.blockReason="" (NOT blocked) and valid content + completion_response = { + "promptFeedback": { + "blockReason": "", + "blockReasonMessage": "", + }, + "candidates": [ + { + "content": {"parts": [{"text": "Hello"}], "role": "model"}, + "finishReason": "STOP", + } + ], + "usageMetadata": { + "promptTokenCount": 10, + "candidatesTokenCount": 5, + "totalTokenCount": 15, + }, + } + + raw_response = MagicMock() + raw_response.json.return_value = completion_response + + result = VertexGeminiConfig().transform_response( + model="gemini-pro", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + # Should NOT be a blocked response — should have valid text + assert len(result.choices) == 1 + assert result.choices[0].finish_reason != "content_filter" + assert result.choices[0].message.content == "Hello" + + +def test_non_streaming_safety_finish_reason_with_function_calls(): + """ + Regression test: when Gemini returns a safety finishReason (SAFETY, RECITATION, etc.) + alongside valid functionCall parts, the response should NOT short-circuit to + content_filter — function calls must be properly extracted. + """ + from unittest.mock import MagicMock + + from litellm import ModelResponse + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + VertexGeminiConfig, + ) + + # finishReason is SAFETY but candidate has functionCall parts + completion_response = { + "candidates": [ + { + "content": { + "parts": [ + { + "functionCall": { + "name": "get_current_weather", + "args": {"location": "Boston, MA"}, + } + } + ], + "role": "model", + }, + "finishReason": "SAFETY", + "safetyRatings": [ + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "NEGLIGIBLE", + "blocked": False, + } + ], + } + ], + "usageMetadata": { + "promptTokenCount": 68, + "candidatesTokenCount": 120, + "totalTokenCount": 188, + }, + } + + raw_response = MagicMock() + raw_response.json.return_value = completion_response + + result = VertexGeminiConfig().transform_response( + model="gemini-pro", + raw_response=raw_response, + model_response=ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + # Must have tool_calls, NOT content_filter + assert len(result.choices) == 1 + assert result.choices[0].finish_reason == "tool_calls", ( + f"finish_reason should be 'tool_calls' when functionCall parts exist, " + f"got '{result.choices[0].finish_reason}'" + ) + assert result.choices[0].message.tool_calls is not None + assert len(result.choices[0].message.tool_calls) == 1 + assert ( + result.choices[0].message.tool_calls[0].function.name + == "get_current_weather" + ) + + def test_vertex_ai_traffic_type_preserved_in_hidden_params_non_streaming(): """Test trafficType is preserved in _hidden_params for non-streaming.""" from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (