add identifier field to BinaryContent class (#2231)

Co-authored-by: Douwe Maan <[email protected]>

Author: GDaamn

pydantic_ai_slim/pydantic_ai/_agent_graph.py

@@ -759,7 +759,7 @@ def process_content(content: Any) -> Any:
                 )
             elif isinstance(content, _messages.MultiModalContentTypes):
                 if isinstance(content, _messages.BinaryContent):
-                    identifier = multi_modal_content_identifier(content.data)
+                    identifier = content.identifier or multi_modal_content_identifier(content.data)
                 else:
                     identifier = multi_modal_content_identifier(content.url)
 

pydantic_ai_slim/pydantic_ai/messages.py

@@ -282,6 +282,14 @@ class BinaryContent:
     media_type: AudioMediaType | ImageMediaType | DocumentMediaType | str
     """The media type of the binary data."""
 
+    identifier: str | None = None
+    """Identifier for the binary content, such as a URL or 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.
 

tests/test_agent.py

@@ -2938,7 +2938,13 @@ def test_binary_content_all_messages_json():
                     {
                         'content': [
                             'Hello',
-                            {'data': 'SGVsbG8=', 'media_type': 'text/plain', 'vendor_metadata': None, 'kind': 'binary'},
+                            {
+                                'data': 'SGVsbG8=',
+                                'media_type': 'text/plain',
+                                'vendor_metadata': None,
+                                'kind': 'binary',
+                                'identifier': None,
+                            },
                         ],
                         'timestamp': IsStr(),
                         'part_kind': 'user-prompt',
@@ -2973,7 +2979,7 @@ def test_binary_content_all_messages_json():
 def test_tool_return_part_binary_content_serialization():
     """Test that ToolReturnPart can properly serialize BinaryContent."""
     png_data = b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc```\x00\x00\x00\x04\x00\x01\xf6\x178\x00\x00\x00\x00IEND\xaeB`\x82'
-    binary_content = BinaryContent(png_data, media_type='image/png')
+    binary_content = BinaryContent(png_data, media_type='image/png', identifier='image_id_1')
 
     tool_return = ToolReturnPart(tool_name='test_tool', content=binary_content, tool_call_id='test_call_123')
 
@@ -2982,10 +2988,12 @@ def test_tool_return_part_binary_content_serialization():
     assert '"kind":"binary"' in response_str
     assert '"media_type":"image/png"' in response_str
     assert '"data":"' in response_str
+    assert '"identifier":"image_id_1"' 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'] == 'image_id_1'
     assert 'data' in response_obj['return_value']
 
 
@@ -3011,6 +3019,50 @@ def get_image() -> BinaryContent:
     assert result.output == 'Image received'
 
 
+def test_tool_returning_binary_content_with_identifier():
+    """Test that a tool returning BinaryContent directly works correctly."""
+
+    def llm(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
+        if len(messages) == 1:
+            return ModelResponse(parts=[ToolCallPart('get_image', {})])
+        else:
+            return ModelResponse(parts=[TextPart('Image received')])
+
+    agent = Agent(FunctionModel(llm))
+
+    @agent.tool_plain
+    def get_image() -> BinaryContent:
+        """Return a simple image."""
+        png_data = b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc```\x00\x00\x00\x04\x00\x01\xf6\x178\x00\x00\x00\x00IEND\xaeB`\x82'
+        return BinaryContent(png_data, media_type='image/png', identifier='image_id_1')
+
+    # This should work without the serialization error
+    result = agent.run_sync('Get an image')
+    assert result.all_messages()[2] == snapshot(
+        ModelRequest(
+            parts=[
+                ToolReturnPart(
+                    tool_name='get_image',
+                    content='See file image_id_1',
+                    tool_call_id=IsStr(),
+                    timestamp=IsNow(tz=timezone.utc),
+                ),
+                UserPromptPart(
+                    content=[
+                        'This is file image_id_1:',
+                        BinaryContent(
+                            data=b'\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc```\x00\x00\x00\x04\x00\x01\xf6\x178\x00\x00\x00\x00IEND\xaeB`\x82',
+                            media_type='image/png',
+                            identifier='image_id_1',
+                        ),
+                    ],
+                    timestamp=IsNow(tz=timezone.utc),
+                ),
+            ]
+        )
+    )
+
+
 def test_instructions_raise_error_when_system_prompt_is_set():
     agent = Agent('test', instructions='An instructions!')