Explicitly request image response modality from Google API when model supports it (#3172)

Author: DouweM

docs/builtin-tools.md

@@ -199,7 +199,7 @@ The [`ImageGenerationTool`][pydantic_ai.builtin_tools.ImageGenerationTool] enabl
 | Provider | Supported | Notes |
 |----------|-----------|-------|
 | OpenAI Responses | ✅ | Full feature support. Only supported by models newer than `gpt-4o`. Metadata about the generated image, like the [`revised_prompt`](https://platform.openai.com/docs/guides/tools-image-generation#revised-prompt) sent to the underlying image model, is available on the [`BuiltinToolReturnPart`][pydantic_ai.messages.BuiltinToolReturnPart] that's available via [`ModelResponse.builtin_tool_calls`][pydantic_ai.messages.ModelResponse.builtin_tool_calls]. |
-| Google | ✅ | No parameter support. Only supported by [image generation models](https://ai.google.dev/gemini-api/docs/image-generation) like `gemini-2.5-flash-image-preview`. These models do not support [structured output](output.md) or [function tools](tools.md). These models will always generate images, even if this built-in tool is not explicitly specified. |
+| Google | ✅ | No parameter support. Only supported by [image generation models](https://ai.google.dev/gemini-api/docs/image-generation) like `gemini-2.5-flash-image`. These models do not support [structured output](output.md) or [function tools](tools.md). These models will always generate images, even if this built-in tool is not explicitly specified. |
 | Anthropic | ❌ | |
 | Groq | ❌ | |
 | Bedrock | ❌ | |
@@ -232,7 +232,7 @@ Image generation with Google [image generation models](https://ai.google.dev/gem
 ```py {title="image_generation_google.py"}
 from pydantic_ai import Agent, BinaryImage
 
-agent = Agent('google-gla:gemini-2.5-flash-image-preview')
+agent = Agent('google-gla:gemini-2.5-flash-image')
 
 result = agent.run_sync('Tell me a two-sentence story about an axolotl with an illustration.')
 print(result.output)

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -73,6 +73,7 @@
         GroundingMetadata,
         HttpOptionsDict,
         MediaResolution,
+        Modality,
         Part,
         PartDict,
         SafetySettingDict,
@@ -415,6 +416,10 @@ async def _build_content_and_config(
         tool_config = self._get_tool_config(model_request_parameters, tools)
         system_instruction, contents = await self._map_messages(messages)
 
+        modalities = [Modality.TEXT.value]
+        if self.profile.supports_image_output:
+            modalities.append(Modality.IMAGE.value)
+
         http_options: HttpOptionsDict = {
             'headers': {'Content-Type': 'application/json', 'User-Agent': get_user_agent()}
         }
@@ -443,6 +448,7 @@ async def _build_content_and_config(
             tool_config=tool_config,
             response_mime_type=response_mime_type,
             response_schema=response_schema,
+            response_modalities=modalities,
         )
         return contents, config
 

tests/models/cassettes/test_google/test_google_image_and_text_output.yaml

@@ -20,7 +20,7 @@ interactions:
         - text: Tell me a two-sentence story about an axolotl with an illustration.
         role: user
       generationConfig: {}
-    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image:generateContent
   response:
     headers:
       alt-svc:
@@ -49,7 +49,7 @@ interactions:
           role: model
         finishReason: STOP
         index: 0
-      modelVersion: gemini-2.5-flash-image-preview
+      modelVersion: gemini-2.5-flash-image
       responseId: ZyjgaJODFaudz7IP2OK9-As
       usageMetadata:
         candidatesTokenCount: 1336

tests/models/cassettes/test_google/test_google_image_generation.yaml

@@ -20,7 +20,7 @@ interactions:
         - text: Generate an image of an axolotl.
         role: user
       generationConfig: {}
-    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image:generateContent
   response:
     headers:
       alt-svc:
@@ -48,7 +48,7 @@ interactions:
           role: model
         finishReason: STOP
         index: 0
-      modelVersion: gemini-2.5-flash-image-preview
+      modelVersion: gemini-2.5-flash-image
       responseId: eMXNaJf4FrDmqtsP2JGRsAQ
       usageMetadata:
         candidatesTokenCount: 1304
@@ -93,7 +93,7 @@ interactions:
         - text: Now give it a sombrero.
         role: user
       generationConfig: {}
-    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image:generateContent
   response:
     headers:
       alt-svc:
@@ -121,7 +121,7 @@ interactions:
           role: model
         finishReason: STOP
         index: 0
-      modelVersion: gemini-2.5-flash-image-preview
+      modelVersion: gemini-2.5-flash-image
       responseId: gMXNaOfYN4iUmtkPpO_k4AQ
       usageMetadata:
         candidatesTokenCount: 1304

tests/models/cassettes/test_google/test_google_image_generation_stream.yaml

@@ -20,7 +20,7 @@ interactions:
         - text: Generate an illustrated two-sentence story about an axolotl.
         role: user
       generationConfig: {}
-    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image:generateContent
   response:
     headers:
       alt-svc:
@@ -49,7 +49,7 @@ interactions:
           role: model
         finishReason: STOP
         index: 0
-      modelVersion: gemini-2.5-flash-image-preview
+      modelVersion: gemini-2.5-flash-image
       responseId: haHNaLDWN5n3qtsPzZTmkAc
       usageMetadata:
         candidatesTokenCount: 1335

tests/models/cassettes/test_google/test_google_image_generation_with_text.yaml

@@ -20,7 +20,7 @@ interactions:
         - text: Tell me a two-sentence story about an axolotl, no image please.
         role: user
       generationConfig: {}
-    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image:generateContent
   response:
     headers:
       alt-svc:
@@ -46,7 +46,7 @@ interactions:
           role: model
         finishReason: STOP
         index: 0
-      modelVersion: gemini-2.5-flash-image-preview
+      modelVersion: gemini-2.5-flash-image
       responseId: tofdaKSRBMSgz7IPx5eRyAU
       usageMetadata:
         candidatesTokenCount: 48
@@ -79,7 +79,7 @@ interactions:
         - text: Generate an image of an axolotl.
         role: user
       generationConfig: {}
-    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image:generateContent
   response:
     headers:
       alt-svc:
@@ -107,7 +107,7 @@ interactions:
           role: model
         finishReason: STOP
         index: 0
-      modelVersion: gemini-2.5-flash-image-preview
+      modelVersion: gemini-2.5-flash-image
       responseId: vIfdaKSMD96ez7IPpr73qAM
       usageMetadata:
         candidatesTokenCount: 1304

tests/models/cassettes/test_google/test_google_image_or_text_output.yaml

@@ -20,7 +20,7 @@ interactions:
         - text: Generate two separate images of axolotls.
         role: user
       generationConfig: {}
-    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image-preview:generateContent
+    uri: https://generativelanguage.googleapis.com/v1beta/models/gemini-2.5-flash-image:generateContent
   response:
     headers:
       alt-svc:
@@ -50,7 +50,7 @@ interactions:
           role: model
         finishReason: STOP
         index: 0
-      modelVersion: gemini-2.5-flash-image-preview
+      modelVersion: gemini-2.5-flash-image
       responseId: R4DdaO-WON7Qz7IPhNrg0QU
       usageMetadata:
         candidatesTokenCount: 1303

tests/models/cassettes/test_google/test_google_multiple_images.yaml

@@ -2590,7 +2590,7 @@ class CityLocation(BaseModel):
 
 
 async def test_google_image_generation(allow_model_requests: None, google_provider: GoogleProvider):
-    m = GoogleModel('gemini-2.5-flash-image-preview', provider=google_provider)
+    m = GoogleModel('gemini-2.5-flash-image', provider=google_provider)
     agent = Agent(m, output_type=BinaryImage)
 
     result = await agent.run('Generate an image of an axolotl.')
@@ -2629,7 +2629,7 @@ async def test_google_image_generation(allow_model_requests: None, google_provid
                     output_tokens=1304,
                     details={'text_prompt_tokens': 10, 'image_candidates_tokens': 1290},
                 ),
-                model_name='gemini-2.5-flash-image-preview',
+                model_name='gemini-2.5-flash-image',
                 timestamp=IsDatetime(),
                 provider_name='google-gla',
                 provider_details={'finish_reason': 'STOP'},
@@ -2673,7 +2673,7 @@ async def test_google_image_generation(allow_model_requests: None, google_provid
                     output_tokens=1304,
                     details={'text_prompt_tokens': 32, 'image_prompt_tokens': 1290, 'image_candidates_tokens': 1290},
                 ),
-                model_name='gemini-2.5-flash-image-preview',
+                model_name='gemini-2.5-flash-image',
                 timestamp=IsDatetime(),
                 provider_name='google-gla',
                 provider_details={'finish_reason': 'STOP'},
@@ -2685,7 +2685,7 @@ async def test_google_image_generation(allow_model_requests: None, google_provid
 
 
 async def test_google_image_generation_stream(allow_model_requests: None, google_provider: GoogleProvider):
-    m = GoogleModel('gemini-2.5-flash-image-preview', provider=google_provider)
+    m = GoogleModel('gemini-2.5-flash-image', provider=google_provider)
     agent = Agent(m, output_type=BinaryImage)
 
     async with agent.run_stream('Generate an image of an axolotl') as result:
@@ -2739,7 +2739,7 @@ async def test_google_image_generation_stream(allow_model_requests: None, google
                     output_tokens=1295,
                     details={'text_prompt_tokens': 10, 'image_candidates_tokens': 1290},
                 ),
-                model_name='gemini-2.5-flash-image-preview',
+                model_name='gemini-2.5-flash-image',
                 timestamp=IsDatetime(),
                 provider_name='google-gla',
                 provider_details={'finish_reason': 'STOP'},
@@ -2768,7 +2768,7 @@ async def test_google_image_generation_stream(allow_model_requests: None, google
 
 
 async def test_google_image_generation_with_text(allow_model_requests: None, google_provider: GoogleProvider):
-    m = GoogleModel('gemini-2.5-flash-image-preview', provider=google_provider)
+    m = GoogleModel('gemini-2.5-flash-image', provider=google_provider)
     agent = Agent(m)
 
     result = await agent.run('Generate an illustrated two-sentence story about an axolotl.')
@@ -2805,7 +2805,7 @@ async def test_google_image_generation_with_text(allow_model_requests: None, goo
                     output_tokens=1335,
                     details={'text_prompt_tokens': 14, 'image_candidates_tokens': 1290},
                 ),
-                model_name='gemini-2.5-flash-image-preview',
+                model_name='gemini-2.5-flash-image',
                 timestamp=IsDatetime(),
                 provider_name='google-gla',
                 provider_details={'finish_reason': 'STOP'},
@@ -2817,8 +2817,8 @@ async def test_google_image_generation_with_text(allow_model_requests: None, goo
 
 
 async def test_google_image_or_text_output(allow_model_requests: None, google_provider: GoogleProvider):
-    m = GoogleModel('gemini-2.5-flash-image-preview', provider=google_provider)
-    # ImageGenerationTool is listed here to indicate just that it doesn't cause any issues, even though it's not necessary with an image-preview model.
+    m = GoogleModel('gemini-2.5-flash-image', provider=google_provider)
+    # ImageGenerationTool is listed here to indicate just that it doesn't cause any issues, even though it's not necessary with an image model.
     agent = Agent(m, output_type=str | BinaryImage, builtin_tools=[ImageGenerationTool()])
 
     result = await agent.run('Tell me a two-sentence story about an axolotl, no image please.')
@@ -2837,7 +2837,7 @@ async def test_google_image_or_text_output(allow_model_requests: None, google_pr
 
 
 async def test_google_image_and_text_output(allow_model_requests: None, google_provider: GoogleProvider):
-    m = GoogleModel('gemini-2.5-flash-image-preview', provider=google_provider)
+    m = GoogleModel('gemini-2.5-flash-image', provider=google_provider)
     agent = Agent(m)
 
     result = await agent.run('Tell me a two-sentence story about an axolotl with an illustration.')
@@ -2860,7 +2860,7 @@ class Animal(BaseModel):
         species: str
         name: str
 
-    model = GoogleModel('gemini-2.5-flash-image-preview', provider=google_provider)
+    model = GoogleModel('gemini-2.5-flash-image', provider=google_provider)
     agent = Agent(model=model, output_type=Animal)
 
     with pytest.raises(UserError, match='Tool output is not supported by this model.'):
@@ -2872,7 +2872,7 @@ class Animal(BaseModel):
         species: str
         name: str
 
-    model = GoogleModel('gemini-2.5-flash-image-preview', provider=google_provider)
+    model = GoogleModel('gemini-2.5-flash-image', provider=google_provider)
     agent = Agent(model=model, output_type=NativeOutput(Animal))
 
     with pytest.raises(UserError, match='Native structured output is not supported by this model.'):
@@ -2886,15 +2886,15 @@ class Animal(BaseModel):
         species: str
         name: str
 
-    model = GoogleModel('gemini-2.5-flash-image-preview', provider=google_provider)
+    model = GoogleModel('gemini-2.5-flash-image', provider=google_provider)
     agent = Agent(model=model, output_type=PromptedOutput(Animal))
 
     with pytest.raises(UserError, match='JSON output is not supported by this model.'):
         await agent.run('Generate an image of an axolotl.')
 
 
 async def test_google_image_generation_with_tools(allow_model_requests: None, google_provider: GoogleProvider):
-    model = GoogleModel('gemini-2.5-flash-image-preview', provider=google_provider)
+    model = GoogleModel('gemini-2.5-flash-image', provider=google_provider)
     agent = Agent(model=model, output_type=BinaryImage)
 
     @agent.tool_plain
@@ -2917,15 +2917,9 @@ async def test_google_image_generation_tool(allow_model_requests: None, google_p
 
 
 async def test_google_vertexai_image_generation(allow_model_requests: None, vertex_provider: GoogleProvider):
-    model = GoogleModel('gemini-2.5-flash-image-preview', provider=vertex_provider)
+    model = GoogleModel('gemini-2.5-flash-image', provider=vertex_provider)
 
     agent = Agent(model, output_type=BinaryImage)
 
     result = await agent.run('Generate an image of an axolotl.')
-    assert result.output == snapshot(
-        BinaryImage(
-            data=IsBytes(),
-            media_type='image/png',
-            identifier='f3edd8',
-        )
-    )
+    assert result.output == snapshot(BinaryImage(data=IsBytes(), media_type='image/png', identifier='b037a4'))