fix: Fix auto-variadic assignment (#10688)

* Try out idea

* Undo some changes

* update function to handle Optional[list] as well

* Add unit tests for auto-variadic + conversion

* add test for optional list

* add reno

* remove unnecessary tests for now

* add another unit test

Author: sjrl

haystack/core/component/types.py

@@ -85,21 +85,18 @@ def __post_init__(self):
         except AttributeError:
             self.is_lazy_variadic = False
             self.is_greedy = False
+
+        # We need to "unpack" the type inside the Variadic annotation, otherwise the pipeline connection api will try
+        # to match `Annotated[type, HAYSTACK_VARIADIC_ANNOTATION]`.
+        #
+        # Note1: Variadic is expressed as an annotation of one single type, so the return value of get_args will
+        # always be a one-item tuple.
+        #
+        # Note2: a pipeline always passes a list of items when a component input is declared as Variadic, so the
+        # type itself always wraps an iterable of the declared type. For example, Variadic[int] is eventually an
+        # alias for Iterable[int]. Since we're interested in getting the inner type `int`, we call `get_args`
+        # twice: the first time to get `list[int]` out of `Variadic`, the second time to get `int` out of `list[int]`.
         if self.is_lazy_variadic or self.is_greedy:
-            # We need to "unpack" the type inside the Variadic annotation,
-            # otherwise the pipeline connection api will try to match
-            # `Annotated[type, HAYSTACK_VARIADIC_ANNOTATION]`.
-            #
-            # Note1: Variadic is expressed as an annotation of one single type,
-            # so the return value of get_args will always be a one-item tuple.
-            #
-            # Note2: a pipeline always passes a list of items when a component
-            # input is declared as Variadic, so the type itself always wraps
-            # an iterable of the declared type. For example, Variadic[int]
-            # is eventually an alias for Iterable[int]. Since we're interested
-            # in getting the inner type `int`, we call `get_args` twice: the
-            # first time to get `list[int]` out of `Variadic`, the second time
-            # to get `int` out of `list[int]`.
             self.type = get_args(get_args(self.type)[0])[0]
 
 

haystack/core/pipeline/base.py

@@ -10,7 +10,7 @@
 from datetime import datetime
 from enum import IntEnum
 from pathlib import Path
-from typing import Any, TextIO, TypeVar
+from typing import Any, TextIO, TypeVar, Union, get_args
 
 import networkx
 
@@ -615,16 +615,7 @@ def connect(self, sender: str, receiver: str) -> "PipelineBase":  # noqa: PLR091
             return self
 
         if receiver_socket.senders:
-            # We automatically set the receiver socket as variadic if:
-            # - it has at least one sender already connected
-            # - it's not already variadic
-            # - its origin type is list
-            if not receiver_socket.is_variadic and _safe_get_origin(receiver_socket.type) == list:
-                receiver_socket.is_lazy_variadic = True
-                # We also disable wrapping inputs into list so the sender outputs matches the type of the receiver
-                # socket.
-                receiver_socket.wrap_input_in_list = False
-
+            receiver_socket = self._make_socket_auto_variadic(receiver_socket=receiver_socket)
             if not receiver_socket.is_variadic:
                 # Only variadic input sockets can receive from multiple senders
                 msg = (
@@ -949,23 +940,48 @@ def validate_input(self, data: dict[str, Any]) -> None:
 
                 # Check if an input is provided more than once for non-variadic sockets
                 if socket.senders and socket_name in component_inputs:
-                    # We automatically set the receiver socket as lazy variadic if:
-                    # - it has at least one sender already connected
-                    # - it's not already variadic
-                    # - its origin type is list
-                    if not socket.is_variadic and _safe_get_origin(socket.type) == list:
-                        socket.is_lazy_variadic = True
-                        # We also disable wrapping inputs into list so the sender outputs matches the type of the
-                        # receiver socket.
-                        socket.wrap_input_in_list = False
-
+                    socket = self._make_socket_auto_variadic(receiver_socket=socket)
                     if not socket.is_variadic:
                         raise ValueError(
                             f"Component '{component_name}' cannot accept multiple inputs to '{socket_name}'. "
                             f"It is already connected to component '{socket.senders[0]}' so it cannot accept "
                             "additional inputs."
                         )
 
+    def _make_socket_auto_variadic(self, receiver_socket: InputSocket) -> InputSocket:
+        """
+        Checks if the receiver socket can be made lazy variadic and modifies it in-place if that's the case.
+
+        We automatically set the receiver socket as lazy variadic if:
+            - it has at least one sender already connected
+            - it's not already variadic
+            - its type is list or Optional[list]
+
+        NOTE: We also disable wrapping inputs into list for these auto-variadic sockets, so the sender outputs match the
+        type of the receiver socket.
+
+        :returns:
+            The potentially modified receiver socket.
+        """
+        # If it's already variadic, we don't change anything
+        if receiver_socket.is_variadic:
+            return receiver_socket
+
+        origin = _safe_get_origin(receiver_socket.type)
+
+        # Unwrap Optional types
+        if origin == Union:
+            non_none_args = [a for a in get_args(receiver_socket.type) if a is not type(None)]
+            if len(non_none_args) == 1:
+                origin = _safe_get_origin(non_none_args[0])
+
+        # If the origin is list, we can make the socket lazy variadic
+        if origin == list:
+            receiver_socket.is_lazy_variadic = True
+            receiver_socket.wrap_input_in_list = False
+
+        return receiver_socket
+
     def _prepare_component_input_data(self, data: dict[str, Any]) -> dict[str, dict[str, Any]]:
         """
         Prepares input data for pipeline components.

releasenotes/notes/fix-auto-variadic-optional-list-483de240f8755021.yaml

@@ -0,0 +1,8 @@
+---
+enhancements:
+  - |
+    Auto variadic sockets now also support ``Optional[list[...]]`` input types, in addition to plain ``list[...]``.
+fixes:
+  - |
+    Fixed smart connection logic to support connecting multiple outputs to a socket whose type is ``Optional[list[...]]`` (e.g. ``list[ChatMessage] | None``).
+    Previously, connecting two ``list[ChatMessage]`` outputs to ``Agent.messages`` would fail after its type was updated from ``list[ChatMessage]`` to ``list[ChatMessage] | None``.

test/core/pipeline/test_pipeline.py

@@ -7,11 +7,19 @@
 
 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 PipelineRuntimeError
 from haystack.core.pipeline import Pipeline
-from haystack.dataclasses.document import Document
+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
@@ -251,3 +259,53 @@ def run(self, document: Document) -> dict[str, Document]:
         # 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:
+                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!"),
+        ]

test/core/pipeline/test_pipeline_base.py

@@ -1949,6 +1949,31 @@ def run(self, numbers: list[int]) -> dict[str, list[int]]:
         assert receiver.__haystack_input__._sockets_dict == {"numbers": inp_socket}  # type: ignore[attr-defined]
         assert receiver.__haystack_input__._sockets_dict["numbers"].senders == ["sender1", "sender2"]  # type: ignore[attr-defined]
 
+    def test_connect_auto_variadic_optional_list(self):
+        @component
+        class ListAcceptor:
+            @component.output_types(result=list[int])
+            def run(self, numbers: list[int] | None = None) -> dict[str, list[int]]:
+                return {"result": numbers or []}
+
+        pipeline = PipelineBase()
+        receiver = ListAcceptor()
+        pipeline.add_component("sender1", ListAcceptor())
+        pipeline.add_component("sender2", ListAcceptor())
+        pipeline.add_component("receiver", receiver)
+
+        pipeline.connect("sender1.result", "receiver.numbers")
+        pipeline.connect("sender2.result", "receiver.numbers")
+
+        # Check that the receiver's input socket is correctly set to lazy variadic with wrap_input_in_list=False
+        inp_socket = InputSocket(
+            name="numbers", type=list[int] | None, senders=["sender1", "sender2"], default_value=None
+        )
+        inp_socket.is_lazy_variadic = True
+        inp_socket.wrap_input_in_list = False
+        assert receiver.__haystack_input__._sockets_dict == {"numbers": inp_socket}  # type: ignore[attr-defined]
+        assert receiver.__haystack_input__._sockets_dict["numbers"].senders == ["sender1", "sender2"]  # type: ignore[attr-defined]
+
     def test_connect_with_conversion_chat_message_to_str(self):
         @component
         class ChatMessageOutput: