From 5c0c7b60f560220586b8de237a96af5afda8d345 Mon Sep 17 00:00:00 2001
From: siddharth <siddharth.github@gmail.com>
Date: Fri, 27 Mar 2026 16:11:17 -0400
Subject: [PATCH] fix(gemini): merge extra_body into generationConfig for image
 preview

---
 .../gemini/image_generation/transformation.py |  9 ++++
 .../gemini/test_gemini_image_generation.py    | 44 +++++++++++++++++++
 2 files changed, 53 insertions(+)
 create mode 100644 tests/test_litellm/llms/gemini/test_gemini_image_generation.py

diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py
index b094fc133d7..c2f468776da 100644
--- a/litellm/llms/gemini/image_generation/transformation.py
+++ b/litellm/llms/gemini/image_generation/transformation.py
@@ -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
diff --git a/tests/test_litellm/llms/gemini/test_gemini_image_generation.py b/tests/test_litellm/llms/gemini/test_gemini_image_generation.py
new file mode 100644
index 00000000000..f5658da5df2
--- /dev/null
+++ b/tests/test_litellm/llms/gemini/test_gemini_image_generation.py
@@ -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"}