fix(gemini): merge extra_body into generationConfig for image preview

This commit is contained in:
siddharth 2026-03-27 16:11:17 -04:00
parent d949085310
commit 5c0c7b60f5
No known key found for this signature in database
GPG key ID: 1C52862D13014F31
2 changed files with 53 additions and 0 deletions

View file

@ -184,6 +184,15 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
"contents": [{"parts": [{"text": prompt}]}],
"generationConfig": {"response_modalities": ["IMAGE", "TEXT"]},
}
extra_body: Optional[dict] = optional_params.pop("extra_body", None)
if extra_body is not None:
for k, v in extra_body.items():
if k in {"cache", "tags"}:
continue
if k in request_body and isinstance(request_body[k], dict) and isinstance(v, dict):
request_body[k].update(v)
else:
request_body[k] = v
return request_body
else:
# For other Imagen models, use the original Imagen format

View file

@ -0,0 +1,44 @@
import pytest
import litellm
from litellm.llms.gemini.image_generation.transformation import GoogleImageGenConfig
def test_gemini_flash_image_generation_extra_body():
"""
Test that extra_body is properly extracted and deep merged into the
generationConfig for Gemini 3.1 Flash Image Preview requests.
"""
config = GoogleImageGenConfig()
model = "gemini-3.1-flash-image-preview"
prompt = "A realistic, high-quality close-up..."
# User's optional parameters, containing the extra_body we want to inject
optional_params = {
"extra_body": {
"generationConfig": {
"imageConfig": {
"imageSize": "2K"
}
}
}
}
litellm_params = {}
headers = {}
request_body = config.transform_image_generation_request(
model=model,
prompt=prompt,
optional_params=optional_params,
litellm_params=litellm_params,
headers=headers
)
assert "generationConfig" in request_body
gen_config = request_body["generationConfig"]
# Should contain original hardcoded item
assert "response_modalities" in gen_config
assert gen_config["response_modalities"] == ["IMAGE", "TEXT"]
# Should contain the extra_body injected item
assert "imageConfig" in gen_config
assert gen_config["imageConfig"] == {"imageSize": "2K"}