diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 08f00fefc11..d11417e7cfa 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/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. diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 771e94fb799..1d417e529c0 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/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