diff --git a/python/packages/core/agent_framework/_serialization.py b/python/packages/core/agent_framework/_serialization.py index 384fa05a5b..fc383ded5f 100644 --- a/python/packages/core/agent_framework/_serialization.py +++ b/python/packages/core/agent_framework/_serialization.py @@ -151,6 +151,43 @@ def _is_serialization_protocol(value: Any) -> TypeGuard[SerializationProtocol]: return callable(getattr(value, "to_dict", None)) and callable(getattr(value, "from_dict", None)) +def _iter_instance_fields(instance: Any) -> dict[str, Any]: + """Return attributes stored in ``__dict__`` or slots.""" + fields = dict(getattr(instance, "__dict__", {})) + for cls in type(instance).__mro__: + slots = cls.__dict__.get("__slots__", ()) + if isinstance(slots, str): + slots = (slots,) + for field_name in slots: + if field_name not in {"__dict__", "__weakref__"} and hasattr(instance, field_name): + fields[field_name] = getattr(instance, field_name) + return fields + + +def get_pickle_state(instance: Any, omitted_fields: set[str]) -> dict[str, Any]: + """Build pickle state while omitting runtime-only fields.""" + state = _iter_instance_fields(instance) + for field_name in omitted_fields: + state.pop(field_name, None) + return state + + +def restore_pickle_state( + instance: Any, + state: dict[str, Any] | tuple[dict[str, Any], dict[str, Any]], + omitted_fields: set[str], +) -> None: + """Restore dict- and slot-backed pickle state.""" + if isinstance(state, tuple): + dict_state, slot_state = state + state = {**dict_state, **slot_state} + for field_name, value in state.items(): + object.__setattr__(instance, field_name, value) + for field_name in omitted_fields: + if field_name in _iter_instance_fields(instance) or hasattr(instance, "__dict__"): + object.__setattr__(instance, field_name, None) + + class SerializationMixin: """Mixin class providing comprehensive serialization and deserialization capabilities. @@ -284,6 +321,15 @@ def __init__(self, **kwargs): DEFAULT_EXCLUDE: ClassVar[set[str]] = set() INJECTABLE: ClassVar[set[str]] = set() _SHALLOW_COPY_FIELDS: ClassVar[set[str]] = {"raw_representation"} + _PICKLE_OMIT_FIELDS: ClassVar[set[str]] = {"raw_representation"} + + def __copy__(self) -> SerializationMixin: + """Create a shallow copy without invoking pickle state hooks.""" + cls = type(self) + result = cls.__new__(cls) + for field_name, value in _iter_instance_fields(self).items(): + object.__setattr__(result, field_name, value) + return result def __deepcopy__(self, memo: dict[int, Any]) -> SerializationMixin: """Create a deep copy, preserving ``_SHALLOW_COPY_FIELDS`` by reference. @@ -296,13 +342,21 @@ def __deepcopy__(self, memo: dict[int, Any]) -> SerializationMixin: cls = type(self) result = cls.__new__(cls) memo[id(self)] = result - for k, v in self.__dict__.items(): + for k, v in _iter_instance_fields(self).items(): if k in cls._SHALLOW_COPY_FIELDS: object.__setattr__(result, k, v) else: object.__setattr__(result, k, copy.deepcopy(v, memo)) return result + def __getstate__(self) -> dict[str, Any]: + """Return pickle state without runtime-only shallow-copy fields.""" + return get_pickle_state(self, self._PICKLE_OMIT_FIELDS) + + def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, Any]]) -> None: + """Restore pickle state and reset runtime-only shallow-copy fields.""" + restore_pickle_state(self, state, self._PICKLE_OMIT_FIELDS) + def to_dict(self, *, exclude: set[str] | None = None, exclude_none: bool = True) -> dict[str, Any]: """Convert the instance and any nested objects to a dictionary. diff --git a/python/packages/core/agent_framework/_types.py b/python/packages/core/agent_framework/_types.py index 9134dc6a04..e4fadb00a3 100644 --- a/python/packages/core/agent_framework/_types.py +++ b/python/packages/core/agent_framework/_types.py @@ -30,7 +30,7 @@ from typing_extensions import TypedDict from ._feature_stage import ExperimentalFeature, experimental -from ._serialization import SerializationMixin +from ._serialization import SerializationMixin, get_pickle_state, restore_pickle_state from .exceptions import AdditionItemMismatch, ContentError if sys.version_info >= (3, 13): @@ -483,6 +483,7 @@ class Content: """ _SHALLOW_COPY_FIELDS: ClassVar[set[str]] = {"raw_representation"} + _PICKLE_OMIT_FIELDS: ClassVar[set[str]] = {"raw_representation"} def __init__( self, @@ -610,6 +611,28 @@ def __deepcopy__(self, memo: dict[int, Any]) -> Content: object.__setattr__(result, k, deepcopy(v, memo)) return result + def __copy__(self) -> Content: + """Create a shallow copy while preserving provider runtime fields.""" + cls = type(self) + result = cls.__new__(cls) + for field_name, value in self.__dict__.items(): + object.__setattr__(result, field_name, value) + return result + + def __getstate__(self) -> dict[str, Any]: + """Return pickle state without runtime-only shallow-copy fields.""" + state = get_pickle_state(self, self._PICKLE_OMIT_FIELDS) + if self.annotations is not None: + state["annotations"] = [ + {key: value for key, value in annotation.items() if key != "raw_representation"} + for annotation in self.annotations + ] + return state + + def __setstate__(self, state: dict[str, Any] | tuple[dict[str, Any], dict[str, Any]]) -> None: + """Restore pickle state and reset runtime-only shallow-copy fields.""" + restore_pickle_state(self, state, self._PICKLE_OMIT_FIELDS) + @classmethod def from_text( cls: type[ContentT], diff --git a/python/packages/core/tests/core/test_serializable_mixin.py b/python/packages/core/tests/core/test_serializable_mixin.py index 03853e8386..d5b14cb001 100644 --- a/python/packages/core/tests/core/test_serializable_mixin.py +++ b/python/packages/core/tests/core/test_serializable_mixin.py @@ -572,6 +572,59 @@ def __init__(self, items: list, opaque: Any = None, additional_properties: dict assert cloned.items is not obj.items assert cloned.items == ["a"] + def test_shallow_copy_preserves_pickle_omitted_fields(self): + """Shallow copies retain runtime fields that pickle omits.""" + + class TestClass(SerializationMixin): + def __init__(self, raw_representation: Any): + self.raw_representation = raw_representation + + raw = object() + cloned = copy.copy(TestClass(raw)) + + assert cloned.raw_representation is raw + + def test_pickle_restores_slot_fields(self): + """Pickle state should include fields declared in slots.""" + class TestClass(SerializationMixin): + __slots__ = ("value",) + + def __init__(self, value: str): + self.value = value + + original = TestClass("value") + restored = TestClass.__new__(TestClass) + restored.__setstate__(original.__getstate__()) + + assert restored.value == "value" + + def test_pickle_restores_legacy_tuple_state(self): + """Pickle restoration should accept the legacy dict-and-slots tuple.""" + + class TestClass(SerializationMixin): + __slots__ = ("value",) + + def __init__(self): + self.value = "new" + + restored = TestClass.__new__(TestClass) + restored.__setstate__(({"other": "dict"}, {"value": "legacy"})) + + assert restored.value == "legacy" + + def test_pickle_omission_is_separate_from_shallow_copy_policy(self): + """Fields shallow-copied by default remain persistent unless explicitly omitted.""" + class TestClass(SerializationMixin): + _PICKLE_OMIT_FIELDS = set() + + def __init__(self, raw_representation: Any): + self.raw_representation = raw_representation + + raw = {"provider": "value"} + state = TestClass(raw).__getstate__() + + assert state["raw_representation"] == raw + def test_dependency_dict_merge_does_not_mutate_input(self): """Test that dict dependency merging does not mutate the caller's input dictionary.""" diff --git a/python/packages/core/tests/core/test_types.py b/python/packages/core/tests/core/test_types.py index a4d5957f58..e936da808a 100644 --- a/python/packages/core/tests/core/test_types.py +++ b/python/packages/core/tests/core/test_types.py @@ -2733,6 +2733,29 @@ def test_content_deepcopy_discards_raw_representation(caplog: pytest.LogCaptureF assert caplog.messages == ["Discarding field 'raw_representation' while deep-copying Content."] +def test_content_pickle_discards_nested_annotation_raw_representation() -> None: + """Pickle should omit provider objects stored on annotations.""" + import pickle + + raw = object() + annotation: Annotation = {"type": "citation", "url": "https://example.com", "raw_representation": raw} + content = Content.from_text("hello", annotations=[annotation]) + + restored = pickle.loads(pickle.dumps(content)) + + assert restored.annotations == [{"type": "citation", "url": "https://example.com"}] + + +def test_content_shallow_copy_preserves_raw_representation() -> None: + """Shallow copies of Content retain provider runtime fields.""" + import copy + + raw = _NonCopyableRaw() + cloned = copy.copy(Content.from_text("hello", raw_representation=raw)) + + assert cloned.raw_representation is raw + + def test_message_deepcopy_preserves_raw_representation(): """Test that deepcopy of Message keeps raw_representation by reference.""" import copy diff --git a/python/packages/core/tests/workflow/test_agent_executor.py b/python/packages/core/tests/workflow/test_agent_executor.py index ccb1e9425b..2cc2ed2ce6 100644 --- a/python/packages/core/tests/workflow/test_agent_executor.py +++ b/python/packages/core/tests/workflow/test_agent_executor.py @@ -1,5 +1,7 @@ # Copyright (c) Microsoft. All rights reserved. +import pickle + from collections.abc import AsyncIterable, Awaitable from typing import Any, Literal, overload @@ -336,6 +338,9 @@ class _NonCopyableRaw: def __deepcopy__(self, memo: dict) -> Any: raise TypeError("Cannot deepcopy this object") + def __reduce__(self) -> Any: + raise TypeError("Cannot pickle this object") + class _AgentWithRawRepr(BaseAgent): """Agent that returns responses with a non-copyable raw_representation.""" @@ -387,6 +392,20 @@ async def test_agent_executor_workflow_with_non_copyable_raw_representation() -> assert agent_responses[0].raw_representation is raw +def test_serialization_mixin_omits_non_pickleable_raw_representation() -> None: + """Pickling framework objects should not include runtime-only raw representations.""" + raw = _NonCopyableRaw() + response = AgentResponse( + messages=[Message("assistant", [Content.from_text(text="reply", raw_representation=raw)])], + raw_representation=raw, + ) + + restored = pickle.loads(pickle.dumps(response)) + + assert restored.raw_representation is None + assert restored.messages[0].contents[0].raw_representation is None + + # --------------------------------------------------------------------------- # Context mode tests # ---------------------------------------------------------------------------