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
This commit is contained in:
Chenglun Hu 2026-05-28 16:03:31 +08:00
parent be658d5d29
commit ff85c5bb88
3 changed files with 78 additions and 0 deletions

View file

@ -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"])

View file

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

View file

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