Merge 80257675c8 into 5103d8a0ee

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(

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))
 

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)

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"}}]},

tests/unit/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py

@@ -540,6 +540,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):
@@ -815,7 +840,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"}}]},