Coverage for haystack/core/pipeline/pipeline.py: 94%
327 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 asyncio
6import contextlib
7from collections.abc import AsyncGenerator, AsyncIterator, Mapping
8from typing import Any, ClassVar, cast
10from haystack import logging, tracing
11from haystack.core.component import Component
12from haystack.core.errors import BreakpointException, PipelineInvalidPipelineSnapshotError, PipelineRuntimeError
13from haystack.core.pipeline.base import (
14 _COMPONENT_INPUT,
15 _COMPONENT_OUTPUT,
16 _COMPONENT_VISITS,
17 ComponentPriority,
18 PipelineBase,
19 _validate_component_output_keys,
20)
21from haystack.core.pipeline.breakpoint import (
22 SnapshotCallback,
23 _create_pipeline_snapshot,
24 _save_pipeline_snapshot,
25 _validate_break_point_against_pipeline,
26 _validate_pipeline_snapshot_against_pipeline,
27)
28from haystack.core.pipeline.utils import _deepcopy_with_exceptions
29from haystack.core.serialization_security import mark_deserialization_internal
30from haystack.dataclasses import AsyncStreamingCallbackT, StreamingCallbackT, StreamingChunk, select_streaming_callback
31from haystack.dataclasses.breakpoints import INTERNAL_INPUTS_FORMAT, Breakpoint, PipelineSnapshot
32from haystack.dataclasses.streaming_chunk import _invoke_streaming_callback
33from haystack.telemetry import pipeline_running
34from haystack.utils import _deserialize_value_with_schema
35from haystack.utils.async_utils import _execute_component_async
36from haystack.utils.misc import _get_output_dir
38logger = logging.getLogger(__name__)
41class _EndOfStream:
42 """Sentinel type indicating no more chunks will arrive on the stream."""
45class PipelineStreamHandle:
46 """
47 Handle returned by `Pipeline.stream()`.
49 Async-iterable over `StreamingChunk`s produced by streaming components in the pipeline. After iteration ends,
50 `result` holds the final pipeline output dict.
52 By default, iteration cleans up automatically: if the consumer abandons iteration, the underlying pipeline task is
53 cancelled. `aclose()` is also available for explicit cleanup.
54 """
56 _END_OF_STREAM: ClassVar[_EndOfStream] = _EndOfStream()
57 _CLEANUP_TIMEOUT_SECONDS: ClassVar[float] = 1.0
59 def __init__(
60 self,
61 queue: asyncio.Queue["StreamingChunk | _EndOfStream"],
62 task: asyncio.Task[dict[str, Any]],
63 cancel_on_abandon: bool = True,
64 ) -> None:
65 self._queue = queue
66 self._task = task
67 self._cancel_on_abandon = cancel_on_abandon
69 async def __aiter__(self) -> AsyncIterator[StreamingChunk]:
70 """
71 Drain the queue and cancel the pipeline task if iteration is abandoned.
73 `__aiter__` is an async generator function: each `async for` call gets a generator and `try/finally` runs on
74 exit. When `cancel_on_abandon` is True (default), abandoned iteration cancels the pipeline task; when False,
75 the task is left running to completion.
76 """
77 try:
78 while True:
79 item = await self._queue.get()
80 if item is self._END_OF_STREAM:
81 await self._task # called to make exceptions surface
82 return
83 yield cast(StreamingChunk, item) # at this point, item is guaranteed to be a StreamingChunk
85 finally:
86 if self._cancel_on_abandon:
87 await self.aclose()
89 @property
90 def result(self) -> dict[str, Any]:
91 """
92 Final pipeline output dict, available only after a successful, complete run.
94 Raises a `RuntimeError` if the pipeline has not finished or was cancelled. If the pipeline failed, re-raises the
95 original exception.
96 """
97 if not self._task.done():
98 raise RuntimeError("Pipeline has not finished; iterate the handle first.")
99 if self._task.cancelled():
100 raise RuntimeError("Pipeline was cancelled; no result available.")
101 exc = self._task.exception()
102 if exc is not None:
103 raise exc
104 return self._task.result()
106 async def aclose(self) -> None:
107 """
108 Cancel the underlying pipeline task.
110 Bounded by `_CLEANUP_TIMEOUT_SECONDS` so that components cannot block cleanup indefinitely.
111 """
112 if not self._task.done():
113 self._task.cancel()
114 with contextlib.suppress(BaseException):
115 await asyncio.wait_for(self._task, timeout=self._CLEANUP_TIMEOUT_SECONDS)
118class Pipeline(PipelineBase):
119 """
120 Orchestration engine that runs components according to the execution graph.
122 Supports both a synchronous run path (`run`) and an asynchronous run path
123 (`run_async`, `run_async_generator`, `stream`).
124 """
126 @staticmethod
127 def _run_component(
128 component_name: str,
129 component: dict[str, Any],
130 inputs: dict[str, Any],
131 component_visits: dict[str, int],
132 parent_span: tracing.Span | None = None,
133 *,
134 break_point: Breakpoint | None = None,
135 ) -> Mapping[str, Any]:
136 """
137 Runs a Component with the given inputs.
139 :param component_name: Name of the Component.
140 :param component: Component with component metadata.
141 :param inputs: Inputs for the Component.
142 :param component_visits: Current state of component visits.
143 :param parent_span: The parent span to use for the newly created span.
144 This is to allow tracing to be correctly linked to the pipeline run.
145 :param break_point: An optional breakpoint. If it targets this component and its visit
146 count matches the current one, a `BreakpointException` is raised before the Component runs.
147 :raises PipelineRuntimeError: If Component doesn't return a dictionary.
148 :return: The output of the Component.
149 """
150 if (
151 isinstance(break_point, Breakpoint)
152 and break_point.component_name == component_name
153 and break_point.visit_count == component_visits[component_name]
154 ):
155 raise BreakpointException.from_triggered_breakpoint(break_point=break_point)
157 instance: Component = component["instance"]
159 with PipelineBase._create_component_span(
160 component_name=component_name, instance=instance, inputs=inputs, parent_span=parent_span
161 ) as span:
162 # deepcopy inputs before passing to the tracer so that even if a tracer mutates them
163 # the component always receives the original unmodified values
164 inputs_copy = _deepcopy_with_exceptions(inputs)
165 span.set_content_tag(_COMPONENT_INPUT, inputs)
166 logger.info("Running component {component_name}", component_name=component_name)
168 try:
169 component_output = instance.run(**inputs_copy)
170 except BreakpointException as error:
171 # Re-raise BreakpointException to preserve the original exception context
172 # This is important when Agent components internally use Pipeline._run_component
173 # and trigger breakpoints that need to bubble up to the main pipeline
174 raise error
176 # Any components that internally use Pipeline._run_component could raise a PipelineRuntimeError with
177 # additional context (e.g. Agent raises an agent snapshot) so we re-raise here instead of wrapping it in
178 # another PipelineRuntimeError
180 except PipelineRuntimeError as runtime_error:
181 raise runtime_error
183 # Catch all other exceptions and wrap them in a PipelineRuntimeError
184 except Exception as error:
185 raise PipelineRuntimeError.from_exception(component_name, instance.__class__, error) from error
187 component_visits[component_name] += 1
189 if not isinstance(component_output, Mapping):
190 raise PipelineRuntimeError.from_invalid_output(component_name, instance.__class__, component_output)
192 _validate_component_output_keys(component_name, component, component_output)
194 span.set_tag(_COMPONENT_VISITS, component_visits[component_name])
195 span.set_content_tag(_COMPONENT_OUTPUT, component_output)
197 return component_output
199 @mark_deserialization_internal
200 def run( # noqa: PLR0915, PLR0912, C901
201 self,
202 data: dict[str, Any],
203 include_outputs_from: set[str] | None = None,
204 *,
205 break_point: Breakpoint | None = None,
206 pipeline_snapshot: PipelineSnapshot | None = None,
207 snapshot_callback: SnapshotCallback | None = None,
208 ) -> dict[str, Any]:
209 """
210 Runs the Pipeline with given input data.
212 `run` executes synchronously and blocks the calling thread until the run completes. In an async context,
213 use `run_async` instead.
215 Usage:
216 ```python
217 from haystack import Pipeline, Document
218 from haystack.components.builders.answer_builder import AnswerBuilder
219 from haystack.components.builders.chat_prompt_builder import ChatPromptBuilder
220 from haystack.components.generators.chat import OpenAIChatGenerator
221 from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
222 from haystack.dataclasses import ChatMessage
223 from haystack.document_stores.in_memory import InMemoryDocumentStore
224 from haystack.utils import Secret
226 # Write documents to InMemoryDocumentStore
227 document_store = InMemoryDocumentStore()
228 document_store.write_documents([
229 Document(content="My name is Jean and I live in Paris."),
230 Document(content="My name is Mark and I live in Berlin."),
231 Document(content="My name is Giorgio and I live in Rome.")
232 ])
234 retriever = InMemoryBM25Retriever(document_store=document_store)
236 prompt_template = \"\"\"
237 Given these documents, answer the question.
238 Documents:
239 {% for doc in documents %}
240 {{ doc.content }}
241 {% endfor %}
242 Question: {{question}}
243 Answer:
244 \"\"\"
246 template = [ChatMessage.from_user(prompt_template)]
247 prompt_builder = ChatPromptBuilder(
248 template=template,
249 required_variables=["question", "documents"],
250 variables=["question", "documents"]
251 )
253 llm = OpenAIChatGenerator()
254 rag_pipeline = Pipeline()
255 rag_pipeline.add_component("retriever", retriever)
256 rag_pipeline.add_component("prompt_builder", prompt_builder)
257 rag_pipeline.add_component("llm", llm)
258 rag_pipeline.connect("retriever", "prompt_builder.documents")
259 rag_pipeline.connect("prompt_builder", "llm")
261 question = "Who lives in Paris?"
262 results = rag_pipeline.run(
263 {
264 "retriever": {"query": question},
265 "prompt_builder": {"question": question},
266 }
267 )
269 print(results["llm"]["replies"][0].text)
270 # Jean lives in Paris
271 ```
273 :param data:
274 A dictionary of inputs for the pipeline's components. Each key is a component name
275 and its value is a dictionary of that component's input parameters:
276 ```
277 data = {
278 "comp1": {"input1": 1, "input2": 2},
279 }
280 ```
281 For convenience, this format is also supported when input names are unique:
282 ```
283 data = {
284 "input1": 1, "input2": 2,
285 }
286 ```
287 :param include_outputs_from:
288 Set of component names whose individual outputs are to be
289 included in the pipeline's output. For components that are
290 invoked multiple times (in a loop), only the last-produced
291 output is included.
293 :param break_point:
294 A breakpoint that pauses execution before the specified component runs by raising a
295 `BreakpointException` carrying a `PipelineSnapshot` of the current pipeline state.
297 :param pipeline_snapshot:
298 A snapshot of a previously interrupted pipeline execution to resume from. Can be combined with
299 `break_point` to step through a pipeline: resume from the snapshot and pause again at the next
300 breakpoint. The `break_point` must target a different component or visit count than the one the
301 snapshot was created at, otherwise it would trigger again before any progress is made.
303 :param snapshot_callback:
304 Optional callback function that is invoked when a pipeline snapshot is created.
305 The callback receives a `PipelineSnapshot` object and can return an optional string
306 (e.g., a file path or identifier).
307 If provided, the callback is used instead of the default file-saving behavior,
308 allowing custom handling of snapshots (e.g., saving to a database, sending to a remote service).
309 If not provided, the default behavior saves snapshots to a JSON file.
311 :returns:
312 A dictionary where each entry corresponds to a component name
313 and its output. If `include_outputs_from` is `None`, this dictionary
314 will only contain the outputs of leaf components, i.e., components
315 without outgoing connections.
317 :raises ValueError:
318 If invalid inputs are provided to the pipeline.
319 :raises PipelineRuntimeError:
320 If the Pipeline contains cycles with unsupported connections that would cause
321 it to get stuck and fail running.
322 Or if a Component fails or returns output in an unsupported type.
323 :raises PipelineMaxComponentRuns:
324 If a Component reaches the maximum number of times it can be run in this Pipeline.
325 :raises PipelineBreakpointException:
326 When a pipeline_breakpoint is triggered. Contains the component name, state, and partial results.
327 """
328 pipeline_running(self) # telemetry
330 if (
331 break_point
332 and pipeline_snapshot
333 and break_point.component_name == pipeline_snapshot.break_point.component_name
334 and break_point.visit_count == pipeline_snapshot.break_point.visit_count
335 ):
336 msg = (
337 "The provided break_point targets the same component and visit count as the break_point of the "
338 "pipeline_snapshot. It would trigger again before the resumed component runs, so the pipeline "
339 "could not make any progress. Provide a break_point with a different component or visit count."
340 )
341 raise PipelineInvalidPipelineSnapshotError(message=msg)
343 # make sure all breakpoints are valid, i.e. reference components in the pipeline
344 if break_point:
345 _validate_break_point_against_pipeline(break_point, self.graph)
347 # warm up the pipeline by running each component's warm_up method
348 self.warm_up()
350 if include_outputs_from is None:
351 include_outputs_from = set()
353 pipeline_outputs: dict[str, Any] = {}
354 # Set when resuming from a snapshot that predates `INTERNAL_INPUTS_FORMAT` and therefore lost the sender of
355 # each input. Cleared as soon as the paused component has run.
356 legacy_resume_component: str | None = None
358 if not pipeline_snapshot:
359 # normalize `data`
360 data = self._prepare_component_input_data(data)
362 # Raise ValueError if input is malformed in some way
363 self.validate_input(data)
365 # We create a list of components in the pipeline sorted by name, so that the algorithm runs
366 # deterministically and independent of insertion order into the pipeline.
367 ordered_component_names = sorted(self.graph.nodes.keys())
369 # We track component visits to decide if a component can run.
370 component_visits = dict.fromkeys(ordered_component_names, 0)
372 inputs = self._convert_to_internal_format(pipeline_inputs=data)
374 else:
375 # Validate the pipeline snapshot against the current pipeline graph
376 _validate_pipeline_snapshot_against_pipeline(pipeline_snapshot, self.graph)
378 # Handle resuming the pipeline from a snapshot
379 component_visits = pipeline_snapshot.pipeline_state.component_visits
380 ordered_component_names = pipeline_snapshot.ordered_component_names
381 data = _deserialize_value_with_schema(pipeline_snapshot.original_input_data)
383 if pipeline_snapshot.pipeline_state.inputs_format == INTERNAL_INPUTS_FORMAT:
384 inputs = _deserialize_value_with_schema(pipeline_snapshot.pipeline_state.inputs)
385 else:
386 # A legacy snapshot lost the sender of each input, so the only thing we can do is treat them as if
387 # they came from outside the pipeline. The paused component then has to be let through once.
388 inputs = self._convert_to_internal_format(
389 pipeline_inputs=_deserialize_value_with_schema(pipeline_snapshot.pipeline_state.inputs)
390 )
391 legacy_resume_component = pipeline_snapshot.break_point.component_name
393 # include_outputs_from from the snapshot when resuming
394 include_outputs_from = pipeline_snapshot.include_outputs_from
396 # also intermediate_outputs from the snapshot when resuming
397 pipeline_outputs = _deserialize_value_with_schema(pipeline_snapshot.pipeline_state.pipeline_outputs)
399 cached_topological_sort = None
400 # We need to access a component's receivers multiple times during a pipeline run.
401 # We store them here for easy access.
402 cached_receivers = {name: self._find_receivers_from(name) for name in ordered_component_names}
404 with tracing.tracer.trace(
405 "haystack.pipeline.run",
406 tags={
407 "haystack.pipeline.input_data": data,
408 "haystack.pipeline.metadata": self.metadata,
409 "haystack.pipeline.max_runs_per_component": self._max_runs_per_component,
410 "haystack.pipeline.execution_mode": "sync",
411 },
412 ) as span:
413 priority_queue = self._fill_queue(ordered_component_names, inputs, component_visits)
415 # check if pipeline is blocked before execution
416 self.validate_pipeline(priority_queue)
418 while True:
419 candidate = self._get_next_runnable_component(priority_queue, component_visits)
421 # If there are no runnable components left, we can exit the loop.
422 # In practice this rarely happens because the queue is constantly refilled even with components that
423 # have already run. They just get a BLOCKED priority since their inputs have already been consumed.
424 if candidate is None:
425 break
427 priority, component_name, component = candidate
429 # If the next component is blocked, we do a check to see if the pipeline is possibly blocked and raise
430 # a warning if it is.
431 if priority == ComponentPriority.BLOCKED:
432 if self._is_pipeline_possibly_blocked(current_pipeline_outputs=pipeline_outputs):
433 # Pipeline is most likely blocked (most likely a configuration issue) so we raise a warning.
434 self._find_components_blocking_pipeline(
435 priority_queue=priority_queue, component_visits=component_visits, inputs=inputs
436 )
437 # We always exit the loop since we cannot run the next component.
438 break
440 if len(priority_queue) > 0 and priority == ComponentPriority.DEFER:
441 component_name, topological_sort = self._tiebreak_waiting_components(
442 component_name=component_name,
443 priority=priority,
444 priority_queue=priority_queue,
445 topological_sort=cached_topological_sort,
446 )
447 cached_topological_sort = topological_sort
448 component = self._get_component_with_graph_metadata_and_visits(
449 component_name, component_visits[component_name]
450 )
452 is_resume = legacy_resume_component == component_name
453 if is_resume:
454 # Only the first execution of the paused component needs the legacy handling. Later visits of the
455 # same component inside a loop receive regular inputs and must consume them normally.
456 legacy_resume_component = None
458 # A snapshot has to store the component's inputs as they were before the component consumed them, so
459 # that resuming re-triggers the component the same way this run did.
460 component_inputs_before_consume = inputs.get(component_name, {})
461 component_inputs = self._consume_component_inputs(
462 component_name=component_name, component=component, inputs=inputs, is_resume=is_resume
463 )
465 # We need to add missing defaults using default values from input sockets because the run signature
466 # might not provide these defaults for components with inputs defined dynamically upon component
467 # initialization
468 component_inputs = self._add_missing_input_defaults(component_inputs, component["input_sockets"])
470 try:
471 component_outputs = self._run_component(
472 component_name=component_name,
473 component=component,
474 inputs=component_inputs, # the inputs to the current component
475 component_visits=component_visits,
476 parent_span=span,
477 break_point=break_point,
478 )
479 except (BreakpointException, PipelineRuntimeError) as error:
480 saved_break_point: Breakpoint
481 if isinstance(error, PipelineRuntimeError):
482 saved_break_point = Breakpoint(
483 component_name=component_name,
484 visit_count=component_visits[component_name],
485 snapshot_file_path=_get_output_dir("pipeline_snapshot"),
486 )
487 else:
488 saved_break_point = error.break_point
490 # Create a snapshot of the state of the pipeline before the error occurred.
491 pipeline_snapshot = _create_pipeline_snapshot(
492 inputs=_deepcopy_with_exceptions(inputs),
493 component_inputs=_deepcopy_with_exceptions(component_inputs_before_consume),
494 break_point=saved_break_point,
495 component_visits=component_visits,
496 original_input_data=data,
497 ordered_component_names=ordered_component_names,
498 include_outputs_from=include_outputs_from,
499 pipeline_outputs=pipeline_outputs,
500 )
502 # Attach the pipeline snapshot to the error before re-raising
503 error.pipeline_snapshot = pipeline_snapshot
504 full_file_path = _save_pipeline_snapshot(
505 pipeline_snapshot=pipeline_snapshot,
506 raise_on_failure=isinstance(error, BreakpointException),
507 snapshot_callback=snapshot_callback,
508 )
509 error.pipeline_snapshot_file_path = full_file_path
510 raise error
512 # Updates global input state with component outputs and returns outputs that should go to
513 # pipeline outputs.
514 component_pipeline_outputs = self._write_component_outputs(
515 component_name=component_name,
516 component_outputs=component_outputs,
517 inputs=inputs,
518 receivers=cached_receivers[component_name],
519 include_outputs_from=include_outputs_from,
520 )
522 if component_pipeline_outputs or component_name in include_outputs_from:
523 pipeline_outputs[component_name] = component_pipeline_outputs
524 if self._is_queue_stale(priority_queue):
525 priority_queue = self._fill_queue(ordered_component_names, inputs, component_visits)
527 if isinstance(break_point, Breakpoint):
528 logger.warning(
529 "The given breakpoint {break_point} was never triggered. This is because:\n"
530 "1. The provided component is not a part of the pipeline execution path.\n"
531 "2. The component did not reach the visit count specified in the pipeline_breakpoint",
532 break_point=break_point,
533 )
535 # Set here so the tag reflects the final outputs.
536 span.set_content_tag("haystack.pipeline.output_data", pipeline_outputs)
538 return pipeline_outputs
540 @staticmethod
541 async def _run_component_async(
542 component_name: str,
543 component: dict[str, Any],
544 component_inputs: dict[str, Any],
545 component_visits: dict[str, int],
546 parent_span: tracing.Span | None = None,
547 *,
548 break_point: Breakpoint | None = None,
549 ) -> Mapping[str, Any]:
550 """
551 Executes a single component asynchronously.
553 If the component supports async execution, it is awaited directly as it will run async;
554 otherwise the component is offloaded to executor.
556 The method also updates the `visits` count of the component, writes outputs to `inputs_state`,
557 and returns pruned outputs that get stored in `pipeline_outputs`.
559 :param component_name: The name of the component.
560 :param component_inputs: Inputs for the component.
561 :returns: Outputs from the component that can be yielded from run_async_generator.
562 """
563 if (
564 isinstance(break_point, Breakpoint)
565 and break_point.component_name == component_name
566 and break_point.visit_count == component_visits[component_name]
567 ):
568 raise BreakpointException.from_triggered_breakpoint(break_point=break_point)
570 instance: Component = component["instance"]
572 with PipelineBase._create_component_span(
573 component_name=component_name, instance=instance, inputs=component_inputs, parent_span=parent_span
574 ) as span:
575 # deepcopy inputs before passing to the tracer so that even if a tracer mutates them
576 # the component always receives the original unmodified values
577 component_inputs_copy = _deepcopy_with_exceptions(component_inputs)
578 span.set_content_tag(_COMPONENT_INPUT, component_inputs)
579 logger.info("Running component {component_name}", component_name=component_name)
581 try:
582 # For sync-only components, _run_component_async dispatches to a thread via asyncio.to_thread,
583 # which copies the current contextvars context — preserving e.g. the active tracing span.
584 outputs = await _execute_component_async(instance, **component_inputs_copy)
585 except Exception as error:
586 raise PipelineRuntimeError.from_exception(component_name, instance.__class__, error) from error
588 component_visits[component_name] += 1
590 if not isinstance(outputs, Mapping):
591 raise PipelineRuntimeError.from_invalid_output(component_name, instance.__class__, outputs)
593 _validate_component_output_keys(component_name, component, outputs)
595 span.set_tag(_COMPONENT_VISITS, component_visits[component_name])
596 span.set_content_tag(_COMPONENT_OUTPUT, outputs)
598 return outputs
600 @staticmethod
601 async def _wait_for_tasks(
602 running_tasks: dict[asyncio.Task, str], scheduled_components: set[str], *, return_when: str
603 ) -> AsyncIterator[dict[str, Any]]:
604 """
605 Waits for running tasks to finish and yields their partial outputs.
607 :param running_tasks: Mapping of in-flight tasks to the name of the component they run. Finished tasks are
608 removed in place.
609 :param scheduled_components: Set of component names that are scheduled but not yet finished. Finished
610 components are discarded in place.
611 :param return_when: Either `asyncio.FIRST_COMPLETED` to wait for a single task or `asyncio.ALL_COMPLETED` to
612 wait for every running task.
613 :returns: An async iterator of partial outputs, one per finished component that produced an output.
614 """
615 if not running_tasks:
616 return
618 done, _pending = await asyncio.wait(running_tasks.keys(), return_when=return_when)
619 for finished in done:
620 finished_component_name = running_tasks.pop(finished)
621 try:
622 partial_result = finished.result()
623 except Exception:
624 # A component failed. Cancel and drain the remaining in-flight tasks so they don't keep running in
625 # the background (and leak) after the run is aborted, then re-raise the original error.
626 await Pipeline._cancel_in_flight_tasks(running_tasks, scheduled_components)
627 raise
628 scheduled_components.discard(finished_component_name)
629 if partial_result:
630 yield {finished_component_name: _deepcopy_with_exceptions(partial_result)}
632 @staticmethod
633 async def _cancel_in_flight_tasks(running_tasks: dict[asyncio.Task, str], scheduled_components: set[str]) -> None:
634 """
635 Cancels all in-flight tasks and waits for the cancellations to settle.
637 Called when a component fails or when the run is abandoned early so that sibling tasks don't keep running in
638 the background after the pipeline run is aborted. Exceptions from the cancelled tasks are suppressed since we
639 are already unwinding.
641 Note: cancellation is only effective for components that run natively async. Sync components are offloaded to
642 a thread via `asyncio.to_thread` and a running thread cannot be interrupted: cancelling its task abandons
643 the await, but the thread keeps running until the component's `run` returns. Its outputs are then discarded
644 (the task never writes them to the pipeline state), so state stays consistent, but side effects (e.g. API
645 calls) still complete and the thread can outlive this cleanup.
647 :param running_tasks: Mapping of in-flight tasks to component names. Cleared in place.
648 :param scheduled_components: Set of scheduled-but-unfinished component names. Cleared in place.
649 """
650 for task in running_tasks:
651 task.cancel()
652 # return_exceptions=True so a failing or cancelled sibling doesn't mask the original error we re-raise.
653 await asyncio.gather(*running_tasks.keys(), return_exceptions=True)
654 for component_name in running_tasks.values():
655 scheduled_components.discard(component_name)
656 running_tasks.clear()
658 async def _run_component_in_isolation(
659 self,
660 *,
661 component_name: str,
662 inputs: dict[str, dict[str, list[dict[str, Any]]]],
663 pipeline_outputs: dict[str, Any],
664 component_visits: dict[str, int],
665 running_tasks: dict[asyncio.Task, str],
666 scheduled_components: set[str],
667 cached_receivers: dict[str, Any],
668 include_outputs_from: set[str],
669 parent_span: tracing.Span | None,
670 ) -> AsyncIterator[dict[str, Any]]:
671 """
672 Runs a component with HIGHEST priority in isolation.
674 We need to run components with HIGHEST priority (i.e. components with a GreedyVariadic input socket) by
675 themselves, without any other components running concurrently. Otherwise, downstream components could produce
676 additional inputs for the GreedyVariadic socket.
678 :param component_name: The name of the component to run.
679 :param inputs: The global input state shared by all components. Mutated in place.
680 :param pipeline_outputs: The accumulated pipeline outputs. Mutated in place.
681 :param component_visits: Current state of component visits. Mutated in place.
682 :param running_tasks: Mapping of in-flight tasks to component names. Drained in place before running.
683 :param scheduled_components: Set of scheduled-but-unfinished component names. Mutated in place.
684 :param cached_receivers: Precomputed mapping of component name to its downstream receivers.
685 :param include_outputs_from: Set of component names whose outputs should always be included in the output.
686 :param parent_span: The parent tracing span for the pipeline run.
687 :returns: An async iterator of partial outputs.
688 """
689 # 1) Wait for all in-flight tasks to finish so the HIGHEST component runs alone.
690 async for partial_outputs in self._wait_for_tasks(
691 running_tasks, scheduled_components, return_when=asyncio.ALL_COMPLETED
692 ):
693 yield partial_outputs
695 if component_name in scheduled_components:
696 # If it's already scheduled for some reason, skip.
697 return
699 # 2) Run the HIGHEST component by itself.
700 scheduled_components.add(component_name)
701 component = self._get_component_with_graph_metadata_and_visits(component_name, component_visits[component_name])
702 component_inputs = self._consume_component_inputs(component_name, component, inputs)
703 component_inputs = self._add_missing_input_defaults(component_inputs, component["input_sockets"])
705 component_outputs = await self._run_component_async(
706 component_name=component_name,
707 component=component,
708 component_inputs=component_inputs,
709 component_visits=component_visits,
710 parent_span=parent_span,
711 )
713 pruned = self._write_component_outputs(
714 component_name=component_name,
715 component_outputs=component_outputs,
716 inputs=inputs,
717 receivers=cached_receivers[component_name],
718 include_outputs_from=include_outputs_from,
719 )
720 if pruned or component_name in include_outputs_from:
721 pipeline_outputs[component_name] = pruned
723 scheduled_components.remove(component_name)
724 if pruned or component_name in include_outputs_from:
725 yield {component_name: _deepcopy_with_exceptions(pruned)}
727 def _schedule_component(
728 self,
729 *,
730 component_name: str,
731 inputs: dict[str, dict[str, list[dict[str, Any]]]],
732 pipeline_outputs: dict[str, Any],
733 component_visits: dict[str, int],
734 running_tasks: dict[asyncio.Task, str],
735 scheduled_components: set[str],
736 ready_sem: asyncio.Semaphore,
737 cached_receivers: dict[str, Any],
738 include_outputs_from: set[str],
739 parent_span: tracing.Span | None,
740 ) -> None:
741 """
742 Schedules a component to run as a background task without waiting for it to finish.
744 Inputs are consumed synchronously here (before the task is created) so that other components scheduled in the
745 same iteration of the scheduling loop observe the updated input state and don't race for the same inputs.
747 :param component_name: The name of the component to schedule.
748 :param inputs: The global input state shared by all components. Mutated in place.
749 :param pipeline_outputs: The accumulated pipeline outputs. Mutated in place by the task once it finishes.
750 :param component_visits: Current state of component visits. Mutated in place by the task once it finishes.
751 :param running_tasks: Mapping of in-flight tasks to component names. The new task is registered here.
752 :param scheduled_components: Set of scheduled-but-unfinished component names. Mutated in place.
753 :param ready_sem: Semaphore bounding how many components run concurrently.
754 :param cached_receivers: Precomputed mapping of component name to its downstream receivers.
755 :param include_outputs_from: Set of component names whose outputs should always be included in the output.
756 :param parent_span: The parent tracing span for the pipeline run.
757 """
758 if component_name in scheduled_components:
759 return # already scheduled, do nothing
761 scheduled_components.add(component_name)
763 component = self._get_component_with_graph_metadata_and_visits(component_name, component_visits[component_name])
764 component_inputs = self._consume_component_inputs(component_name, component, inputs)
765 component_inputs = self._add_missing_input_defaults(component_inputs, component["input_sockets"])
767 async def _runner() -> Mapping[str, Any]:
768 async with ready_sem:
769 component_outputs = await self._run_component_async(
770 component_name=component_name,
771 component=component,
772 component_inputs=component_inputs,
773 component_visits=component_visits,
774 parent_span=parent_span,
775 )
777 pruned = self._write_component_outputs(
778 component_name=component_name,
779 component_outputs=component_outputs,
780 inputs=inputs,
781 receivers=cached_receivers[component_name],
782 include_outputs_from=include_outputs_from,
783 )
784 if pruned or component_name in include_outputs_from:
785 pipeline_outputs[component_name] = pruned
787 scheduled_components.remove(component_name)
788 return pruned
790 task = asyncio.create_task(_runner())
791 running_tasks[task] = component_name
793 @mark_deserialization_internal
794 async def run_async_generator( # noqa: PLR0915,C901
795 self, data: dict[str, Any], include_outputs_from: set[str] | None = None, concurrency_limit: int = 4
796 ) -> AsyncGenerator[dict[str, Any], None]:
797 """
798 Executes the pipeline step by step asynchronously, yielding partial outputs when any component finishes.
800 Usage:
801 ```python
802 from haystack import Document
803 from haystack.components.builders import ChatPromptBuilder
804 from haystack.dataclasses import ChatMessage
805 from haystack.utils import Secret
806 from haystack.document_stores.in_memory import InMemoryDocumentStore
807 from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
808 from haystack.components.generators.chat import OpenAIChatGenerator
809 from haystack.components.builders.prompt_builder import PromptBuilder
810 from haystack import Pipeline
811 import asyncio
813 # Write documents to InMemoryDocumentStore
814 document_store = InMemoryDocumentStore()
815 document_store.write_documents([
816 Document(content="My name is Jean and I live in Paris."),
817 Document(content="My name is Mark and I live in Berlin."),
818 Document(content="My name is Giorgio and I live in Rome.")
819 ])
821 prompt_template = [
822 ChatMessage.from_user(
823 '''
824 Given these documents, answer the question.
825 Documents:
826 {% for doc in documents %}
827 {{ doc.content }}
828 {% endfor %}
829 Question: {{question}}
830 Answer:
831 ''')
832 ]
834 # Create and connect pipeline components
835 retriever = InMemoryBM25Retriever(document_store=document_store)
836 prompt_builder = ChatPromptBuilder(template=prompt_template)
837 llm = OpenAIChatGenerator()
839 rag_pipeline = Pipeline()
840 rag_pipeline.add_component("retriever", retriever)
841 rag_pipeline.add_component("prompt_builder", prompt_builder)
842 rag_pipeline.add_component("llm", llm)
843 rag_pipeline.connect("retriever", "prompt_builder.documents")
844 rag_pipeline.connect("prompt_builder", "llm")
846 # Prepare input data
847 question = "Who lives in Paris?"
848 data = {
849 "retriever": {"query": question},
850 "prompt_builder": {"question": question},
851 }
854 # Process results as they become available
855 async def process_results():
856 async for partial_output in rag_pipeline.run_async_generator(
857 data=data,
858 include_outputs_from={"retriever", "llm"}
859 ):
860 # Each partial_output contains the results from a completed component
861 if "retriever" in partial_output:
862 print("Retrieved documents:", len(partial_output["retriever"]["documents"]))
863 if "llm" in partial_output:
864 print("Generated answer:", partial_output["llm"]["replies"][0])
867 asyncio.run(process_results())
868 ```
870 :param data: Initial input data to the pipeline.
871 :param concurrency_limit: The maximum number of components that are allowed to run concurrently.
872 :param include_outputs_from:
873 Set of component names whose individual outputs are to be
874 included in the pipeline's output. For components that are
875 invoked multiple times (in a loop), only the last-produced
876 output is included.
877 :return: An async iterator containing partial (and final) outputs.
879 :raises ValueError:
880 If invalid inputs are provided to the pipeline, or if `concurrency_limit` is less than 1.
881 :raises PipelineMaxComponentRuns:
882 If a component exceeds the maximum number of allowed executions within the pipeline.
883 :raises PipelineRuntimeError:
884 If the Pipeline contains cycles with unsupported connections that would cause
885 it to get stuck and fail running.
886 Or if a Component fails or returns output in an unsupported type.
887 """
888 if concurrency_limit < 1:
889 raise ValueError("concurrency_limit must be greater than or equal to 1.")
891 pipeline_running(self) # telemetry
893 # warm up the pipeline by running each component's warm_up_async (or warm_up) method
894 await self.warm_up_async()
896 if include_outputs_from is None:
897 include_outputs_from = set()
899 pipeline_outputs: dict[str, Any] = {}
901 # Normalize `data` and raise ValueError if the input is malformed in some way.
902 data = self._prepare_component_input_data(data)
904 # Raise ValueError if input is malformed in some way
905 self.validate_input(data)
907 # We create a list of components in the pipeline sorted by name, so that the algorithm runs
908 # deterministically and independent of insertion order into the pipeline.
909 ordered_component_names = sorted(self.graph.nodes.keys())
911 # We track component visits to decide if a component can run.
912 component_visits = dict.fromkeys(ordered_component_names, 0)
914 cached_topological_sort = None
915 # We need to access a component's receivers multiple times during a pipeline run.
916 # We store them here for easy access.
917 cached_receivers = {name: self._find_receivers_from(name) for name in ordered_component_names}
919 # Ephemeral concurrency state shared (and mutated in place) by the scheduling helpers below.
920 ready_sem = asyncio.Semaphore(concurrency_limit)
921 running_tasks: dict[asyncio.Task, str] = {}
922 # A set of component names that have been scheduled but not finished.
923 scheduled_components: set[str] = set()
925 with tracing.tracer.trace(
926 "haystack.pipeline.run",
927 tags={
928 "haystack.pipeline.input_data": data,
929 "haystack.pipeline.metadata": self.metadata,
930 "haystack.pipeline.max_runs_per_component": self._max_runs_per_component,
931 "haystack.pipeline.execution_mode": "async",
932 },
933 ) as parent_span:
934 inputs = self._convert_to_internal_format(pipeline_inputs=data)
936 # check if pipeline is blocked before execution
937 self.validate_pipeline(self._fill_queue(ordered_component_names, inputs, component_visits))
939 try:
940 while True:
941 # We rebuild the priority queue every iteration: each iteration waits for one or more concurrent
942 # tasks to finish, which mutates `inputs` and can change many components' priorities at once, so
943 # we rebuild to give every scheduling decision an up-to-date view.
944 priority_queue = self._fill_queue(ordered_component_names, inputs, component_visits)
945 candidate = self._get_next_runnable_component(priority_queue, component_visits)
947 # If we can't make progress with the queue but tasks are running, we wait for one to finish and
948 # retry to potentially unblock the priority queue.
949 if (candidate is None or candidate[0] == ComponentPriority.BLOCKED) and running_tasks:
950 async for partial_outputs in self._wait_for_tasks(
951 running_tasks, scheduled_components, return_when=asyncio.FIRST_COMPLETED
952 ):
953 yield partial_outputs
954 continue
956 # If there are no runnable components left and nothing is running, we can exit the loop.
957 if candidate is None and not running_tasks:
958 break
960 priority, component_name, component = candidate # type: ignore
962 # If the next component is blocked, we do a check to see if the pipeline is possibly blocked and
963 # raise a warning if it is.
964 if priority == ComponentPriority.BLOCKED and not running_tasks:
965 if self._is_pipeline_possibly_blocked(current_pipeline_outputs=pipeline_outputs):
966 # Pipeline is most likely blocked (most likely a configuration issue) so we raise a warning.
967 self._find_components_blocking_pipeline(
968 priority_queue=priority_queue, component_visits=component_visits, inputs=inputs
969 )
970 # We always exit the loop since we cannot run the next component.
971 break
973 # If the next component is already scheduled, we wait for a task to finish to make progress.
974 if component_name in scheduled_components:
975 async for partial_outputs in self._wait_for_tasks(
976 running_tasks, scheduled_components, return_when=asyncio.FIRST_COMPLETED
977 ):
978 yield partial_outputs
979 continue
981 if priority == ComponentPriority.HIGHEST:
982 # A HIGHEST priority component must run alone, so we hand off to the isolation helper.
983 async for partial_outputs in self._run_component_in_isolation(
984 component_name=component_name,
985 inputs=inputs,
986 pipeline_outputs=pipeline_outputs,
987 component_visits=component_visits,
988 running_tasks=running_tasks,
989 scheduled_components=scheduled_components,
990 cached_receivers=cached_receivers,
991 include_outputs_from=include_outputs_from,
992 parent_span=parent_span,
993 ):
994 yield partial_outputs
995 continue
997 if priority == ComponentPriority.READY:
998 # Schedule this component, then schedule as many additional READY components as concurrency
999 # allows.
1000 self._schedule_component(
1001 component_name=component_name,
1002 inputs=inputs,
1003 pipeline_outputs=pipeline_outputs,
1004 component_visits=component_visits,
1005 running_tasks=running_tasks,
1006 scheduled_components=scheduled_components,
1007 ready_sem=ready_sem,
1008 cached_receivers=cached_receivers,
1009 include_outputs_from=include_outputs_from,
1010 parent_span=parent_span,
1011 )
1013 # Possibly schedule more READY tasks if concurrency not fully used
1014 while len(priority_queue) > 0 and not ready_sem.locked():
1015 peek_priority, peek_name = priority_queue.peek()
1016 if peek_priority != ComponentPriority.READY:
1017 # We stop scheduling: the next component is BLOCKED (can't run), HIGHEST (must run
1018 # alone), or DEFER (waiting for more inputs - we only schedule it once it becomes
1019 # READY).
1020 break
1021 priority_queue.pop()
1022 self._schedule_component(
1023 component_name=peek_name,
1024 inputs=inputs,
1025 pipeline_outputs=pipeline_outputs,
1026 component_visits=component_visits,
1027 running_tasks=running_tasks,
1028 scheduled_components=scheduled_components,
1029 ready_sem=ready_sem,
1030 cached_receivers=cached_receivers,
1031 include_outputs_from=include_outputs_from,
1032 parent_span=parent_span,
1033 )
1035 # We only schedule components with priority DEFER when no other tasks are running.
1036 elif priority == ComponentPriority.DEFER and not running_tasks:
1037 if len(priority_queue) > 0:
1038 component_name, cached_topological_sort = self._tiebreak_waiting_components(
1039 component_name=component_name,
1040 priority=priority,
1041 priority_queue=priority_queue,
1042 topological_sort=cached_topological_sort,
1043 )
1045 self._schedule_component(
1046 component_name=component_name,
1047 inputs=inputs,
1048 pipeline_outputs=pipeline_outputs,
1049 component_visits=component_visits,
1050 running_tasks=running_tasks,
1051 scheduled_components=scheduled_components,
1052 ready_sem=ready_sem,
1053 cached_receivers=cached_receivers,
1054 include_outputs_from=include_outputs_from,
1055 parent_span=parent_span,
1056 )
1058 # To make progress, we wait for one task to complete before restarting the loop.
1059 async for partial_outputs in self._wait_for_tasks(
1060 running_tasks, scheduled_components, return_when=asyncio.FIRST_COMPLETED
1061 ):
1062 yield partial_outputs
1064 # Safety net: drain any leftover tasks once the scheduling loop has finished. With the current loop
1065 # both `break` paths require `running_tasks` to be empty, so this is a no-op. We keep it so that a
1066 # future change adding a `break` that leaves tasks in flight doesn't lose outputs.
1067 async for partial_outputs in self._wait_for_tasks(
1068 running_tasks, scheduled_components, return_when=asyncio.ALL_COMPLETED
1069 ):
1070 yield partial_outputs
1072 # Set here so the tag reflects the final outputs.
1073 parent_span.set_content_tag("haystack.pipeline.output_data", pipeline_outputs)
1075 # Yield the final pipeline outputs.
1076 yield pipeline_outputs
1077 finally:
1078 # If iteration is abandoned early (e.g. the consumer stops iterating the generator and closes it) or
1079 # the run is cancelled, cancel any tasks still in flight so they don't leak.
1080 # This is a no-op on normal completion and on a component error, since no tasks are left running by then
1081 await self._cancel_in_flight_tasks(running_tasks, scheduled_components)
1083 @mark_deserialization_internal
1084 async def run_async(
1085 self, data: dict[str, Any], include_outputs_from: set[str] | None = None, concurrency_limit: int = 4
1086 ) -> dict[str, Any]:
1087 """
1088 Provides an asynchronous interface to run the pipeline with provided input data.
1090 This method allows the pipeline to be integrated into an asynchronous workflow, enabling non-blocking
1091 execution of pipeline components.
1093 Usage:
1094 ```python
1095 import asyncio
1097 from haystack import Document
1098 from haystack.components.builders import ChatPromptBuilder
1099 from haystack.components.generators.chat import OpenAIChatGenerator
1100 from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
1101 from haystack import Pipeline
1102 from haystack.dataclasses import ChatMessage
1103 from haystack.document_stores.in_memory import InMemoryDocumentStore
1105 # Write documents to InMemoryDocumentStore
1106 document_store = InMemoryDocumentStore()
1107 document_store.write_documents([
1108 Document(content="My name is Jean and I live in Paris."),
1109 Document(content="My name is Mark and I live in Berlin."),
1110 Document(content="My name is Giorgio and I live in Rome.")
1111 ])
1113 prompt_template = [
1114 ChatMessage.from_user(
1115 '''
1116 Given these documents, answer the question.
1117 Documents:
1118 {% for doc in documents %}
1119 {{ doc.content }}
1120 {% endfor %}
1121 Question: {{question}}
1122 Answer:
1123 ''')
1124 ]
1126 retriever = InMemoryBM25Retriever(document_store=document_store)
1127 prompt_builder = ChatPromptBuilder(template=prompt_template)
1128 llm = OpenAIChatGenerator()
1130 rag_pipeline = Pipeline()
1131 rag_pipeline.add_component("retriever", retriever)
1132 rag_pipeline.add_component("prompt_builder", prompt_builder)
1133 rag_pipeline.add_component("llm", llm)
1134 rag_pipeline.connect("retriever", "prompt_builder.documents")
1135 rag_pipeline.connect("prompt_builder", "llm")
1137 # Ask a question
1138 question = "Who lives in Paris?"
1140 async def run_inner(data, include_outputs_from):
1141 return await rag_pipeline.run_async(data=data, include_outputs_from=include_outputs_from)
1143 data = {
1144 "retriever": {"query": question},
1145 "prompt_builder": {"question": question},
1146 }
1148 results = asyncio.run(run_inner(data, include_outputs_from={"retriever", "llm"}))
1150 print(results["llm"]["replies"])
1151 # [ChatMessage(_role=<ChatRole.ASSISTANT: 'assistant'>, _content=[TextContent(text='Jean lives in Paris.')],
1152 # _name=None, _meta={'model': 'gpt-5-mini', 'index': 0, 'finish_reason': 'stop', 'usage':
1153 # {'completion_tokens': 6, 'prompt_tokens': 69, 'total_tokens': 75,
1154 # 'completion_tokens_details': CompletionTokensDetails(accepted_prediction_tokens=0,
1155 # audio_tokens=0, reasoning_tokens=0, rejected_prediction_tokens=0), 'prompt_tokens_details':
1156 # PromptTokensDetails(audio_tokens=0, cached_tokens=0)}})]
1157 ```
1159 :param data:
1160 A dictionary of inputs for the pipeline's components. Each key is a component name
1161 and its value is a dictionary of that component's input parameters:
1162 ```
1163 data = {
1164 "comp1": {"input1": 1, "input2": 2},
1165 }
1166 ```
1167 For convenience, this format is also supported when input names are unique:
1168 ```
1169 data = {
1170 "input1": 1, "input2": 2,
1171 }
1172 ```
1173 :param include_outputs_from:
1174 Set of component names whose individual outputs are to be
1175 included in the pipeline's output. For components that are
1176 invoked multiple times (in a loop), only the last-produced
1177 output is included.
1178 :param concurrency_limit: The maximum number of components that should be allowed to run concurrently.
1179 :returns:
1180 A dictionary where each entry corresponds to a component name
1181 and its output. If `include_outputs_from` is `None`, this dictionary
1182 will only contain the outputs of leaf components, i.e., components
1183 without outgoing connections.
1185 :raises ValueError:
1186 If invalid inputs are provided to the pipeline, or if `concurrency_limit` is less than 1.
1187 :raises PipelineRuntimeError:
1188 If the Pipeline contains cycles with unsupported connections that would cause
1189 it to get stuck and fail running.
1190 Or if a Component fails or returns output in an unsupported type.
1191 :raises PipelineMaxComponentRuns:
1192 If a Component reaches the maximum number of times it can be run in this Pipeline.
1193 """
1194 final: dict[str, Any] = {}
1195 async for partial in self.run_async_generator(
1196 data=data, concurrency_limit=concurrency_limit, include_outputs_from=include_outputs_from
1197 ):
1198 final = partial
1199 return final or {}
1201 @mark_deserialization_internal
1202 def stream(
1203 self,
1204 data: dict[str, Any],
1205 *,
1206 streaming_components: list[str] | None = None,
1207 include_outputs_from: set[str] | None = None,
1208 concurrency_limit: int = 4,
1209 cancel_on_abandon: bool = True,
1210 ) -> PipelineStreamHandle:
1211 """
1212 Run the pipeline and return a handle that streams `StreamingChunk`s as they arrive.
1214 Iterate the handle with `async for` to consume chunks; after iteration ends, `handle.result` holds the final
1215 pipeline output dict (same as `run_async`). By default, if iteration is abandoned, the underlying pipeline task
1216 is cancelled automatically. Pass `cancel_on_abandon=False` to instead let the pipeline run to completion.
1218 For every async-capable component that exposes a `streaming_callback` input socket, a forwarder is injected at
1219 runtime that pushes chunks onto the handle's queue. If a `streaming_callback` is provided at component init or
1220 at runtime (inside `data`, e.g. `data={"llm": {"streaming_callback": cb}}`), it is also invoked for each chunk.
1221 Async callbacks are preferred; a sync callback is accepted but will run synchronously on the event loop and
1222 may block it.
1224 Usage:
1225 ```python
1226 import asyncio
1228 from haystack.components.builders import ChatPromptBuilder
1229 from haystack.components.generators.chat import OpenAIChatGenerator
1230 from haystack import Pipeline
1231 from haystack.dataclasses import ChatMessage
1233 pipe = Pipeline()
1234 pipe.add_component(
1235 "prompt_builder",
1236 ChatPromptBuilder(template=[ChatMessage.from_user("Tell me about {{topic}}")]),
1237 )
1238 pipe.add_component("llm", OpenAIChatGenerator())
1239 pipe.connect("prompt_builder.prompt", "llm.messages")
1241 async def main():
1242 handle = pipe.stream(data={"prompt_builder": {"topic": "Italy"}})
1243 async for chunk in handle:
1244 print(chunk.content, end="", flush=True)
1245 return handle.result
1247 result = asyncio.run(main())
1248 print(result["llm"]["replies"])
1249 ```
1251 :param data:
1252 A dictionary of inputs for the pipeline's components. Each key is a component name
1253 and its value is a dictionary of that component's input parameters:
1254 ```
1255 data = {
1256 "comp1": {"input1": 1, "input2": 2},
1257 }
1258 ```
1259 For convenience, this format is also supported when input names are unique:
1260 ```
1261 data = {
1262 "input1": 1, "input2": 2,
1263 }
1264 ```
1265 :param streaming_components: Names of components to stream from. If `None` (default), every streaming-capable
1266 component is forwarded. If a list, only the listed components are forwarded; unknown names or names of
1267 components that do not support streaming raise `ValueError`.
1268 :param include_outputs_from:
1269 Set of component names whose individual outputs are to be
1270 included in the pipeline's output. For components that are
1271 invoked multiple times (in a loop), only the last-produced
1272 output is included.
1273 :param concurrency_limit: The maximum number of components that should be allowed to run concurrently.
1274 :param cancel_on_abandon: If `True` (default), the underlying pipeline task is cancelled when iteration is
1275 abandoned. If `False`, the pipeline runs to completion even when the consumer stops reading.
1276 :returns:
1277 A `PipelineStreamHandle` that is async-iterable over `StreamingChunk`s. After iteration ends,
1278 `handle.result` holds the final pipeline output dict (same shape as `run_async`).
1280 :raises ValueError:
1281 If `streaming_components` contains unknown component names or components that do not support streaming,
1282 or if invalid inputs are provided to the pipeline, or if `concurrency_limit` is less than 1.
1283 :raises PipelineRuntimeError:
1284 Surfaced during iteration. If the Pipeline contains cycles with unsupported connections that would cause
1285 it to get stuck and fail running, or if a Component fails or returns output in an unsupported type.
1286 :raises PipelineMaxComponentRuns:
1287 Surfaced during iteration. If a Component reaches the maximum number of times it can be run in this
1288 Pipeline.
1289 """
1290 streaming_capable = {
1291 name
1292 for name in self.graph.nodes
1293 if getattr(self.graph.nodes[name]["instance"], "__haystack_supports_async__", False)
1294 and "streaming_callback" in self.graph.nodes[name]["instance"].__haystack_input__
1295 }
1296 if streaming_components is not None:
1297 requested = set(streaming_components)
1298 unknown = requested - set(self.graph.nodes)
1299 non_streaming = requested - unknown - streaming_capable
1300 if unknown:
1301 raise ValueError(f"Unknown components in streaming_components: {sorted(unknown)}")
1302 if non_streaming:
1303 raise ValueError(f"These components do not support streaming: {sorted(non_streaming)}")
1305 queue: asyncio.Queue[StreamingChunk | _EndOfStream] = asyncio.Queue()
1307 def make_forwarder(user_callback: StreamingCallbackT | None) -> AsyncStreamingCallbackT:
1308 async def forwarder(chunk: StreamingChunk) -> None:
1309 await queue.put(chunk)
1310 if user_callback is not None:
1311 await _invoke_streaming_callback(user_callback, chunk)
1313 return forwarder
1315 new_data: dict[str, Any] = self._prepare_component_input_data(data)
1316 for name in streaming_capable:
1317 if streaming_components is not None and name not in streaming_components:
1318 continue
1319 instance = self.graph.nodes[name]["instance"]
1320 comp_inputs = new_data.setdefault(name, {})
1322 user_callback = select_streaming_callback(
1323 init_callback=getattr(instance, "streaming_callback", None),
1324 runtime_callback=comp_inputs.get("streaming_callback"),
1325 requires_async=True,
1326 )
1327 comp_inputs["streaming_callback"] = make_forwarder(user_callback)
1329 async def runner() -> dict[str, Any]:
1330 try:
1331 return await self.run_async(
1332 new_data, include_outputs_from=include_outputs_from, concurrency_limit=concurrency_limit
1333 )
1334 finally:
1335 await queue.put(PipelineStreamHandle._END_OF_STREAM)
1337 task = asyncio.create_task(runner())
1338 return PipelineStreamHandle(queue=queue, task=task, cancel_on_abandon=cancel_on_abandon)