diff --git a/haystack/components/agents/state/state.py b/haystack/components/agents/state/state.py index 8c0daa0ca08..7db24cdcd47 100644 --- a/haystack/components/agents/state/state.py +++ b/haystack/components/agents/state/state.py @@ -7,7 +7,8 @@ from typing import Any, get_args from haystack.dataclasses import ChatMessage -from haystack.utils import _deserialize_value_with_schema, _serialize_value_with_schema +from haystack.utils import _deserialize_value_with_schema +from haystack.utils.base_serialization import _serialize_with_field_fallback from haystack.utils.callable_serialization import deserialize_callable, serialize_callable from haystack.utils.type_serialization import deserialize_type, serialize_type @@ -194,7 +195,9 @@ def to_dict(self) -> dict[str, Any]: """ serialized = {} serialized["schema"] = _schema_to_dict(self.schema) - serialized["data"] = _serialize_value_with_schema(self._data) + # Field-level fallback so a single non-serializable value omits only that field instead of + # failing the whole State serialization. + serialized["data"] = _serialize_with_field_fallback(self._data, description="the agent's State data") return serialized @classmethod diff --git a/haystack/core/pipeline/breakpoint.py b/haystack/core/pipeline/breakpoint.py index a3d6edd9726..932c6190ca6 100644 --- a/haystack/core/pipeline/breakpoint.py +++ b/haystack/core/pipeline/breakpoint.py @@ -13,9 +13,8 @@ from haystack import logging from haystack.core.errors import PipelineInvalidPipelineSnapshotError -from haystack.core.pipeline.utils import _deepcopy_with_exceptions from haystack.dataclasses.breakpoints import Breakpoint, PipelineSnapshot, PipelineState -from haystack.utils.base_serialization import _serialize_value_with_schema +from haystack.utils.base_serialization import _serialize_with_field_fallback logger = logging.getLogger(__name__) @@ -292,54 +291,3 @@ def _transform_json_structure(data: dict[str, Any] | list[Any] | Any) -> Any: # For other data types, just return the value as is. return data - - -def _serialize_with_field_fallback(payload: Any, *, description: str) -> dict[str, Any]: - """ - Serialize a payload and, on failure, retry field-by-field to preserve resumable fields. - - If the whole payload serializes, the result is returned as-is. Otherwise, and if the payload is a - mapping, each top-level field is serialized individually and only the failing fields are omitted. - When the payload is not a mapping, or when every field fails to serialize, the helper returns a - structurally valid empty-object payload so that the downstream ``_deserialize_value_with_schema`` - can still load it back instead of raising ``DeserializationError`` on a bare ``{}``. - - :param payload: The value to serialize. - :param description: Short human-readable label used in warning messages, for example - ``"the agent's chat_generator inputs"`` or ``"the inputs of the current pipeline state"``. - :returns: A dict of the form ``{"serialization_schema": ..., "serialized_data": ...}``. - """ - try: - return _serialize_value_with_schema(_deepcopy_with_exceptions(payload)) - except Exception as error: - logger.warning( - "Failed to serialize {description}. " - "Haystack will omit only the non-serializable fields when possible. Error: {e}", - description=description, - e=error, - ) - - serialized_properties: dict[str, Any] = {} - serialized_data: dict[str, Any] = {} - - if isinstance(payload, dict): - for field_name, value in payload.items(): - try: - serialized_value = _serialize_value_with_schema(_deepcopy_with_exceptions(value)) - except Exception as field_error: - logger.warning( - "Failed to serialize the '{field_name}' field of {description}. " - "The field will be omitted from the snapshot. Error: {e}", - field_name=field_name, - description=description, - e=field_error, - ) - continue - - serialized_properties[field_name] = serialized_value["serialization_schema"] - serialized_data[field_name] = serialized_value["serialized_data"] - - return { - "serialization_schema": {"type": "object", "properties": serialized_properties}, - "serialized_data": serialized_data, - } diff --git a/haystack/utils/base_serialization.py b/haystack/utils/base_serialization.py index 30fb81c9f2b..0db98b6d39b 100644 --- a/haystack/utils/base_serialization.py +++ b/haystack/utils/base_serialization.py @@ -8,7 +8,7 @@ import pydantic from haystack import logging -from haystack.core.errors import DeserializationError +from haystack.core.errors import DeserializationError, SerializationError from haystack.core.serialization import generate_qualified_class_name, import_class_by_name from haystack.utils import deserialize_callable, serialize_callable @@ -60,7 +60,7 @@ def _serialize_value_with_schema(payload: Any) -> dict[str, Any]: # noqa: PLR09 return {"serialization_schema": {"type": "object", "properties": schema}, "serialized_data": data} # Handle array case - iterate through elements - if isinstance(payload, (list, tuple, set)): + if isinstance(payload, (list, tuple, set, frozenset)): # Serialize each item in the array serialized_list = [] for item in payload: @@ -76,9 +76,12 @@ def _serialize_value_with_schema(payload: Any) -> dict[str, Any]: # noqa: PLR09 else: base_schema = {"type": "array", "items": {}} - # Add JSON Schema properties to infer sets and tuples - if isinstance(payload, set): + # Add JSON Schema properties to infer sets, frozensets and tuples + if isinstance(payload, (set, frozenset)): base_schema["uniqueItems"] = True + # `frozen` distinguishes frozenset from set on deserialization + if isinstance(payload, frozenset): + base_schema["frozen"] = True elif isinstance(payload, tuple): base_schema["minItems"] = len(payload) base_schema["maxItems"] = len(payload) @@ -101,19 +104,71 @@ def _serialize_value_with_schema(payload: Any) -> dict[str, Any]: # noqa: PLR09 type_name = generate_qualified_class_name(type(payload)) return {"serialization_schema": {"type": type_name}, "serialized_data": payload.name} - # Handle arbitrary objects with __dict__ - if hasattr(payload, "__dict__"): - type_name = generate_qualified_class_name(type(payload)) - schema = {"type": type_name} - serialized_data = {} - for key, value in vars(payload).items(): - serialized_value = _serialize_value_with_schema(value) - serialized_data[key] = serialized_value["serialized_data"] - return {"serialization_schema": schema, "serialized_data": serialized_data} - # Handle primitives - schema = {"type": _primitive_schema_type(payload)} - return {"serialization_schema": schema, "serialized_data": payload} + if payload is None or isinstance(payload, (bool, int, float, str)): + return {"serialization_schema": {"type": _primitive_schema_type(payload)}, "serialized_data": payload} + + # Unsupported type: raise so callers can omit the field (see `_serialize_with_field_fallback`) + # instead of embedding a live, non-portable object in the output. + raise SerializationError( + f"Cannot serialize value of type '{type(payload).__name__}'. Supported values are primitives " + f"(str, int, float, bool, None), lists/tuples/sets/frozensets/dicts of supported values, Enums, " + f"callables, pydantic models, and objects exposing a 'to_dict' method." + ) + + +def _serialize_with_field_fallback(payload: Any, *, description: str) -> dict[str, Any]: + """ + Serialize a payload and, on failure, retry field-by-field to preserve serializable fields. + + If the whole payload serializes, the result is returned as-is. Otherwise, and if the payload is a + mapping, each top-level field is serialized individually and only the failing fields are omitted. + When the payload is not a mapping, or when every field fails to serialize, the helper returns a + structurally valid empty-object payload so that the downstream `_deserialize_value_with_schema` + can still load it back instead of raising `DeserializationError` on a bare `{}`. + + :param payload: The value to serialize. + :param description: Short human-readable label used in warning messages, for example + `"the agent's State data"` or `"the inputs of the current pipeline state"`. + :returns: A dict of the form `{"serialization_schema": ..., "serialized_data": ...}`. + """ + # Local import to avoid a circular import at module load time. + from haystack.core.pipeline.utils import _deepcopy_with_exceptions + + try: + return _serialize_value_with_schema(_deepcopy_with_exceptions(payload)) + except Exception as error: + logger.warning( + "Failed to serialize {description}. " + "Haystack will omit only the non-serializable fields when possible. Error: {e}", + description=description, + e=error, + ) + + serialized_properties: dict[str, Any] = {} + serialized_data: dict[str, Any] = {} + + if isinstance(payload, dict): + for field_name, value in payload.items(): + try: + serialized_value = _serialize_value_with_schema(_deepcopy_with_exceptions(value)) + except Exception as field_error: + logger.warning( + "Failed to serialize the '{field_name}' field of {description}. " + "The field will be omitted from the snapshot. Error: {e}", + field_name=field_name, + description=description, + e=field_error, + ) + continue + + serialized_properties[field_name] = serialized_value["serialization_schema"] + serialized_data[field_name] = serialized_value["serialized_data"] + + return { + "serialization_schema": {"type": "object", "properties": serialized_properties}, + "serialized_data": serialized_data, + } def _primitive_schema_type(value: Any) -> str: @@ -183,10 +238,10 @@ def _deserialize_value_with_schema(serialized: dict[str, Any]) -> Any: _deserialize_value_with_schema({"serialization_schema": schema["items"], "serialized_data": item}) for item in data ] - final_array: list | set | tuple - # Is a set if uniqueItems is True + final_array: list | set | frozenset | tuple + # Is a set or frozenset if uniqueItems is True if schema.get("uniqueItems") is True: - final_array = set(deserialized_items) + final_array = frozenset(deserialized_items) if schema.get("frozen") is True else set(deserialized_items) # Is a tuple if minItems and maxItems are set elif schema.get("minItems") is not None and schema.get("maxItems") is not None: final_array = tuple(deserialized_items) diff --git a/releasenotes/notes/fix-serialization-passthrough-9f1c2a4b7e6d3f08.yaml b/releasenotes/notes/fix-serialization-passthrough-9f1c2a4b7e6d3f08.yaml new file mode 100644 index 00000000000..230bbab0116 --- /dev/null +++ b/releasenotes/notes/fix-serialization-passthrough-9f1c2a4b7e6d3f08.yaml @@ -0,0 +1,13 @@ +--- +fixes: + - | + Fixed the schema-aware serialization helper used for pipeline snapshots and ``Agent`` ``State`` + (``_serialize_value_with_schema``) so it no longer silently passes unsupported objects through as + if they were serialized. Values such as ``datetime``, ``bytes``, ``complex`` and arbitrary objects + without a ``to_dict`` method were previously stored unchanged and mislabeled as strings, which broke + JSON storage and round-tripping of snapshots. Unsupported values now raise a ``SerializationError``, + and the callers that build snapshots (pipeline breakpoints and ``State.to_dict``) catch it to omit + only the offending field while keeping the rest of the payload resumable. + - | + Added support for serializing and deserializing ``frozenset`` values in ``_serialize_value_with_schema``. + A ``frozenset`` now round-trips back to a ``frozenset`` instead of being dropped. diff --git a/test/components/agents/test_state_class.py b/test/components/agents/test_state_class.py index f98bd53defe..8426c6f1148 100644 --- a/test/components/agents/test_state_class.py +++ b/test/components/agents/test_state_class.py @@ -3,7 +3,9 @@ # SPDX-License-Identifier: Apache-2.0 import inspect +import logging from dataclasses import dataclass +from datetime import datetime from typing import Dict, Generic, List, Optional, TypeVar, Union import pytest @@ -512,6 +514,16 @@ def test_state_to_dict_typing_list(self): }, } + def test_state_to_dict_omits_non_serializable_field(self, caplog): + # A non-serializable value must not crash serialization: only that field is omitted, + # and a warning is logged. + state = State({"good": {"type": int}, "bad": {"type": object}}, {"good": 1, "bad": datetime(2024, 1, 1)}) + with caplog.at_level(logging.WARNING): + state_dict = state.to_dict() + assert state_dict["data"]["serialized_data"] == {"good": 1} + assert "bad" not in state_dict["data"]["serialization_schema"]["properties"] + assert any("Failed to serialize the 'bad' field of the agent's State data" in msg for msg in caplog.messages) + def test_state_from_dict_typing_list(self): state_dict = { "schema": { diff --git a/test/core/pipeline/test_breakpoint.py b/test/core/pipeline/test_breakpoint.py index 1f8d38927c6..956595da3d9 100644 --- a/test/core/pipeline/test_breakpoint.py +++ b/test/core/pipeline/test_breakpoint.py @@ -349,19 +349,24 @@ def to_dict(self): def test_save_pipeline_snapshot_raises_on_failure(tmp_path, caplog, monkeypatch): monkeypatch.setenv(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "true") + # Point the snapshot directory below an existing file so creating it fails with a filesystem + # error, exercising the raise_on_failure contract. + blocking_file = tmp_path / "not_a_dir" + blocking_file.write_text("i am a file") + snapshot_path = blocking_file / "snapshots" + snapshot = _create_pipeline_snapshot( inputs={}, component_inputs={}, - break_point=Breakpoint(component_name="comp2", snapshot_file_path=str(tmp_path)), + break_point=Breakpoint(component_name="comp2", snapshot_file_path=str(snapshot_path)), component_visits={"comp1": 1, "comp2": 0}, original_input_data={}, ordered_component_names=["comp1", "comp2"], include_outputs_from={"comp1"}, - # We use a non-serializable type (bytes) directly in pipeline outputs to trigger the error - pipeline_outputs={"comp1": {"result": b"test"}}, + pipeline_outputs={"comp1": {"result": "test"}}, ) - with pytest.raises(TypeError): + with pytest.raises(OSError): _save_pipeline_snapshot(snapshot) with caplog.at_level(logging.ERROR): diff --git a/test/utils/test_base_serialization.py b/test/utils/test_base_serialization.py index 9a5827e12b6..67f48071c00 100644 --- a/test/utils/test_base_serialization.py +++ b/test/utils/test_base_serialization.py @@ -2,16 +2,24 @@ # # SPDX-License-Identifier: Apache-2.0 +from datetime import datetime from enum import Enum import pydantic import pytest -from haystack.core.errors import DeserializationError +from haystack.core.errors import DeserializationError, SerializationError from haystack.dataclasses import ChatMessage, Document, GeneratedAnswer from haystack.utils.base_serialization import _deserialize_value_with_schema, _serialize_value_with_schema +class PlainObject: + """An arbitrary object without a ``to_dict`` method (only a `__dict__`).""" + + def __init__(self, value): + self.value = value + + class CustomModel(pydantic.BaseModel): id: int name: str @@ -22,16 +30,6 @@ class CustomEnum(Enum): TWO = "two" -class CustomClass: - def to_dict(self): - return {"key": "value", "more": False} - - @classmethod - def from_dict(cls, data): - assert data == {"key": "value", "more": False} - return cls() - - def simple_calc_function(x: int) -> int: return x * 2 @@ -112,6 +110,19 @@ def test_serializing_and_deserializing_empty_structures(value, result): "serialized_data": [1, 2, 3], }, ), + # frozenset + ( + frozenset({1, 2, 3}), + { + "serialization_schema": { + "type": "array", + "items": {"type": "integer"}, + "uniqueItems": True, + "frozen": True, + }, + "serialized_data": [1, 2, 3], + }, + ), # tuple ( (1, 2, 3), @@ -221,33 +232,66 @@ def test_serializing_and_deserializing_empty_structures(value, result): ) def test_serialize_and_deserialize_sequence_types(value, result): assert _serialize_value_with_schema(value) == result - assert _deserialize_value_with_schema(result) == value + deserialized = _deserialize_value_with_schema(result) + assert deserialized == value + # `frozenset({...}) == set({...})` is True, so check the exact type to catch container regressions. + assert type(deserialized) is type(value) -def test_serialize_and_deserialize_nested_dicts(): - data = {"key1": {"nested1": "value1", "nested2": {"deep": "value2"}}} - expected = { - "serialization_schema": { - "type": "object", - "properties": { - "key1": { +@pytest.mark.parametrize( + "value,result", + [ + pytest.param( + {"key1": {"nested1": "value1", "nested2": {"deep": "value2"}}}, + { + "serialization_schema": { "type": "object", "properties": { - "nested1": {"type": "string"}, - "nested2": {"type": "object", "properties": {"deep": {"type": "string"}}}, + "key1": { + "type": "object", + "properties": { + "nested1": {"type": "string"}, + "nested2": {"type": "object", "properties": {"deep": {"type": "string"}}}, + }, + } }, - } + }, + "serialized_data": {"key1": {"nested1": "value1", "nested2": {"deep": "value2"}}}, }, - }, - "serialized_data": {"key1": {"nested1": "value1", "nested2": {"deep": "value2"}}}, - } - assert _serialize_value_with_schema(data) == expected - assert _deserialize_value_with_schema(expected) == data + id="nested-dicts", + ), + pytest.param( + simple_calc_function, + { + "serialization_schema": {"type": "typing.Callable"}, + "serialized_data": "test_base_serialization.simple_calc_function", + }, + id="callable", + ), + pytest.param( + CustomEnum.ONE, + {"serialization_schema": {"type": "test_base_serialization.CustomEnum"}, "serialized_data": "ONE"}, + id="enum", + ), + pytest.param( + CustomModel(id=1, name="Test"), + { + "serialization_schema": {"type": "test_base_serialization.CustomModel"}, + "serialized_data": {"id": 1, "name": "Test"}, + }, + id="pydantic-model", + ), + ], +) +def test_serialize_and_deserialize_complex_types(value, result): + assert _serialize_value_with_schema(value) == result + assert _deserialize_value_with_schema(result) == value def test_serialize_and_deserialize_value_with_schema_with_various_types(): data = { "numbers": 1, + "key_name": None, "messages": [ChatMessage.from_user(text="Hello, world!"), ChatMessage.from_assistant(text="Hello, world!")], "user_id": "123", "dict_of_lists": {"numbers": [1, 2, 3]}, @@ -267,6 +311,7 @@ def test_serialize_and_deserialize_value_with_schema_with_various_types(): "type": "object", "properties": { "numbers": {"type": "integer"}, + "key_name": {"type": "null"}, "messages": {"type": "array", "items": {"type": "haystack.dataclasses.chat_message.ChatMessage"}}, "user_id": {"type": "string"}, "dict_of_lists": { @@ -286,6 +331,7 @@ def test_serialize_and_deserialize_value_with_schema_with_various_types(): }, "serialized_data": { "numbers": 1, + "key_name": None, "messages": [ {"role": "user", "meta": {}, "name": None, "content": [{"text": "Hello, world!"}]}, {"role": "assistant", "meta": {}, "name": None, "content": [{"text": "Hello, world!"}]}, @@ -327,87 +373,54 @@ def test_serialize_and_deserialize_value_with_schema_with_various_types(): assert _deserialize_value_with_schema(expected) == data -def test_serializing_and_deserializing_custom_class_type(): - custom_type = CustomClass() - data = {"numbers": 1, "custom_type": custom_type} - serialized_data = _serialize_value_with_schema(data) - assert serialized_data == { - "serialization_schema": { - "properties": { - "custom_type": {"type": "test_base_serialization.CustomClass"}, - "numbers": {"type": "integer"}, - }, - "type": "object", - }, - "serialized_data": {"numbers": 1, "custom_type": {"key": "value", "more": False}}, - } - - deserialized_data = _deserialize_value_with_schema(serialized_data) - assert deserialized_data["numbers"] == 1 - assert isinstance(deserialized_data["custom_type"], CustomClass) - - -def test_serialize_and_deserialize_value_with_callable(): - expected = { - "serialization_schema": {"type": "typing.Callable"}, - "serialized_data": "test_base_serialization.simple_calc_function", - } - assert _serialize_value_with_schema(simple_calc_function) == expected - assert _deserialize_value_with_schema(expected) == simple_calc_function - - -def test_serialize_and_deserialize_value_with_enum(): - data = CustomEnum.ONE - expected = {"serialization_schema": {"type": "test_base_serialization.CustomEnum"}, "serialized_data": "ONE"} - assert _serialize_value_with_schema(data) == expected - assert _deserialize_value_with_schema(expected) == data - - -def test_deserialize_value_with_wrong_value(): - with pytest.raises(DeserializationError, match="Value 'NOT_VALID' is not a valid member of Enum"): - _deserialize_value_with_schema( - {"serialization_schema": {"type": "test_base_serialization.CustomEnum"}, "serialized_data": "NOT_VALID"} - ) - - -def test_deserialize_value_with_schema_class_not_importable(): - with pytest.raises( - DeserializationError, match="Class 'test_base_serialization.NonExistentClass' not correctly imported" - ): - _deserialize_value_with_schema( - {"serialization_schema": {"type": "test_base_serialization.NonExistentClass"}, "serialized_data": {}} - ) - - -def test_deserialize_value_with_schema_class_name_without_module(): - with pytest.raises(DeserializationError, match="Class 'NonExistentClass' not correctly imported"): - _deserialize_value_with_schema({"serialization_schema": {"type": "NonExistentClass"}, "serialized_data": {}}) - - -def test_serialize_and_deserialize_pydantic_model(): - model_instance = CustomModel(id=1, name="Test") - serialized = _serialize_value_with_schema(model_instance) - expected_serialized = { - "serialization_schema": {"type": "test_base_serialization.CustomModel"}, - "serialized_data": {"id": 1, "name": "Test"}, - } - assert serialized == expected_serialized +class TestErrorHandling: + @pytest.mark.parametrize( + "value", + [ + pytest.param(datetime(2024, 1, 1), id="datetime"), + pytest.param(b"some bytes", id="bytes"), + pytest.param(3 + 4j, id="complex"), + pytest.param(PlainObject(1), id="object-without-to_dict"), + ], + ) + def test_serialize_unsupported_type_raises(self, value): + with pytest.raises(SerializationError, match="Cannot serialize value of type"): + _serialize_value_with_schema(value) + + def test_serialize_unsupported_nested_value_raises(self): + # An unsupported value nested inside a supported container must not be silently passed through. + with pytest.raises(SerializationError, match="Cannot serialize value of type"): + _serialize_value_with_schema({"good": 1, "bad": datetime(2024, 1, 1)}) + + def test_deserialize_value_with_wrong_value(self): + with pytest.raises(DeserializationError, match="Value 'NOT_VALID' is not a valid member of Enum"): + _deserialize_value_with_schema( + {"serialization_schema": {"type": "test_base_serialization.CustomEnum"}, "serialized_data": "NOT_VALID"} + ) - deserialized = _deserialize_value_with_schema(expected_serialized) - assert isinstance(deserialized, CustomModel) - assert deserialized.id == 1 - assert deserialized.name == "Test" + def test_deserialize_value_with_schema_class_not_importable(self): + with pytest.raises( + DeserializationError, match="Class 'test_base_serialization.NonExistentClass' not correctly imported" + ): + _deserialize_value_with_schema( + {"serialization_schema": {"type": "test_base_serialization.NonExistentClass"}, "serialized_data": {}} + ) + def test_deserialize_value_with_schema_class_name_without_module(self): + with pytest.raises(DeserializationError, match="Class 'NonExistentClass' not correctly imported"): + _deserialize_value_with_schema( + {"serialization_schema": {"type": "NonExistentClass"}, "serialized_data": {}} + ) -def test_deserialize_pydantic_model_with_invalid_data(): - with pytest.raises( - DeserializationError, - match="Failed to deserialize data '{'id': 'not_an_integer', 'name': 'Test'}' into " - "Pydantic model 'test_base_serialization.CustomModel'", - ): - _deserialize_value_with_schema( - { - "serialization_schema": {"type": "test_base_serialization.CustomModel"}, - "serialized_data": {"id": "not_an_integer", "name": "Test"}, - } - ) + def test_deserialize_pydantic_model_with_invalid_data(self): + with pytest.raises( + DeserializationError, + match="Failed to deserialize data '{'id': 'not_an_integer', 'name': 'Test'}' into " + "Pydantic model 'test_base_serialization.CustomModel'", + ): + _deserialize_value_with_schema( + { + "serialization_schema": {"type": "test_base_serialization.CustomModel"}, + "serialized_data": {"id": "not_an_integer", "name": "Test"}, + } + )