fix: send the right data format to gemini image input (#991)

Author: Kludex

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -320,18 +320,24 @@ async def _map_user_prompt(part: UserPromptPart) -> list[_GeminiPartUnion]:
                     content.append({'text': item})
                 elif isinstance(item, BinaryContent):
                     base64_encoded = base64.b64encode(item.data).decode('utf-8')
-                    content.append(_GeminiInlineDataPart(data=base64_encoded, mime_type=item.media_type))
+                    content.append(
+                        _GeminiInlineDataPart(inline_data={'data': base64_encoded, 'mime_type': item.media_type})
+                    )
                 elif isinstance(item, (AudioUrl, ImageUrl)):
                     try:
-                        content.append(_GeminiFileDataData(file_uri=item.url, mime_type=item.media_type))
+                        content.append(
+                            _GeminiFileDataPart(file_data={'file_uri': item.url, 'mime_type': item.media_type})
+                        )
                     except ValueError:
                         # Download the file if can't find the mime type.
                         client = cached_async_http_client()
                         response = await client.get(item.url, follow_redirects=True)
                         response.raise_for_status()
                         base64_encoded = base64.b64encode(response.content).decode('utf-8')
                         content.append(
-                            _GeminiInlineDataPart(data=base64_encoded, mime_type=response.headers['Content-Type'])
+                            _GeminiInlineDataPart(
+                                inline_data={'data': base64_encoded, 'mime_type': response.headers['Content-Type']}
+                            )
                         )
                 else:
                     assert_never(item)
@@ -528,20 +534,28 @@ class _GeminiTextPart(TypedDict):
     text: str
 
 
+class _GeminiInlineData(TypedDict):
+    data: str
+    mime_type: Annotated[str, pydantic.Field(alias='mimeType')]
+
+
 class _GeminiInlineDataPart(TypedDict):
     """See <https://ai.google.dev/api/caching#Blob>."""
 
-    data: str
-    mime_type: Annotated[str, pydantic.Field(alias='mimeType')]
+    inline_data: Annotated[_GeminiInlineData, pydantic.Field(alias='inlineData')]
 
 
-class _GeminiFileDataData(TypedDict):
+class _GeminiFileData(TypedDict):
     """See <https://ai.google.dev/api/caching#FileData>."""
 
     file_uri: Annotated[str, pydantic.Field(alias='fileUri')]
     mime_type: Annotated[str, pydantic.Field(alias='mimeType')]
 
 
+class _GeminiFileDataPart(TypedDict):
+    file_data: Annotated[_GeminiFileData, pydantic.Field(alias='fileData')]
+
+
 class _GeminiFunctionCallPart(TypedDict):
     function_call: Annotated[_GeminiFunctionCall, pydantic.Field(alias='functionCall')]
 
@@ -617,7 +631,7 @@ def _part_discriminator(v: Any) -> str:
         Annotated[_GeminiFunctionCallPart, pydantic.Tag('function_call')],
         Annotated[_GeminiFunctionResponsePart, pydantic.Tag('function_response')],
         Annotated[_GeminiInlineDataPart, pydantic.Tag('inline_data')],
-        Annotated[_GeminiFileDataData, pydantic.Tag('file_data')],
+        Annotated[_GeminiFileDataPart, pydantic.Tag('file_data')],
     ],
     pydantic.Discriminator(_part_discriminator),
 ]

tests/models/cassettes/test_gemini/test_image_as_binary_content_input.yaml

@@ -998,7 +998,7 @@ async def test_image_as_binary_content_input(
     agent = Agent(m)
 
     result = await agent.run(['What is the name of this fruit?', image_content])
-    assert result.data == snapshot('The fruit in the image is a Kiwi.')
+    assert result.data == snapshot('The fruit in the image is a kiwi.')
 
 
 @pytest.mark.vcr()
@@ -1009,4 +1009,4 @@ async def test_image_url_input(allow_model_requests: None, gemini_api_key: str)
     image_url = ImageUrl(url='https://goo.gle/instrument-img')
 
     result = await agent.run(['What is the name of this fruit?', image_url])
-    assert result.data == snapshot('The image shows an organ, not a fruit.')
+    assert result.data == snapshot('This is not a fruit, it is an organ console.')