mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
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:
parent
dddbaa803d
commit
80257675c8
5 changed files with 98 additions and 2 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"}}]},
|
||||
|
|
|
|||
|
|
@ -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"}}]},
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue