From bfabb39fc6b847a50877f7362089c5b6161b78b2 Mon Sep 17 00:00:00 2001 From: Kris Xia Date: Sat, 31 Jan 2026 12:01:31 +0800 Subject: [PATCH] fix: Fix Vertex AI Gemini streaming content_filter handling - Add promptFeedback.blockReason check in chunk_parser - Return content_filter finish_reason when blocked - Extract content filter logic into _check_prompt_level_content_filter() method - Update unit tests to reflect simplified implementation Signed-off-by: Kris Xia --- .../vertex_and_google_ai_studio_gemini.py | 55 ++++++++ ...test_vertex_and_google_ai_studio_gemini.py | 120 ++++++++++++++++++ 2 files changed, 175 insertions(+) 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 b5a6949f27..04ae4b6beb 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 @@ -1732,6 +1732,52 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): else: return "stop" + @staticmethod + def _check_prompt_level_content_filter( + processed_chunk: GenerateContentResponseBody, + response_id: Optional[str], + ) -> Optional["ModelResponseStream"]: + """ + Check if prompt is blocked due to content filtering at the prompt level. + + This handles the case where Vertex AI blocks the prompt before generation begins, + indicated by promptFeedback.blockReason being present. + + Args: + processed_chunk: The parsed response chunk from Vertex AI + response_id: The response ID from the chunk + + Returns: + ModelResponseStream with content_filter finish_reason if blocked, None otherwise. + + Note: + This is consistent with non-streaming _handle_blocked_response() behavior. + Candidate-level content filtering (SAFETY, RECITATION, etc.) is handled + separately via _process_candidates() → _check_finish_reason(). + """ + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + # Check if prompt is blocked due to content filtering + prompt_feedback = processed_chunk.get("promptFeedback") + if prompt_feedback and "blockReason" in prompt_feedback: + verbose_logger.debug( + f"Prompt blocked due to: {prompt_feedback.get('blockReason')} - {prompt_feedback.get('blockReasonMessage')}" + ) + + # Create a content_filter response (consistent with non-streaming _handle_blocked_response) + choice = StreamingChoices( + finish_reason="content_filter", + index=0, + delta=Delta(content=None, role="assistant"), + logprobs=None, + enhancements=None, + ) + + model_response = ModelResponseStream(choices=[choice], id=response_id) + return model_response + + return None + @staticmethod def _calculate_web_search_requests(grounding_metadata: List[dict]) -> Optional[int]: web_search_requests: Optional[int] = None @@ -2813,6 +2859,15 @@ class ModelResponseIterator: processed_chunk = GenerateContentResponseBody(**chunk) # type: ignore response_id = processed_chunk.get("responseId") model_response = ModelResponseStream(choices=[], id=response_id) + + # Check if prompt is blocked due to content filtering + blocked_response = VertexGeminiConfig._check_prompt_level_content_filter( + processed_chunk=processed_chunk, + response_id=response_id, + ) + if blocked_response is not None: + model_response = blocked_response + usage: Optional[Usage] = None _candidates: Optional[List[Candidates]] = processed_chunk.get("candidates") grounding_metadata: List[dict] = [] 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 cb3b51acd6..75fd597ffa 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 @@ -3218,3 +3218,123 @@ def test_video_metadata_only_for_gemini_3(): assert file_part_3 is not None assert "media_resolution" in file_part_3, "Gemini 3 should have media_resolution" assert "video_metadata" in file_part_3, "Gemini 3 should have video_metadata" + + + +def test_chunk_parser_handles_prompt_feedback_block(): + """Test chunk_parser correctly handles promptFeedback.blockReason""" + from unittest.mock import Mock + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + # Arrange - mock a blocked response + blocked_chunk = { + "promptFeedback": { + "blockReason": "PROHIBITED_CONTENT", + "blockReasonMessage": "The prompt is blocked due to prohibited contents" + }, + "responseId": "test_response_id", + "modelVersion": "gemini-3-pro-preview" + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj + ) + + # Act + result = streaming_obj.chunk_parser(blocked_chunk) + + # Assert + assert result is not None, "Result should not be None" + assert len(result.choices) == 1, "Should have exactly one choice" + assert result.choices[0].finish_reason == "content_filter", f"finish_reason should be content_filter, got {result.choices[0].finish_reason}" + assert result.choices[0].delta.content is None, "content should be None" + + +def test_chunk_parser_handles_prompt_feedback_safety_block(): + """Test chunk_parser handles different blockReason types (SAFETY)""" + from unittest.mock import Mock + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + # Arrange - mock a SAFETY blocked response + blocked_chunk = { + "promptFeedback": { + "blockReason": "SAFETY", + "blockReasonMessage": "The prompt is blocked due to safety concerns" + }, + "responseId": "test_safety_response_id", + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj + ) + + # Act + result = streaming_obj.chunk_parser(blocked_chunk) + + # Assert + assert result is not None + assert len(result.choices) == 1 + assert result.choices[0].finish_reason == "content_filter" + + +def test_chunk_parser_handles_prompt_feedback_block_with_usage(): + """Test chunk_parser correctly extracts usageMetadata when promptFeedback.blockReason is present""" + from unittest.mock import Mock + from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( + ModelResponseIterator, + ) + + # Arrange - 模拟一个包含 usageMetadata 的 blocked response + blocked_chunk = { + "promptFeedback": { + "blockReason": "PROHIBITED_CONTENT", + "blockReasonMessage": "The prompt is blocked due to prohibited contents" + }, + "responseId": "test_response_id_with_usage", + "modelVersion": "gemini-3-pro-preview", + "usageMetadata": { + "promptTokenCount": 8175, + "candidatesTokenCount": 0, + "totalTokenCount": 8175 + } + } + + logging_obj = Mock() + logging_obj.optional_params = {} + + streaming_obj = ModelResponseIterator( + streaming_response=iter([]), + sync_stream=True, + logging_obj=logging_obj + ) + + # Act + result = streaming_obj.chunk_parser(blocked_chunk) + + # Assert - 验证 content_filter 响应和 usage 都被正确处理 + assert result is not None, "Result should not be None" + assert len(result.choices) == 1, "Should have exactly one choice" + assert result.choices[0].finish_reason == "content_filter", f"finish_reason should be content_filter, got {result.choices[0].finish_reason}" + assert result.choices[0].delta.content is None, "content should be None" + + # 验证 usage 信息被正确提取 + assert hasattr(result, "usage"), "result should have usage attribute" + assert result.usage is not None, "usage should not be None" + assert result.usage.prompt_tokens == 8175, f"prompt_tokens should be 8175, got {result.usage.prompt_tokens}" + assert result.usage.completion_tokens == 0, f"completion_tokens should be 0, got {result.usage.completion_tokens}" + assert result.usage.total_tokens == 8175, f"total_tokens should be 8175, got {result.usage.total_tokens}" +