diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 17dba475e3e..4b8572e884a 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -1535,8 +1535,17 @@ def from_dict(cls: type[ContentT], data: Mapping[str, Any]) -> ContentT: # Special handling for DataContent with data and media_type if content_type == "data" and "data" in remaining and "media_type" in remaining: - # Use from_data() to properly create the DataContent with URI - return cls.from_data(remaining["data"], remaining["media_type"]) + # Use from_data() to properly create the DataContent with URI. The three + # fields popped above are passed on explicitly: this branch returns before + # the constructor below, so leaving them out dropped on a data content what + # every other content type keeps. + return cls.from_data( + remaining["data"], + remaining["media_type"], + annotations=annotations, + additional_properties=additional_properties, + raw_representation=raw_representation, + ) # Handle nested Content objects (e.g., function_call in function_approval_request) if (function_call := remaining.get("function_call")) and isinstance(function_call, dict): diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index a6797b30449..585a654a3df 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -2607,6 +2607,70 @@ def test_content_roundtrip_preserves_compaction_annotation_dict() -> None: assert annotation[GROUP_TOKEN_COUNT_KEY] is None +def test_content_from_dict_data_keeps_annotations_and_metadata() -> None: + """The `data` + `media_type` branch must not lose what every other type keeps. + + `from_dict` pops `annotations`, `additional_properties` and `raw_representation` + before dispatching, then returns early for a data content. `from_data` accepts + all three, but they were not being passed, so they were silently dropped -- and + this runs inside `Message.from_dict` and `ChatResponse.from_dict`. + """ + content = Content.from_dict({ + "type": "data", + "data": b"hello", + "media_type": "text/plain", + "annotations": [{"type": "citation", "title": "t"}], + "additional_properties": {"origin": "upload"}, + "raw_representation": "provider-blob", + }) + + assert content.type == "data" + assert content.uri == "data:text/plain;base64,aGVsbG8=" + assert content.annotations == [{"type": "citation", "title": "t"}] + assert content.additional_properties == {"origin": "upload"} + assert content.raw_representation == "provider-blob" + + +def test_content_from_dict_data_matches_other_types_on_metadata() -> None: + """Same mapping shape, two content types, same treatment of the extra fields.""" + common = { + "additional_properties": {"origin": "upload"}, + "raw_representation": "provider-blob", + } + data_content = Content.from_dict({ + "type": "data", + "data": b"hello", + "media_type": "text/plain", + **common, + }) + uri_content = Content.from_dict({ + "type": "uri", + "uri": "https://example.com/x.png", + "media_type": "image/png", + **common, + }) + + assert data_content.additional_properties == uri_content.additional_properties + assert data_content.raw_representation == uri_content.raw_representation + + +def test_message_from_dict_keeps_data_content_metadata() -> None: + """Through the reachable path: rebuilding a message from its mapping.""" + message = Message.from_dict({ + "role": "user", + "contents": [ + { + "type": "data", + "data": b"hello", + "media_type": "text/plain", + "additional_properties": {"origin": "upload"}, + } + ], + }) + + assert message.contents[0].additional_properties == {"origin": "upload"} + + def test_content_from_dict_via_json() -> None: """Test Content.from_dict with data parsed from a JSON string.""" data = json.loads(json.dumps({"type": "text", "text": "Hello world"}))