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

python/packages/core/agent_framework/_tools.py

@@ -48,7 +48,7 @@
 from opentelemetry import trace
 from opentelemetry.metrics import Histogram, NoOpHistogram
 from opentelemetry.util.types import AttributeValue
-from pydantic import BaseModel, Field, TypeAdapter, ValidationError, create_model
+from pydantic import BaseModel, Field, PlainSerializer, TypeAdapter, ValidationError, WrapSerializer, create_model
 from pydantic_core import to_jsonable_python
 
 from ._serialization import SerializationMixin
@@ -500,6 +500,39 @@ 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 _has_field_serializer(model_type: type[BaseModel], name: str) -> bool:
+    """Whether a model field has its own serializer, from a decorator or an annotation."""
+    decorators = model_type.__pydantic_decorators__.field_serializers.values()
+    if any(name in decorator.info.fields or "*" in decorator.info.fields for decorator in decorators):
+        return True
+    field = model_type.model_fields.get(name)
+    return field is not None and any(
+        isinstance(metadata, (PlainSerializer, WrapSerializer)) for metadata in field.metadata
+    )
+
+
+def _escape_invalid_utf8(value: Any) -> Any:
+    """Replace bytes that aren't valid UTF-8 with their str() form, inside models, dicts and sequences too."""
+    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_invalid_utf8(item) for name, item in value.__dict__.items()})
+    if isinstance(value, dict):
+        dict_value = cast(dict[object, object], value)
+        return {key: _escape_invalid_utf8(item) for key, item in dict_value.items()}
+    if type(value) in (list, tuple, set, frozenset):
+        return type(value)(_escape_invalid_utf8(item) for item in value)
+    return value
 
 
 def _normalize_tool_description_format(  # pyright: ignore[reportUnusedFunction]
@@ -1372,13 +1405,80 @@ 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, falling back to python mode only for the values that need it.
+
+        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``), each field is dumped in JSON mode on its own and only a failing field falls
+        back. Models inside a failing field, directly or in a list or dict, are dumped the same way,
+        so their JSON serializers and excluded fields still apply. A failing field with its own
+        serializer is dumped in JSON mode again with its invalid UTF-8 bytes escaped, so the
+        serializer still runs, and is left out if that fails too. 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
+        # 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):
+            return model.model_dump()
+        model_type = type(model)
+        dumped: dict[str, Any] = {}
+        for name in python_dump:
+            try:
+                dumped.update(model.model_dump(mode="json", include={name}, **_MODEL_DUMP_JSON_FALLBACK))
+                continue
+            except ValueError:
+                pass
+            value = getattr(model, name, None)
+            if _has_field_serializer(model_type, name):
+                # The python-mode value would skip the field's serializer, and with it any redaction it does.
+                escaped = model.model_copy(update={name: _escape_invalid_utf8(value)})
+                try:
+                    dumped.update(escaped.model_dump(mode="json", include={name}, **_MODEL_DUMP_JSON_FALLBACK))
+                except ValueError:
+                    logger.warning(
+                        "Leaving %s.%s out of the tool result: its serializer's output can't be encoded as JSON.",
+                        model_type.__name__,
+                        name,
+                    )
+                continue
+            for key, python_value in model.model_dump(include={name}).items():
+                if isinstance(value, (BaseModel, dict, list, tuple, set, frozenset)):
+                    dumped[key] = FunctionTool._dump_fallback_value(value)
+                else:
+                    dumped[key] = python_value
+        return dumped
+
+    @staticmethod
+    def _dump_fallback_value(value: Any) -> Any:
+        """Dump a field value JSON mode could not encode, keeping JSON mode for what inside it can be."""
+        if isinstance(value, BaseModel):
+            return FunctionTool._dump_model(value)
+        if isinstance(value, dict):
+            dict_value = cast(dict[object, object], value)
+            return {key: FunctionTool._dump_fallback_value(item) for key, item in dict_value.items()}
+        if isinstance(value, (list, tuple, set, frozenset)):
+            items = cast(Iterable[object], value)
+            return [FunctionTool._dump_fallback_value(item) for item in items]
+        try:
+            return to_jsonable_python(value)
+        except ValueError:
+            # parse_result's json.dumps(default=str) turns it into a string
+            return value
+
     @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,177 @@ 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]
+
+
 # region FunctionTool.parse_result with MCP TextContent-like objects