mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-10 22:41:41 +00:00
fix(gemini): merge extra_body into generationConfig for image preview
This commit is contained in:
parent
d949085310
commit
5c0c7b60f5
2 changed files with 53 additions and 0 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"}
|
||||
Loading…
Add table
Reference in a new issue