Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 65 additions & 41 deletions pydantic_ai_slim/pydantic_ai/messages.py
Original file line number Diff line number Diff line change
Expand Up @@ -114,20 +114,6 @@ class FileUrl(ABC):

_: KW_ONLY

identifier: str
"""The identifier of the file, such as a unique ID. generating one from the url if not explicitly set.

This identifier can be provided to the model in a message to allow it to refer to this file in a tool call argument,
and the tool can look up the file in question by iterating over the message history and finding the matching `FileUrl`.

This identifier is only automatically passed to the model when the `FileUrl` is returned by a tool.
If you're passing the `FileUrl` as a user message, it's up to you to include a separate text part with the identifier,
e.g. "This is file <identifier>:" preceding the `FileUrl`.

It's also included in inline-text delimiters for providers that require inlining text documents, so the model can
distinguish multiple files.
"""

force_download: bool = False
"""If the model supports it:

Expand All @@ -147,27 +133,48 @@ class FileUrl(ABC):
compare=False, default=None
)

_identifier: Annotated[str | None, pydantic.Field(alias='identifier', default=None, exclude=True)] = field(
compare=False, default=None
)

def __init__(
self,
url: str,
*,
force_download: bool = False,
vendor_metadata: dict[str, Any] | None = None,
media_type: str | None = None,
identifier: str | None = None,
force_download: bool = False,
vendor_metadata: dict[str, Any] | None = None,
) -> None:
self.url = url
self._media_type = media_type
self._identifier = identifier
self.force_download = force_download
self.vendor_metadata = vendor_metadata
self._media_type = media_type
self.identifier = identifier or _multi_modal_content_identifier(url)

@pydantic.computed_field
@property
def media_type(self) -> str:
"""Return the media type of the file, based on the URL or the provided `media_type`."""
return self._media_type or self._infer_media_type()

@pydantic.computed_field
@property
def identifier(self) -> str:
"""The identifier of the file, such as a unique ID.

This identifier can be provided to the model in a message to allow it to refer to this file in a tool call argument,
and the tool can look up the file in question by iterating over the message history and finding the matching `FileUrl`.

This identifier is only automatically passed to the model when the `FileUrl` is returned by a tool.
If you're passing the `FileUrl` as a user message, it's up to you to include a separate text part with the identifier,
e.g. "This is file <identifier>:" preceding the `FileUrl`.

It's also included in inline-text delimiters for providers that require inlining text documents, so the model can
distinguish multiple files.
"""
return self._identifier or _multi_modal_content_identifier(self.url)

@abstractmethod
def _infer_media_type(self) -> str:
"""Infer the media type of the file based on the URL."""
Expand Down Expand Up @@ -198,20 +205,21 @@ def __init__(
self,
url: str,
*,
media_type: str | None = None,
identifier: str | None = None,
force_download: bool = False,
vendor_metadata: dict[str, Any] | None = None,
media_type: str | None = None,
kind: Literal['video-url'] = 'video-url',
identifier: str | None = None,
# Required for inline-snapshot which expects all dataclass `__init__` methods to take all field names as kwargs.
_media_type: str | None = None,
_identifier: str | None = None,
) -> None:
super().__init__(
url=url,
force_download=force_download,
vendor_metadata=vendor_metadata,
media_type=media_type or _media_type,
identifier=identifier,
identifier=identifier or _identifier,
)
self.kind = kind

Expand Down Expand Up @@ -273,20 +281,21 @@ def __init__(
self,
url: str,
*,
media_type: str | None = None,
identifier: str | None = None,
force_download: bool = False,
vendor_metadata: dict[str, Any] | None = None,
media_type: str | None = None,
kind: Literal['audio-url'] = 'audio-url',
identifier: str | None = None,
# Required for inline-snapshot which expects all dataclass `__init__` methods to take all field names as kwargs.
_media_type: str | None = None,
_identifier: str | None = None,
) -> None:
super().__init__(
url=url,
force_download=force_download,
vendor_metadata=vendor_metadata,
media_type=media_type or _media_type,
identifier=identifier,
identifier=identifier or _identifier,
)
self.kind = kind

Expand Down Expand Up @@ -335,20 +344,21 @@ def __init__(
self,
url: str,
*,
media_type: str | None = None,
identifier: str | None = None,
force_download: bool = False,
vendor_metadata: dict[str, Any] | None = None,
media_type: str | None = None,
kind: Literal['image-url'] = 'image-url',
identifier: str | None = None,
# Required for inline-snapshot which expects all dataclass `__init__` methods to take all field names as kwargs.
_media_type: str | None = None,
_identifier: str | None = None,
) -> None:
super().__init__(
url=url,
force_download=force_download,
vendor_metadata=vendor_metadata,
media_type=media_type or _media_type,
identifier=identifier,
identifier=identifier or _identifier,
)
self.kind = kind

Expand Down Expand Up @@ -392,20 +402,21 @@ def __init__(
self,
url: str,
*,
media_type: str | None = None,
identifier: str | None = None,
force_download: bool = False,
vendor_metadata: dict[str, Any] | None = None,
media_type: str | None = None,
kind: Literal['document-url'] = 'document-url',
identifier: str | None = None,
# Required for inline-snapshot which expects all dataclass `__init__` methods to take all field names as kwargs.
_media_type: str | None = None,
_identifier: str | None = None,
) -> None:
super().__init__(
url=url,
force_download=force_download,
vendor_metadata=vendor_metadata,
media_type=media_type or _media_type,
identifier=identifier,
identifier=identifier or _identifier,
)
self.kind = kind

Expand Down Expand Up @@ -460,16 +471,6 @@ class BinaryContent:
media_type: AudioMediaType | ImageMediaType | DocumentMediaType | str
"""The media type of the binary data."""

identifier: str
"""Identifier for the binary content, such as a unique ID.
This identifier can be provided to the model in a message to allow it to refer to this file in a tool call argument,
and the tool can look up the file in question by iterating over the message history and finding the matching `BinaryContent`.

This identifier is only automatically passed to the model when the `BinaryContent` is returned by a tool.
If you're passing the `BinaryContent` as a user message, it's up to you to include a separate text part with the identifier,
e.g. "This is file <identifier>:" preceding the `BinaryContent`.
"""

vendor_metadata: dict[str, Any] | None = None
"""Vendor-specific metadata for the file.

Expand All @@ -478,6 +479,10 @@ class BinaryContent:
- `OpenAIChatModel`, `OpenAIResponsesModel`: `BinaryContent.vendor_metadata['detail']` is used as `detail` setting for images
"""

_identifier: Annotated[str | None, pydantic.Field(alias='identifier', default=None, exclude=True)] = field(
compare=False, default=None
)

kind: Literal['binary'] = 'binary'
"""Type identifier, this is available on all parts as a discriminator."""

Expand All @@ -489,10 +494,12 @@ def __init__(
identifier: str | None = None,
vendor_metadata: dict[str, Any] | None = None,
kind: Literal['binary'] = 'binary',
# Required for inline-snapshot which expects all dataclass `__init__` methods to take all field names as kwargs.
_identifier: str | None = None,
) -> None:
self.data = data
self.media_type = media_type
self.identifier = identifier or _multi_modal_content_identifier(data)
self._identifier = identifier or _identifier
self.vendor_metadata = vendor_metadata
self.kind = kind

Expand All @@ -518,6 +525,23 @@ def from_data_uri(cls, data_uri: str) -> Self:
media_type, data = data_uri[len(prefix) :].split(';base64,', 1)
return cls(data=base64.b64decode(data), media_type=media_type)

@pydantic.computed_field
@property
def identifier(self) -> str:
"""Identifier for the binary content, such as a unique ID.

This identifier can be provided to the model in a message to allow it to refer to this file in a tool call argument,
and the tool can look up the file in question by iterating over the message history and finding the matching `BinaryContent`.

This identifier is only automatically passed to the model when the `BinaryContent` is returned by a tool.
If you're passing the `BinaryContent` as a user message, it's up to you to include a separate text part with the identifier,
e.g. "This is file <identifier>:" preceding the `BinaryContent`.

It's also included in inline-text delimiters for providers that require inlining text documents, so the model can
distinguish multiple files.
"""
return self._identifier or _multi_modal_content_identifier(self.data)

@property
def data_uri(self) -> str:
"""Convert the `BinaryContent` to a data URI."""
Expand Down
23 changes: 11 additions & 12 deletions tests/test_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -3537,18 +3537,17 @@ def test_tool_return_part_binary_content_serialization():

tool_return = ToolReturnPart(tool_name='test_tool', content=binary_content, tool_call_id='test_call_123')

response_str = tool_return.model_response_str()

assert '"kind":"binary"' in response_str
assert '"media_type":"image/png"' in response_str
assert '"data":"' in response_str
assert '"identifier":"14a01a"' in response_str

response_obj = tool_return.model_response_object()
assert response_obj['return_value']['kind'] == 'binary'
assert response_obj['return_value']['media_type'] == 'image/png'
assert response_obj['return_value']['identifier'] == '14a01a'
assert 'data' in response_obj['return_value']
assert tool_return.model_response_object() == snapshot(
{
'return_value': {
'data': 'iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADElEQVR4nGNgYGAAAAAEAAH2FzgAAAAASUVORK5CYII=',
'media_type': 'image/png',
'vendor_metadata': None,
'_identifier': None,
'kind': 'binary',
}
}
)


def test_tool_returning_binary_content_directly():
Expand Down
69 changes: 69 additions & 0 deletions tests/test_messages.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

import pytest
from inline_snapshot import snapshot
from pydantic import TypeAdapter

from pydantic_ai import (
AudioUrl,
Expand Down Expand Up @@ -515,3 +516,71 @@ def test_model_response_convenience_methods():
)
]
)


def test_image_url_validation_with_optional_identifier():
image_url_ta = TypeAdapter(ImageUrl)
image = image_url_ta.validate_python({'url': 'https://example.com/image.jpg'})
assert image.url == snapshot('https://example.com/image.jpg')
assert image.identifier == snapshot('39cfc4')
assert image.media_type == snapshot('image/jpeg')
assert image_url_ta.dump_python(image) == snapshot(
{
'url': 'https://example.com/image.jpg',
'force_download': False,
'vendor_metadata': None,
'kind': 'image-url',
'media_type': 'image/jpeg',
'identifier': '39cfc4',
}
)

image = image_url_ta.validate_python(
{'url': 'https://example.com/image.jpg', 'identifier': 'foo', 'media_type': 'image/png'}
)
assert image.url == snapshot('https://example.com/image.jpg')
assert image.identifier == snapshot('foo')
assert image.media_type == snapshot('image/png')
assert image_url_ta.dump_python(image) == snapshot(
{
'url': 'https://example.com/image.jpg',
'force_download': False,
'vendor_metadata': None,
'kind': 'image-url',
'media_type': 'image/png',
'identifier': 'foo',
}
)


def test_binary_content_validation_with_optional_identifier():
binary_content_ta = TypeAdapter(BinaryContent)
binary_content = binary_content_ta.validate_python({'data': b'fake', 'media_type': 'image/jpeg'})
assert binary_content.data == b'fake'
assert binary_content.identifier == snapshot('c053ec')
assert binary_content.media_type == snapshot('image/jpeg')
assert binary_content_ta.dump_python(binary_content) == snapshot(
{
'data': b'fake',
'vendor_metadata': None,
'kind': 'binary',
'media_type': 'image/jpeg',
'identifier': 'c053ec',
}
)

binary_content = binary_content_ta.validate_python(
{'data': b'fake', 'identifier': 'foo', 'media_type': 'image/png'}
)
assert binary_content.data == b'fake'
assert binary_content.identifier == snapshot('foo')
assert binary_content.media_type == snapshot('image/png')
assert binary_content_ta.dump_python(binary_content) == snapshot(
{
'data': b'fake',
'vendor_metadata': None,
'kind': 'binary',
'media_type': 'image/png',
'identifier': 'foo',
}
)