Coverage for haystack/core/pipeline/pipeline.py: 94%
331 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-12 09:42 +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 OpenAIResponsesChatGenerator
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 = OpenAIResponsesChatGenerator()
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 if (
329 break_point
330 and pipeline_snapshot
331 and break_point.component_name == pipeline_snapshot.break_point.component_name
332 and break_point.visit_count == pipeline_snapshot.break_point.visit_count
333 ):
334 msg = (
335 "The provided break_point targets the same component and visit count as the break_point of the "
336 "pipeline_snapshot. It would trigger again before the resumed component runs, so the pipeline "
337 "could not make any progress. Provide a break_point with a different component or visit count."
338 )
339 raise PipelineInvalidPipelineSnapshotError(message=msg)
341 # make sure all breakpoints are valid, i.e. reference components in the pipeline
342 if break_point:
343 _validate_break_point_against_pipeline(break_point, self.graph)
345 # warm up the pipeline by running each component's warm_up method
346 self.warm_up()
348 # after warm-up, so toolsets that load their tools in warm_up() report them
349 pipeline_running(self) # telemetry
351 if include_outputs_from is None:
352 include_outputs_from = set()
354 pipeline_outputs: dict[str, Any] = {}
355 # Set when resuming from a snapshot that predates `INTERNAL_INPUTS_FORMAT` and therefore lost the sender of
356 # each input. Cleared as soon as the paused component has run.
357 legacy_resume_component: str | None = None
359 if not pipeline_snapshot:
360 # normalize `data`
361 data = self._prepare_component_input_data(data)
363 # Raise ValueError if input is malformed in some way
364 self.validate_input(data)
366 # We create a list of components in the pipeline sorted by name, so that the algorithm runs
367 # deterministically and independent of insertion order into the pipeline.
368 ordered_component_names = sorted(self.graph.nodes.keys())
370 # We track component visits to decide if a component can run.
371 component_visits = dict.fromkeys(ordered_component_names, 0)
373 inputs = self._convert_to_internal_format(pipeline_inputs=data)
375 else:
376 # Validate the pipeline snapshot against the current pipeline graph
377 _validate_pipeline_snapshot_against_pipeline(pipeline_snapshot, self.graph)
379 # Handle resuming the pipeline from a snapshot
380 component_visits = pipeline_snapshot.pipeline_state.component_visits
381 ordered_component_names = pipeline_snapshot.ordered_component_names
382 data = _deserialize_value_with_schema(pipeline_snapshot.original_input_data)
384 if pipeline_snapshot.pipeline_state.inputs_format == INTERNAL_INPUTS_FORMAT:
385 inputs = _deserialize_value_with_schema(pipeline_snapshot.pipeline_state.inputs)
386 else:
387 # A legacy snapshot lost the sender of each input, so the only thing we can do is treat them as if
388 # they came from outside the pipeline. The paused component then has to be let through once.
389 inputs = self._convert_to_internal_format(
390 pipeline_inputs=_deserialize_value_with_schema(pipeline_snapshot.pipeline_state.inputs)
391 )
392 legacy_resume_component = pipeline_snapshot.break_point.component_name
394 # include_outputs_from from the snapshot when resuming
395 include_outputs_from = pipeline_snapshot.include_outputs_from
397 # also intermediate_outputs from the snapshot when resuming
398 pipeline_outputs = _deserialize_value_with_schema(pipeline_snapshot.pipeline_state.pipeline_outputs)
400 cached_topological_sort = None
401 # We need to access a component's receivers multiple times during a pipeline run.
402 # We store them here for easy access.
403 cached_receivers = {name: self._find_receivers_from(name) for name in ordered_component_names}
405 with tracing.tracer.trace(
406 "haystack.pipeline.run",
407 tags={
408 "haystack.pipeline.input_data": data,
409 "haystack.pipeline.metadata": self.metadata,
410 "haystack.pipeline.max_runs_per_component": self._max_runs_per_component,
411 "haystack.pipeline.execution_mode": "sync",
412 },
413 ) as span:
414 priority_queue = self._fill_queue(ordered_component_names, inputs, component_visits)
416 # check if pipeline is blocked before execution
417 self.validate_pipeline(priority_queue)
419 while True:
420 candidate = self._get_next_runnable_component(priority_queue, component_visits)
422 # If there are no runnable components left, we can exit the loop.
423 # In practice this rarely happens because the queue is constantly refilled even with components that
424 # have already run. They just get a BLOCKED priority since their inputs have already been consumed.
425 if candidate is None:
426 break
428 priority, component_name, component = candidate
430 # If the next component is blocked, we do a check to see if the pipeline is possibly blocked and raise
431 # a warning if it is.
432 if priority == ComponentPriority.BLOCKED:
433 if self._is_pipeline_possibly_blocked(current_pipeline_outputs=pipeline_outputs):
434 # Pipeline is most likely blocked (most likely a configuration issue) so we raise a warning.
435 self._find_components_blocking_pipeline(
436 priority_queue=priority_queue, component_visits=component_visits, inputs=inputs
437 )
438 # We always exit the loop since we cannot run the next component.
439 break
441 if len(priority_queue) > 0 and priority == ComponentPriority.DEFER:
442 component_name, topological_sort = self._tiebreak_waiting_components(
443 component_name=component_name,
444 priority=priority,
445 priority_queue=priority_queue,
446 topological_sort=cached_topological_sort,
447 )
448 cached_topological_sort = topological_sort
449 component = self._get_component_with_graph_metadata_and_visits(
450 component_name, component_visits[component_name]
451 )
453 is_resume = legacy_resume_component == component_name
454 if is_resume:
455 # Only the first execution of the paused component needs the legacy handling. Later visits of the
456 # same component inside a loop receive regular inputs and must consume them normally.
457 legacy_resume_component = None
459 # A snapshot has to store the component's inputs as they were before the component consumed them, so
460 # that resuming re-triggers the component the same way this run did.
461 component_inputs_before_consume = inputs.get(component_name, {})
462 component_inputs = self._consume_component_inputs(
463 component_name=component_name, component=component, inputs=inputs, is_resume=is_resume
464 )
466 # We need to add missing defaults using default values from input sockets because the run signature
467 # might not provide these defaults for components with inputs defined dynamically upon component
468 # initialization
469 component_inputs = self._add_missing_input_defaults(component_inputs, component["input_sockets"])
471 try:
472 component_outputs = self._run_component(
473 component_name=component_name,
474 component=component,
475 inputs=component_inputs, # the inputs to the current component
476 component_visits=component_visits,
477 parent_span=span,
478 break_point=break_point,
479 )
480 except (BreakpointException, PipelineRuntimeError) as error:
481 saved_break_point: Breakpoint
482 if isinstance(error, PipelineRuntimeError):
483 saved_break_point = Breakpoint(
484 component_name=component_name,
485 visit_count=component_visits[component_name],
486 snapshot_file_path=_get_output_dir("pipeline_snapshot"),
487 )
488 else:
489 saved_break_point = error.break_point
491 # Create a snapshot of the state of the pipeline before the error occurred.
492 pipeline_snapshot = _create_pipeline_snapshot(
493 inputs=_deepcopy_with_exceptions(inputs),
494 component_inputs=_deepcopy_with_exceptions(component_inputs_before_consume),
495 break_point=saved_break_point,
496 component_visits=component_visits,
497 original_input_data=data,
498 ordered_component_names=ordered_component_names,
499 include_outputs_from=include_outputs_from,
500 pipeline_outputs=pipeline_outputs,
501 )
503 # Attach the pipeline snapshot to the error before re-raising
504 error.pipeline_snapshot = pipeline_snapshot
505 full_file_path = _save_pipeline_snapshot(
506 pipeline_snapshot=pipeline_snapshot,
507 raise_on_failure=isinstance(error, BreakpointException),
508 snapshot_callback=snapshot_callback,
509 )
510 error.pipeline_snapshot_file_path = full_file_path
511 raise error
513 # Updates global input state with component outputs and returns outputs that should go to
514 # pipeline outputs.
515 component_pipeline_outputs = self._write_component_outputs(
516 component_name=component_name,
517 component_outputs=component_outputs,
518 inputs=inputs,
519 receivers=cached_receivers[component_name],
520 include_outputs_from=include_outputs_from,
521 )
523 if component_pipeline_outputs or component_name in include_outputs_from:
524 pipeline_outputs[component_name] = component_pipeline_outputs
525 if self._is_queue_stale(priority_queue):
526 priority_queue = self._fill_queue(ordered_component_names, inputs, component_visits)
528 if isinstance(break_point, Breakpoint):
529 logger.warning(
530 "The given breakpoint {break_point} was never triggered. This is because:\n"
531 "1. The provided component is not a part of the pipeline execution path.\n"
532 "2. The component did not reach the visit count specified in the pipeline_breakpoint",
533 break_point=break_point,
534 )
536 # Set here so the tag reflects the final outputs.
537 span.set_content_tag("haystack.pipeline.output_data", pipeline_outputs)
539 return pipeline_outputs
541 @staticmethod
542 async def _run_component_async(
543 component_name: str,
544 component: dict[str, Any],
545 component_inputs: dict[str, Any],
546 component_visits: dict[str, int],
547 parent_span: tracing.Span | None = None,
548 *,
549 break_point: Breakpoint | None = None,
550 ) -> Mapping[str, Any]:
551 """
552 Executes a single component asynchronously.
554 If the component supports async execution, it is awaited directly as it will run async;
555 otherwise the component is offloaded to executor.
557 The method also updates the `visits` count of the component, writes outputs to `inputs_state`,
558 and returns pruned outputs that get stored in `pipeline_outputs`.
560 :param component_name: The name of the component.
561 :param component_inputs: Inputs for the component.
562 :returns: Outputs from the component that can be yielded from run_async_generator.
563 """
564 if (
565 isinstance(break_point, Breakpoint)
566 and break_point.component_name == component_name
567 and break_point.visit_count == component_visits[component_name]
568 ):
569 raise BreakpointException.from_triggered_breakpoint(break_point=break_point)
571 instance: Component = component["instance"]
573 with PipelineBase._create_component_span(
574 component_name=component_name, instance=instance, inputs=component_inputs, parent_span=parent_span
575 ) as span:
576 # deepcopy inputs before passing to the tracer so that even if a tracer mutates them
577 # the component always receives the original unmodified values
578 component_inputs_copy = _deepcopy_with_exceptions(component_inputs)
579 span.set_content_tag(_COMPONENT_INPUT, component_inputs)
580 logger.info("Running component {component_name}", component_name=component_name)
582 try:
583 # For sync-only components, _run_component_async dispatches to a thread via asyncio.to_thread,
584 # which copies the current contextvars context — preserving e.g. the active tracing span.
585 outputs = await _execute_component_async(instance, **component_inputs_copy)
586 except BreakpointException as error:
587 # Re-raise BreakpointException to preserve the original exception context, matching
588 # the sync _run_component: a breakpoint triggered by a nested component (e.g. an Agent)
589 # must bubble up to the main pipeline instead of being wrapped.
590 raise error
592 # A component that internally uses Pipeline._run_component_async could raise a PipelineRuntimeError
593 # carrying additional context (e.g. an agent snapshot); re-raise it instead of wrapping it in
594 # another PipelineRuntimeError, matching the sync _run_component.
595 except PipelineRuntimeError as runtime_error:
596 raise runtime_error
598 except Exception as error:
599 raise PipelineRuntimeError.from_exception(component_name, instance.__class__, error) from error
601 component_visits[component_name] += 1
603 if not isinstance(outputs, Mapping):
604 raise PipelineRuntimeError.from_invalid_output(component_name, instance.__class__, outputs)
606 _validate_component_output_keys(component_name, component, outputs)
608 span.set_tag(_COMPONENT_VISITS, component_visits[component_name])
609 span.set_content_tag(_COMPONENT_OUTPUT, outputs)
611 return outputs
613 @staticmethod
614 async def _wait_for_tasks(
615 running_tasks: dict[asyncio.Task, str], scheduled_components: set[str], *, return_when: str
616 ) -> AsyncIterator[dict[str, Any]]:
617 """
618 Waits for running tasks to finish and yields their partial outputs.
620 :param running_tasks: Mapping of in-flight tasks to the name of the component they run. Finished tasks are
621 removed in place.
622 :param scheduled_components: Set of component names that are scheduled but not yet finished. Finished
623 components are discarded in place.
624 :param return_when: Either `asyncio.FIRST_COMPLETED` to wait for a single task or `asyncio.ALL_COMPLETED` to
625 wait for every running task.
626 :returns: An async iterator of partial outputs, one per finished component that produced an output.
627 """
628 if not running_tasks:
629 return
631 done, _pending = await asyncio.wait(running_tasks.keys(), return_when=return_when)
632 for finished in done:
633 finished_component_name = running_tasks.pop(finished)
634 try:
635 partial_result = finished.result()
636 except Exception:
637 # A component failed. Cancel and drain the remaining in-flight tasks so they don't keep running in
638 # the background (and leak) after the run is aborted, then re-raise the original error.
639 await Pipeline._cancel_in_flight_tasks(running_tasks, scheduled_components)
640 raise
641 scheduled_components.discard(finished_component_name)
642 if partial_result:
643 yield {finished_component_name: _deepcopy_with_exceptions(partial_result)}
645 @staticmethod
646 async def _cancel_in_flight_tasks(running_tasks: dict[asyncio.Task, str], scheduled_components: set[str]) -> None:
647 """
648 Cancels all in-flight tasks and waits for the cancellations to settle.
650 Called when a component fails or when the run is abandoned early so that sibling tasks don't keep running in
651 the background after the pipeline run is aborted. Exceptions from the cancelled tasks are suppressed since we
652 are already unwinding.
654 Note: cancellation is only effective for components that run natively async. Sync components are offloaded to
655 a thread via `asyncio.to_thread` and a running thread cannot be interrupted: cancelling its task abandons
656 the await, but the thread keeps running until the component's `run` returns. Its outputs are then discarded
657 (the task never writes them to the pipeline state), so state stays consistent, but side effects (e.g. API
658 calls) still complete and the thread can outlive this cleanup.
660 :param running_tasks: Mapping of in-flight tasks to component names. Cleared in place.
661 :param scheduled_components: Set of scheduled-but-unfinished component names. Cleared in place.
662 """
663 for task in running_tasks:
664 task.cancel()
665 # return_exceptions=True so a failing or cancelled sibling doesn't mask the original error we re-raise.
666 await asyncio.gather(*running_tasks.keys(), return_exceptions=True)
667 for component_name in running_tasks.values():
668 scheduled_components.discard(component_name)
669 running_tasks.clear()
671 async def _run_component_in_isolation(
672 self,
673 *,
674 component_name: str,
675 inputs: dict[str, dict[str, list[dict[str, Any]]]],
676 pipeline_outputs: dict[str, Any],
677 component_visits: dict[str, int],
678 running_tasks: dict[asyncio.Task, str],
679 scheduled_components: set[str],
680 cached_receivers: dict[str, Any],
681 include_outputs_from: set[str],
682 parent_span: tracing.Span | None,
683 ) -> AsyncIterator[dict[str, Any]]:
684 """
685 Runs a component with HIGHEST priority in isolation.
687 We need to run components with HIGHEST priority (i.e. components with a GreedyVariadic input socket) by
688 themselves, without any other components running concurrently. Otherwise, downstream components could produce
689 additional inputs for the GreedyVariadic socket.
691 :param component_name: The name of the component to run.
692 :param inputs: The global input state shared by all components. Mutated in place.
693 :param pipeline_outputs: The accumulated pipeline outputs. Mutated in place.
694 :param component_visits: Current state of component visits. Mutated in place.
695 :param running_tasks: Mapping of in-flight tasks to component names. Drained in place before running.
696 :param scheduled_components: Set of scheduled-but-unfinished component names. Mutated in place.
697 :param cached_receivers: Precomputed mapping of component name to its downstream receivers.
698 :param include_outputs_from: Set of component names whose outputs should always be included in the output.
699 :param parent_span: The parent tracing span for the pipeline run.
700 :returns: An async iterator of partial outputs.
701 """
702 # 1) Wait for all in-flight tasks to finish so the HIGHEST component runs alone.
703 async for partial_outputs in self._wait_for_tasks(
704 running_tasks, scheduled_components, return_when=asyncio.ALL_COMPLETED
705 ):
706 yield partial_outputs
708 if component_name in scheduled_components:
709 # If it's already scheduled for some reason, skip.
710 return
712 # 2) Run the HIGHEST component by itself.
713 scheduled_components.add(component_name)
714 component = self._get_component_with_graph_metadata_and_visits(component_name, component_visits[component_name])
715 component_inputs = self._consume_component_inputs(component_name, component, inputs)
716 component_inputs = self._add_missing_input_defaults(component_inputs, component["input_sockets"])
718 component_outputs = await self._run_component_async(
719 component_name=component_name,
720 component=component,
721 component_inputs=component_inputs,
722 component_visits=component_visits,
723 parent_span=parent_span,
724 )
726 pruned = self._write_component_outputs(
727 component_name=component_name,
728 component_outputs=component_outputs,
729 inputs=inputs,
730 receivers=cached_receivers[component_name],
731 include_outputs_from=include_outputs_from,
732 )
733 if pruned or component_name in include_outputs_from:
734 pipeline_outputs[component_name] = pruned
736 scheduled_components.remove(component_name)
737 if pruned or component_name in include_outputs_from:
738 yield {component_name: _deepcopy_with_exceptions(pruned)}
740 def _schedule_component(
741 self,
742 *,
743 component_name: str,
744 inputs: dict[str, dict[str, list[dict[str, Any]]]],
745 pipeline_outputs: dict[str, Any],
746 component_visits: dict[str, int],
747 running_tasks: dict[asyncio.Task, str],
748 scheduled_components: set[str],
749 ready_sem: asyncio.Semaphore,
750 cached_receivers: dict[str, Any],
751 include_outputs_from: set[str],
752 parent_span: tracing.Span | None,
753 ) -> None:
754 """
755 Schedules a component to run as a background task without waiting for it to finish.
757 Inputs are consumed synchronously here (before the task is created) so that other components scheduled in the
758 same iteration of the scheduling loop observe the updated input state and don't race for the same inputs.
760 :param component_name: The name of the component to schedule.
761 :param inputs: The global input state shared by all components. Mutated in place.
762 :param pipeline_outputs: The accumulated pipeline outputs. Mutated in place by the task once it finishes.
763 :param component_visits: Current state of component visits. Mutated in place by the task once it finishes.
764 :param running_tasks: Mapping of in-flight tasks to component names. The new task is registered here.
765 :param scheduled_components: Set of scheduled-but-unfinished component names. Mutated in place.
766 :param ready_sem: Semaphore bounding how many components run concurrently.
767 :param cached_receivers: Precomputed mapping of component name to its downstream receivers.
768 :param include_outputs_from: Set of component names whose outputs should always be included in the output.
769 :param parent_span: The parent tracing span for the pipeline run.
770 """
771 if component_name in scheduled_components:
772 return # already scheduled, do nothing
774 scheduled_components.add(component_name)
776 component = self._get_component_with_graph_metadata_and_visits(component_name, component_visits[component_name])
777 component_inputs = self._consume_component_inputs(component_name, component, inputs)
778 component_inputs = self._add_missing_input_defaults(component_inputs, component["input_sockets"])
780 async def _runner() -> Mapping[str, Any]:
781 async with ready_sem:
782 component_outputs = await self._run_component_async(
783 component_name=component_name,
784 component=component,
785 component_inputs=component_inputs,
786 component_visits=component_visits,
787 parent_span=parent_span,
788 )
790 pruned = self._write_component_outputs(
791 component_name=component_name,
792 component_outputs=component_outputs,
793 inputs=inputs,
794 receivers=cached_receivers[component_name],
795 include_outputs_from=include_outputs_from,
796 )
797 if pruned or component_name in include_outputs_from:
798 pipeline_outputs[component_name] = pruned
800 scheduled_components.remove(component_name)
801 return pruned
803 task = asyncio.create_task(_runner())
804 running_tasks[task] = component_name
806 @mark_deserialization_internal
807 async def run_async_generator( # noqa: PLR0915,C901
808 self, data: dict[str, Any], include_outputs_from: set[str] | None = None, concurrency_limit: int = 4
809 ) -> AsyncGenerator[dict[str, Any], None]:
810 """
811 Executes the pipeline step by step asynchronously, yielding partial outputs when any component finishes.
813 Usage:
814 ```python
815 from haystack import Document
816 from haystack.components.builders import ChatPromptBuilder
817 from haystack.dataclasses import ChatMessage
818 from haystack.utils import Secret
819 from haystack.document_stores.in_memory import InMemoryDocumentStore
820 from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
821 from haystack.components.generators.chat import OpenAIResponsesChatGenerator
822 from haystack.components.builders.prompt_builder import PromptBuilder
823 from haystack import Pipeline
824 import asyncio
826 # Write documents to InMemoryDocumentStore
827 document_store = InMemoryDocumentStore()
828 document_store.write_documents([
829 Document(content="My name is Jean and I live in Paris."),
830 Document(content="My name is Mark and I live in Berlin."),
831 Document(content="My name is Giorgio and I live in Rome.")
832 ])
834 prompt_template = [
835 ChatMessage.from_user(
836 '''
837 Given these documents, answer the question.
838 Documents:
839 {% for doc in documents %}
840 {{ doc.content }}
841 {% endfor %}
842 Question: {{question}}
843 Answer:
844 ''')
845 ]
847 # Create and connect pipeline components
848 retriever = InMemoryBM25Retriever(document_store=document_store)
849 prompt_builder = ChatPromptBuilder(template=prompt_template)
850 llm = OpenAIResponsesChatGenerator()
852 rag_pipeline = Pipeline()
853 rag_pipeline.add_component("retriever", retriever)
854 rag_pipeline.add_component("prompt_builder", prompt_builder)
855 rag_pipeline.add_component("llm", llm)
856 rag_pipeline.connect("retriever", "prompt_builder.documents")
857 rag_pipeline.connect("prompt_builder", "llm")
859 # Prepare input data
860 question = "Who lives in Paris?"
861 data = {
862 "retriever": {"query": question},
863 "prompt_builder": {"question": question},
864 }
867 # Process results as they become available
868 async def process_results():
869 async for partial_output in rag_pipeline.run_async_generator(
870 data=data,
871 include_outputs_from={"retriever", "llm"}
872 ):
873 # Each partial_output contains the results from a completed component
874 if "retriever" in partial_output:
875 print("Retrieved documents:", len(partial_output["retriever"]["documents"]))
876 if "llm" in partial_output:
877 print("Generated answer:", partial_output["llm"]["replies"][0])
880 asyncio.run(process_results())
881 ```
883 :param data: Initial input data to the pipeline.
884 :param concurrency_limit: The maximum number of components that are allowed to run concurrently.
885 :param include_outputs_from:
886 Set of component names whose individual outputs are to be
887 included in the pipeline's output. For components that are
888 invoked multiple times (in a loop), only the last-produced
889 output is included.
890 :return: An async iterator containing partial (and final) outputs.
892 :raises ValueError:
893 If invalid inputs are provided to the pipeline, or if `concurrency_limit` is less than 1.
894 :raises PipelineMaxComponentRuns:
895 If a component exceeds the maximum number of allowed executions within the pipeline.
896 :raises PipelineRuntimeError:
897 If the Pipeline contains cycles with unsupported connections that would cause
898 it to get stuck and fail running.
899 Or if a Component fails or returns output in an unsupported type.
900 """
901 if concurrency_limit < 1:
902 raise ValueError("concurrency_limit must be greater than or equal to 1.")
904 # warm up the pipeline by running each component's warm_up_async (or warm_up) method
905 await self.warm_up_async()
907 # after warm-up, so toolsets that load their tools in warm_up() report them
908 pipeline_running(self) # telemetry
910 if include_outputs_from is None:
911 include_outputs_from = set()
913 pipeline_outputs: dict[str, Any] = {}
915 # Normalize `data` and raise ValueError if the input is malformed in some way.
916 data = self._prepare_component_input_data(data)
918 # Raise ValueError if input is malformed in some way
919 self.validate_input(data)
921 # We create a list of components in the pipeline sorted by name, so that the algorithm runs
922 # deterministically and independent of insertion order into the pipeline.
923 ordered_component_names = sorted(self.graph.nodes.keys())
925 # We track component visits to decide if a component can run.
926 component_visits = dict.fromkeys(ordered_component_names, 0)
928 cached_topological_sort = None
929 # We need to access a component's receivers multiple times during a pipeline run.
930 # We store them here for easy access.
931 cached_receivers = {name: self._find_receivers_from(name) for name in ordered_component_names}
933 # Ephemeral concurrency state shared (and mutated in place) by the scheduling helpers below.
934 ready_sem = asyncio.Semaphore(concurrency_limit)
935 running_tasks: dict[asyncio.Task, str] = {}
936 # A set of component names that have been scheduled but not finished.
937 scheduled_components: set[str] = set()
939 with tracing.tracer.trace(
940 "haystack.pipeline.run",
941 tags={
942 "haystack.pipeline.input_data": data,
943 "haystack.pipeline.metadata": self.metadata,
944 "haystack.pipeline.max_runs_per_component": self._max_runs_per_component,
945 "haystack.pipeline.execution_mode": "async",
946 },
947 ) as parent_span:
948 inputs = self._convert_to_internal_format(pipeline_inputs=data)
950 # check if pipeline is blocked before execution
951 self.validate_pipeline(self._fill_queue(ordered_component_names, inputs, component_visits))
953 try:
954 while True:
955 # We rebuild the priority queue every iteration: each iteration waits for one or more concurrent
956 # tasks to finish, which mutates `inputs` and can change many components' priorities at once, so
957 # we rebuild to give every scheduling decision an up-to-date view.
958 priority_queue = self._fill_queue(ordered_component_names, inputs, component_visits)
959 candidate = self._get_next_runnable_component(priority_queue, component_visits)
961 # If we can't make progress with the queue but tasks are running, we wait for one to finish and
962 # retry to potentially unblock the priority queue.
963 if (candidate is None or candidate[0] == ComponentPriority.BLOCKED) and running_tasks:
964 async for partial_outputs in self._wait_for_tasks(
965 running_tasks, scheduled_components, return_when=asyncio.FIRST_COMPLETED
966 ):
967 yield partial_outputs
968 continue
970 # If there are no runnable components left and nothing is running, we can exit the loop.
971 if candidate is None and not running_tasks:
972 break
974 priority, component_name, component = candidate # type: ignore
976 # If the next component is blocked, we do a check to see if the pipeline is possibly blocked and
977 # raise a warning if it is.
978 if priority == ComponentPriority.BLOCKED and not running_tasks:
979 if self._is_pipeline_possibly_blocked(current_pipeline_outputs=pipeline_outputs):
980 # Pipeline is most likely blocked (most likely a configuration issue) so we raise a warning.
981 self._find_components_blocking_pipeline(
982 priority_queue=priority_queue, component_visits=component_visits, inputs=inputs
983 )
984 # We always exit the loop since we cannot run the next component.
985 break
987 # If the next component is already scheduled, we wait for a task to finish to make progress.
988 if component_name in scheduled_components:
989 async for partial_outputs in self._wait_for_tasks(
990 running_tasks, scheduled_components, return_when=asyncio.FIRST_COMPLETED
991 ):
992 yield partial_outputs
993 continue
995 if priority == ComponentPriority.HIGHEST:
996 # A HIGHEST priority component must run alone, so we hand off to the isolation helper.
997 async for partial_outputs in self._run_component_in_isolation(
998 component_name=component_name,
999 inputs=inputs,
1000 pipeline_outputs=pipeline_outputs,
1001 component_visits=component_visits,
1002 running_tasks=running_tasks,
1003 scheduled_components=scheduled_components,
1004 cached_receivers=cached_receivers,
1005 include_outputs_from=include_outputs_from,
1006 parent_span=parent_span,
1007 ):
1008 yield partial_outputs
1009 continue
1011 if priority == ComponentPriority.READY:
1012 # Schedule this component, then schedule as many additional READY components as concurrency
1013 # allows.
1014 self._schedule_component(
1015 component_name=component_name,
1016 inputs=inputs,
1017 pipeline_outputs=pipeline_outputs,
1018 component_visits=component_visits,
1019 running_tasks=running_tasks,
1020 scheduled_components=scheduled_components,
1021 ready_sem=ready_sem,
1022 cached_receivers=cached_receivers,
1023 include_outputs_from=include_outputs_from,
1024 parent_span=parent_span,
1025 )
1027 # Possibly schedule more READY tasks if concurrency not fully used
1028 while len(priority_queue) > 0 and not ready_sem.locked():
1029 peek_priority, peek_name = priority_queue.peek()
1030 if peek_priority != ComponentPriority.READY:
1031 # We stop scheduling: the next component is BLOCKED (can't run), HIGHEST (must run
1032 # alone), or DEFER (waiting for more inputs - we only schedule it once it becomes
1033 # READY).
1034 break
1035 priority_queue.pop()
1036 self._schedule_component(
1037 component_name=peek_name,
1038 inputs=inputs,
1039 pipeline_outputs=pipeline_outputs,
1040 component_visits=component_visits,
1041 running_tasks=running_tasks,
1042 scheduled_components=scheduled_components,
1043 ready_sem=ready_sem,
1044 cached_receivers=cached_receivers,
1045 include_outputs_from=include_outputs_from,
1046 parent_span=parent_span,
1047 )
1049 # We only schedule components with priority DEFER when no other tasks are running.
1050 elif priority == ComponentPriority.DEFER and not running_tasks:
1051 if len(priority_queue) > 0:
1052 component_name, cached_topological_sort = self._tiebreak_waiting_components(
1053 component_name=component_name,
1054 priority=priority,
1055 priority_queue=priority_queue,
1056 topological_sort=cached_topological_sort,
1057 )
1059 self._schedule_component(
1060 component_name=component_name,
1061 inputs=inputs,
1062 pipeline_outputs=pipeline_outputs,
1063 component_visits=component_visits,
1064 running_tasks=running_tasks,
1065 scheduled_components=scheduled_components,
1066 ready_sem=ready_sem,
1067 cached_receivers=cached_receivers,
1068 include_outputs_from=include_outputs_from,
1069 parent_span=parent_span,
1070 )
1072 # To make progress, we wait for one task to complete before restarting the loop.
1073 async for partial_outputs in self._wait_for_tasks(
1074 running_tasks, scheduled_components, return_when=asyncio.FIRST_COMPLETED
1075 ):
1076 yield partial_outputs
1078 # Safety net: drain any leftover tasks once the scheduling loop has finished. With the current loop
1079 # both `break` paths require `running_tasks` to be empty, so this is a no-op. We keep it so that a
1080 # future change adding a `break` that leaves tasks in flight doesn't lose outputs.
1081 async for partial_outputs in self._wait_for_tasks(
1082 running_tasks, scheduled_components, return_when=asyncio.ALL_COMPLETED
1083 ):
1084 yield partial_outputs
1086 # Set here so the tag reflects the final outputs.
1087 parent_span.set_content_tag("haystack.pipeline.output_data", pipeline_outputs)
1089 # Yield the final pipeline outputs.
1090 yield pipeline_outputs
1091 finally:
1092 # If iteration is abandoned early (e.g. the consumer stops iterating the generator and closes it) or
1093 # the run is cancelled, cancel any tasks still in flight so they don't leak.
1094 # This is a no-op on normal completion and on a component error, since no tasks are left running by then
1095 await self._cancel_in_flight_tasks(running_tasks, scheduled_components)
1097 @mark_deserialization_internal
1098 async def run_async(
1099 self, data: dict[str, Any], include_outputs_from: set[str] | None = None, concurrency_limit: int = 4
1100 ) -> dict[str, Any]:
1101 """
1102 Provides an asynchronous interface to run the pipeline with provided input data.
1104 This method allows the pipeline to be integrated into an asynchronous workflow, enabling non-blocking
1105 execution of pipeline components.
1107 Usage:
1108 ```python
1109 import asyncio
1111 from haystack import Document
1112 from haystack.components.builders import ChatPromptBuilder
1113 from haystack.components.generators.chat import OpenAIResponsesChatGenerator
1114 from haystack.components.retrievers.in_memory import InMemoryBM25Retriever
1115 from haystack import Pipeline
1116 from haystack.dataclasses import ChatMessage
1117 from haystack.document_stores.in_memory import InMemoryDocumentStore
1119 # Write documents to InMemoryDocumentStore
1120 document_store = InMemoryDocumentStore()
1121 document_store.write_documents([
1122 Document(content="My name is Jean and I live in Paris."),
1123 Document(content="My name is Mark and I live in Berlin."),
1124 Document(content="My name is Giorgio and I live in Rome.")
1125 ])
1127 prompt_template = [
1128 ChatMessage.from_user(
1129 '''
1130 Given these documents, answer the question.
1131 Documents:
1132 {% for doc in documents %}
1133 {{ doc.content }}
1134 {% endfor %}
1135 Question: {{question}}
1136 Answer:
1137 ''')
1138 ]
1140 retriever = InMemoryBM25Retriever(document_store=document_store)
1141 prompt_builder = ChatPromptBuilder(template=prompt_template)
1142 llm = OpenAIResponsesChatGenerator()
1144 rag_pipeline = Pipeline()
1145 rag_pipeline.add_component("retriever", retriever)
1146 rag_pipeline.add_component("prompt_builder", prompt_builder)
1147 rag_pipeline.add_component("llm", llm)
1148 rag_pipeline.connect("retriever", "prompt_builder.documents")
1149 rag_pipeline.connect("prompt_builder", "llm")
1151 # Ask a question
1152 question = "Who lives in Paris?"
1154 async def run_inner(data, include_outputs_from):
1155 return await rag_pipeline.run_async(data=data, include_outputs_from=include_outputs_from)
1157 data = {
1158 "retriever": {"query": question},
1159 "prompt_builder": {"question": question},
1160 }
1162 results = asyncio.run(run_inner(data, include_outputs_from={"retriever", "llm"}))
1164 print(results["llm"]["replies"][0].text)
1165 # >> Jean lives in Paris.
1166 ```
1168 :param data:
1169 A dictionary of inputs for the pipeline's components. Each key is a component name
1170 and its value is a dictionary of that component's input parameters:
1171 ```
1172 data = {
1173 "comp1": {"input1": 1, "input2": 2},
1174 }
1175 ```
1176 For convenience, this format is also supported when input names are unique:
1177 ```
1178 data = {
1179 "input1": 1, "input2": 2,
1180 }
1181 ```
1182 :param include_outputs_from:
1183 Set of component names whose individual outputs are to be
1184 included in the pipeline's output. For components that are
1185 invoked multiple times (in a loop), only the last-produced
1186 output is included.
1187 :param concurrency_limit: The maximum number of components that should be allowed to run concurrently.
1188 :returns:
1189 A dictionary where each entry corresponds to a component name
1190 and its output. If `include_outputs_from` is `None`, this dictionary
1191 will only contain the outputs of leaf components, i.e., components
1192 without outgoing connections.
1194 :raises ValueError:
1195 If invalid inputs are provided to the pipeline, or if `concurrency_limit` is less than 1.
1196 :raises PipelineRuntimeError:
1197 If the Pipeline contains cycles with unsupported connections that would cause
1198 it to get stuck and fail running.
1199 Or if a Component fails or returns output in an unsupported type.
1200 :raises PipelineMaxComponentRuns:
1201 If a Component reaches the maximum number of times it can be run in this Pipeline.
1202 """
1203 final: dict[str, Any] = {}
1204 async for partial in self.run_async_generator(
1205 data=data, concurrency_limit=concurrency_limit, include_outputs_from=include_outputs_from
1206 ):
1207 final = partial
1208 return final or {}
1210 @mark_deserialization_internal
1211 def stream(
1212 self,
1213 data: dict[str, Any],
1214 *,
1215 streaming_components: list[str] | None = None,
1216 include_outputs_from: set[str] | None = None,
1217 concurrency_limit: int = 4,
1218 cancel_on_abandon: bool = True,
1219 ) -> PipelineStreamHandle:
1220 """
1221 Run the pipeline and return a handle that streams `StreamingChunk`s as they arrive.
1223 Iterate the handle with `async for` to consume chunks; after iteration ends, `handle.result` holds the final
1224 pipeline output dict (same as `run_async`). By default, if iteration is abandoned, the underlying pipeline task
1225 is cancelled automatically. Pass `cancel_on_abandon=False` to instead let the pipeline run to completion.
1227 For every async-capable component that exposes a `streaming_callback` input socket, a forwarder is injected at
1228 runtime that pushes chunks onto the handle's queue. If a `streaming_callback` is provided at component init or
1229 at runtime (inside `data`, e.g. `data={"llm": {"streaming_callback": cb}}`), it is also invoked for each chunk.
1230 Async callbacks are preferred; a sync callback is accepted but will run synchronously on the event loop and
1231 may block it.
1233 Usage:
1234 ```python
1235 import asyncio
1237 from haystack.components.builders import ChatPromptBuilder
1238 from haystack.components.generators.chat import OpenAIResponsesChatGenerator
1239 from haystack import Pipeline
1240 from haystack.dataclasses import ChatMessage
1242 pipe = Pipeline()
1243 pipe.add_component(
1244 "prompt_builder",
1245 ChatPromptBuilder(template=[ChatMessage.from_user("Tell me about {{topic}}")]),
1246 )
1247 pipe.add_component("llm", OpenAIResponsesChatGenerator())
1248 pipe.connect("prompt_builder.prompt", "llm.messages")
1250 async def main():
1251 handle = pipe.stream(data={"prompt_builder": {"topic": "Italy"}})
1252 async for chunk in handle:
1253 print(chunk.content, end="", flush=True)
1254 return handle.result
1256 result = asyncio.run(main())
1257 print(result["llm"]["replies"])
1258 ```
1260 :param data:
1261 A dictionary of inputs for the pipeline's components. Each key is a component name
1262 and its value is a dictionary of that component's input parameters:
1263 ```
1264 data = {
1265 "comp1": {"input1": 1, "input2": 2},
1266 }
1267 ```
1268 For convenience, this format is also supported when input names are unique:
1269 ```
1270 data = {
1271 "input1": 1, "input2": 2,
1272 }
1273 ```
1274 :param streaming_components: Names of components to stream from. If `None` (default), every streaming-capable
1275 component is forwarded. If a list, only the listed components are forwarded; unknown names or names of
1276 components that do not support streaming raise `ValueError`.
1277 :param include_outputs_from:
1278 Set of component names whose individual outputs are to be
1279 included in the pipeline's output. For components that are
1280 invoked multiple times (in a loop), only the last-produced
1281 output is included.
1282 :param concurrency_limit: The maximum number of components that should be allowed to run concurrently.
1283 :param cancel_on_abandon: If `True` (default), the underlying pipeline task is cancelled when iteration is
1284 abandoned. If `False`, the pipeline runs to completion even when the consumer stops reading.
1285 :returns:
1286 A `PipelineStreamHandle` that is async-iterable over `StreamingChunk`s. After iteration ends,
1287 `handle.result` holds the final pipeline output dict (same shape as `run_async`).
1289 :raises ValueError:
1290 If `streaming_components` contains unknown component names or components that do not support streaming,
1291 or if invalid inputs are provided to the pipeline, or if `concurrency_limit` is less than 1.
1292 :raises PipelineRuntimeError:
1293 Surfaced during iteration. If the Pipeline contains cycles with unsupported connections that would cause
1294 it to get stuck and fail running, or if a Component fails or returns output in an unsupported type.
1295 :raises PipelineMaxComponentRuns:
1296 Surfaced during iteration. If a Component reaches the maximum number of times it can be run in this
1297 Pipeline.
1298 """
1299 streaming_capable = {
1300 name
1301 for name in self.graph.nodes
1302 if getattr(self.graph.nodes[name]["instance"], "__haystack_supports_async__", False)
1303 and "streaming_callback" in self.graph.nodes[name]["instance"].__haystack_input__
1304 }
1305 if streaming_components is not None:
1306 requested = set(streaming_components)
1307 unknown = requested - set(self.graph.nodes)
1308 non_streaming = requested - unknown - streaming_capable
1309 if unknown:
1310 raise ValueError(f"Unknown components in streaming_components: {sorted(unknown)}")
1311 if non_streaming:
1312 raise ValueError(f"These components do not support streaming: {sorted(non_streaming)}")
1314 queue: asyncio.Queue[StreamingChunk | _EndOfStream] = asyncio.Queue()
1316 def make_forwarder(user_callback: StreamingCallbackT | None) -> AsyncStreamingCallbackT:
1317 async def forwarder(chunk: StreamingChunk) -> None:
1318 await queue.put(chunk)
1319 if user_callback is not None:
1320 await _invoke_streaming_callback(user_callback, chunk)
1322 return forwarder
1324 new_data: dict[str, Any] = self._prepare_component_input_data(data)
1325 for name in streaming_capable:
1326 if streaming_components is not None and name not in streaming_components:
1327 continue
1328 instance = self.graph.nodes[name]["instance"]
1329 comp_inputs = new_data.setdefault(name, {})
1331 user_callback = select_streaming_callback(
1332 init_callback=getattr(instance, "streaming_callback", None),
1333 runtime_callback=comp_inputs.get("streaming_callback"),
1334 requires_async=True,
1335 )
1336 comp_inputs["streaming_callback"] = make_forwarder(user_callback)
1338 async def runner() -> dict[str, Any]:
1339 try:
1340 return await self.run_async(
1341 new_data, include_outputs_from=include_outputs_from, concurrency_limit=concurrency_limit
1342 )
1343 finally:
1344 await queue.put(PipelineStreamHandle._END_OF_STREAM)
1346 task = asyncio.create_task(runner())
1347 return PipelineStreamHandle(queue=queue, task=task, cancel_on_abandon=cancel_on_abandon)