Python: fix(core): serialize Pydantic tool results in JSON mode

python/packages/core/agent_framework/_tools.py

@@ -500,6 +500,38 @@ def _annotation_includes_function_invocation_context(annotation: Any) -> bool:
 _ToolDescriptionFormat: TypeAlias = _ToolParameterFormat | Mapping[str, _ToolParameterFormat]
 _NormalizedToolDescriptionFormat: TypeAlias = _ToolParameterFormat | dict[str, _ToolParameterFormat]
 _TOOL_DESCRIPTION_FORMAT_ERROR = "tool_description_format must be 'compact', 'json', or a tool-name mapping."
+# model_dump's fallback (str for values JSON mode doesn't know) only exists on newer Pydantic 2 versions.
+_MODEL_DUMP_JSON_FALLBACK: dict[str, Any] = (
+    {"fallback": str} if "fallback" in inspect.signature(BaseModel.model_dump).parameters else {}
+)
+
+
+def _escape_unencodable(value: Any) -> Any:
+    """Copy a value with the parts JSON mode can't encode replaced, inside models, dicts and sequences too.
+
+    Bytes that aren't valid UTF-8 become their str() form. On Pydantic versions without ``fallback``,
+    values JSON mode doesn't know become str() too. Models are copied, not rebuilt, so their class stays
+    the same and the parent's declared field types still decide how they are serialized.
+    """
+    if isinstance(value, bytes):
+        try:
+            value.decode("utf-8")
+        except UnicodeDecodeError:
+            return str(value).encode()
+        return value
+    if isinstance(value, BaseModel):
+        return value.model_copy(update={name: _escape_unencodable(item) for name, item in value.__dict__.items()})
+    if isinstance(value, dict):
+        dict_value = cast(dict[object, object], value)
+        return {key: _escape_unencodable(item) for key, item in dict_value.items()}
+    if type(value) in (list, tuple, set, frozenset):
+        return type(value)(_escape_unencodable(item) for item in value)
+    if not _MODEL_DUMP_JSON_FALLBACK:
+        try:
+            to_jsonable_python(value)
+        except ValueError:
+            return str(value)
+    return value
 
 
 def _normalize_tool_description_format(  # pyright: ignore[reportUnusedFunction]
@@ -1372,13 +1404,50 @@ def _make_dumpable(value: Any) -> Any:
         if isinstance(value, Content):
             return value.to_dict(exclude={"raw_representation", "additional_properties"})
         if isinstance(value, BaseModel):
-            return value.model_dump()
+            return FunctionTool._dump_model(value)
         if hasattr(value, "to_dict"):
             return value.to_dict()
         if hasattr(value, "text") and isinstance(value.text, str):
             return value.text
         return value
 
+    @staticmethod
+    def _dump_model(model: BaseModel) -> Any:
+        """Dump a Pydantic model in JSON mode.
+
+        JSON mode applies enum, datetime and UUID encoding and JSON-only field serializers. When it
+        fails (bytes that aren't valid UTF-8, or an arbitrary type on Pydantic versions without
+        ``fallback``), those values are escaped in a copy of the model and Pydantic dumps the copy, so
+        the declared field types, serializers and exclusions still decide what is written. Computed
+        fields are dumped from the original model. A field that still can't be encoded is left out.
+        Excluded fields stay out either way, unlike the str(result) fallback in invoke.
+        """
+        try:
+            return model.model_dump(mode="json", **_MODEL_DUMP_JSON_FALLBACK)
+        except ValueError:
+            pass
+        escaped = _escape_unencodable(model)
+        # The python-mode dump gives the serialized field names without the excluded ones. The class's
+        # model_computed_fields can't be used for that, it's a property on Pydantic 2.0.
+        python_dump = model.model_dump(by_alias=False)
+        if not isinstance(python_dump, dict):
+            try:
+                return escaped.model_dump(mode="json", **_MODEL_DUMP_JSON_FALLBACK)
+            except ValueError:
+                logger.warning("Leaving %s out of the tool result: it can't be encoded as JSON.", type(model).__name__)
+                return None
+        fields = type(model).model_fields
+        dumped: dict[str, Any] = {}
+        for name in python_dump:
+            source = escaped if name in fields else model
+            try:
+                dumped.update(source.model_dump(mode="json", include={name}, **_MODEL_DUMP_JSON_FALLBACK))
+            except ValueError:
+                logger.warning(
+                    "Leaving %s.%s out of the tool result: it can't be encoded as JSON.", type(model).__name__, name
+                )
+        return dumped
+
     @staticmethod
     def parse_result(result: Any) -> list[Content]:
         """Convert a raw function return value to a list of Content items.

python/packages/core/tests/core/test_types.py

@@ -4527,6 +4527,202 @@ def test_parse_result_nested_pydantic_model():
     assert "18.0" in parsed[0].text or "18" in parsed[0].text  # type: ignore[operator]  # pyrefly: ignore[not-iterable]  # ty: ignore[unsupported-operator]
 
 
+def test_parse_result_pydantic_model_uses_json_mode():
+    """Enum and datetime fields are serialized the way Pydantic writes them to JSON."""
+    from enum import Enum
+
+    class Status(Enum):
+        OPEN = "open"
+
+    class Ticket(BaseModel):
+        status: Status
+        created: datetime
+
+    parsed = FunctionTool.parse_result(Ticket(status=Status.OPEN, created=datetime(2024, 1, 1, 12)))
+
+    assert json.loads(parsed[0].text) == {"status": "open", "created": "2024-01-01T12:00:00"}  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+
+def test_parse_result_pydantic_model_with_arbitrary_type():
+    """Only the field Pydantic cannot serialize to JSON falls back to str; the others stay in JSON mode."""
+    from pydantic import field_serializer
+
+    class Handle:
+        def __str__(self) -> str:
+            return "handle-1"
+
+    class Result(BaseModel):
+        model_config = {"arbitrary_types_allowed": True}
+
+        handle: Handle
+        created: datetime
+        token: str
+
+        @field_serializer("token", when_used="json")
+        def _redact(self, value: str) -> str:
+            return "***"
+
+    parsed = FunctionTool.parse_result(Result(handle=Handle(), created=datetime(2024, 1, 1, 12), token="secret"))
+
+    assert json.loads(parsed[0].text) == {"handle": "handle-1", "created": "2024-01-01T12:00:00", "token": "***"}  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+
+def test_parse_result_pydantic_model_with_non_utf8_bytes_falls_back_per_field():
+    """Only the bytes fall back; JSON-only serializers still run and excluded fields stay out."""
+    from pydantic import Field, field_serializer
+
+    class Result(BaseModel):
+        data: bytes
+        created: datetime
+        token: str
+        api_key: str = Field(exclude=True)
+
+        @field_serializer("token", when_used="json")
+        def _redact(self, value: str) -> str:
+            return "***"
+
+    result = Result(data=b"\xff\xfe", created=datetime(2024, 1, 1, 12), token="secret", api_key="key")
+    parsed = FunctionTool.parse_result(result)
+
+    assert json.loads(parsed[0].text) == {  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        "data": str(b"\xff\xfe"),
+        "created": "2024-01-01T12:00:00",
+        "token": "***",
+    }
+
+
+def test_parse_result_nested_pydantic_model_with_non_utf8_bytes():
+    """A nested model with bytes JSON mode cannot encode keeps its JSON-only serializers."""
+    from pydantic import field_serializer
+
+    class Inner(BaseModel):
+        data: bytes
+        token: str
+
+        @field_serializer("token", when_used="json")
+        def _redact(self, value: str) -> str:
+            return "***"
+
+    class Outer(BaseModel):
+        inner: Inner
+        created: datetime
+
+    parsed = FunctionTool.parse_result(
+        Outer(inner=Inner(data=b"\xff\xfe", token="secret"), created=datetime(2024, 1, 1, 12))
+    )
+
+    assert json.loads(parsed[0].text) == {  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        "inner": {"data": str(b"\xff\xfe"), "token": "***"},
+        "created": "2024-01-01T12:00:00",
+    }
+
+
+def test_parse_result_pydantic_model_with_non_utf8_bytes_keeps_computed_fields():
+    """The per-field fallback includes computed fields and leaves excluded fields out."""
+    from pydantic import Field, computed_field
+
+    class Result(BaseModel):
+        data: bytes
+        api_key: str = Field(exclude=True)
+
+        @computed_field  # type: ignore[prop-decorator]
+        @property
+        def size(self) -> int:
+            return len(self.data)
+
+    parsed = FunctionTool.parse_result(Result(data=b"\xff\xfe", api_key="key"))
+
+    assert json.loads(parsed[0].text) == {"data": str(b"\xff\xfe"), "size": 2}  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+
+def test_parse_result_pydantic_models_in_list_and_dict_with_non_utf8_bytes():
+    """Models inside a list or dict field that falls back keep their JSON-only serializers and exclusions."""
+    from pydantic import Field, field_serializer
+
+    class Inner(BaseModel):
+        data: bytes
+        created: datetime
+        token: str
+        api_key: str = Field(exclude=True)
+
+        @field_serializer("token", when_used="json")
+        def _redact(self, value: str) -> str:
+            return "***"
+
+    class Outer(BaseModel):
+        items: list[Inner]
+        by_name: dict[str, Inner]
+
+    inner = Inner(data=b"\xff\xfe", created=datetime(2024, 1, 1, 12), token="secret", api_key="key")
+    parsed = FunctionTool.parse_result(Outer(items=[inner], by_name={"a": inner}))
+
+    expected_inner = {"data": str(b"\xff\xfe"), "created": "2024-01-01T12:00:00", "token": "***"}
+    assert json.loads(parsed[0].text) == {"items": [expected_inner], "by_name": {"a": expected_inner}}  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+
+def test_parse_result_pydantic_field_serializer_output_with_non_utf8_bytes():
+    """A JSON-only serializer whose output has bytes JSON mode cannot encode still runs; only the bytes fall back."""
+    from pydantic import field_serializer
+
+    class Result(BaseModel):
+        info: dict[str, Any]
+        created: datetime
+
+        @field_serializer("info", when_used="json")
+        def _redact(self, value: dict[str, Any]) -> dict[str, Any]:
+            return {**value, "token": "***"}
+
+    result = Result(info={"token": "secret", "data": b"\xff\xfe"}, created=datetime(2024, 1, 1, 12))
+    parsed = FunctionTool.parse_result(result)
+
+    assert json.loads(parsed[0].text) == {  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        "info": {"token": "***", "data": str(b"\xff\xfe")},
+        "created": "2024-01-01T12:00:00",
+    }
+
+
+def test_parse_result_pydantic_field_serializer_output_that_cannot_be_encoded_is_left_out():
+    """A field whose serializer output JSON mode cannot encode is left out instead of sent without the serializer."""
+    from pydantic import field_serializer
+
+    class Result(BaseModel):
+        token: str
+        created: datetime
+
+        @field_serializer("token", when_used="json")
+        def _redact(self, value: str) -> bytes:
+            return b"\xff***"
+
+    parsed = FunctionTool.parse_result(Result(token="secret", created=datetime(2024, 1, 1, 12)))
+
+    assert json.loads(parsed[0].text) == {"created": "2024-01-01T12:00:00"}  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+
+
+def test_parse_result_pydantic_fallback_keeps_declared_field_types():
+    """A subclass in a field declared as its base keeps the subclass-only fields out when the fallback is used."""
+
+    class Public(BaseModel):
+        data: bytes
+
+    class Secret(Public):
+        api_key: str
+
+    class Envelope(BaseModel):
+        item: Public
+        items: list[Public]
+        by_name: dict[str, Public]
+
+    secret = Secret(data=b"\xff\xfe", api_key="key")
+    parsed = FunctionTool.parse_result(Envelope(item=secret, items=[secret], by_name={"a": secret}))
+
+    expected_item = {"data": str(b"\xff\xfe")}
+    assert json.loads(parsed[0].text) == {  # type: ignore[arg-type]  # pyrefly: ignore[bad-argument-type]  # ty: ignore[invalid-argument-type]
+        "item": expected_item,
+        "items": [expected_item],
+        "by_name": {"a": expected_item},
+    }
+
+
 # region FunctionTool.parse_result with MCP TextContent-like objects