Coverage for haystack/core/pipeline/breakpoint.py: 90%
106 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
5import json
6import os
7from collections.abc import Callable
8from datetime import datetime
9from pathlib import Path
10from typing import Any
12from networkx import MultiDiGraph
14from haystack import logging
15from haystack.core.errors import PipelineInvalidPipelineSnapshotError
16from haystack.dataclasses.breakpoints import INTERNAL_INPUTS_FORMAT, Breakpoint, PipelineSnapshot, PipelineState
17from haystack.utils.base_serialization import _serialize_with_field_fallback
19logger = logging.getLogger(__name__)
21# Environment variable to control pipeline snapshot file saving (enabled by default)
22HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED = "HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED"
24# Type alias for snapshot callback function
25# The callback receives a PipelineSnapshot and optionally returns a file path string
26SnapshotCallback = Callable[[PipelineSnapshot], str | None]
29def _is_snapshot_save_enabled() -> bool:
30 """
31 Check if pipeline snapshot file saving is enabled via environment variable.
33 The environment variable HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED controls whether
34 pipeline snapshots are saved to files. By default (when the variable is not set),
35 saving is disabled. Only "true" and "1" (case-insensitive) enable saving; any other value disables it.
37 :returns: True if snapshot saving is enabled, False otherwise.
38 """
39 value = os.environ.get(HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED, "false").lower()
40 return value in ("true", "1")
43def _validate_break_point_against_pipeline(break_point: Breakpoint, graph: MultiDiGraph) -> None:
44 """
45 Validates the breakpoints passed to the pipeline.
47 Makes sure the breakpoint contains a valid components registered in the pipeline.
49 :param break_point: a breakpoint to validate
50 """
51 if break_point.component_name not in graph.nodes:
52 raise ValueError(f"break_point {break_point} is not a registered component in the pipeline")
55def _validate_pipeline_snapshot_against_pipeline(pipeline_snapshot: PipelineSnapshot, graph: MultiDiGraph) -> None:
56 """
57 Validates that the pipeline_snapshot contains valid configuration for the current pipeline.
59 Raises a PipelineInvalidPipelineSnapshotError if any component in pipeline_snapshot is not part of the
60 target pipeline.
62 :param pipeline_snapshot: The saved state to validate.
63 """
65 pipeline_state = pipeline_snapshot.pipeline_state
66 valid_components = set(graph.nodes.keys())
68 # Check if the ordered_component_names are valid components in the pipeline
69 invalid_ordered_components = set(pipeline_snapshot.ordered_component_names) - valid_components
70 if invalid_ordered_components:
71 raise PipelineInvalidPipelineSnapshotError(
72 f"Invalid pipeline snapshot: components {invalid_ordered_components} in 'ordered_component_names' "
73 f"are not part of the current pipeline."
74 )
76 # Check if the original_input_data is valid components in the pipeline
77 serialized_input_data = pipeline_snapshot.original_input_data["serialized_data"]
78 invalid_input_data = set(serialized_input_data.keys()) - valid_components
79 if invalid_input_data:
80 raise PipelineInvalidPipelineSnapshotError(
81 f"Invalid pipeline snapshot: components {invalid_input_data} in 'input_data' "
82 f"are not part of the current pipeline."
83 )
85 # Validate 'component_visits'
86 invalid_component_visits = set(pipeline_state.component_visits.keys()) - valid_components
87 if invalid_component_visits:
88 raise PipelineInvalidPipelineSnapshotError(
89 f"Invalid pipeline snapshot: components {invalid_component_visits} in 'component_visits' "
90 f"are not part of the current pipeline."
91 )
93 component_name = pipeline_snapshot.break_point.component_name
94 visit_count = pipeline_snapshot.pipeline_state.component_visits[component_name]
96 logger.info(
97 "Resuming pipeline from {component} with visit count {visits}", component=component_name, visits=visit_count
98 )
101def load_pipeline_snapshot(file_path: str | Path) -> PipelineSnapshot:
102 """
103 Load a saved pipeline snapshot.
105 :param file_path: Path to the pipeline_snapshot file.
106 :returns:
107 Dict containing the loaded pipeline_snapshot.
108 """
110 file_path = Path(file_path)
112 try:
113 with open(file_path, encoding="utf-8") as f:
114 pipeline_snapshot_dict = json.load(f)
115 except FileNotFoundError as e:
116 raise FileNotFoundError(f"File not found: {file_path}") from e
117 except json.JSONDecodeError as e:
118 raise json.JSONDecodeError(f"Invalid JSON file {file_path}: {str(e)}", e.doc, e.pos) from e
119 except OSError as e:
120 raise OSError(f"Error reading {file_path}: {str(e)}") from e
122 try:
123 pipeline_snapshot = PipelineSnapshot.from_dict(pipeline_snapshot_dict)
124 except ValueError as e:
125 raise ValueError(f"Invalid pipeline snapshot from {file_path}: {str(e)}") from e
127 logger.info("Successfully loaded the pipeline snapshot from: {file_path}", file_path=file_path)
128 return pipeline_snapshot
131def _save_pipeline_snapshot(
132 pipeline_snapshot: PipelineSnapshot,
133 raise_on_failure: bool = True,
134 snapshot_callback: SnapshotCallback | None = None,
135) -> str | None:
136 """
137 Save the pipeline snapshot dictionary to a JSON file, or invoke a custom callback.
139 If a `snapshot_callback` is provided, it will be called with the pipeline snapshot instead of saving to a file.
140 This allows users to customize how snapshots are handled (e.g., saving to a database, sending to a remote service).
142 When no callback is provided, the default behavior saves to a JSON file:
143 - The filename is generated based on the component name, visit count, and timestamp.
144 - The component name is taken from the break point's `component_name`.
145 - The visit count is taken from the pipeline state's `component_visits` for the component name.
146 - The timestamp is taken from the pipeline snapshot's `timestamp` or the current time if not available.
147 - The file path is taken from the break point's `snapshot_file_path`.
148 - If the `snapshot_file_path` is None, the function will return without saving.
150 The default file saving behavior is disabled. To enable it, set the environment variable
151 `HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED` to "true" or "1". When disabled,
152 the function will return None without saving to a file (custom callbacks are still invoked).
154 :param pipeline_snapshot: The pipeline snapshot to save.
155 :param raise_on_failure: If True, raises an exception if saving fails. If False, logs the error and returns.
156 :param snapshot_callback: Optional callback function that receives the PipelineSnapshot.
157 If provided, the callback is invoked instead of the default file-saving behavior.
158 The callback should return an optional string (e.g., a file path or identifier) or None.
160 :returns:
161 The full path to the saved JSON file (or the value returned by the callback), or None if
162 `snapshot_file_path` is None, no callback is provided, or snapshot saving is disabled.
163 :raises:
164 Exception: If saving the JSON snapshot fails (when raise_on_failure is True).
165 """
166 # If a callback is provided, use it instead of the default file-saving behavior
167 if snapshot_callback is not None:
168 try:
169 result = snapshot_callback(pipeline_snapshot)
170 logger.info("Pipeline snapshot handled by custom callback.")
171 return result
172 except Exception as error:
173 logger.exception("Failed to handle pipeline snapshot with custom callback. Error: {error}", error=error)
174 if raise_on_failure:
175 raise
176 return None
178 # Check if snapshot saving is enabled via environment variable (enabled by default)
179 if not _is_snapshot_save_enabled():
180 logger.debug("Pipeline snapshot file saving is disabled via HAYSTACK_PIPELINE_SNAPSHOT_SAVE_ENABLED env var.")
181 return None
183 break_point = pipeline_snapshot.break_point
184 snapshot_file_path = break_point.snapshot_file_path
186 if snapshot_file_path is None:
187 return None
189 dt = pipeline_snapshot.timestamp or datetime.now()
190 snapshot_dir = Path(snapshot_file_path)
192 component_name = break_point.component_name
193 visit_nr = pipeline_snapshot.pipeline_state.component_visits.get(component_name, 0)
194 timestamp = dt.strftime("%Y_%m_%d_%H_%M_%S")
195 file_name = f"{component_name}_{visit_nr}_{timestamp}.json"
196 full_path = snapshot_dir / file_name
198 try:
199 snapshot_dir.mkdir(parents=True, exist_ok=True)
200 with open(full_path, "w") as f_out:
201 json.dump(pipeline_snapshot.to_dict(), f_out, indent=2)
202 logger.info(
203 "Pipeline snapshot saved to '{full_path}'. You can use this file to debug or resume the pipeline.",
204 full_path=full_path,
205 )
206 except Exception as error:
207 logger.exception("Failed to save pipeline snapshot to '{full_path}'. Error: {e}", full_path=full_path, e=error)
208 if raise_on_failure:
209 raise
211 return str(full_path)
214def _create_pipeline_snapshot(
215 *,
216 inputs: dict[str, Any],
217 component_inputs: dict[str, Any],
218 break_point: Breakpoint,
219 component_visits: dict[str, int],
220 original_input_data: dict[str, Any],
221 ordered_component_names: list[str],
222 include_outputs_from: set[str],
223 pipeline_outputs: dict[str, Any],
224) -> PipelineSnapshot:
225 """
226 Create a snapshot of the pipeline at the point where the breakpoint was triggered.
228 :param inputs: The current pipeline snapshot inputs.
229 :param component_inputs: The inputs of the component that triggered the breakpoint, as they were before the
230 component consumed them.
231 :param break_point: The breakpoint that triggered the snapshot.
232 :param component_visits: The visit count of the component that triggered the breakpoint.
233 :param original_input_data: The original input data.
234 :param ordered_component_names: The ordered component names.
235 :param include_outputs_from: Set of component names whose outputs should be included in the pipeline results.
236 :param pipeline_outputs: The current outputs of the pipeline.
237 :returns:
238 A PipelineSnapshot containing the state of the pipeline at the point of the breakpoint.
239 """
240 component_name = break_point.component_name
242 transformed_original_input_data = _transform_json_structure(original_input_data)
244 serialized_inputs = _serialize_with_field_fallback(
245 {**inputs, component_name: component_inputs}, description="the inputs of the current pipeline state"
246 )
247 serialized_original_input_data = _serialize_with_field_fallback(
248 transformed_original_input_data, description="original input data for `pipeline.run`"
249 )
250 serialized_pipeline_outputs = _serialize_with_field_fallback(
251 pipeline_outputs, description="outputs of the current pipeline state"
252 )
254 return PipelineSnapshot(
255 pipeline_state=PipelineState(
256 inputs=serialized_inputs,
257 component_visits=component_visits,
258 pipeline_outputs=serialized_pipeline_outputs,
259 inputs_format=INTERNAL_INPUTS_FORMAT,
260 ),
261 timestamp=datetime.now(),
262 break_point=break_point,
263 original_input_data=serialized_original_input_data,
264 ordered_component_names=ordered_component_names,
265 include_outputs_from=include_outputs_from,
266 )
269def _transform_json_structure(data: dict[str, Any] | list[Any] | Any) -> Any:
270 """
271 Transforms a JSON structure by removing the 'sender' key and moving the 'value' to the top level.
273 For example:
274 "key": [{"sender": null, "value": "some value"}] -> "key": "some value"
276 :param data: The JSON structure to transform.
277 :returns: The transformed structure.
278 """
279 if isinstance(data, dict):
280 # If this dict has both 'sender' and 'value', return just the value
281 if "value" in data and "sender" in data:
282 return data["value"]
283 # Otherwise, recursively process each key-value pair
284 return {k: _transform_json_structure(v) for k, v in data.items()}
286 if isinstance(data, list):
287 # First, transform each item in the list.
288 transformed = [_transform_json_structure(item) for item in data]
289 # If the original list has exactly one element and that element was a dict
290 # with 'sender' and 'value', then unwrap the list.
291 if len(data) == 1 and isinstance(data[0], dict) and "value" in data[0] and "sender" in data[0]:
292 return transformed[0]
293 return transformed
295 # For other data types, just return the value as is.
296 return data