diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 5ba5b86459a..9742dc01e32 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -24,7 +24,20 @@ from copy import deepcopy from datetime import datetime from inspect import isawaitable, ismethod -from typing import TYPE_CHECKING, Any, ClassVar, Final, Generic, Literal, NamedTuple, NewType, TypeAlias, cast, overload +from typing import ( + TYPE_CHECKING, + Any, + ClassVar, + Final, + Generic, + Literal, + NamedTuple, + NewType, + TypeAlias, + cast, + get_args, + overload, +) from typing_extensions import Required, TypedDict @@ -442,6 +455,8 @@ def _restore_compaction_annotation_in_additional_properties( "code_interpreter_tool_result", }) +_CONTENT_TYPES: Final[frozenset[str]] = frozenset(get_args(ContentType)) + class TextSpanRegion(TypedDict, total=False): """TypedDict representation of a text span region annotation.""" @@ -1906,6 +1921,13 @@ def from_dict(cls: type[ContentT], data: Mapping[str, Any]) -> ContentT: remaining["outputs"] = [cls.from_dict(item) if isinstance(item, dict) else item for item in output_items] # type: ignore[reportUnknownVariableType] if (content_items := remaining.get("items")) and isinstance(content_items, list): remaining["items"] = [cls.from_dict(item) if isinstance(item, dict) else item for item in content_items] # type: ignore[reportUnknownVariableType] + # A single nested Content (e.g., outputs in image_generation_tool_result) and Content items in a tool output + # (e.g., output in mcp_server_tool_result). These fields can also hold plain values such as tool JSON. + # A list in outputs was converted above and passes through unchanged. + if (single_output := remaining.get("outputs")) is not None: + remaining["outputs"] = cls._from_serialized_value(single_output) + if (tool_output := remaining.get("output")) and isinstance(tool_output, list): + remaining["output"] = [cls._from_serialized_value(item) for item in cast(list[Any], tool_output)] return cls( type=content_type, @@ -1915,6 +1937,30 @@ def from_dict(cls: type[ContentT], data: Mapping[str, Any]) -> ContentT: **remaining, ) + @classmethod + def _from_serialized_value(cls: type[ContentT], value: Any) -> Any: + """Rebuild a Content from its serialized mapping, or return the value unchanged. + + A mapping is rebuilt only when it has a known content type and the rebuilt Content serializes back to the + same mapping, with or without ``None`` values. Any other value, such as tool JSON that only shares a + content type name, is kept as it is. A mapping identical to a serialized Content (including its + ``additional_properties``) cannot be told apart from one and is rebuilt. + """ + if not isinstance(value, dict): + return value + mapping = cast(dict[str, Any], value) + content_type = mapping.get("type") + if not isinstance(content_type, str) or content_type not in _CONTENT_TYPES: + return mapping + try: + content = cls.from_dict(mapping) + if content.to_dict() == mapping or content.to_dict(exclude_none=False) == mapping: + return content + except Exception: + # The mapping is not a serialized Content (for example tool JSON with other fields), so keep it. + return mapping + return mapping + def __add__(self, other: Content) -> Content: """Concatenate or merge two Content instances.""" if not isinstance(other, Content): diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index ee936e3f9c2..8ae2b371da3 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -4705,6 +4705,77 @@ def test_function_result_items_roundtrip_via_dict(): assert restored.items[1].type == "data" +@pytest.mark.parametrize("serialization", ["dict", "dict_with_none", "json"]) +def test_tool_result_nested_content_roundtrips(serialization: str): + """Test that nested Content in image generation and MCP server tool results survives a round trip.""" + image = Content.from_uri("data:image/png;base64,iVBORw0KGgo=") + output = [ + Content.from_text("42", annotations=[{"type": "citation", "url": "https://example.com/source"}]), + Content.from_uri("https://example.com/report.pdf", media_type="application/pdf"), + Content.from_hosted_file("file_1"), + Content.from_function_result(call_id="call_1", result=[Content.from_text("done")]), + ] + message = Message( + "assistant", + [ + Content.from_image_generation_tool_result(image_id="ig_1", outputs=image), + Content.from_mcp_server_tool_result("mcp_1", output=output), + ], + ) + + if serialization == "dict": + restored = Message.from_dict(message.to_dict()) + elif serialization == "dict_with_none": + restored = Message.from_dict(message.to_dict(exclude_none=False)) + else: + restored = Message.from_json(message.to_json()) + restored_image, restored_mcp = restored.contents + + assert isinstance(restored_image.outputs, Content) + assert restored_image.outputs == image + assert restored_mcp.output is not None + assert all(isinstance(item, Content) for item in restored_mcp.output) + assert restored_mcp.output == output + + +@pytest.mark.parametrize( + "tool_json", + [ + {"type": "text", "value": 1}, + {"type": "text", "text": "x"}, + {"type": "data", "data": "AAAA", "media_type": "image/png"}, + {"type": "text", "text": "x", "annotations": {"audience": ["user"]}}, + {"type": "text", "function_call": {"a": 1}}, + ], + ids=["unknown_field", "text_shape", "data_with_str_data", "mcp_annotations", "invalid_nested_content"], +) +def test_tool_result_json_with_content_type_name_roundtrips(tool_json: dict[str, Any]): + """Test that tool JSON whose type is a content type name is kept as it is in a round trip.""" + mcp_result = Content.from_mcp_server_tool_result("mcp_1", output=[Content.from_text("42"), tool_json]) + image_result = Content.from_image_generation_tool_result(image_id="ig_1", outputs=tool_json) + + restored_mcp = Content.from_dict(mcp_result.to_dict()) + restored_image = Content.from_dict(image_result.to_dict()) + + assert restored_mcp.output == [Content.from_text("42"), tool_json] + assert isinstance(restored_mcp.output[1], dict) + assert restored_image.outputs == tool_json + assert restored_mcp.to_dict()["output"][1] == tool_json + assert restored_image.to_dict()["outputs"] == tool_json + + +def test_tool_result_plain_outputs_roundtrip_via_dict(): + """Test that output values that are not serialized Content stay unchanged in a round trip.""" + output = ["plain", {"type": "object", "value": 1}, {"type": ["x"]}, {"answer": 42}] + mcp_result = Content.from_mcp_server_tool_result("mcp_1", output=output) + image_result = Content.from_image_generation_tool_result( + image_id="ig_1", outputs={"url": "https://example.com/a.png"} + ) + + assert Content.from_dict(mcp_result.to_dict()).output == output + assert Content.from_dict(image_result.to_dict()).outputs == {"url": "https://example.com/a.png"} + + def test_from_function_result_with_non_content_list(): """Test Content.from_function_result with a list of non-Content objects falls back to str.""" result = Content.from_function_result(call_id="test-789", result=["hello", "world"])