From 6ec3e2837020fd4148a07a3eddfd97cf64c5fbab Mon Sep 17 00:00:00 2001 From: David Date: Mon, 21 Jul 2025 23:16:33 -0400 Subject: [PATCH 1/4] Support setting FileUrl media_type explicitely at construction time --- pydantic_ai_slim/pydantic_ai/messages.py | 26 +++++++++++++++--------- tests/test_messages.py | 21 +++++++++++-------- 2 files changed, 29 insertions(+), 18 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/messages.py b/pydantic_ai_slim/pydantic_ai/messages.py index 731ccb5ca..056695604 100644 --- a/pydantic_ai_slim/pydantic_ai/messages.py +++ b/pydantic_ai_slim/pydantic_ai/messages.py @@ -92,6 +92,9 @@ class FileUrl(ABC): url: str """The URL of the file.""" + _media_type: str | None = field(default=None, repr=False) + """Optional override for the media type of the file, in case it cannot be inferred from the URL.""" + force_download: bool = False """If the model supports it: @@ -106,11 +109,18 @@ class FileUrl(ABC): - `GoogleModel`: `VideoUrl.vendor_metadata` is used as `video_metadata`: https://ai.google.dev/gemini-api/docs/video-understanding#customize-video-processing """ - @property @abstractmethod - def media_type(self) -> str: + def _infer_media_type_from_url(self) -> str: """Return the media type of the file, based on the url.""" + @property + def media_type(self) -> str: + """Return the media type of the file, based on the url or the provided `_media_type`.""" + if self._media_type is not None: + return self._media_type + else: + return self._infer_media_type_from_url() + @property @abstractmethod def format(self) -> str: @@ -129,8 +139,7 @@ class VideoUrl(FileUrl): kind: Literal['video-url'] = 'video-url' """Type identifier, this is available on all parts as a discriminator.""" - @property - def media_type(self) -> VideoMediaType: + def _infer_media_type_from_url(self) -> VideoMediaType: """Return the media type of the video, based on the url.""" if self.url.endswith('.mkv'): return 'video/x-matroska' @@ -180,8 +189,7 @@ class AudioUrl(FileUrl): kind: Literal['audio-url'] = 'audio-url' """Type identifier, this is available on all parts as a discriminator.""" - @property - def media_type(self) -> AudioMediaType: + def _infer_media_type_from_url(self) -> AudioMediaType: """Return the media type of the audio file, based on the url. References: @@ -218,8 +226,7 @@ class ImageUrl(FileUrl): kind: Literal['image-url'] = 'image-url' """Type identifier, this is available on all parts as a discriminator.""" - @property - def media_type(self) -> ImageMediaType: + def _infer_media_type_from_url(self) -> ImageMediaType: """Return the media type of the image, based on the url.""" if self.url.endswith(('.jpg', '.jpeg')): return 'image/jpeg' @@ -251,8 +258,7 @@ class DocumentUrl(FileUrl): kind: Literal['document-url'] = 'document-url' """Type identifier, this is available on all parts as a discriminator.""" - @property - def media_type(self) -> str: + def _infer_media_type_from_url(self) -> str: """Return the media type of the document, based on the url.""" type_, _ = guess_type(self.url) if type_ is None: diff --git a/tests/test_messages.py b/tests/test_messages.py index d95aca32a..d1b3b90ac 100644 --- a/tests/test_messages.py +++ b/tests/test_messages.py @@ -17,16 +17,20 @@ def test_image_url(): assert image_url.media_type == 'image/jpeg' assert image_url.format == 'jpeg' + image_url = ImageUrl(url='https://example.com/image', _media_type='image/jpeg') + assert image_url.media_type == 'image/jpeg' + assert image_url.format == 'jpeg' -def test_video_url(): - with pytest.raises(ValueError, match='Unknown video file extension: https://example.com/video.potato'): - video_url = VideoUrl(url='https://example.com/video.potato') - video_url.media_type +def test_video_url(): video_url = VideoUrl(url='https://example.com/video.mp4') assert video_url.media_type == 'video/mp4' assert video_url.format == 'mp4' + video_url = VideoUrl(url='https://example.com/video', _media_type='video/mp4') + assert video_url.media_type == 'video/mp4' + assert video_url.format == 'mp4' + @pytest.mark.parametrize( 'url,is_youtube', @@ -45,14 +49,14 @@ def test_youtube_video_url(url: str, is_youtube: bool): def test_document_url(): - with pytest.raises(ValueError, match='Unknown document file extension: https://example.com/document.potato'): - document_url = DocumentUrl(url='https://example.com/document.potato') - document_url.media_type - document_url = DocumentUrl(url='https://example.com/document.pdf') assert document_url.media_type == 'application/pdf' assert document_url.format == 'pdf' + document_url = DocumentUrl(url='https://example.com/document', _media_type='application/pdf') + assert document_url.media_type == 'application/pdf' + assert document_url.format == 'pdf' + @pytest.mark.parametrize( 'media_type, format', @@ -129,6 +133,7 @@ def test_binary_content_document(media_type: str, format: str): pytest.param(AudioUrl('foobar.flac'), 'audio/flac', 'flac', id='flac'), pytest.param(AudioUrl('foobar.aiff'), 'audio/aiff', 'aiff', id='aiff'), pytest.param(AudioUrl('foobar.aac'), 'audio/aac', 'aac', id='aac'), + pytest.param(AudioUrl('foobar', 'audio/mpeg'), 'audio/mpeg', 'mp3', id='mp3'), ], ) def test_audio_url(audio_url: AudioUrl, media_type: str, format: str): From 0a94e5c75d66327969a389a4b2ffdfd1c46b74e3 Mon Sep 17 00:00:00 2001 From: David Date: Mon, 21 Jul 2025 23:18:49 -0400 Subject: [PATCH 2/4] Support passing Gemini File API urls for google-gla --- pydantic_ai_slim/pydantic_ai/models/google.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/pydantic_ai_slim/pydantic_ai/models/google.py b/pydantic_ai_slim/pydantic_ai/models/google.py index 9ec1260d4..ad5da243c 100644 --- a/pydantic_ai_slim/pydantic_ai/models/google.py +++ b/pydantic_ai_slim/pydantic_ai/models/google.py @@ -411,7 +411,12 @@ async def _map_user_prompt(self, part: UserPromptPart) -> list[PartDict]: file_data_dict['video_metadata'] = item.vendor_metadata content.append(file_data_dict) # type: ignore elif isinstance(item, FileUrl): - if self.system == 'google-gla' or item.force_download: + if item.force_download or ( + # google-gla does not support passing file urls directly, except for youtube videos + # (see above) and files uploaded to the file API (which cannot be downloaded anyway) + self.system == 'google-gla' + and not item.url.startswith(r'https://generativelanguage.googleapis.com/v1beta/files') + ): downloaded_item = await download_item(item, data_format='base64') inline_data = {'data': downloaded_item['data'], 'mime_type': downloaded_item['data_type']} content.append({'inline_data': inline_data}) # type: ignore From 390218f82bb0f41eeb724b0c936447d41db1a735 Mon Sep 17 00:00:00 2001 From: David Date: Tue, 22 Jul 2025 14:57:28 -0400 Subject: [PATCH 3/4] Make FileUrl media type override "non private" init parameter --- pydantic_ai_slim/pydantic_ai/messages.py | 24 +++++++++++++----------- tests/test_messages.py | 6 +++--- 2 files changed, 16 insertions(+), 14 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/messages.py b/pydantic_ai_slim/pydantic_ai/messages.py index 056695604..baf3f6752 100644 --- a/pydantic_ai_slim/pydantic_ai/messages.py +++ b/pydantic_ai_slim/pydantic_ai/messages.py @@ -3,7 +3,7 @@ import base64 from abc import ABC, abstractmethod from collections.abc import Sequence -from dataclasses import dataclass, field, replace +from dataclasses import InitVar, dataclass, field, replace from datetime import datetime from mimetypes import guess_type from typing import TYPE_CHECKING, Annotated, Any, Literal, Union, cast, overload @@ -92,7 +92,7 @@ class FileUrl(ABC): url: str """The URL of the file.""" - _media_type: str | None = field(default=None, repr=False) + media_type_: InitVar[str | None] = None """Optional override for the media type of the file, in case it cannot be inferred from the URL.""" force_download: bool = False @@ -109,17 +109,19 @@ class FileUrl(ABC): - `GoogleModel`: `VideoUrl.vendor_metadata` is used as `video_metadata`: https://ai.google.dev/gemini-api/docs/video-understanding#customize-video-processing """ + _media_type: str | None = field(init=False, repr=False) + + def __post_init__(self, media_type_: str | None = None) -> None: + self._media_type = media_type_ + @abstractmethod - def _infer_media_type_from_url(self) -> str: + def _infer_media_type(self) -> str: """Return the media type of the file, based on the url.""" @property def media_type(self) -> str: """Return the media type of the file, based on the url or the provided `_media_type`.""" - if self._media_type is not None: - return self._media_type - else: - return self._infer_media_type_from_url() + return self._media_type or self._infer_media_type() @property @abstractmethod @@ -139,7 +141,7 @@ class VideoUrl(FileUrl): kind: Literal['video-url'] = 'video-url' """Type identifier, this is available on all parts as a discriminator.""" - def _infer_media_type_from_url(self) -> VideoMediaType: + def _infer_media_type(self) -> VideoMediaType: """Return the media type of the video, based on the url.""" if self.url.endswith('.mkv'): return 'video/x-matroska' @@ -189,7 +191,7 @@ class AudioUrl(FileUrl): kind: Literal['audio-url'] = 'audio-url' """Type identifier, this is available on all parts as a discriminator.""" - def _infer_media_type_from_url(self) -> AudioMediaType: + def _infer_media_type(self) -> AudioMediaType: """Return the media type of the audio file, based on the url. References: @@ -226,7 +228,7 @@ class ImageUrl(FileUrl): kind: Literal['image-url'] = 'image-url' """Type identifier, this is available on all parts as a discriminator.""" - def _infer_media_type_from_url(self) -> ImageMediaType: + def _infer_media_type(self) -> ImageMediaType: """Return the media type of the image, based on the url.""" if self.url.endswith(('.jpg', '.jpeg')): return 'image/jpeg' @@ -258,7 +260,7 @@ class DocumentUrl(FileUrl): kind: Literal['document-url'] = 'document-url' """Type identifier, this is available on all parts as a discriminator.""" - def _infer_media_type_from_url(self) -> str: + def _infer_media_type(self) -> str: """Return the media type of the document, based on the url.""" type_, _ = guess_type(self.url) if type_ is None: diff --git a/tests/test_messages.py b/tests/test_messages.py index d1b3b90ac..f0aa7f073 100644 --- a/tests/test_messages.py +++ b/tests/test_messages.py @@ -17,7 +17,7 @@ def test_image_url(): assert image_url.media_type == 'image/jpeg' assert image_url.format == 'jpeg' - image_url = ImageUrl(url='https://example.com/image', _media_type='image/jpeg') + image_url = ImageUrl(url='https://example.com/image', media_type_='image/jpeg') assert image_url.media_type == 'image/jpeg' assert image_url.format == 'jpeg' @@ -27,7 +27,7 @@ def test_video_url(): assert video_url.media_type == 'video/mp4' assert video_url.format == 'mp4' - video_url = VideoUrl(url='https://example.com/video', _media_type='video/mp4') + video_url = VideoUrl(url='https://example.com/video', media_type_='video/mp4') assert video_url.media_type == 'video/mp4' assert video_url.format == 'mp4' @@ -53,7 +53,7 @@ def test_document_url(): assert document_url.media_type == 'application/pdf' assert document_url.format == 'pdf' - document_url = DocumentUrl(url='https://example.com/document', _media_type='application/pdf') + document_url = DocumentUrl(url='https://example.com/document', media_type_='application/pdf') assert document_url.media_type == 'application/pdf' assert document_url.format == 'pdf' From b610d25b8478d6b113ab84c021bd72b2197fe786 Mon Sep 17 00:00:00 2001 From: David Date: Wed, 23 Jul 2025 22:26:25 -0400 Subject: [PATCH 4/4] Define custom __init__ for FIleUrl classes to match media_type override parameter name with property name --- pydantic_ai_slim/pydantic_ai/messages.py | 72 ++++++++++++++++++++---- tests/test_messages.py | 8 +-- 2 files changed, 65 insertions(+), 15 deletions(-) diff --git a/pydantic_ai_slim/pydantic_ai/messages.py b/pydantic_ai_slim/pydantic_ai/messages.py index baf3f6752..4f98d995a 100644 --- a/pydantic_ai_slim/pydantic_ai/messages.py +++ b/pydantic_ai_slim/pydantic_ai/messages.py @@ -3,7 +3,7 @@ import base64 from abc import ABC, abstractmethod from collections.abc import Sequence -from dataclasses import InitVar, dataclass, field, replace +from dataclasses import dataclass, field, replace from datetime import datetime from mimetypes import guess_type from typing import TYPE_CHECKING, Annotated, Any, Literal, Union, cast, overload @@ -85,16 +85,13 @@ def otel_event(self, settings: InstrumentationSettings) -> Event: __repr__ = _utils.dataclasses_no_defaults_repr -@dataclass(repr=False) +@dataclass(init=False, repr=False) class FileUrl(ABC): """Abstract base class for any URL-based file.""" url: str """The URL of the file.""" - media_type_: InitVar[str | None] = None - """Optional override for the media type of the file, in case it cannot be inferred from the URL.""" - force_download: bool = False """If the model supports it: @@ -111,8 +108,17 @@ class FileUrl(ABC): _media_type: str | None = field(init=False, repr=False) - def __post_init__(self, media_type_: str | None = None) -> None: - self._media_type = media_type_ + def __init__( + self, + url: str, + force_download: bool = False, + vendor_metadata: dict[str, Any] | None = None, + media_type: str | None = None, + ) -> None: + self.url = url + self.vendor_metadata = vendor_metadata + self.force_download = force_download + self._media_type = media_type @abstractmethod def _infer_media_type(self) -> str: @@ -131,7 +137,7 @@ def format(self) -> str: __repr__ = _utils.dataclasses_no_defaults_repr -@dataclass(repr=False) +@dataclass(init=False, repr=False) class VideoUrl(FileUrl): """A URL to a video.""" @@ -141,6 +147,17 @@ class VideoUrl(FileUrl): kind: Literal['video-url'] = 'video-url' """Type identifier, this is available on all parts as a discriminator.""" + def __init__( + self, + url: str, + force_download: bool = False, + vendor_metadata: dict[str, Any] | None = None, + media_type: str | None = None, + kind: Literal['video-url'] = 'video-url', + ) -> None: + super().__init__(url=url, force_download=force_download, vendor_metadata=vendor_metadata, media_type=media_type) + self.kind = kind + def _infer_media_type(self) -> VideoMediaType: """Return the media type of the video, based on the url.""" if self.url.endswith('.mkv'): @@ -181,7 +198,7 @@ def format(self) -> VideoFormat: return _video_format_lookup[self.media_type] -@dataclass(repr=False) +@dataclass(init=False, repr=False) class AudioUrl(FileUrl): """A URL to an audio file.""" @@ -191,6 +208,17 @@ class AudioUrl(FileUrl): kind: Literal['audio-url'] = 'audio-url' """Type identifier, this is available on all parts as a discriminator.""" + def __init__( + self, + url: str, + force_download: bool = False, + vendor_metadata: dict[str, Any] | None = None, + media_type: str | None = None, + kind: Literal['audio-url'] = 'audio-url', + ) -> None: + super().__init__(url=url, force_download=force_download, vendor_metadata=vendor_metadata, media_type=media_type) + self.kind = kind + def _infer_media_type(self) -> AudioMediaType: """Return the media type of the audio file, based on the url. @@ -218,7 +246,7 @@ def format(self) -> AudioFormat: return _audio_format_lookup[self.media_type] -@dataclass(repr=False) +@dataclass(init=False, repr=False) class ImageUrl(FileUrl): """A URL to an image.""" @@ -228,6 +256,17 @@ class ImageUrl(FileUrl): kind: Literal['image-url'] = 'image-url' """Type identifier, this is available on all parts as a discriminator.""" + def __init__( + self, + url: str, + force_download: bool = False, + vendor_metadata: dict[str, Any] | None = None, + media_type: str | None = None, + kind: Literal['image-url'] = 'image-url', + ) -> None: + super().__init__(url=url, force_download=force_download, vendor_metadata=vendor_metadata, media_type=media_type) + self.kind = kind + def _infer_media_type(self) -> ImageMediaType: """Return the media type of the image, based on the url.""" if self.url.endswith(('.jpg', '.jpeg')): @@ -250,7 +289,7 @@ def format(self) -> ImageFormat: return _image_format_lookup[self.media_type] -@dataclass(repr=False) +@dataclass(init=False, repr=False) class DocumentUrl(FileUrl): """The URL of the document.""" @@ -260,6 +299,17 @@ class DocumentUrl(FileUrl): kind: Literal['document-url'] = 'document-url' """Type identifier, this is available on all parts as a discriminator.""" + def __init__( + self, + url: str, + force_download: bool = False, + vendor_metadata: dict[str, Any] | None = None, + media_type: str | None = None, + kind: Literal['document-url'] = 'document-url', + ) -> None: + super().__init__(url=url, force_download=force_download, vendor_metadata=vendor_metadata, media_type=media_type) + self.kind = kind + def _infer_media_type(self) -> str: """Return the media type of the document, based on the url.""" type_, _ = guess_type(self.url) diff --git a/tests/test_messages.py b/tests/test_messages.py index f0aa7f073..e1f3a0b9f 100644 --- a/tests/test_messages.py +++ b/tests/test_messages.py @@ -17,7 +17,7 @@ def test_image_url(): assert image_url.media_type == 'image/jpeg' assert image_url.format == 'jpeg' - image_url = ImageUrl(url='https://example.com/image', media_type_='image/jpeg') + image_url = ImageUrl(url='https://example.com/image', media_type='image/jpeg') assert image_url.media_type == 'image/jpeg' assert image_url.format == 'jpeg' @@ -27,7 +27,7 @@ def test_video_url(): assert video_url.media_type == 'video/mp4' assert video_url.format == 'mp4' - video_url = VideoUrl(url='https://example.com/video', media_type_='video/mp4') + video_url = VideoUrl(url='https://example.com/video', media_type='video/mp4') assert video_url.media_type == 'video/mp4' assert video_url.format == 'mp4' @@ -53,7 +53,7 @@ def test_document_url(): assert document_url.media_type == 'application/pdf' assert document_url.format == 'pdf' - document_url = DocumentUrl(url='https://example.com/document', media_type_='application/pdf') + document_url = DocumentUrl(url='https://example.com/document', media_type='application/pdf') assert document_url.media_type == 'application/pdf' assert document_url.format == 'pdf' @@ -133,7 +133,7 @@ def test_binary_content_document(media_type: str, format: str): pytest.param(AudioUrl('foobar.flac'), 'audio/flac', 'flac', id='flac'), pytest.param(AudioUrl('foobar.aiff'), 'audio/aiff', 'aiff', id='aiff'), pytest.param(AudioUrl('foobar.aac'), 'audio/aac', 'aac', id='aac'), - pytest.param(AudioUrl('foobar', 'audio/mpeg'), 'audio/mpeg', 'mp3', id='mp3'), + pytest.param(AudioUrl('foobar', media_type='audio/mpeg'), 'audio/mpeg', 'mp3', id='mp3'), ], ) def test_audio_url(audio_url: AudioUrl, media_type: str, format: str):