diff --git a/python/packages/core/agent_framework/_tools.py b/python/packages/core/agent_framework/_tools.py index 08f00fefc1..169d8c1c51 100644 --- a/python/packages/core/agent_framework/_tools.py +++ b/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. diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index 771e94fb79..8616dbb857 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/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