diff --git a/litellm/llms/gemini/image_generation/transformation.py b/litellm/llms/gemini/image_generation/transformation.py index 63b835df9d0..73aef15e4c7 100644 --- a/litellm/llms/gemini/image_generation/transformation.py +++ b/litellm/llms/gemini/image_generation/transformation.py @@ -255,9 +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, + 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 89ed9f1a8a5..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,9 +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, + 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 f063418f92e..95cc2f9f2fb 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2129,6 +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. + provider_specific_fields: Provider-specific fields not part of OpenAI spec. https://platform.openai.com/docs/api-reference/images/object """ @@ -2136,9 +2137,12 @@ class ImageObject(OpenAIImage): b64_json: Optional[str] = None url: Optional[str] = None revised_prompt: Optional[str] = None + provider_specific_fields: Optional[Dict[str, Any]] = None - def __init__(self, b64_json=None, url=None, revised_prompt=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 + 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 b91438b3cac..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 @@ -230,6 +230,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].provider_specific_fields["thought_signature"] == "test_signature_abc123" + class TestVertexAIImagenImageGenerationConfig: def setup_method(self):