From ff85c5bb88f31c0a7d6eac4e89b294cd83ee3c18 Mon Sep 17 00:00:00 2001 From: Chenglun Hu Date: Thu, 28 May 2026 16:03:31 +0800 Subject: [PATCH] fix(gemini): surface finishReason on image-gen safety blocks Image-gen transforms iterated candidates[].content.parts[] for inlineData and ignored finishReason, so an IMAGE_SAFETY / IMAGE_PROHIBITED_CONTENT block returned an empty ImageResponse with no exception. The chat path already surfaces these (get_flagged_finish_reasons + ContentPolicyViolationError); mirror that for image gen by raising on a flagged finishReason when no image data was produced. Covers gemini/ and vertex_ai/ via one shared helper. Fixes #28989 --- .../gemini/image_generation/transformation.py | 34 +++++++++++++++++++ .../vertex_gemini_transformation.py | 11 ++++++ ...rtex_ai_image_generation_transformation.py | 33 ++++++++++++++++++ 3 files changed, 78 insertions(+) diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index dcdec46edca..cbd1d85e7dc 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -31,6 +31,35 @@ else: LiteLLMLoggingObj = Any +def raise_if_image_gen_flagged( + response_data: dict, + model: str, + raw_response: httpx.Response, + llm_provider: str = "gemini", +) -> None: + """ + Gemini image-gen returns a candidate with a flagged finishReason (e.g. + IMAGE_SAFETY / IMAGE_PROHIBITED_CONTENT) and no inlineData on a refusal. + The chat path surfaces these; the image path used to drop them silently + and return an empty ImageResponse. Raise so callers can tell a refusal + apart from an unrelated failure. Reasons reuse the central finish-reason + map (content_filter == flagged). + """ + from litellm.exceptions import ContentPolicyViolationError + from litellm.litellm_core_utils.core_helpers import _FINISH_REASON_MAP + + for candidate in response_data.get("candidates", []): + finish_reason = candidate.get("finishReason") + if finish_reason and _FINISH_REASON_MAP.get(finish_reason) == "content_filter": + raise ContentPolicyViolationError( + message=f"Gemini image generation blocked with finishReason={finish_reason}", + model=model, + llm_provider=llm_provider, + response=raw_response, + provider_specific_fields={"finish_reason": finish_reason}, + ) + + class GoogleImageGenConfig(BaseImageGenerationConfig): DEFAULT_BASE_URL: str = "https://generativelanguage.googleapis.com/v1beta" @@ -213,6 +242,11 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): ) ) + # A safety/prohibited block returns a candidate with finishReason and + # no inlineData — surface it instead of returning empty data. + if not model_response.data: + raise_if_image_gen_flagged(response_data, model, raw_response) + # Extract usage metadata for Gemini models if "usageMetadata" in response_data: model_response.usage = transform_gemini_image_usage(response_data["usageMetadata"]) diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py index 572725ac789..e6c98688de1 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -321,6 +321,17 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): ) ) + # A safety/prohibited block returns a candidate with finishReason and + # no inlineData — surface it instead of returning empty data. + if not model_response.data: + from litellm.llms.gemini.image_generation.transformation import ( + raise_if_image_gen_flagged, + ) + + raise_if_image_gen_flagged( + response_data, model, raw_response, llm_provider="vertex_ai" + ) + if usage_metadata := response_data.get("usageMetadata", None): model_response.usage = self._transform_image_usage(usage_metadata) diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 8c72bdee525..d03a89dfe66 100644 --- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -443,6 +443,39 @@ class TestVertexAIGeminiImageGenerationConfig: assert result.usage.web_search_requests == 2 + @pytest.mark.parametrize( + "finish_reason", ["IMAGE_SAFETY", "IMAGE_PROHIBITED_CONTENT", "SAFETY"] + ) + def test_transform_image_generation_response_raises_on_safety_block( + self, finish_reason + ): + """A safety/prohibited block returns a candidate with finishReason and no + inlineData; it must raise ContentPolicyViolationError, not return empty data.""" + from litellm.exceptions import ContentPolicyViolationError + + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [{"finishReason": finish_reason, "content": {"parts": []}}] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + with pytest.raises(ContentPolicyViolationError) as exc: + self.config.transform_image_generation_response( + model="gemini-2.5-flash-image", + raw_response=mock_response, + model_response=ImageResponse(), + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + assert finish_reason in str(exc.value) + assert exc.value.llm_provider == "vertex_ai" + class TestVertexAIImagenImageGenerationConfig: def setup_method(self):