From 80257675c899e1bbe46b96fc0eff831fbc0caf86 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 | 9 +++++ .../test_gemini_image_usage.py | 30 ++++++++++++++++ ..._gemini_image_generation_transformation.py | 1 - ...rtex_ai_image_generation_transformation.py | 26 +++++++++++++- 5 files changed, 98 insertions(+), 2 deletions(-) diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 1f986b11af2..cf04aff8c7f 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -43,6 +43,35 @@ _JSON_CONTAINER: Final[TypeAdapter[Mapping[str, object] | str | Sequence[object] ) +def raise_if_image_gen_flagged( + response_object: Mapping[str, object], + 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 go through the central finish-reason + map (content_filter == flagged). + """ + from litellm.exceptions import ContentPolicyViolationError + from litellm.litellm_core_utils.core_helpers import map_finish_reason + + for candidate in _JSON_OBJECTS.validate_python(response_object.get("candidates", [])): + finish_reason = candidate.get("finishReason") + if isinstance(finish_reason, str) and map_finish_reason(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" @@ -226,6 +255,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_object, model, raw_response) + # Extract usage metadata for Gemini models if "usageMetadata" in response_object: model_response.usage = transform_gemini_image_usage( 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 f6df772dd56..9a16004f50a 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -333,6 +333,15 @@ 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_object, model, raw_response, llm_provider="vertex_ai") + if usage_metadata := response_object.get("usageMetadata", None): model_response.usage = self._transform_image_usage(_JSON_DICT.validate_python(usage_metadata)) diff --git a/tests/unit/llms/gemini/image_generation/test_gemini_image_usage.py b/tests/unit/llms/gemini/image_generation/test_gemini_image_usage.py index 456626cab8f..119b74f0e90 100644 --- a/tests/unit/llms/gemini/image_generation/test_gemini_image_usage.py +++ b/tests/unit/llms/gemini/image_generation/test_gemini_image_usage.py @@ -259,6 +259,36 @@ def test_gemini_image_generation_accumulates_multiple_image_prompt_token_details litellm.model_cost = previous_model_cost +@pytest.mark.parametrize("finish_reason", ["IMAGE_SAFETY", "IMAGE_PROHIBITED_CONTENT", "SAFETY"]) +def test_gemini_image_generation_raises_on_safety_block(finish_reason): + """A safety/prohibited block returns a candidate with finishReason and no + inlineData; GoogleImageGenConfig must raise ContentPolicyViolationError, not + return empty data.""" + import httpx + + 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 = {} + + config = GoogleImageGenConfig() + with pytest.raises(ContentPolicyViolationError) as exc: + config.transform_image_generation_response( + model="gemini/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 == "gemini" + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py b/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py index 138aafcbfcd..7637757cd85 100644 --- a/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py +++ b/tests/unit/llms/gemini/test_gemini_image_generation_transformation.py @@ -482,7 +482,6 @@ _GEMINI_IMAGE_MODEL = "gemini-3.1-flash-image-preview" {"candidates": ""}, {"candidates": {}}, {"candidates": [{}]}, - {"candidates": [{"finishReason": "SAFETY"}]}, {"candidates": [{"content": {}}]}, {"candidates": [{"content": {"parts": ""}}]}, {"candidates": [{"content": {"parts": "plain text"}}]}, diff --git a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py index 5ac16283094..e97a582cfdb 100644 --- a/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py +++ b/tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py @@ -448,6 +448,31 @@ 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 = {} + + 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): @@ -724,7 +749,6 @@ def test_gemini_image_generation_response_rejects_non_object_usage_metadata_with {"candidates": ""}, {"candidates": {}}, {"candidates": [{}]}, - {"candidates": [{"finishReason": "SAFETY"}]}, {"candidates": [{"content": {}}]}, {"candidates": [{"content": {"parts": ""}}]}, {"candidates": [{"content": {"parts": "plain text"}}]},