chore: import upstream snapshot with attribution

haystack/core/component/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from haystack.core.component.component import Component, component
+from haystack.core.component.types import InputSocket, OutputSocket
+
+__all__ = ["component", "Component", "InputSocket", "OutputSocket"]

haystack/core/component/component.py

@@ -0,0 +1,644 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""
+Attributes:
+
+    component: Marks a class as a component. Any class decorated with `@component` can be used by a Pipeline.
+
+All components must follow the contract below. This docstring is the source of truth for components contract.
+
+<hr>
+
+`@component` decorator
+
+All component classes must be decorated with the `@component` decorator. This allows Haystack to discover them.
+
+<hr>
+
+`__init__(self, **kwargs)`
+
+Optional method.
+
+Components may have an `__init__` method where they define:
+
+- `self.init_parameters = {same parameters that the __init__ method received}`:
+    In this dictionary you can store any state the components wish to be persisted when they are saved.
+    These values will be given to the `__init__` method of a new instance when the pipeline is loaded.
+    Note that by default the `@component` decorator saves the arguments automatically.
+    However, if a component sets their own `init_parameters` manually in `__init__()`, that will be used instead.
+    Note: all of the values contained here **must be JSON serializable**. Serialize them manually if needed.
+
+Components should take only "basic" Python types as parameters of their `__init__` function, or iterables and
+dictionaries containing only such values. Anything else (objects, functions, etc) will raise an exception at init
+time. If there's the need for such values, consider serializing them to a string.
+
+If you need to accept classes or callables, accept either a string import path or the callable itself. Resolve strings
+to objects in `__init__`, and serialize objects back to importable strings in `to_dict()` so that `from_dict()` can load
+them (for example, store `"module_path.symbol_name"` and load it via `importlib`). This keeps init parameters JSON
+serializable for pipeline save/load. See `haystack.testing.sample_components.accumulate.Accumulate` for a reference
+implementation.
+
+The `__init__` must be extremely lightweight, because it's a frequent operation during the construction and
+validation of the pipeline. If a component has some heavy state to initialize (models, backends, etc...) refer to
+the `warm_up()` method.
+
+<hr>
+
+`warm_up(self)`
+
+Optional method.
+
+This method is called by Pipeline before the graph execution. Make sure to avoid double-initializations,
+because Pipeline will not keep track of which components it called `warm_up()` on.
+
+<hr>
+
+`run(self, data)`
+
+Mandatory method.
+
+This is the method where the main functionality of the component should be carried out. It's called by
+`Pipeline.run()`.
+
+When the component should run, Pipeline will call this method with an instance of the dataclass returned by the
+method decorated with `@component.input`. This dataclass contains:
+
+- all the input values coming from other components connected to it,
+- if any is missing, the corresponding value defined in `self.defaults`, if it exists.
+
+`run()` must return a single instance of the dataclass declared through the method decorated with
+`@component.output`.
+
+"""
+
+import inspect
+import typing
+from collections.abc import Callable, Coroutine, Iterator, Mapping
+from contextlib import contextmanager
+from contextvars import ContextVar
+from copy import deepcopy
+from dataclasses import dataclass
+from types import new_class
+from typing import Any, ParamSpec, Protocol, TypeVar, overload, runtime_checkable
+
+from haystack import logging
+from haystack.core.errors import ComponentError
+
+from .sockets import Sockets
+from .types import InputSocket, OutputSocket, _empty
+
+logger = logging.getLogger(__name__)
+
+RunParamsT = ParamSpec("RunParamsT")
+RunReturnT = TypeVar("RunReturnT", bound=Mapping[str, Any] | Coroutine[Any, Any, Mapping[str, Any]])
+
+
+@dataclass
+class PreInitHookPayload:
+    """
+    Payload for the hook called before a component instance is initialized.
+
+    :param callback:
+        Receives the following inputs: component class and init parameter keyword args.
+    :param in_progress:
+        Flag to indicate if the hook is currently being executed.
+        Used to prevent it from being called recursively (if the component's constructor
+        instantiates another component).
+    """
+
+    callback: Callable
+    in_progress: bool = False
+
+
+_COMPONENT_PRE_INIT_HOOK: ContextVar[PreInitHookPayload | None] = ContextVar("component_pre_init_hook", default=None)
+
+
+@contextmanager
+def _hook_component_init(callback: Callable) -> Iterator[None]:
+    """
+    Context manager to set a callback that will be invoked before a component's constructor is called.
+
+    The callback receives the component class and the init parameters (as keyword arguments) and can modify the init
+    parameters in place.
+
+    :param callback:
+        Callback function to invoke.
+    """
+    token = _COMPONENT_PRE_INIT_HOOK.set(PreInitHookPayload(callback))
+    try:
+        yield
+    finally:
+        _COMPONENT_PRE_INIT_HOOK.reset(token)
+
+
+@runtime_checkable
+class Component(Protocol):
+    """
+    Note this is only used by type checking tools.
+
+    In order to implement the `Component` protocol, custom components need to
+    have a `run` method. The signature of the method and its return value
+    won't be checked, i.e. classes with the following methods:
+
+        def run(self, param: str) -> dict[str, Any]:
+            ...
+
+    and
+
+        def run(self, **kwargs):
+            ...
+
+    will be both considered as respecting the protocol. This makes the type
+    checking much weaker, but we have other places where we ensure code is
+    dealing with actual Components.
+
+    The protocol is runtime checkable so it'll be possible to assert:
+
+        isinstance(MyComponent, Component)
+    """
+
+    # The following expression defines a run method compatible with any input signature.
+    # Its type is equivalent to Callable[..., dict[str, Any]].
+    # See https://typing.python.org/en/latest/spec/callables.html#meaning-of-in-callable.
+    #
+    # Using `run: Callable[..., dict[str, Any]]` directly leads to type errors: the protocol would expect a settable
+    # attribute `run`, while the actual implementation is a read-only method.
+    # For example:
+    # from haystack import Pipeline, component
+    # @component
+    # class MyComponent:
+    #     @component.output_types(out=str)
+    #     def run(self):
+    #         return {"out": "Hello, world!"}
+    # pipeline = Pipeline()
+    # pipeline.add_component("my_component", MyComponent())
+    #
+    # mypy raises:
+    # error: Argument 2 to "add_component" of "PipelineBase" has incompatible type "MyComponent"; expected "Component"
+    # [arg-type]
+    # note: Protocol member Component.run expected settable variable, got read-only attribute
+
+    def run(self, *args: Any, **kwargs: Any) -> Mapping[str, Any]:  # noqa: D102
+        ...
+
+
+class ComponentMeta(type):
+    @staticmethod
+    def _positional_to_kwargs(cls_type: type, args: tuple[Any, ...]) -> dict[str, Any]:
+        """
+        Convert positional arguments to keyword arguments based on the signature of the `__init__` method.
+        """
+        init_signature = inspect.signature(cls_type.__init__)  # type:ignore[misc]
+        init_params = {name: info for name, info in init_signature.parameters.items() if name != "self"}
+
+        out = {}
+        for arg, (name, info) in zip(args, init_params.items(), strict=False):
+            if info.kind == inspect.Parameter.VAR_POSITIONAL:
+                raise ComponentError(
+                    "Pre-init hooks do not support components with variadic positional args in their init method"
+                )
+
+            assert info.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.POSITIONAL_ONLY)
+            out[name] = arg
+        return out
+
+    @staticmethod
+    def _parse_and_set_output_sockets(instance: Any) -> None:
+        has_async_run = hasattr(instance, "run_async")
+
+        # If `component.set_output_types()` was called in the component constructor,
+        # `__haystack_output__` is already populated, no need to do anything.
+        if not hasattr(instance, "__haystack_output__"):
+            # If that's not the case, we need to populate `__haystack_output__`
+            #
+            # If either of the run methods were decorated, they'll have a field assigned that
+            # stores the output specification. If both run methods were decorated, we ensure that
+            # outputs are the same. We deepcopy the content of the cache to transfer ownership from
+            # the class method to the actual instance, so that different instances of the same class
+            # won't share this data.
+
+            run_output_types = getattr(instance.run, "_output_types_cache", {})
+            async_run_output_types = getattr(instance.run_async, "_output_types_cache", {}) if has_async_run else {}
+
+            if has_async_run and run_output_types != async_run_output_types:
+                raise ComponentError("Output type specifications of 'run' and 'run_async' methods must be the same")
+            output_types_cache = run_output_types
+
+            instance.__haystack_output__ = Sockets(instance, deepcopy(output_types_cache), OutputSocket)
+
+    @staticmethod
+    def _parse_and_set_input_sockets(component_cls: type, instance: Any) -> None:
+        def inner(method: Callable[..., Any], sockets: Sockets) -> inspect.Signature:
+            from inspect import Parameter
+
+            run_signature = inspect.signature(method)
+            try:
+                # TypeError is raised if the argument is not of a type that can contain annotations
+                run_hints = typing.get_type_hints(method)
+            except TypeError:
+                run_hints = None
+
+            for param_name, param_info in run_signature.parameters.items():
+                if param_name == "self" or param_info.kind in (Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD):
+                    continue
+
+                # We prefer the type annotation from inspect.signature, but if it's a string we need to resolve it
+                # using the hints. The type annotation can be a string if the component is using postponed evaluation
+                # of annotations.
+                annotation = param_info.annotation
+                if isinstance(annotation, str) and run_hints is not None:
+                    annotation = run_hints.get(param_name, annotation)
+
+                socket_kwargs = {"name": param_name, "type": annotation}
+                if param_info.default != Parameter.empty:
+                    socket_kwargs["default_value"] = param_info.default
+
+                new_socket = InputSocket(**socket_kwargs)
+
+                # Also ensure that new sockets don't override existing ones.
+                existing_socket = sockets.get(param_name)
+                if existing_socket is not None and existing_socket != new_socket:
+                    raise ComponentError(
+                        "set_input_types()/set_input_type() cannot override the parameters of the 'run' method"
+                    )
+
+                sockets[param_name] = new_socket
+
+            return run_signature
+
+        # Create the sockets if set_input_types() wasn't called in the constructor.
+        if not hasattr(instance, "__haystack_input__"):
+            instance.__haystack_input__ = Sockets(instance, {}, InputSocket)
+
+        inner(getattr(component_cls, "run"), instance.__haystack_input__)  # noqa: B009
+
+        # Ensure that the sockets are the same for the async method, if it exists.
+        async_run = getattr(component_cls, "run_async", None)
+        if async_run is not None:
+            run_sockets = Sockets(instance, {}, InputSocket)
+            async_run_sockets = Sockets(instance, {}, InputSocket)
+
+            # Can't use the sockets from above as they might contain
+            # values set with set_input_types().
+            run_sig = inner(getattr(component_cls, "run"), run_sockets)  # noqa: B009
+            async_run_sig = inner(async_run, async_run_sockets)
+
+            if async_run_sockets != run_sockets or run_sig != async_run_sig:
+                sig_diff = _compare_run_methods_signatures(run_sig, async_run_sig)
+                raise ComponentError(
+                    f"Parameters of 'run' and 'run_async' methods must be the same.\nDifferences found:\n{sig_diff}"
+                )
+
+    def __call__(cls, *args: Any, **kwargs: Any) -> Any:
+        """
+        This method is called when clients instantiate a Component and runs before __new__ and __init__.
+        """
+        # This will call __new__ then __init__, giving us back the Component instance
+        pre_init_hook = _COMPONENT_PRE_INIT_HOOK.get()
+        if pre_init_hook is None or pre_init_hook.in_progress:
+            instance = super().__call__(*args, **kwargs)
+        else:
+            try:
+                pre_init_hook.in_progress = True
+                named_positional_args = ComponentMeta._positional_to_kwargs(cls, args)
+                assert set(named_positional_args.keys()).intersection(kwargs.keys()) == set(), (
+                    "positional and keyword arguments overlap"
+                )
+                kwargs.update(named_positional_args)
+                pre_init_hook.callback(cls, kwargs)
+                instance = super().__call__(**kwargs)
+            finally:
+                pre_init_hook.in_progress = False
+
+        # Before returning, we have the chance to modify the newly created
+        # Component instance, so we take the chance and set up the I/O sockets
+        has_async_run = hasattr(instance, "run_async")
+        if has_async_run and not inspect.iscoroutinefunction(instance.run_async):
+            raise ComponentError(f"Method 'run_async' of component '{cls.__name__}' must be a coroutine")
+        instance.__haystack_supports_async__ = has_async_run
+
+        ComponentMeta._parse_and_set_input_sockets(cls, instance)
+        ComponentMeta._parse_and_set_output_sockets(instance)
+
+        # Since a Component can't be used in multiple Pipelines at the same time
+        # we need to know if it's already owned by a Pipeline when adding it to one.
+        # We use this flag to check that.
+        instance.__haystack_added_to_pipeline__ = None
+
+        return instance
+
+
+def _component_repr(component: Component) -> str:
+    """
+    All Components override their __repr__ method with this one.
+
+    It prints the component name and the input/output sockets.
+    """
+    result = object.__repr__(component)
+    if pipeline := getattr(component, "__haystack_added_to_pipeline__", None):
+        # This Component has been added in a Pipeline, let's get the name from there.
+        result += f"\n{pipeline.get_component_name(component)}"
+
+    # We're explicitly ignoring the type here because we're sure that the component
+    # has the __haystack_input__ and __haystack_output__ attributes at this point
+    return (
+        f"{result}\n{getattr(component, '__haystack_input__', '<invalid_input_sockets>')}"
+        f"\n{getattr(component, '__haystack_output__', '<invalid_output_sockets>')}"
+    )
+
+
+def _component_run_has_kwargs(component_cls: type) -> bool:
+    run_method = getattr(component_cls, "run", None)
+    if run_method is None:
+        return False
+    return any(
+        param.kind == inspect.Parameter.VAR_KEYWORD for param in inspect.signature(run_method).parameters.values()
+    )
+
+
+def _compare_run_methods_signatures(run_sig: inspect.Signature, async_run_sig: inspect.Signature) -> str:
+    """
+    Builds a detailed error message with the differences between the signatures of the run and run_async methods.
+
+    :param run_sig: The signature of the run method
+    :param async_run_sig: The signature of the run_async method
+
+    :returns:
+        A detailed error message if signatures don't match, empty string if they do
+    """
+    differences = []
+    run_params = list(run_sig.parameters.items())
+    async_params = list(async_run_sig.parameters.items())
+
+    if len(run_params) != len(async_params):
+        differences.append(
+            f"Different number of parameters: run has {len(run_params)}, run_async has {len(async_params)}"
+        )
+
+    for (run_name, run_param), (async_name, async_param) in zip(run_params, async_params, strict=False):
+        if run_name != async_name:
+            differences.append(f"Parameter name mismatch: {run_name} vs {async_name}")
+
+        if run_param.annotation != async_param.annotation:
+            differences.append(
+                f"Parameter '{run_name}' type mismatch: {run_param.annotation} vs {async_param.annotation}"
+            )
+
+        if run_param.default != async_param.default:
+            differences.append(
+                f"Parameter '{run_name}' default value mismatch: {run_param.default} vs {async_param.default}"
+            )
+
+        if run_param.kind != async_param.kind:
+            differences.append(
+                f"Parameter '{run_name}' kind (POSITIONAL, KEYWORD, etc.) mismatch: "
+                f"{run_param.kind} vs {async_param.kind}"
+            )
+
+    return "\n".join(differences)
+
+
+T = TypeVar("T", bound=Component)
+
+
+class _Component:
+    """
+    See module's docstring.
+
+    Args:
+        cls: the class that should be used as a component.
+
+    Returns:
+        A class that can be recognized as a component.
+
+    Raises:
+        ComponentError: if the class provided has no `run()` method or otherwise doesn't respect the component contract.
+    """
+
+    def __init__(self) -> None:
+        self.registry: dict[str, type] = {}
+
+    def set_input_type(
+        self,
+        instance: Component,
+        name: str,
+        type: Any,  # noqa: A002
+        default: Any = _empty,
+    ) -> None:
+        """
+        Add a single input socket to the component instance.
+
+        Replaces any existing input socket with the same name.
+
+        :param instance: Component instance where the input type will be added.
+        :param name: name of the input socket.
+        :param type: type of the input socket.
+        :param default: default value of the input socket, defaults to _empty
+        """
+        if not _component_run_has_kwargs(instance.__class__):
+            raise ComponentError(
+                "Cannot set input types on a component that doesn't have a kwargs parameter in the 'run' method"
+            )
+
+        if not hasattr(instance, "__haystack_input__"):
+            instance.__haystack_input__ = Sockets(instance, {}, InputSocket)  # type: ignore
+        instance.__haystack_input__[name] = InputSocket(name=name, type=type, default_value=default)  # type: ignore
+
+    def set_input_types(self, instance: Any, **types: type[Any]) -> None:
+        """
+        Method that specifies the input types when 'kwargs' is passed to the run method.
+
+        Use as:
+
+        ```python
+        @component
+        class MyComponent:
+
+            def __init__(self, value: int) -> None:
+                component.set_input_types(self, value_1=str, value_2=str)
+                ...
+
+            @component.output_types(output_1=int, output_2=str)
+            def run(self, **kwargs):
+                return {"output_1": kwargs["value_1"], "output_2": ""}
+        ```
+
+        Note that if the `run()` method also specifies some parameters, those will take precedence.
+
+        For example:
+
+        ```python
+        @component
+        class MyComponent:
+
+            def __init__(self, value: int) -> None:
+                component.set_input_types(self, value_1=str, value_2=str)
+                ...
+
+            @component.output_types(output_1=int, output_2=str)
+            def run(self, value_0: str, value_1: Optional[str] = None, **kwargs):
+                return {"output_1": kwargs["value_1"], "output_2": ""}
+        ```
+
+        would add a mandatory `value_0` parameters, make the `value_1`
+        parameter optional with a default None, and keep the `value_2`
+        parameter mandatory as specified in `set_input_types`.
+
+        """
+        if not _component_run_has_kwargs(instance.__class__):
+            raise ComponentError(
+                "Cannot set input types on a component that doesn't have a kwargs parameter in the 'run' method"
+            )
+
+        instance.__haystack_input__ = Sockets(
+            instance, {name: InputSocket(name=name, type=type_) for name, type_ in types.items()}, InputSocket
+        )
+
+    def set_output_types(self, instance: Any, **types: type[Any]) -> None:
+        """
+        Method that specifies the output types when the 'run' method is not decorated with 'component.output_types'.
+
+        Use as:
+
+        ```python
+        @component
+        class MyComponent:
+
+            def __init__(self, value: int) -> None:
+                component.set_output_types(self, output_1=int, output_2=str)
+                ...
+
+            # no decorators here
+            def run(self, value: int):
+                return {"output_1": 1, "output_2": "2"}
+
+            # also no decorators here
+            async def run_async(self, value: int):
+                return {"output_1": 1, "output_2": "2"}
+        ```
+        """
+        has_run_decorator = hasattr(instance.run, "_output_types_cache")
+        has_run_async_decorator = hasattr(instance, "run_async") and hasattr(instance.run_async, "_output_types_cache")
+        if has_run_decorator or has_run_async_decorator:
+            raise ComponentError(
+                "Cannot call `set_output_types` on a component that already has the 'output_types' decorator on its "
+                "`run` or `run_async` methods."
+            )
+
+        instance.__haystack_output__ = Sockets(
+            instance, {name: OutputSocket(name=name, type=type_) for name, type_ in types.items()}, OutputSocket
+        )
+
+    def output_types(
+        self, **types: Any
+    ) -> Callable[[Callable[RunParamsT, RunReturnT]], Callable[RunParamsT, RunReturnT]]:
+        """
+        Decorator factory that specifies the output types of a component.
+
+        Use as:
+        ```python
+        @component
+        class MyComponent:
+            @component.output_types(output_1=int, output_2=str)
+            def run(self, value: int):
+                return {"output_1": 1, "output_2": "2"}
+        ```
+        """
+
+        def output_types_decorator(run_method: Callable[RunParamsT, RunReturnT]) -> Callable[RunParamsT, RunReturnT]:
+            """
+            Decorator that sets the output types of the decorated method.
+
+            This happens at class creation time, and since we don't have the decorated
+            class available here, we temporarily store the output types as an attribute of
+            the decorated method. The ComponentMeta metaclass will use this data to create
+            sockets at instance creation time.
+            """
+            method_name = run_method.__name__
+            if method_name not in ("run", "run_async"):
+                raise ComponentError("'output_types' decorator can only be used on 'run' and 'run_async' methods")
+
+            setattr(  # noqa: B010
+                run_method,
+                "_output_types_cache",
+                {name: OutputSocket(name=name, type=type_) for name, type_ in types.items()},
+            )
+            return run_method
+
+        return output_types_decorator
+
+    def _component(self, cls: type[T]) -> type[T]:
+        """
+        Decorator validating the structure of the component and registering it in the components registry.
+        """
+        logger.debug("Registering {component} as a component", component=cls)
+
+        # Check for required methods and fail as soon as possible
+        if not hasattr(cls, "run"):
+            raise ComponentError(f"{cls.__name__} must have a 'run()' method. See the docs for more information.")
+
+        def copy_class_namespace(namespace: dict[str, Any]) -> None:
+            """
+            This is the callback that `typing.new_class` will use to populate the newly created class.
+
+            Simply copy the whole namespace from the decorated class.
+            """
+            for key, val in dict(cls.__dict__).items():
+                # __dict__ and __weakref__ are class-bound, we should let Python recreate them.
+                if key in ("__dict__", "__weakref__"):
+                    continue
+                namespace[key] = val
+
+        # Recreate the decorated component class so it uses our metaclass.
+        # We must explicitly redefine the type of the class to make sure language servers
+        # and type checkers understand that the class is of the correct type.
+        new_cls: type[T] = new_class(cls.__name__, cls.__bases__, {"metaclass": ComponentMeta}, copy_class_namespace)
+
+        # Save the component in the class registry (for deserialization)
+        class_path = f"{new_cls.__module__}.{new_cls.__name__}"
+        if class_path in self.registry:
+            # Corner case, but it may occur easily in notebooks when re-running cells.
+            logger.debug(
+                "Component {component} is already registered. Previous imported from '{module_name}', \
+                new imported from '{new_module_name}'",
+                component=class_path,
+                module_name=self.registry[class_path],
+                new_module_name=new_cls,
+            )
+        self.registry[class_path] = new_cls
+        logger.debug("Registered Component {component}", component=new_cls)
+
+        # Override the __repr__ method with a default one
+        # mypy is not happy that:
+        # 1) we are assigning a method to a class
+        # 2) _component_repr has a different type (Callable[[Component], str]) than the expected
+        # __repr__ method (Callable[[object], str])
+        new_cls.__repr__ = _component_repr  # type: ignore[assignment]
+
+        return new_cls
+
+    # Call signature when the decorator is used without parens (@component).
+    @overload
+    def __call__(self, cls: type[T]) -> type[T]: ...
+
+    # Overload allowing the decorator to be used with parens (@component()).
+    @overload
+    def __call__(self) -> Callable[[type[T]], type[T]]: ...
+
+    def __call__(self, cls: type[T] | None = None) -> type[T] | Callable[[type[T]], type[T]]:
+        # We must wrap the call to the decorator in a function for it to work
+        # correctly with or without parens
+        def wrap(cls: type[T]) -> type[T]:
+            return self._component(cls)
+
+        if cls:
+            # Decorator is called without parens
+            return wrap(cls)
+
+        # Decorator is called with parens
+        return wrap
+
+
+component = _Component()

haystack/core/component/sockets.py

@@ -0,0 +1,143 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any
+
+from haystack.core.type_utils import _type_name
+
+from .types import InputSocket, OutputSocket
+
+SocketsDict = dict[str, InputSocket | OutputSocket]
+SocketsIOType = type[InputSocket] | type[OutputSocket]
+
+
+class Sockets:  # noqa: PLW1641
+    """
+    Represents the inputs or outputs of a `Component`.
+
+    Depending on the type passed to the constructor, it will represent either the inputs or the outputs of
+    the `Component`.
+
+    Usage:
+    ```python
+    from typing import Any
+    from haystack.components.builders.prompt_builder import PromptBuilder
+    from haystack.core.component.sockets import Sockets
+    from haystack.core.component.types import InputSocket, OutputSocket
+
+
+    prompt_template = \"""
+    Given these documents, answer the question.\nDocuments:
+    {% for doc in documents %}
+        {{ doc.content }}
+    {% endfor %}
+
+    \nQuestion: {{question}}
+    \nAnswer:
+    \"""
+
+    prompt_builder = PromptBuilder(template=prompt_template)
+    sockets = {"question": InputSocket("question", Any), "documents": InputSocket("documents", Any)}
+    inputs = Sockets(component=prompt_builder, sockets_dict=sockets, sockets_io_type=InputSocket)
+    inputs
+    # >> Inputs:
+    # >>   - question: Any
+    # >>   - documents: Any
+
+    inputs.question
+    # >> InputSocket(name='question', type=typing.Any, default_value=<class 'haystack.core.component.types._empty'>, ...
+    ```
+    """
+
+    # We're using a forward declaration here to avoid a circular import.
+    def __init__(
+        self,
+        component: "Component",  # type: ignore[name-defined] # noqa: F821
+        sockets_dict: SocketsDict,
+        sockets_io_type: SocketsIOType,
+    ) -> None:
+        """
+        Create a new Sockets object.
+
+        We don't do any enforcement on the types of the sockets here, the `sockets_type` is only used for
+        the `__repr__` method.
+        We could do without it and use the type of a random value in the `sockets` dict, but that wouldn't
+        work for components that have no sockets at all. Either input or output.
+
+        :param component:
+            The component that these sockets belong to.
+        :param sockets_dict:
+            A dictionary of sockets.
+        :param sockets_io_type:
+            The type of the sockets.
+        """
+        self._sockets_io_type = sockets_io_type
+        self._component = component
+        self._sockets_dict = sockets_dict
+        self.__dict__.update(sockets_dict)
+
+    def __eq__(self, value: object) -> bool:
+        if not isinstance(value, Sockets):
+            return False
+
+        return (
+            self._sockets_io_type == value._sockets_io_type
+            and self._component == value._component
+            and self._sockets_dict == value._sockets_dict
+        )
+
+    def __setitem__(self, key: str, socket: InputSocket | OutputSocket) -> None:
+        """
+        Adds a new socket to this Sockets object.
+
+        This eases a bit updating the list of sockets after Sockets has been created.
+        That should happen only in the `component` decorator.
+        """
+        self._sockets_dict[key] = socket
+        self.__dict__[key] = socket
+
+    def __contains__(self, key: str) -> bool:
+        return key in self._sockets_dict
+
+    def get(self, key: str, default: InputSocket | OutputSocket | None = None) -> InputSocket | OutputSocket | None:
+        """
+        Get a socket from the Sockets object.
+
+        :param key:
+            The name of the socket to get.
+        :param default:
+            The value to return if the key is not found.
+        :returns:
+            The socket with the given key or `default` if the key is not found.
+        """
+        return self._sockets_dict.get(key, default)
+
+    def _component_name(self) -> str:
+        if pipeline := self._component.__haystack_added_to_pipeline__:
+            # This Component has been added in a Pipeline, let's get the name from there.
+            return pipeline.get_component_name(self._component)
+
+        # This Component has not been added to a Pipeline yet, so we can't know its name.
+        # Let's use default __repr__. We don't call repr() directly as Components have a custom
+        # __repr__ method and that would lead to infinite recursion since we call Sockets.__repr__ in it.
+        return object.__repr__(self._component)
+
+    def __getattribute__(self, name: Any) -> Any:
+        try:
+            sockets = object.__getattribute__(self, "_sockets")
+            if name in sockets:
+                return sockets[name]
+        except AttributeError:
+            pass
+
+        return object.__getattribute__(self, name)
+
+    def __repr__(self) -> str:
+        result = ""
+        if self._sockets_io_type == InputSocket:
+            result = "Inputs:\n"
+        elif self._sockets_io_type == OutputSocket:
+            result = "Outputs:\n"
+
+        return result + "\n".join([f"  - {n}: {_type_name(s.type)}" for n, s in self._sockets_dict.items()])

haystack/core/component/types.py

@@ -0,0 +1,137 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from types import UnionType
+from typing import Annotated, Any, TypeAlias, TypedDict, TypeVar, get_args
+
+from haystack.core.errors import ComponentError
+
+HAYSTACK_VARIADIC_ANNOTATION = "__haystack__variadic_t"
+HAYSTACK_GREEDY_VARIADIC_ANNOTATION = "__haystack__greedy_variadic_t"
+
+# # Generic type variable used in the Variadic container
+T = TypeVar("T")
+
+
+# Variadic is a custom annotation type we use to mark input types.
+# This type doesn't do anything else than "marking" the contained
+# type so it can be used in the `InputSocket` creation where we
+# check that its annotation equals to HAYSTACK_VARIADIC_ANNOTATION
+Variadic: TypeAlias = Annotated[Iterable[T], HAYSTACK_VARIADIC_ANNOTATION]
+
+# GreedyVariadic type is similar to Variadic.
+# The only difference is the way it's treated by the Pipeline when input is received
+# in a socket with this type.
+# Instead of waiting for other inputs to be received, Components that have a GreedyVariadic
+# input will be run right after receiving the first input.
+# Even if there are multiple connections to that socket.
+GreedyVariadic: TypeAlias = Annotated[Iterable[T], HAYSTACK_GREEDY_VARIADIC_ANNOTATION]
+
+
+class _empty:
+    """Custom object for marking InputSocket.default_value as not set."""
+
+
+@dataclass
+class InputSocket:
+    """
+    Represents an input of a `Component`.
+
+    :param name:
+        The name of the input.
+    :param type:
+        The type of the input.
+    :param default_value:
+        The default value of the input. If not set, the input is mandatory.
+    :param is_lazy_variadic:
+        Whether the input is a lazy variadic or not.
+    :param is_greedy:
+        Whether the input is a greedy variadic or not.
+    :param senders:
+        The list of components that send data to this input.
+    :param wrap_input_in_list:
+        Whether to wrap the input in a list before passing it to the component.
+        Only applies to lazy variadic inputs so when is_lazy_variadic is True.
+    """
+
+    name: str
+    type: type | UnionType
+    default_value: Any = _empty
+    is_lazy_variadic: bool = field(init=False)
+    is_greedy: bool = field(init=False)
+    senders: list[str] = field(default_factory=list)
+    wrap_input_in_list: bool = True
+
+    @property
+    def is_variadic(self) -> bool:
+        """Check if the input is variadic."""
+        return self.is_greedy or self.is_lazy_variadic
+
+    @property
+    def is_mandatory(self) -> bool:
+        """Check if the input is mandatory."""
+        return self.default_value == _empty
+
+    def __post_init__(self) -> None:
+        try:
+            # __metadata__ is a tuple
+            self.is_lazy_variadic = (
+                hasattr(self.type, "__metadata__") and self.type.__metadata__[0] == HAYSTACK_VARIADIC_ANNOTATION
+            )
+            self.is_greedy = (
+                hasattr(self.type, "__metadata__") and self.type.__metadata__[0] == HAYSTACK_GREEDY_VARIADIC_ANNOTATION
+            )
+        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:
+            outer_args = get_args(self.type)
+            inner_type = outer_args[0]
+            inner_args = get_args(inner_type)
+            if not inner_args:
+                raise ComponentError(
+                    f"Variadic input '{self.name}' must have a type argument, e.g. Variadic[int]. "
+                    f"Got bare {inner_type!r} without a type argument."
+                )
+            self.type = inner_args[0]
+
+
+class InputSocketTypeDescriptor(TypedDict):
+    """
+    Describes the type of `InputSocket`.
+    """
+
+    type: type | UnionType
+    is_mandatory: bool
+
+
+@dataclass
+class OutputSocket:
+    """
+    Represents an output of a `Component`.
+
+    :param name:
+        The name of the output.
+    :param type:
+        The type of the output.
+    :param receivers:
+        The list of components that receive the output of this component.
+    """
+
+    name: str
+    type: type
+    receivers: list[str] = field(default_factory=list)