deepset-ai--haystack/test/core/pipeline/test_pipeline.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import logging
from concurrent.futures import ThreadPoolExecutor

import pytest

from haystack.components.agents import Agent
from haystack.components.joiners import BranchJoiner
from haystack.core.component import component
from haystack.core.errors import PipelineConnectError, PipelineRuntimeError
from haystack.core.pipeline import Pipeline
from haystack.dataclasses import ChatMessage, Document


@component
class MockChatGenerator:
    @component.output_types(replies=list[ChatMessage])
    def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
        return {"replies": [ChatMessage.from_assistant("Hello, world!")]}


@component
class StringProducer:
    def __init__(self, value: str = "Hello"):
        self.value = value

    @component.output_types(text=str)
    def run(self) -> dict[str, str]:
        return {"text": self.value}


@component
class ListStrProducer:
    def __init__(self, values: list[str] | None = None):
        self.values = values or ["Hello", "Hi"]

    @component.output_types(texts=list[str])
    def run(self) -> dict[str, list[str]]:
        return {"texts": self.values}


@component
class ListStrAcceptor:
    @component.output_types(result=list[str])
    def run(self, texts: list[str]) -> dict[str, list[str]]:
        return {"result": texts}


@component
class ChatMessageProducer:
    def __init__(self, value: str = "Hello"):
        self.value = value

    @component.output_types(message=ChatMessage)
    def run(self) -> dict[str, ChatMessage]:
        return {"message": ChatMessage.from_user(self.value)}


@component
class ListChatMessageProducer:
    def __init__(self, values: list[str] | None = None):
        self.values = values or ["Hello", "Hi"]

    @component.output_types(messages=list[ChatMessage])
    def run(self) -> dict[str, list[ChatMessage]]:
        return {"messages": [ChatMessage.from_user(v) for v in self.values]}


@component
class ListChatMessageAcceptor:
    @component.output_types(result=list[ChatMessage])
    def run(self, messages: list[ChatMessage]) -> dict[str, list[ChatMessage]]:
        return {"result": messages}


@component
class WrongOutput:
    @component.output_types(output=str)
    def run(self, value: str) -> dict[str, str]:
        return "not_a_dict"  # type: ignore[return-value]


class TestPipeline:
    """
    This class contains only unit tests for the Pipeline class.

    It doesn't test Pipeline.run(), that is done separately in a different way.
    """

    def test_pipeline_thread_safety(self, waiting_component, spying_tracer):
        # Initialize pipeline with synchronous components
        pp = Pipeline()
        pp.add_component("wait", waiting_component())

        run_data = [{"wait_for": 0.001}, {"wait_for": 0.002}]

        # Use ThreadPoolExecutor to run pipeline calls in parallel
        with ThreadPoolExecutor(max_workers=len(run_data)) as executor:
            # Submit pipeline runs to the executor
            futures = [executor.submit(pp.run, data) for data in run_data]

            # Wait for all futures to complete
            for future in futures:
                future.result()

        # Verify component visits using tracer
        component_spans = [sp for sp in spying_tracer.spans if sp.operation_name == "haystack.component.run"]

        for span in component_spans:
            assert span.tags["haystack.component.visits"] == 1

    def test_prepare_component_inputs(self):
        pp = Pipeline()
        component_name = "joiner_1"
        pp.add_component(component_name, BranchJoiner(type_=str))
        pp.add_component("joiner_2", BranchJoiner(type_=str))
        pp.connect(component_name, "joiner_2")
        inputs = {"joiner_1": {"value": [{"sender": None, "value": "test_value"}]}}
        comp_dict = pp._get_component_with_graph_metadata_and_visits(component_name, 0)

        _ = pp._consume_component_inputs(component_name=component_name, component=comp_dict, inputs=inputs)
        # We remove input in greedy variadic sockets, even if they are from the user
        assert inputs == {"joiner_1": {}}

    def test__run_component_success(self):
        """Test successful component execution"""
        pp = Pipeline()
        component_name = "joiner_1"
        pp.add_component(component_name, BranchJoiner(type_=str))
        pp.add_component("joiner_2", BranchJoiner(type_=str))
        pp.connect(component_name, "joiner_2")

        outputs = pp._run_component(
            component_name=component_name,
            component=pp._get_component_with_graph_metadata_and_visits(component_name, 0),
            inputs={"value": ["test_value"]},
            component_visits={component_name: 0, "joiner_2": 0},
        )
        assert outputs == {"value": "test_value"}

    def test__run_component_fail(self):
        """Test error when component doesn't return a dictionary"""
        pp = Pipeline()
        pp.add_component("wrong", WrongOutput())

        with pytest.raises(PipelineRuntimeError) as exc_info:
            pp._run_component(
                component_name="wrong",
                component=pp._get_component_with_graph_metadata_and_visits("wrong", 0),
                inputs={"value": "test_value"},
                component_visits={"wrong": 0},
            )
        assert "Expected a dict" in str(exc_info.value)

    def test_run_component_error(self):
        """Test error when component fails to run"""

        @component
        class ErroringComponent:
            @component.output_types(output=str)
            def run(self):
                raise ValueError("Test error")

        pp = Pipeline()
        pp.add_component("erroring_component", ErroringComponent())

        with pytest.raises(PipelineRuntimeError) as exc_info:
            pp._run_component(
                component_name="erroring_component",
                component=pp._get_component_with_graph_metadata_and_visits("erroring_component", 0),
                inputs={"wrong": {"value": [{"sender": None, "value": "test_value"}]}},
                component_visits={"erroring_component": 0},
            )
        assert "Component name: 'erroring_component'" in str(exc_info.value)

    def test_component_with_empty_dict_as_output_appears_in_results(self):
        """Test that components that return an empty dict as output appear in results as an empty dict"""

        @component
        class Producer:
            def __init__(self, prefix: str):
                self.prefix = prefix

            @component.output_types(value=str | None)
            def run(self, text: str | None) -> dict[str, str | None]:
                return {"value": f"{self.prefix}: {text}"}

        @component
        class EmptyProcessor:
            @component.output_types()
            def run(self, sources: list[str]) -> dict:
                # Returns empty dict when sources is empty
                return {}

        @component
        class Combiner:
            @component.output_types(combined=str)
            def run(self, input_a: str | None, input_b: str | None) -> dict[str, str]:
                if input_a is None:
                    input_a = ""
                if input_b is None:
                    input_b = ""
                return {"combined": f"{input_a} | {input_b}"}

        pp = Pipeline()
        pp.add_component("producer_a", Producer("A"))
        pp.add_component("producer_b", Producer("B"))
        pp.add_component("empty_processor", EmptyProcessor())
        pp.add_component("combiner", Combiner())

        pp.connect("producer_a.value", "combiner.input_a")
        pp.connect("producer_b.value", "combiner.input_b")

        result = pp.run(
            {"producer_a": {"text": "hello"}, "producer_b": {"text": "world"}, "empty_processor": {"sources": []}},
            include_outputs_from={"producer_a", "empty_processor", "combiner"},
        )

        # Producer A should appear in results because it's in include_outputs_from
        assert "producer_a" in result
        assert result["producer_a"] == {"value": "A: hello"}
        # Producer B should NOT appear since it's not in include_outputs_from
        assert "producer_b" not in result
        # Combiner should appear in results
        assert "combiner" in result
        assert result["combiner"] == {"combined": "A: hello | B: world"}
        # Empty processor should appear in results even though it returns an empty dict
        # because it's in include_outputs_from
        assert "empty_processor" in result
        assert result["empty_processor"] == {}

    def test__run_component_warns_on_extra_output_keys(self, caplog):
        """Test that a warning is raised when a component returns undeclared output keys."""
        caplog.set_level(logging.WARNING)

        @component
        class ExtraKeyComponent:
            @component.output_types(output=str)
            def run(self, value: str) -> dict[str, str]:
                return {"output": value, "extra_key": "unexpected"}

        pp = Pipeline()
        pp.add_component("extra", ExtraKeyComponent())

        pp._run_component(
            component_name="extra",
            component=pp._get_component_with_graph_metadata_and_visits("extra", 0),
            inputs={"value": "test"},
            component_visits={"extra": 0},
        )
        assert "returned output keys" in caplog.text
        assert "extra_key" in caplog.text
        assert "not declared" in caplog.text

    def test__run_component_no_warning_on_correct_output_keys(self, caplog):
        """Test that no warning is raised when a component returns the correct output keys."""
        caplog.set_level(logging.WARNING)

        @component
        class CorrectComponent:
            @component.output_types(output=str)
            def run(self, value: str) -> dict[str, str]:
                return {"output": value}

        pp = Pipeline()
        pp.add_component("correct", CorrectComponent())

        pp._run_component(
            component_name="correct",
            component=pp._get_component_with_graph_metadata_and_visits("correct", 0),
            inputs={"value": "test"},
            component_visits={"correct": 0},
        )
        assert "returned output keys" not in caplog.text
        assert "did not produce output keys" not in caplog.text

    def test_pipeline_is_possibly_blocked_warning_message(self, caplog):
        """
        Test that the pipeline raises a warning when it is possibly blocked due to missing inputs.

        The situation below looks a little contrived, but it has happened in practice that users create pipelines
        and accidentally made a mistake in their component code.
        """
        caplog.set_level(logging.WARNING)

        @component
        class MisconfiguredComponent:
            # Here we purposely declare other_output which is not actually returned by the run() method
            @component.output_types(output=str, other_output=str)
            def run(self, required_input: str) -> dict[str, str]:
                return {"output": "test"}

        @component
        class SimpleComponentTwoInputs:
            @component.output_types(output=str)
            def run(self, required_input: str, second_required_input: str) -> dict[str, str]:
                return {"output": "test"}

        pp = Pipeline()
        pp.add_component("first", MisconfiguredComponent())
        pp.add_component("second", SimpleComponentTwoInputs())

        # NOTE: We connect both outputs from the first component to the second component, but the first component
        # doesn't actually produce other_output, so the second component will be blocked due to missing input.
        pp.connect("first.output", "second.required_input")
        pp.connect("first.other_output", "second.second_required_input")

        pp.run({"first": {"required_input": "test"}})
        assert "Cannot run pipeline - the pipeline appears to be blocked." in caplog.text
        assert " - 'second' (SimpleComponentTwoInputs)" in caplog.text

    def test_pipeline_ensure_inputs_are_deep_copied(self):
        """
        Test to ensure that pipeline deep copies the inputs before passing them to components.

        This is important to prevent unintended side effects when components modify their inputs especially when
        the output from one component is passed to multiple other components.

        Some other notes about how this situation can arise in practice:
        - When a component returns a mutable object (like a Document) and that output is passed to multiple other
          components.
        - This doesn't happen when using output types like strings or integers, because they are not shared by
          reference so we will only commonly see this for objects like our dataclasses.
        """

        @component
        class SimpleComponent:
            @component.output_types(output=Document)
            def run(self, document: Document) -> dict[str, Document]:
                # Creates a new document to avoid modifying in place
                new_document = Document(content=document.content)
                return {"output": new_document}

        @component
        class ModifyingComponent:
            @component.output_types(output=Document)
            def run(self, document: Document) -> dict[str, Document]:
                # Modifies the incoming document inplace
                document.content = "modified"
                return {"output": document}

        pp = Pipeline()
        pp.add_component("first", SimpleComponent())
        pp.add_component("modifier", ModifyingComponent())
        # It's important that the following component has a name lower down the alphabetical order than "modifier",
        # since the pipeline runs components in a first-in-first-out manner based on ordered_component_names which is
        # sorted alphabetically.
        pp.add_component("second", SimpleComponent())

        pp.connect("first.output", "modifier.document")
        pp.connect("first.output", "second.document")

        result = pp.run({"first": {"document": Document(content="original")}})

        assert result["modifier"]["output"].content == "modified"
        # Without deep copying the inputs, the second component would also see the modified document and produce
        # "modified" instead of "original"
        assert result["second"]["output"].content == "original"

    def test_pipeline_does_not_corrupt_outputs(self):
        """
        Test that a component's output collected via include_outputs_from is not corrupted when a downstream
        component receives and mutates the same data in-place.
        """

        @component
        class Producer:
            @component.output_types(doc=Document)
            def run(self) -> dict:
                return {"doc": Document(content="original")}

        @component
        class Mutator:
            @component.output_types(doc=Document)
            def run(self, doc: Document) -> dict:
                # Modifies the incoming document inplace
                doc.content = "mutated"
                return {"doc": doc}

        pipe = Pipeline()
        pipe.add_component("producer", Producer())
        pipe.add_component("mutator", Mutator())
        pipe.connect("producer.doc", "mutator.doc")

        result = pipe.run({}, include_outputs_from={"producer"})

        assert result["producer"]["doc"].content == "original"
        assert result["mutator"]["doc"].content == "mutated"

    def test_auto_variadic_connection_to_agent(self):
        @component
        class MessageProducer:
            @component.output_types(messages=list[ChatMessage])
            def run(self) -> dict[str, list[ChatMessage]]:
                return {"messages": [ChatMessage.from_user("Hello, world!")]}

        p = Pipeline()
        p.add_component("message_producer", MessageProducer())
        p.add_component("message_producer2", MessageProducer())
        p.add_component("agent", Agent(chat_generator=MockChatGenerator()))
        p.connect("message_producer", "agent.messages")
        p.connect("message_producer2", "agent.messages")

        result = p.run({})
        assert result["agent"]["messages"] == [
            ChatMessage.from_user("Hello, world!"),
            ChatMessage.from_user("Hello, world!"),
            ChatMessage.from_assistant("Hello, world!"),
        ]

    def test_run_auto_variadic_str_to_list_str(self):
        """Two str producers connected to a list[str] input are auto-joined and flattened at runtime."""
        p = Pipeline()
        p.add_component("producer1", StringProducer("hello"))
        p.add_component("producer2", StringProducer("world"))
        p.add_component("receiver", ListStrAcceptor())
        p.connect("producer1.text", "receiver.texts")
        p.connect("producer2.text", "receiver.texts")
        result = p.run({})
        assert result["receiver"]["result"] == ["hello", "world"]

    def test_run_auto_variadic_str_and_list_str_to_list_str(self):
        """A str producer and a list[str] producer connected to a list[str] input are auto-joined at runtime."""
        p = Pipeline()
        p.add_component("str_producer", StringProducer("hello"))
        p.add_component("list_producer", ListStrProducer(["world", "!"]))
        p.add_component("receiver", ListStrAcceptor())
        p.connect("str_producer.text", "receiver.texts")
        p.connect("list_producer.texts", "receiver.texts")
        result = p.run({})
        assert result["receiver"]["result"] == ["world", "!", "hello"]

    def test_run_auto_variadic_chat_message_to_list_str(self):
        """Two ChatMessage producers connected to a list[str] input are converted and auto-joined at runtime."""
        p = Pipeline()
        p.add_component("producer1", ChatMessageProducer("hello"))
        p.add_component("producer2", ChatMessageProducer("world"))
        p.add_component("receiver", ListStrAcceptor())
        p.connect("producer1.message", "receiver.texts")
        p.connect("producer2.message", "receiver.texts")
        result = p.run({})
        assert result["receiver"]["result"] == ["hello", "world"]

    def test_run_auto_variadic_str_and_chat_message_to_list_str(self):
        """A str producer and a ChatMessage producer connected to a list[str] input are auto-joined at runtime."""
        p = Pipeline()
        p.add_component("str_producer", StringProducer("hello"))
        p.add_component("chat_producer", ChatMessageProducer("world"))
        p.add_component("receiver", ListStrAcceptor())
        p.connect("str_producer.text", "receiver.texts")
        p.connect("chat_producer.message", "receiver.texts")
        result = p.run({})
        assert result["receiver"]["result"] == ["world", "hello"]

    def test_run_auto_variadic_chat_message_to_list_chat_message(self):
        """Two ChatMessage producers connected to a list[ChatMessage] input are auto-joined at runtime."""
        p = Pipeline()
        p.add_component("producer1", ChatMessageProducer("hello"))
        p.add_component("producer2", ChatMessageProducer("world"))
        p.add_component("receiver", ListChatMessageAcceptor())
        p.connect("producer1.message", "receiver.messages")
        p.connect("producer2.message", "receiver.messages")
        result = p.run({})
        assert [m.text for m in result["receiver"]["result"]] == ["hello", "world"]

    def test_run_auto_variadic_str_to_list_chat_message(self):
        """Two str producers connected to a list[ChatMessage] input are converted and auto-joined at runtime."""
        p = Pipeline()
        p.add_component("producer1", StringProducer("hello"))
        p.add_component("producer2", StringProducer("world"))
        p.add_component("receiver", ListChatMessageAcceptor())
        p.connect("producer1.text", "receiver.messages")
        p.connect("producer2.text", "receiver.messages")
        result = p.run({})
        assert [m.text for m in result["receiver"]["result"]] == ["hello", "world"]

    def test_run_auto_variadic_str_and_chat_message_to_list_chat_message(self):
        """A str and a ChatMessage producer connected to a list[ChatMessage] input are auto-joined at runtime."""
        p = Pipeline()
        p.add_component("str_producer", StringProducer("hello"))
        p.add_component("chat_producer", ChatMessageProducer("world"))
        p.add_component("receiver", ListChatMessageAcceptor())
        p.connect("str_producer.text", "receiver.messages")
        p.connect("chat_producer.message", "receiver.messages")
        result = p.run({})
        assert [m.text for m in result["receiver"]["result"]] == ["world", "hello"]

    def test_run_auto_variadic_chat_message_and_list_chat_message_to_list_chat_message(self):
        """A ChatMessage and a list[ChatMessage] producer connected to list[ChatMessage] are auto-joined at runtime."""
        p = Pipeline()
        p.add_component("chat_producer", ChatMessageProducer("hello"))
        p.add_component("list_producer", ListChatMessageProducer(["world", "!"]))
        p.add_component("receiver", ListChatMessageAcceptor())
        p.connect("chat_producer.message", "receiver.messages")
        p.connect("list_producer.messages", "receiver.messages")
        result = p.run({})
        assert [m.text for m in result["receiver"]["result"]] == ["hello", "world", "!"]

    def test_connect_rejects_list_of_documents_to_single_document(self):
        @component
        class DocsProducer:
            @component.output_types(docs=list[Document])
            def run(self) -> dict[str, list[Document]]:
                return {"docs": [Document(content="a"), Document(content="b"), Document(content="c")]}

        @component
        class DocConsumer:
            @component.output_types(out=Document)
            def run(self, doc: Document) -> dict[str, Document]:
                return {"out": doc}

        p = Pipeline()
        p.add_component("producer", DocsProducer())
        p.add_component("consumer", DocConsumer())
        with pytest.raises(PipelineConnectError):
            p.connect("producer.docs", "consumer.doc")

    def test_run_raises_when_multi_element_list_is_unwrapped_at_runtime(self):
        @component
        class MultiStrProducer:
            @component.output_types(texts=list[str])
            def run(self) -> dict[str, list[str]]:
                return {"texts": ["first", "second", "third"]}

        @component
        class SingleStrConsumer:
            @component.output_types(out=str)
            def run(self, text: str) -> dict[str, str]:
                return {"out": text}

        p = Pipeline()
        p.add_component("producer", MultiStrProducer())
        p.add_component("consumer", SingleStrConsumer())
        p.connect("producer.texts", "consumer.text")

        with pytest.raises(PipelineRuntimeError, match="Cannot unwrap a list of 3 items"):
            p.run({})

    def test_run_single_element_list_unwrap_still_works(self):
        @component
        class SingleStrProducer:
            @component.output_types(texts=list[str])
            def run(self) -> dict[str, list[str]]:
                return {"texts": ["only-one"]}

        @component
        class SingleStrConsumer:
            @component.output_types(out=str)
            def run(self, text: str) -> dict[str, str]:
                return {"out": text}

        p = Pipeline()
        p.add_component("producer", SingleStrProducer())
        p.add_component("consumer", SingleStrConsumer())
        p.connect("producer.texts", "consumer.text")
        assert p.run({}) == {"consumer": {"out": "only-one"}}