From 619b65adb30f0e809ae0d20d7a270be9ab5777f5 Mon Sep 17 00:00:00 2001 From: Chesars Date: Thu, 27 Nov 2025 17:40:06 -0300 Subject: [PATCH 1/2] fix(image-gen): add thought_signature to ImageObject for Gemini 3 Pro Fixes #17184 - Gemini 3 Pro image preview model returns a thoughtSignature field required for interactive image editing. This change: - Adds thought_signature field to ImageObject class - Updates Gemini and Vertex AI transformations to extract thoughtSignature - Adds test for thought_signature in response transformation --- .../gemini/image_generation/transformation.py | 1 + .../vertex_gemini_transformation.py | 1 + litellm/types/utils.py | 5 ++- ...rtex_ai_image_generation_transformation.py | 41 +++++++++++++++++++ 4 files changed, 47 insertions(+), 1 deletion(-) diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index e79414394fa..cb9d38f938c 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -232,6 +232,7 @@ class GoogleImageGenConfig(BaseImageGenerationConfig): model_response.data.append(ImageObject( b64_json=inline_data["data"], url=None, + thought_signature=part.get("thoughtSignature"), )) else: # Original Imagen format - predictions with generated images 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 149e0850bf0..74e8b28d511 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -258,6 +258,7 @@ class VertexAIGeminiImageGenerationConfig(BaseImageGenerationConfig, VertexLLM): model_response.data.append(ImageObject( b64_json=inline_data["data"], url=None, + thought_signature=part.get("thoughtSignature"), )) return model_response diff --git a/litellm/types/utils.py b/litellm/types/utils.py index b0d081d8f87..3b789ee8140 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1708,6 +1708,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). https://platform.openai.com/docs/api-reference/images/object """ @@ -1715,9 +1716,11 @@ class ImageObject(OpenAIImage): b64_json: Optional[str] = None url: Optional[str] = None revised_prompt: Optional[str] = None + thought_signature: Optional[str] = None - def __init__(self, b64_json=None, url=None, revised_prompt=None, **kwargs): + def __init__(self, b64_json=None, url=None, revised_prompt=None, thought_signature=None, **kwargs): super().__init__(b64_json=b64_json, url=url, revised_prompt=revised_prompt) # type: ignore + self.thought_signature = thought_signature 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 9f33400594b..06b4893e058 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 @@ -209,6 +209,47 @@ class TestVertexAIGeminiImageGenerationConfig: assert result.data[0].b64_json == "image1" assert result.data[1].b64_json == "image2" + def test_transform_image_generation_response_signature(self): + """Test response transformation includes thoughtSignature for Gemini 3 Pro""" + mock_response = MagicMock(spec=httpx.Response) + mock_response.status_code = 200 + mock_response.json.return_value = { + "candidates": [ + { + "content": { + "parts": [ + { + "inlineData": { + "mimeType": "image/png", + "data": "base64_encoded_image_data", + }, + "thoughtSignature": "test_signature_abc123", + } + ] + } + } + ] + } + mock_response.headers = {} + + from litellm.types.utils import ImageResponse + + model_response = ImageResponse() + result = self.config.transform_image_generation_response( + model="gemini-3-pro-image-preview", + raw_response=mock_response, + model_response=model_response, + logging_obj=MagicMock(), + request_data={}, + optional_params={}, + litellm_params={}, + encoding=None, + ) + + assert len(result.data) == 1 + assert result.data[0].b64_json == "base64_encoded_image_data" + assert result.data[0].thought_signature == "test_signature_abc123" + class TestVertexAIImagenImageGenerationConfig: def setup_method(self): From 979c14f9e59f86a9052643bb965a6c0af883d2c1 Mon Sep 17 00:00:00 2001 From: Chesars Date: Fri, 28 Nov 2025 17:16:43 -0300 Subject: [PATCH 2/2] 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 cb9d38f938c..27a9bf3e32e 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -229,10 +229,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, )) else: # Original Imagen format - predictions with generated images 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 74e8b28d511..08689f56ffe 100644 --- a/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py +++ b/litellm/llms/vertex_ai/image_generation/vertex_gemini_transformation.py @@ -255,10 +255,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, )) return model_response diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 3b789ee8140..6bbe87ca455 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1708,7 +1708,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 """ @@ -1716,11 +1716,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 06b4893e058..096835d6b80 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 @@ -248,7 +248,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: