Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 5 additions & 2 deletions haystack/components/agents/state/state.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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
Expand Down
54 changes: 1 addition & 53 deletions haystack/core/pipeline/breakpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand Down Expand Up @@ -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,
}
93 changes: 74 additions & 19 deletions haystack/utils/base_serialization.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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:
Expand All @@ -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)
Expand All @@ -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:
Expand Down Expand Up @@ -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)
Expand Down
Original file line number Diff line number Diff line change
@@ -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.
12 changes: 12 additions & 0 deletions test/components/agents/test_state_class.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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": {
Expand Down
13 changes: 9 additions & 4 deletions test/core/pipeline/test_breakpoint.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
Loading
Loading