Coverage for haystack/core/component/component.py: 99%
179 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-21 13:53 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5"""
6Attributes:
8 component: Marks a class as a component. Any class decorated with `@component` can be used by a Pipeline.
10All components must follow the contract below. This docstring is the source of truth for components contract.
12<hr>
14`@component` decorator
16All component classes must be decorated with the `@component` decorator. This allows Haystack to discover them.
18<hr>
20`__init__(self, **kwargs)`
22Optional method.
24Components may have an `__init__` method where they define:
26- `self.init_parameters = {same parameters that the __init__ method received}`:
27 In this dictionary you can store any state the components wish to be persisted when they are saved.
28 These values will be given to the `__init__` method of a new instance when the pipeline is loaded.
29 Note that by default the `@component` decorator saves the arguments automatically.
30 However, if a component sets their own `init_parameters` manually in `__init__()`, that will be used instead.
31 Note: all of the values contained here **must be JSON serializable**. Serialize them manually if needed.
33Components should take only "basic" Python types as parameters of their `__init__` function, or iterables and
34dictionaries containing only such values. Anything else (objects, functions, etc) will raise an exception at init
35time. If there's the need for such values, consider serializing them to a string.
37If you need to accept classes or callables, accept either a string import path or the callable itself. Resolve strings
38to objects in `__init__`, and serialize objects back to importable strings in `to_dict()` so that `from_dict()` can load
39them (for example, store `"module_path.symbol_name"` and load it via `importlib`). This keeps init parameters JSON
40serializable for pipeline save/load. See `haystack.testing.sample_components.accumulate.Accumulate` for a reference
41implementation.
43The `__init__` must be extremely lightweight, because it's a frequent operation during the construction and
44validation of the pipeline. If a component has some heavy state to initialize (models, backends, etc...) refer to
45the `warm_up()` method.
47<hr>
49`warm_up(self)`
51Optional method.
53This method is called by Pipeline before the graph execution. Make sure to avoid double-initializations,
54because Pipeline will not keep track of which components it called `warm_up()` on.
56<hr>
58`run(self, data)`
60Mandatory method.
62This is the method where the main functionality of the component should be carried out. It's called by
63`Pipeline.run()`.
65When the component should run, Pipeline will call this method with an instance of the dataclass returned by the
66method decorated with `@component.input`. This dataclass contains:
68- all the input values coming from other components connected to it,
69- if any is missing, the corresponding value defined in `self.defaults`, if it exists.
71`run()` must return a single instance of the dataclass declared through the method decorated with
72`@component.output`.
74"""
76import inspect
77from collections.abc import Callable, Coroutine, Iterator, Mapping
78from contextlib import contextmanager
79from contextvars import ContextVar
80from copy import deepcopy
81from dataclasses import dataclass
82from types import new_class
83from typing import Any, ParamSpec, Protocol, TypeVar, overload, runtime_checkable
85from haystack import logging
86from haystack.core.errors import ComponentError
87from haystack.core.type_utils import _resolve_parameter_types
89from .sockets import Sockets
90from .types import InputSocket, OutputSocket, _empty
92logger = logging.getLogger(__name__)
94RunParamsT = ParamSpec("RunParamsT")
95RunReturnT = TypeVar("RunReturnT", bound=Mapping[str, Any] | Coroutine[Any, Any, Mapping[str, Any]])
98@dataclass
99class PreInitHookPayload:
100 """
101 Payload for the hook called before a component instance is initialized.
103 :param callback:
104 Receives the following inputs: component class and init parameter keyword args.
105 :param in_progress:
106 Flag to indicate if the hook is currently being executed.
107 Used to prevent it from being called recursively (if the component's constructor
108 instantiates another component).
109 """
111 callback: Callable
112 in_progress: bool = False
115_COMPONENT_PRE_INIT_HOOK: ContextVar[PreInitHookPayload | None] = ContextVar("component_pre_init_hook", default=None)
118@contextmanager
119def _hook_component_init(callback: Callable) -> Iterator[None]:
120 """
121 Context manager to set a callback that will be invoked before a component's constructor is called.
123 The callback receives the component class and the init parameters (as keyword arguments) and can modify the init
124 parameters in place.
126 :param callback:
127 Callback function to invoke.
128 """
129 token = _COMPONENT_PRE_INIT_HOOK.set(PreInitHookPayload(callback))
130 try:
131 yield
132 finally:
133 _COMPONENT_PRE_INIT_HOOK.reset(token)
136@runtime_checkable
137class Component(Protocol):
138 """
139 Note this is only used by type checking tools.
141 In order to implement the `Component` protocol, custom components need to
142 have a `run` method. The signature of the method and its return value
143 won't be checked, i.e. classes with the following methods:
145 def run(self, param: str) -> dict[str, Any]:
146 ...
148 and
150 def run(self, **kwargs):
151 ...
153 will be both considered as respecting the protocol. This makes the type
154 checking much weaker, but we have other places where we ensure code is
155 dealing with actual Components.
157 The protocol is runtime checkable so it'll be possible to assert:
159 isinstance(MyComponent, Component)
160 """
162 # The following expression defines a run method compatible with any input signature.
163 # Its type is equivalent to Callable[..., dict[str, Any]].
164 # See https://typing.python.org/en/latest/spec/callables.html#meaning-of-in-callable.
165 #
166 # Using `run: Callable[..., dict[str, Any]]` directly leads to type errors: the protocol would expect a settable
167 # attribute `run`, while the actual implementation is a read-only method.
168 # For example:
169 # from haystack import Pipeline, component
170 # @component
171 # class MyComponent:
172 # @component.output_types(out=str)
173 # def run(self):
174 # return {"out": "Hello, world!"}
175 # pipeline = Pipeline()
176 # pipeline.add_component("my_component", MyComponent())
177 #
178 # mypy raises:
179 # error: Argument 2 to "add_component" of "PipelineBase" has incompatible type "MyComponent"; expected "Component"
180 # [arg-type]
181 # note: Protocol member Component.run expected settable variable, got read-only attribute
183 def run(self, *args: Any, **kwargs: Any) -> Mapping[str, Any]: # noqa: D102
184 ...
187class ComponentMeta(type):
188 @staticmethod
189 def _positional_to_kwargs(cls_type: type, args: tuple[Any, ...]) -> dict[str, Any]:
190 """
191 Convert positional arguments to keyword arguments based on the signature of the `__init__` method.
192 """
193 init_signature = inspect.signature(cls_type.__init__) # type:ignore[misc]
194 init_params = {name: info for name, info in init_signature.parameters.items() if name != "self"}
196 out = {}
197 for arg, (name, info) in zip(args, init_params.items(), strict=False):
198 if info.kind == inspect.Parameter.VAR_POSITIONAL:
199 raise ComponentError(
200 "Pre-init hooks do not support components with variadic positional args in their init method"
201 )
203 assert info.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD, inspect.Parameter.POSITIONAL_ONLY)
204 out[name] = arg
205 return out
207 @staticmethod
208 def _parse_and_set_output_sockets(instance: Any) -> None:
209 has_async_run = hasattr(instance, "run_async")
211 # If `component.set_output_types()` was called in the component constructor,
212 # `__haystack_output__` is already populated, no need to do anything.
213 if not hasattr(instance, "__haystack_output__"):
214 # If that's not the case, we need to populate `__haystack_output__`
215 #
216 # If either of the run methods were decorated, they'll have a field assigned that
217 # stores the output specification. If both run methods were decorated, we ensure that
218 # outputs are the same. We deepcopy the content of the cache to transfer ownership from
219 # the class method to the actual instance, so that different instances of the same class
220 # won't share this data.
222 run_output_types = getattr(instance.run, "_output_types_cache", {})
223 async_run_output_types = getattr(instance.run_async, "_output_types_cache", {}) if has_async_run else {}
225 if has_async_run and run_output_types != async_run_output_types:
226 raise ComponentError("Output type specifications of 'run' and 'run_async' methods must be the same")
227 output_types_cache = run_output_types
229 instance.__haystack_output__ = Sockets(instance, deepcopy(output_types_cache), OutputSocket)
231 @staticmethod
232 def _parse_and_set_input_sockets(component_cls: type, instance: Any) -> None:
233 def inner(method: Callable[..., Any], sockets: Sockets) -> inspect.Signature:
234 from inspect import Parameter
236 run_signature = inspect.signature(method)
237 # Resolves the annotations of components using postponed evaluation of annotations, where they are stored
238 # as strings.
239 param_types = _resolve_parameter_types(method)
241 for param_name, param_info in run_signature.parameters.items():
242 if param_name == "self" or param_info.kind in (Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD):
243 continue
245 socket_kwargs = {"name": param_name, "type": param_types[param_name]}
246 if param_info.default != Parameter.empty:
247 socket_kwargs["default_value"] = param_info.default
249 new_socket = InputSocket(**socket_kwargs)
251 # Also ensure that new sockets don't override existing ones.
252 existing_socket = sockets.get(param_name)
253 if existing_socket is not None and existing_socket != new_socket:
254 raise ComponentError(
255 "set_input_types()/set_input_type() cannot override the parameters of the 'run' method"
256 )
258 sockets[param_name] = new_socket
260 return run_signature
262 # Create the sockets if set_input_types() wasn't called in the constructor.
263 if not hasattr(instance, "__haystack_input__"):
264 instance.__haystack_input__ = Sockets(instance, {}, InputSocket)
266 inner(getattr(component_cls, "run"), instance.__haystack_input__) # noqa: B009
268 # Ensure that the sockets are the same for the async method, if it exists.
269 async_run = getattr(component_cls, "run_async", None)
270 if async_run is not None:
271 run_sockets = Sockets(instance, {}, InputSocket)
272 async_run_sockets = Sockets(instance, {}, InputSocket)
274 # Can't use the sockets from above as they might contain
275 # values set with set_input_types().
276 run_sig = inner(getattr(component_cls, "run"), run_sockets) # noqa: B009
277 async_run_sig = inner(async_run, async_run_sockets)
279 if async_run_sockets != run_sockets or run_sig != async_run_sig:
280 sig_diff = _compare_run_methods_signatures(run_sig, async_run_sig)
281 raise ComponentError(
282 f"Parameters of 'run' and 'run_async' methods must be the same.\nDifferences found:\n{sig_diff}"
283 )
285 def __call__(cls, *args: Any, **kwargs: Any) -> Any:
286 """
287 This method is called when clients instantiate a Component and runs before __new__ and __init__.
288 """
289 # This will call __new__ then __init__, giving us back the Component instance
290 pre_init_hook = _COMPONENT_PRE_INIT_HOOK.get()
291 if pre_init_hook is None or pre_init_hook.in_progress:
292 instance = super().__call__(*args, **kwargs)
293 else:
294 try:
295 pre_init_hook.in_progress = True
296 named_positional_args = ComponentMeta._positional_to_kwargs(cls, args)
297 assert set(named_positional_args.keys()).intersection(kwargs.keys()) == set(), (
298 "positional and keyword arguments overlap"
299 )
300 kwargs.update(named_positional_args)
301 pre_init_hook.callback(cls, kwargs)
302 instance = super().__call__(**kwargs)
303 finally:
304 pre_init_hook.in_progress = False
306 # Before returning, we have the chance to modify the newly created
307 # Component instance, so we take the chance and set up the I/O sockets
308 has_async_run = hasattr(instance, "run_async")
309 if has_async_run and not inspect.iscoroutinefunction(instance.run_async):
310 raise ComponentError(f"Method 'run_async' of component '{cls.__name__}' must be a coroutine")
311 instance.__haystack_supports_async__ = has_async_run
313 ComponentMeta._parse_and_set_input_sockets(cls, instance)
314 ComponentMeta._parse_and_set_output_sockets(instance)
316 # Since a Component can't be used in multiple Pipelines at the same time
317 # we need to know if it's already owned by a Pipeline when adding it to one.
318 # We use this flag to check that.
319 instance.__haystack_added_to_pipeline__ = None
321 return instance
324def _component_repr(component: Component) -> str:
325 """
326 All Components override their __repr__ method with this one.
328 It prints the component name and the input/output sockets.
329 """
330 result = object.__repr__(component)
331 if pipeline := getattr(component, "__haystack_added_to_pipeline__", None):
332 # This Component has been added in a Pipeline, let's get the name from there.
333 result += f"\n{pipeline.get_component_name(component)}"
335 # We're explicitly ignoring the type here because we're sure that the component
336 # has the __haystack_input__ and __haystack_output__ attributes at this point
337 return (
338 f"{result}\n{getattr(component, '__haystack_input__', '<invalid_input_sockets>')}"
339 f"\n{getattr(component, '__haystack_output__', '<invalid_output_sockets>')}"
340 )
343def _component_run_has_kwargs(component_cls: type) -> bool:
344 run_method = getattr(component_cls, "run", None)
345 if run_method is None:
346 return False
347 return any(
348 param.kind == inspect.Parameter.VAR_KEYWORD for param in inspect.signature(run_method).parameters.values()
349 )
352def _compare_run_methods_signatures(run_sig: inspect.Signature, async_run_sig: inspect.Signature) -> str:
353 """
354 Builds a detailed error message with the differences between the signatures of the run and run_async methods.
356 :param run_sig: The signature of the run method
357 :param async_run_sig: The signature of the run_async method
359 :returns:
360 A detailed error message if signatures don't match, empty string if they do
361 """
362 differences = []
363 run_params = list(run_sig.parameters.items())
364 async_params = list(async_run_sig.parameters.items())
366 if len(run_params) != len(async_params):
367 differences.append(
368 f"Different number of parameters: run has {len(run_params)}, run_async has {len(async_params)}"
369 )
371 for (run_name, run_param), (async_name, async_param) in zip(run_params, async_params, strict=False):
372 if run_name != async_name:
373 differences.append(f"Parameter name mismatch: {run_name} vs {async_name}")
375 if run_param.annotation != async_param.annotation:
376 differences.append(
377 f"Parameter '{run_name}' type mismatch: {run_param.annotation} vs {async_param.annotation}"
378 )
380 if run_param.default != async_param.default:
381 differences.append(
382 f"Parameter '{run_name}' default value mismatch: {run_param.default} vs {async_param.default}"
383 )
385 if run_param.kind != async_param.kind:
386 differences.append(
387 f"Parameter '{run_name}' kind (POSITIONAL, KEYWORD, etc.) mismatch: "
388 f"{run_param.kind} vs {async_param.kind}"
389 )
391 return "\n".join(differences)
394T = TypeVar("T", bound=Component)
397class _Component:
398 """
399 See module's docstring.
401 Args:
402 cls: the class that should be used as a component.
404 Returns:
405 A class that can be recognized as a component.
407 Raises:
408 ComponentError: if the class provided has no `run()` method or otherwise doesn't respect the component contract.
409 """
411 def __init__(self) -> None:
412 self.registry: dict[str, type] = {}
414 def set_input_type(
415 self,
416 instance: Component,
417 name: str,
418 type: Any, # noqa: A002
419 default: Any = _empty,
420 ) -> None:
421 """
422 Add a single input socket to the component instance.
424 Replaces any existing input socket with the same name.
426 :param instance: Component instance where the input type will be added.
427 :param name: name of the input socket.
428 :param type: type of the input socket.
429 :param default: default value of the input socket, defaults to _empty
430 """
431 if not _component_run_has_kwargs(instance.__class__):
432 raise ComponentError(
433 "Cannot set input types on a component that doesn't have a kwargs parameter in the 'run' method"
434 )
436 if not hasattr(instance, "__haystack_input__"):
437 instance.__haystack_input__ = Sockets(instance, {}, InputSocket) # type: ignore
438 instance.__haystack_input__[name] = InputSocket(name=name, type=type, default_value=default) # type: ignore
440 def set_input_types(self, instance: Any, **types: type[Any]) -> None:
441 """
442 Method that specifies the input types when 'kwargs' is passed to the run method.
444 Use as:
446 ```python
447 @component
448 class MyComponent:
450 def __init__(self, value: int) -> None:
451 component.set_input_types(self, value_1=str, value_2=str)
452 ...
454 @component.output_types(output_1=int, output_2=str)
455 def run(self, **kwargs):
456 return {"output_1": kwargs["value_1"], "output_2": ""}
457 ```
459 Note that if the `run()` method also specifies some parameters, those will take precedence.
461 For example:
463 ```python
464 @component
465 class MyComponent:
467 def __init__(self, value: int) -> None:
468 component.set_input_types(self, value_1=str, value_2=str)
469 ...
471 @component.output_types(output_1=int, output_2=str)
472 def run(self, value_0: str, value_1: Optional[str] = None, **kwargs):
473 return {"output_1": kwargs["value_1"], "output_2": ""}
474 ```
476 would add a mandatory `value_0` parameters, make the `value_1`
477 parameter optional with a default None, and keep the `value_2`
478 parameter mandatory as specified in `set_input_types`.
480 """
481 if not _component_run_has_kwargs(instance.__class__):
482 raise ComponentError(
483 "Cannot set input types on a component that doesn't have a kwargs parameter in the 'run' method"
484 )
486 instance.__haystack_input__ = Sockets(
487 instance, {name: InputSocket(name=name, type=type_) for name, type_ in types.items()}, InputSocket
488 )
490 def set_output_types(self, instance: Any, **types: type[Any]) -> None:
491 """
492 Method that specifies the output types when the 'run' method is not decorated with 'component.output_types'.
494 Use as:
496 ```python
497 @component
498 class MyComponent:
500 def __init__(self, value: int) -> None:
501 component.set_output_types(self, output_1=int, output_2=str)
502 ...
504 # no decorators here
505 def run(self, value: int):
506 return {"output_1": 1, "output_2": "2"}
508 # also no decorators here
509 async def run_async(self, value: int):
510 return {"output_1": 1, "output_2": "2"}
511 ```
512 """
513 has_run_decorator = hasattr(instance.run, "_output_types_cache")
514 has_run_async_decorator = hasattr(instance, "run_async") and hasattr(instance.run_async, "_output_types_cache")
515 if has_run_decorator or has_run_async_decorator:
516 raise ComponentError(
517 "Cannot call `set_output_types` on a component that already has the 'output_types' decorator on its "
518 "`run` or `run_async` methods."
519 )
521 instance.__haystack_output__ = Sockets(
522 instance, {name: OutputSocket(name=name, type=type_) for name, type_ in types.items()}, OutputSocket
523 )
525 def output_types(
526 self, **types: Any
527 ) -> Callable[[Callable[RunParamsT, RunReturnT]], Callable[RunParamsT, RunReturnT]]:
528 """
529 Decorator factory that specifies the output types of a component.
531 Use as:
532 ```python
533 @component
534 class MyComponent:
535 @component.output_types(output_1=int, output_2=str)
536 def run(self, value: int):
537 return {"output_1": 1, "output_2": "2"}
538 ```
539 """
541 def output_types_decorator(run_method: Callable[RunParamsT, RunReturnT]) -> Callable[RunParamsT, RunReturnT]:
542 """
543 Decorator that sets the output types of the decorated method.
545 This happens at class creation time, and since we don't have the decorated
546 class available here, we temporarily store the output types as an attribute of
547 the decorated method. The ComponentMeta metaclass will use this data to create
548 sockets at instance creation time.
549 """
550 method_name = run_method.__name__
551 if method_name not in ("run", "run_async"):
552 raise ComponentError("'output_types' decorator can only be used on 'run' and 'run_async' methods")
554 setattr( # noqa: B010
555 run_method,
556 "_output_types_cache",
557 {name: OutputSocket(name=name, type=type_) for name, type_ in types.items()},
558 )
559 return run_method
561 return output_types_decorator
563 def _component(self, cls: type[T]) -> type[T]:
564 """
565 Decorator validating the structure of the component and registering it in the components registry.
566 """
567 logger.debug("Registering {component} as a component", component=cls)
569 # Check for required methods and fail as soon as possible
570 if not hasattr(cls, "run"):
571 raise ComponentError(f"{cls.__name__} must have a 'run()' method. See the docs for more information.")
573 def copy_class_namespace(namespace: dict[str, Any]) -> None:
574 """
575 This is the callback that `typing.new_class` will use to populate the newly created class.
577 Simply copy the whole namespace from the decorated class.
578 """
579 for key, val in dict(cls.__dict__).items():
580 # __dict__ and __weakref__ are class-bound, we should let Python recreate them.
581 if key in ("__dict__", "__weakref__"):
582 continue
583 namespace[key] = val
585 # Recreate the decorated component class so it uses our metaclass.
586 # We must explicitly redefine the type of the class to make sure language servers
587 # and type checkers understand that the class is of the correct type.
588 new_cls: type[T] = new_class(cls.__name__, cls.__bases__, {"metaclass": ComponentMeta}, copy_class_namespace)
590 # Save the component in the class registry (for deserialization)
591 class_path = f"{new_cls.__module__}.{new_cls.__name__}"
592 if class_path in self.registry:
593 # Corner case, but it may occur easily in notebooks when re-running cells.
594 logger.debug(
595 "Component {component} is already registered. Previous imported from '{module_name}', \
596 new imported from '{new_module_name}'",
597 component=class_path,
598 module_name=self.registry[class_path],
599 new_module_name=new_cls,
600 )
601 self.registry[class_path] = new_cls
602 logger.debug("Registered Component {component}", component=new_cls)
604 # Override the __repr__ method with a default one
605 # mypy is not happy that:
606 # 1) we are assigning a method to a class
607 # 2) _component_repr has a different type (Callable[[Component], str]) than the expected
608 # __repr__ method (Callable[[object], str])
609 new_cls.__repr__ = _component_repr # type: ignore[assignment]
611 return new_cls
613 # Call signature when the decorator is used without parens (@component).
614 @overload
615 def __call__(self, cls: type[T]) -> type[T]: ...
617 # Overload allowing the decorator to be used with parens (@component()).
618 @overload
619 def __call__(self) -> Callable[[type[T]], type[T]]: ...
621 def __call__(self, cls: type[T] | None = None) -> type[T] | Callable[[type[T]], type[T]]:
622 # We must wrap the call to the decorator in a function for it to work
623 # correctly with or without parens
624 def wrap(cls: type[T]) -> type[T]:
625 return self._component(cls)
627 if cls:
628 # Decorator is called without parens
629 return wrap(cls)
631 # Decorator is called with parens
632 return wrap
635component = _Component()