From ff85c5bb88f31c0a7d6eac4e89b294cd83ee3c18 Mon Sep 17 00:00:00 2001 From: Chenglun Hu Date: Thu, 28 May 2026 16:03:31 +0800 Subject: [PATCH 1/3] 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): From dc1a9c7d258b5b290003a91161b24dbf741a7b32 Mon Sep 17 00:00:00 2001 From: Chenglun Hu Date: Thu, 28 May 2026 19:31:35 +0800 Subject: [PATCH 2/3] test(gemini): cover GoogleImageGenConfig safety-block path greptile flagged that the safety-block raise was only tested on the vertex config; add the matching test for the gemini (AI Studio) path. --- .../test_gemini_image_usage.py | 34 +++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/llm_translation/test_gemini_image_usage.py b/tests/llm_translation/test_gemini_image_usage.py index 096f9c4796c..8fc12408e5e 100644 --- a/tests/llm_translation/test_gemini_image_usage.py +++ b/tests/llm_translation/test_gemini_image_usage.py @@ -280,3 +280,37 @@ def test_gemini_image_generation_accumulates_multiple_image_prompt_token_details else: os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = previous_local_model_cost_map 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" From 295f73f9dc567b2d841afb73a48199bbbea16400 Mon Sep 17 00:00:00 2001 From: Chenglun Hu Date: Thu, 16 Jul 2026 13:05:07 +0800 Subject: [PATCH 3/3] chore: ruff format --- .../vertex_gemini_transformation.py | 4 +- .../test_gemini_image_usage.py | 90 ++++++------------- ...rtex_ai_image_generation_transformation.py | 12 +-- 3 files changed, 33 insertions(+), 73 deletions(-) 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 e6c98688de1..836b9ea41ae 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -328,9 +328,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): raise_if_image_gen_flagged, ) - raise_if_image_gen_flagged( - response_data, model, raw_response, llm_provider="vertex_ai" - ) + 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/llm_translation/test_gemini_image_usage.py b/tests/llm_translation/test_gemini_image_usage.py index 8fc12408e5e..2460d6f5cd3 100644 --- a/tests/llm_translation/test_gemini_image_usage.py +++ b/tests/llm_translation/test_gemini_image_usage.py @@ -1,7 +1,7 @@ """ Test for Gemini image generation usage metadata extraction. -This test verifies the fix for issue #18323 where image_generation() +This test verifies the fix for issue #18323 where image_generation() was returning usage=0 while completion() returned proper token usage. """ @@ -57,9 +57,7 @@ def test_gemini_image_generation_usage_metadata(model_name: str): }, } - with patch( - "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post: # Mock successful HTTP response mock_http_response = MagicMock() mock_http_response.json.return_value = mock_response_data @@ -87,56 +85,38 @@ def test_gemini_image_generation_usage_metadata(model_name: str): # but it should still have the ImageUsage fields (input_tokens, output_tokens, etc.) # Validate token counts match the mock response - assert hasattr( - response.usage, "input_tokens" - ), "Usage should have input_tokens attribute" - assert hasattr( - response.usage, "output_tokens" - ), "Usage should have output_tokens attribute" - assert hasattr( - response.usage, "total_tokens" - ), "Usage should have total_tokens attribute" + assert hasattr(response.usage, "input_tokens"), "Usage should have input_tokens attribute" + assert hasattr(response.usage, "output_tokens"), "Usage should have output_tokens attribute" + assert hasattr(response.usage, "total_tokens"), "Usage should have total_tokens attribute" - assert ( - response.usage.input_tokens == 35 - ), f"Expected input_tokens=35, got {response.usage.input_tokens}" - assert ( - response.usage.output_tokens == 1716 - ), f"Expected output_tokens=1716, got {response.usage.output_tokens}" - assert ( - response.usage.total_tokens == 1751 - ), f"Expected total_tokens=1751, got {response.usage.total_tokens}" + assert response.usage.input_tokens == 35, f"Expected input_tokens=35, got {response.usage.input_tokens}" + assert response.usage.output_tokens == 1716, f"Expected output_tokens=1716, got {response.usage.output_tokens}" + assert response.usage.total_tokens == 1751, f"Expected total_tokens=1751, got {response.usage.total_tokens}" # Validate input tokens details - assert hasattr( - response.usage, "input_tokens_details" - ), "Usage should have input_tokens_details attribute" - assert ( - response.usage.input_tokens_details is not None - ), "Input tokens details should not be None" + assert hasattr(response.usage, "input_tokens_details"), "Usage should have input_tokens_details attribute" + assert response.usage.input_tokens_details is not None, "Input tokens details should not be None" # input_tokens_details might be a dict or an object if isinstance(response.usage.input_tokens_details, dict): - assert ( - response.usage.input_tokens_details["text_tokens"] == 35 - ), f"Expected text_tokens=35, got {response.usage.input_tokens_details['text_tokens']}" - assert ( - response.usage.input_tokens_details["image_tokens"] == 0 - ), f"Expected image_tokens=0, got {response.usage.input_tokens_details['image_tokens']}" + assert response.usage.input_tokens_details["text_tokens"] == 35, ( + f"Expected text_tokens=35, got {response.usage.input_tokens_details['text_tokens']}" + ) + assert response.usage.input_tokens_details["image_tokens"] == 0, ( + f"Expected image_tokens=0, got {response.usage.input_tokens_details['image_tokens']}" + ) else: - assert ( - response.usage.input_tokens_details.text_tokens == 35 - ), f"Expected text_tokens=35, got {response.usage.input_tokens_details.text_tokens}" - assert ( - response.usage.input_tokens_details.image_tokens == 0 - ), f"Expected image_tokens=0, got {response.usage.input_tokens_details.image_tokens}" + assert response.usage.input_tokens_details.text_tokens == 35, ( + f"Expected text_tokens=35, got {response.usage.input_tokens_details.text_tokens}" + ) + assert response.usage.input_tokens_details.image_tokens == 0, ( + f"Expected image_tokens=0, got {response.usage.input_tokens_details.image_tokens}" + ) # Verify the usage is not all zeros (the bug we're fixing) assert response.usage.total_tokens > 0, "Total tokens should be greater than 0" assert response.usage.input_tokens > 0, "Input tokens should be greater than 0" - assert ( - response.usage.output_tokens > 0 - ), "Output tokens should be greater than 0" + assert response.usage.output_tokens > 0, "Output tokens should be greater than 0" def test_gemini_image_generation_without_usage_metadata(): @@ -162,9 +142,7 @@ def test_gemini_image_generation_without_usage_metadata(): ] } - with patch( - "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post: # Mock successful HTTP response mock_http_response = MagicMock() mock_http_response.json.return_value = mock_response_data @@ -197,13 +175,9 @@ def test_gemini_imagen_models_no_usage_extraction(): """ # Mock response data for Imagen models (different format) - mock_response_data = { - "predictions": [{"bytesBase64Encoded": "test_base64_image_data"}] - } + mock_response_data = {"predictions": [{"bytesBase64Encoded": "test_base64_image_data"}]} - with patch( - "litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post" - ) as mock_post: + with patch("litellm.llms.custom_httpx.llm_http_handler.HTTPHandler.post") as mock_post: # Mock successful HTTP response mock_http_response = MagicMock() mock_http_response.json.return_value = mock_response_data @@ -267,9 +241,7 @@ def test_gemini_image_generation_accumulates_multiple_image_prompt_token_details model_info = litellm.get_model_info(model=model, custom_llm_provider="gemini") expected_image_tokens = 190 expected_total_prompt_tokens = 200 - expected_prompt_cost = ( - expected_total_prompt_tokens * model_info["input_cost_per_token"] - ) + expected_prompt_cost = expected_total_prompt_tokens * model_info["input_cost_per_token"] assert parsed_usage.input_tokens_details.image_tokens == expected_image_tokens assert parsed_usage.input_tokens_details.text_tokens == 10 @@ -282,9 +254,7 @@ 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"] -) +@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 @@ -295,9 +265,7 @@ def test_gemini_image_generation_raises_on_safety_block(finish_reason): mock_response = MagicMock(spec=httpx.Response) mock_response.status_code = 200 - mock_response.json.return_value = { - "candidates": [{"finishReason": finish_reason, "content": {"parts": []}}] - } + mock_response.json.return_value = {"candidates": [{"finishReason": finish_reason, "content": {"parts": []}}]} mock_response.headers = {} config = GoogleImageGenConfig() 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 d03a89dfe66..676be2a0f2b 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,21 +443,15 @@ 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 - ): + @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.json.return_value = {"candidates": [{"finishReason": finish_reason, "content": {"parts": []}}]} mock_response.headers = {} from litellm.types.utils import ImageResponse