haystack/test/core/pipeline/test_looping_pipeline.py
Silvano Cerza d4f6531c52
feat: Refactor Pipeline.run() (#6729)
* First rough implementation of refactored run

* Further improve run logic

* Properly handle variadic input in run

* Further work

* Enhance names and add more documentation

* Fix issue with output distribution

* This works

* Enhance run comments

* Mark Multiplexer as greedy

* Remove MergeLoop in favour of Multiplexer in tests

* Remove FirstIntSelector in favour of Multiplexer

* Handle corner when waiting for input is stuck

* Remove unused import

* Handle mutable input data in run and misbehaving components

* Handle run input validation

* Test validation

* Fix pylint

* Fix mypy

* Call warm_up in run to fix tests
2024-01-18 17:53:47 +01:00

87 lines
3.2 KiB
Python

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0
import logging
from haystack.components.others import Multiplexer
from haystack.core.pipeline import Pipeline
from haystack.testing.sample_components import Accumulate, AddFixedValue, Threshold
logging.basicConfig(level=logging.DEBUG)
def test_pipeline():
accumulator = Accumulate()
pipeline = Pipeline(max_loops_allowed=10)
pipeline.add_component("add_one", AddFixedValue(add=1))
pipeline.add_component("multiplexer", Multiplexer(type_=int))
pipeline.add_component("below_10", Threshold(threshold=10))
pipeline.add_component("accumulator", accumulator)
pipeline.add_component("add_two", AddFixedValue(add=2))
pipeline.connect("add_one.result", "multiplexer")
pipeline.connect("multiplexer.value", "below_10.value")
pipeline.connect("below_10.below", "accumulator.value")
pipeline.connect("accumulator.value", "multiplexer")
pipeline.connect("below_10.above", "add_two.value")
results = pipeline.run({"add_one": {"value": 3}})
assert results == {"add_two": {"result": 18}}
assert accumulator.state == 16
def test_pipeline_direct_io_loop():
accumulator = Accumulate()
pipeline = Pipeline(max_loops_allowed=10)
pipeline.add_component("multiplexer", Multiplexer(type_=int))
pipeline.add_component("below_10", Threshold(threshold=10))
pipeline.add_component("accumulator", accumulator)
pipeline.connect("multiplexer.value", "below_10.value")
pipeline.connect("below_10.below", "accumulator.value")
pipeline.connect("accumulator.value", "multiplexer")
results = pipeline.run({"multiplexer": {"value": 4}})
assert results == {"below_10": {"above": 16}}
assert accumulator.state == 16
def test_pipeline_fixed_merger_input():
accumulator = Accumulate()
pipeline = Pipeline(max_loops_allowed=10)
pipeline.add_component("multiplexer", Multiplexer(type_=int))
pipeline.add_component("below_10", Threshold(threshold=10))
pipeline.add_component("accumulator", accumulator)
pipeline.add_component("add_two", AddFixedValue(add=2))
pipeline.connect("multiplexer.value", "below_10.value")
pipeline.connect("below_10.below", "accumulator.value")
pipeline.connect("accumulator.value", "multiplexer")
pipeline.connect("below_10.above", "add_two.value")
results = pipeline.run({"multiplexer": {"value": 4}})
assert results == {"add_two": {"result": 18}}
assert accumulator.state == 16
def test_pipeline_variadic_merger_input():
accumulator = Accumulate()
pipeline = Pipeline(max_loops_allowed=10)
pipeline.add_component("multiplexer", Multiplexer(type_=int))
pipeline.add_component("below_10", Threshold(threshold=10))
pipeline.add_component("accumulator", accumulator)
pipeline.add_component("add_two", AddFixedValue(add=2))
pipeline.connect("multiplexer", "below_10.value")
pipeline.connect("below_10.below", "accumulator.value")
pipeline.connect("accumulator.value", "multiplexer.value")
pipeline.connect("below_10.above", "add_two.value")
results = pipeline.run({"multiplexer": {"value": 4}})
assert results == {"add_two": {"result": 18}}
assert accumulator.state == 16