Coverage for haystack/utils/base_serialization.py: 93%

136 statements  

« prev     ^ index     » next       coverage.py v7.15.4, created at 2026-08-21 13:53 +0000

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from enum import Enum 

6from typing import Any 

7 

8import pydantic 

9 

10from haystack import logging 

11from haystack.core.errors import DeserializationError, SerializationError 

12from haystack.core.serialization import generate_qualified_class_name, import_class_by_name 

13from haystack.utils import deserialize_callable, serialize_callable 

14 

15logger = logging.getLogger(__name__) 

16 

17_PRIMITIVE_TO_SCHEMA_MAP = {type(None): "null", bool: "boolean", int: "integer", float: "number", str: "string"} 

18 

19 

20def _serialize_value_with_schema(payload: Any) -> dict[str, Any]: # noqa: PLR0911 

21 """ 

22 Serializes a value into a schema-aware format suitable for storage or transmission. 

23 

24 The output format separates the schema information from the actual data, making it easier 

25 to deserialize complex nested structures correctly. 

26 

27 The function handles: 

28 - Objects with to_dict() methods (e.g. dataclasses) 

29 - Objects with __dict__ attributes 

30 - Dictionaries 

31 - Lists, tuples, and sets, including ones holding mixed types. A mixed-type array records one 

32 schema per position under the JSON Schema `prefixItems` keyword, while a homogeneous array 

33 keeps the single shared `items` schema. 

34 - Primitive types (str, int, float, bool, None) 

35 

36 This is for runtime values (Agent/State data, pipeline inputs/outputs at a breakpoint), not 

37 Component definitions — see `default_to_dict` in `core/serialization.py` for that other format. 

38 Don't merge the two; they're not interchangeable. 

39 

40 :param payload: The value to serialize (can be any type) 

41 :returns: The serialized dict representation of the given value. Contains two keys: 

42 - "serialization_schema": Contains type information for each field. 

43 - "serialized_data": Contains the actual data in a simplified format. 

44 

45 """ 

46 # Handle pydantic 

47 if isinstance(payload, pydantic.BaseModel): 

48 type_name = generate_qualified_class_name(type(payload)) 

49 return {"serialization_schema": {"type": type_name}, "serialized_data": payload.model_dump()} 

50 

51 # Handle dictionary case - iterate through fields 

52 if isinstance(payload, dict): 

53 schema: dict[str, Any] = {} 

54 data: dict[str, Any] = {} 

55 

56 for field, val in payload.items(): 

57 # Recursively serialize each field 

58 serialized_value = _serialize_value_with_schema(val) 

59 schema[field] = serialized_value["serialization_schema"] 

60 data[field] = serialized_value["serialized_data"] 

61 

62 return {"serialization_schema": {"type": "object", "properties": schema}, "serialized_data": data} 

63 

64 # Handle array case - iterate through elements 

65 if isinstance(payload, (list, tuple, set, frozenset)): 

66 # Serialize each item in the array, keeping its own schema 

67 serialized_list = [] 

68 item_schemas: list[Any] = [] 

69 base_schema: dict[str, Any] 

70 for item in payload: 

71 serialized_value = _serialize_value_with_schema(item) 

72 serialized_list.append(serialized_value["serialized_data"]) 

73 item_schemas.append(serialized_value["serialization_schema"]) 

74 

75 if not item_schemas: 

76 base_schema = {"type": "array", "items": {}} 

77 elif all(item_schema == item_schemas[0] for item_schema in item_schemas): 

78 # Homogeneous array: keep the historical `items` envelope so older readers keep working 

79 base_schema = {"type": "array", "items": item_schemas[0]} 

80 else: 

81 # Mixed-type array: describe every position with the JSON Schema `prefixItems` keyword 

82 base_schema = {"type": "array", "prefixItems": item_schemas} 

83 

84 # Add JSON Schema properties to infer sets, frozensets and tuples 

85 if isinstance(payload, (set, frozenset)): 

86 base_schema["uniqueItems"] = True 

87 # `frozen` distinguishes frozenset from set on deserialization 

88 if isinstance(payload, frozenset): 

89 base_schema["frozen"] = True 

90 elif isinstance(payload, tuple): 

91 base_schema["minItems"] = len(payload) 

92 base_schema["maxItems"] = len(payload) 

93 

94 return {"serialization_schema": base_schema, "serialized_data": serialized_list} 

95 

96 # Handle Haystack style objects (e.g. dataclasses and Components) 

97 if hasattr(payload, "to_dict") and callable(payload.to_dict): 

98 type_name = generate_qualified_class_name(type(payload)) 

99 schema = {"type": type_name} 

100 return {"serialization_schema": schema, "serialized_data": payload.to_dict()} 

101 

102 # Handle callable functions serialization 

103 if callable(payload) and not isinstance(payload, type): 

104 serialized = serialize_callable(payload) 

105 return {"serialization_schema": {"type": "typing.Callable"}, "serialized_data": serialized} 

106 

107 # Handle Enums 

108 if isinstance(payload, Enum): 

109 type_name = generate_qualified_class_name(type(payload)) 

110 return {"serialization_schema": {"type": type_name}, "serialized_data": payload.name} 

111 

112 # Handle primitives 

113 if payload is None or isinstance(payload, (bool, int, float, str)): 

114 return {"serialization_schema": {"type": _primitive_schema_type(payload)}, "serialized_data": payload} 

115 

116 # Unsupported type: raise so callers can omit the field (see `_serialize_with_field_fallback`) 

117 # instead of embedding a live, non-portable object in the output. 

118 raise SerializationError( 

119 f"Cannot serialize value of type '{type(payload).__name__}'. Supported values are primitives " 

120 f"(str, int, float, bool, None), lists/tuples/sets/frozensets/dicts of supported values, Enums, " 

121 f"callables, pydantic models, and objects exposing a 'to_dict' method." 

122 ) 

123 

124 

125def _serialize_with_field_fallback(payload: Any, *, description: str) -> dict[str, Any]: 

126 """ 

127 Serialize a payload and, on failure, retry field-by-field to preserve serializable fields. 

128 

129 If the whole payload serializes, the result is returned as-is. Otherwise, and if the payload is a 

130 mapping, each top-level field is serialized individually and only the failing fields are omitted. 

131 When the payload is not a mapping, or when every field fails to serialize, the helper returns a 

132 structurally valid empty-object payload so that the downstream `_deserialize_value_with_schema` 

133 can still load it back instead of raising `DeserializationError` on a bare `{}`. 

134 

135 :param payload: The value to serialize. 

136 :param description: Short human-readable label used in warning messages, for example 

137 `"the agent's State data"` or `"the inputs of the current pipeline state"`. 

138 :returns: A dict of the form `{"serialization_schema": ..., "serialized_data": ...}`. 

139 """ 

140 # Local import to avoid a circular import at module load time. 

141 from haystack.core.pipeline.utils import _deepcopy_with_exceptions 

142 

143 try: 

144 return _serialize_value_with_schema(_deepcopy_with_exceptions(payload)) 

145 except Exception as error: 

146 logger.warning( 

147 "Failed to serialize {description}. " 

148 "Haystack will omit only the non-serializable fields when possible. Error: {e}", 

149 description=description, 

150 e=error, 

151 ) 

152 

153 serialized_properties: dict[str, Any] = {} 

154 serialized_data: dict[str, Any] = {} 

155 

156 if isinstance(payload, dict): 

157 for field_name, value in payload.items(): 

158 try: 

159 serialized_value = _serialize_value_with_schema(_deepcopy_with_exceptions(value)) 

160 except Exception as field_error: 

161 logger.warning( 

162 "Failed to serialize the '{field_name}' field of {description}. " 

163 "The field will be omitted from the snapshot. Error: {e}", 

164 field_name=field_name, 

165 description=description, 

166 e=field_error, 

167 ) 

168 continue 

169 

170 serialized_properties[field_name] = serialized_value["serialization_schema"] 

171 serialized_data[field_name] = serialized_value["serialized_data"] 

172 

173 return { 

174 "serialization_schema": {"type": "object", "properties": serialized_properties}, 

175 "serialized_data": serialized_data, 

176 } 

177 

178 

179def _primitive_schema_type(value: Any) -> str: 

180 """ 

181 Helper function to determine the schema type for primitive values. 

182 """ 

183 for py_type, schema_value in _PRIMITIVE_TO_SCHEMA_MAP.items(): 

184 if isinstance(value, py_type): 

185 return schema_value 

186 logger.warning( 

187 "Unsupported primitive type '{value_type}', falling back to 'string'", value_type=type(value).__name__ 

188 ) 

189 return "string" # fallback 

190 

191 

192def _deserialize_value_with_schema(serialized: dict[str, Any]) -> Any: 

193 """ 

194 Deserializes a value with schema information back to its original form. 

195 

196 Takes a dict of the form: 

197 { 

198 "serialization_schema": {"type": "integer"} or {"type": "object", "properties": {...}}, 

199 "serialized_data": <the actual data> 

200 } 

201 

202 For array types, `prefixItems` (one schema per position, emitted for mixed-type arrays) takes 

203 precedence over `items` (a single shared schema, emitted for homogeneous arrays). 

204 

205 :param serialized: The serialized dict with schema and data. 

206 :returns: The deserialized value in its original form. 

207 """ 

208 

209 if not serialized or "serialization_schema" not in serialized or "serialized_data" not in serialized: 

210 raise DeserializationError( 

211 f"Invalid format of passed serialized payload. Expected a dictionary with keys " 

212 f"'serialization_schema' and 'serialized_data'. Got: {serialized}" 

213 ) 

214 schema = serialized["serialization_schema"] 

215 data = serialized["serialized_data"] 

216 

217 schema_type = schema.get("type") 

218 

219 if not schema_type: 

220 # for backward compatibility till Haystack 2.16 we use legacy implementation 

221 raise DeserializationError( 

222 "Missing 'type' key in 'serialization_schema'. This likely indicates that you're using a serialized " 

223 "State object created with a version of Haystack older than 2.15.0. " 

224 "Support for the old serialization format is removed in Haystack 2.16.0. " 

225 "Please upgrade to the new serialization format to ensure forward compatibility." 

226 ) 

227 

228 # Handle object case (dictionary with properties) 

229 if schema_type == "object": 

230 properties = schema["properties"] 

231 result: dict[str, Any] = {} 

232 for field, raw_value in data.items(): 

233 field_schema = properties[field] 

234 # Recursively deserialize each field - avoid creating temporary dict 

235 result[field] = _deserialize_value_with_schema( 

236 {"serialization_schema": field_schema, "serialized_data": raw_value} 

237 ) 

238 return result 

239 

240 # Handle array case 

241 if schema_type == "array": 

242 # Mixed-type arrays carry one schema per position under `prefixItems`, 

243 # homogeneous ones carry a single shared schema under `items`. 

244 prefix_items = schema.get("prefixItems") 

245 if prefix_items is not None: 

246 if len(prefix_items) != len(data): 

247 raise DeserializationError( 

248 f"Invalid array payload: 'prefixItems' declares {len(prefix_items)} element schemas " 

249 f"but 'serialized_data' holds {len(data)} elements." 

250 ) 

251 item_schemas: list[Any] = list(prefix_items) 

252 else: 

253 item_schemas = [schema["items"]] * len(data) 

254 

255 # Deserialize each item with the schema of its own position 

256 deserialized_items = [ 

257 _deserialize_value_with_schema({"serialization_schema": item_schema, "serialized_data": item}) 

258 for item_schema, item in zip(item_schemas, data, strict=True) 

259 ] 

260 final_array: list | set | frozenset | tuple 

261 # Is a set or frozenset if uniqueItems is True 

262 if schema.get("uniqueItems") is True: 

263 final_array = frozenset(deserialized_items) if schema.get("frozen") is True else set(deserialized_items) 

264 # Is a tuple if minItems and maxItems are set 

265 elif schema.get("minItems") is not None and schema.get("maxItems") is not None: 

266 final_array = tuple(deserialized_items) 

267 else: 

268 # Otherwise, it's a list 

269 final_array = list(deserialized_items) 

270 return final_array 

271 

272 # Handle primitive types 

273 if schema_type in _PRIMITIVE_TO_SCHEMA_MAP.values(): 

274 return data 

275 

276 # Handle callable functions 

277 if schema_type == "typing.Callable": 

278 return deserialize_callable(data) 

279 

280 # Handle custom class types 

281 return _deserialize_value({"type": schema_type, "data": data}) 

282 

283 

284def _deserialize_value(value: dict[str, Any]) -> Any: 

285 """ 

286 Helper function to deserialize values from their envelope format {"type": T, "data": D}. 

287 

288 This handles: 

289 - Custom classes (with a from_dict method) 

290 - Enums 

291 - Fallback for arbitrary classes (sets attributes on a blank instance) 

292 

293 :param value: The value to deserialize 

294 :returns: 

295 The deserialized value 

296 :raises DeserializationError: 

297 If the type cannot be imported or the value is not valid for the type. 

298 """ 

299 # 1) Envelope case 

300 value_type = value["type"] 

301 payload = value["data"] 

302 

303 # Custom class where value_type is a qualified class name 

304 # ValueError covers type names without a module prefix, which import_class_by_name cannot split 

305 try: 

306 cls = import_class_by_name(value_type) 

307 except (ImportError, ValueError) as e: 

308 raise DeserializationError(f"Class '{value_type}' not correctly imported") from e 

309 

310 # try from_dict (e.g. Haystack dataclasses and Components) 

311 if hasattr(cls, "from_dict") and callable(cls.from_dict): 

312 return cls.from_dict(payload) 

313 

314 # handle pydantic models 

315 if issubclass(cls, pydantic.BaseModel): 

316 try: 

317 return cls.model_validate(payload) 

318 except Exception as e: 

319 raise DeserializationError( 

320 f"Failed to deserialize data '{payload}' into Pydantic model '{value_type}'" 

321 ) from e 

322 

323 # handle enum types 

324 if issubclass(cls, Enum): 

325 try: 

326 return cls[payload] 

327 except Exception as e: 

328 raise DeserializationError(f"Value '{payload}' is not a valid member of Enum '{value_type}'") from e 

329 

330 # fallback: set attributes on a blank instance 

331 deserialized_payload = {k: _deserialize_value(v) for k, v in payload.items()} 

332 instance = cls.__new__(cls) 

333 for attr_name, attr_value in deserialized_payload.items(): 

334 setattr(instance, attr_name, attr_value) 

335 return instance