From 11fd92c21b097426f6d4e0116d002c77766497f2 Mon Sep 17 00:00:00 2001
From: Chesars <cesarponce19544@gmail.com>
Date: Fri, 28 Nov 2025 17:16:43 -0300
Subject: [PATCH] refactor(image-gen): move thought_signature to
 provider_specific_fields

Per review feedback, thought_signature should not be a root-level
param on ImageObject as it's not OpenAI compatible. Moved to
provider_specific_fields dict to match the pattern used in chat
completions (Message, Delta, Choices, etc).
---
 litellm/llms/gemini/image_generation/transformation.py   | 3 ++-
 .../image_generation/vertex_gemini_transformation.py     | 3 ++-
 litellm/types/utils.py                                   | 9 +++++----
 .../test_vertex_ai_image_generation_transformation.py    | 2 +-
 4 files changed, 10 insertions(+), 7 deletions(-)

diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py
index 46cbf6ae877..73aef15e4c7 100644
--- a/litellm/llms/gemini/image_generation/transformation.py
+++ b/litellm/llms/gemini/image_generation/transformation.py
@@ -255,10 +255,11 @@ class GoogleImageGenConfig(BaseImageGenerationConfig):
                     if "inlineData" in part:
                         inline_data = part["inlineData"]
                         if "data" in inline_data:
+                            thought_sig = part.get("thoughtSignature")
                             model_response.data.append(ImageObject(
                                 b64_json=inline_data["data"],
                                 url=None,
-                                thought_signature=part.get("thoughtSignature"),
+                                provider_specific_fields={"thought_signature": thought_sig} if thought_sig else None,
                             ))
             
             # Extract usage metadata for Gemini models
diff --git a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py
index f266f860712..ba3df88be14 100644
--- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py
+++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py
@@ -295,10 +295,11 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM):
                 if "inlineData" in part:
                     inline_data = part["inlineData"]
                     if "data" in inline_data:
+                        thought_sig = part.get("thoughtSignature")
                         model_response.data.append(ImageObject(
                             b64_json=inline_data["data"],
                             url=None,
-                            thought_signature=part.get("thoughtSignature"),
+                            provider_specific_fields={"thought_signature": thought_sig} if thought_sig else None,
                         ))
 
         if usage_metadata := response_data.get("usageMetadata", None):
diff --git a/litellm/types/utils.py b/litellm/types/utils.py
index 077277f4830..09c944e1fe8 100644
--- a/litellm/types/utils.py
+++ b/litellm/types/utils.py
@@ -2129,7 +2129,7 @@ class ImageObject(OpenAIImage):
     b64_json: The base64-encoded JSON of the generated image, if response_format is b64_json.
     url: The URL of the generated image, if response_format is url (default).
     revised_prompt: The prompt that was used to generate the image, if there was any revision to the prompt.
-    thought_signature: The thought signature returned by Gemini image generation models (used for interactive image editing).
+    provider_specific_fields: Provider-specific fields not part of OpenAI spec.
 
     https://platform.openai.com/docs/api-reference/images/object
     """
@@ -2137,11 +2137,12 @@ class ImageObject(OpenAIImage):
     b64_json: Optional[str] = None
     url: Optional[str] = None
     revised_prompt: Optional[str] = None
-    thought_signature: Optional[str] = None
+    provider_specific_fields: Optional[Dict[str, Any]] = None
 
-    def __init__(self, b64_json=None, url=None, revised_prompt=None, thought_signature=None, **kwargs):
+    def __init__(self, b64_json=None, url=None, revised_prompt=None, provider_specific_fields=None, **kwargs):
         super().__init__(b64_json=b64_json, url=url, revised_prompt=revised_prompt)  # type: ignore
-        self.thought_signature = thought_signature
+        if provider_specific_fields:
+            self.provider_specific_fields = provider_specific_fields
 
     def __contains__(self, key):
         # Define custom behavior for the 'in' operator
diff --git a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py
index b4f5053955e..6736eaffebd 100644
--- a/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py
+++ b/tests/test_litellm/llms/vertex_ai/image_generation/test_vertex_ai_image_generation_transformation.py
@@ -269,7 +269,7 @@ class TestVertexAIGeminiImageGenerationConfig:
 
         assert len(result.data) == 1
         assert result.data[0].b64_json == "base64_encoded_image_data"
-        assert result.data[0].thought_signature == "test_signature_abc123"
+        assert result.data[0].provider_specific_fields["thought_signature"] == "test_signature_abc123"
 
 
 class TestVertexAIImagenImageGenerationConfig: