scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-06 08:16:43 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 78

Merge 295f73f9dc into 67c7b97fd2

Browse source
This commit is contained in:
hcl 2026-08-27 16:08:13 -05:00 • committed by GitHub
commit fc9fcaef4d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 127 additions and 55 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
litellm/llms/gemini/image_generation/transformation.py
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"])

9
litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py
View file

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

112
tests/llm_translation/test_gemini_image_usage.py
View file

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

27
tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py
View file

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