diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 3943c0a7dae..14a53d2ad02 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 799307f98c7..f2ef94fdc72 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -320,6 +320,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_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 096f9c4796c..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 @@ -280,3 +252,33 @@ 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" 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 54607cc5284..6c8e8d7d177 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 @@ -441,6 +441,33 @@ 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):