Coverage for haystack/core/component/types.py: 93%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from collections.abc import Iterable 

6from dataclasses import dataclass, field 

7from types import UnionType 

8from typing import Annotated, Any, TypeAlias, TypedDict, TypeVar, get_args 

9 

10from haystack.core.errors import ComponentError 

11 

12HAYSTACK_VARIADIC_ANNOTATION = "__haystack__variadic_t" 

13HAYSTACK_GREEDY_VARIADIC_ANNOTATION = "__haystack__greedy_variadic_t" 

14 

15# # Generic type variable used in the Variadic container 

16T = TypeVar("T") 

17 

18 

19# Variadic is a custom annotation type we use to mark input types. 

20# This type doesn't do anything else than "marking" the contained 

21# type so it can be used in the `InputSocket` creation where we 

22# check that its annotation equals to HAYSTACK_VARIADIC_ANNOTATION 

23Variadic: TypeAlias = Annotated[Iterable[T], HAYSTACK_VARIADIC_ANNOTATION] 

24 

25# GreedyVariadic type is similar to Variadic. 

26# The only difference is the way it's treated by the Pipeline when input is received 

27# in a socket with this type. 

28# Instead of waiting for other inputs to be received, Components that have a GreedyVariadic 

29# input will be run right after receiving the first input. 

30# Even if there are multiple connections to that socket. 

31GreedyVariadic: TypeAlias = Annotated[Iterable[T], HAYSTACK_GREEDY_VARIADIC_ANNOTATION] 

32 

33 

34class _empty: 

35 """Custom object for marking InputSocket.default_value as not set.""" 

36 

37 

38@dataclass 

39class InputSocket: 

40 """ 

41 Represents an input of a `Component`. 

42 

43 :param name: 

44 The name of the input. 

45 :param type: 

46 The type of the input. 

47 :param default_value: 

48 The default value of the input. If not set, the input is mandatory. 

49 :param is_lazy_variadic: 

50 Whether the input is a lazy variadic or not. 

51 :param is_greedy: 

52 Whether the input is a greedy variadic or not. 

53 :param senders: 

54 The list of components that send data to this input. 

55 :param wrap_input_in_list: 

56 Whether to wrap the input in a list before passing it to the component. 

57 Only applies to lazy variadic inputs so when is_lazy_variadic is True. 

58 """ 

59 

60 name: str 

61 type: type | UnionType 

62 default_value: Any = _empty 

63 is_lazy_variadic: bool = field(init=False) 

64 is_greedy: bool = field(init=False) 

65 senders: list[str] = field(default_factory=list) 

66 wrap_input_in_list: bool = True 

67 

68 @property 

69 def is_variadic(self) -> bool: 

70 """Check if the input is variadic.""" 

71 return self.is_greedy or self.is_lazy_variadic 

72 

73 @property 

74 def is_mandatory(self) -> bool: 

75 """Check if the input is mandatory.""" 

76 # Identity, so a default with a custom `__eq__`, such as a DataFrame, is never compared to the sentinel. 

77 return self.default_value is _empty 

78 

79 def __post_init__(self) -> None: 

80 try: 

81 # __metadata__ is a tuple 

82 self.is_lazy_variadic = ( 

83 hasattr(self.type, "__metadata__") and self.type.__metadata__[0] == HAYSTACK_VARIADIC_ANNOTATION 

84 ) 

85 self.is_greedy = ( 

86 hasattr(self.type, "__metadata__") and self.type.__metadata__[0] == HAYSTACK_GREEDY_VARIADIC_ANNOTATION 

87 ) 

88 except AttributeError: 

89 self.is_lazy_variadic = False 

90 self.is_greedy = False 

91 

92 # We need to "unpack" the type inside the Variadic annotation, otherwise the pipeline connection api will try 

93 # to match `Annotated[type, HAYSTACK_VARIADIC_ANNOTATION]`. 

94 # 

95 # Note1: Variadic is expressed as an annotation of one single type, so the return value of get_args will 

96 # always be a one-item tuple. 

97 # 

98 # Note2: a pipeline always passes a list of items when a component input is declared as Variadic, so the 

99 # type itself always wraps an iterable of the declared type. For example, Variadic[int] is eventually an 

100 # alias for Iterable[int]. Since we're interested in getting the inner type `int`, we call `get_args` 

101 # twice: the first time to get `list[int]` out of `Variadic`, the second time to get `int` out of `list[int]`. 

102 if self.is_lazy_variadic or self.is_greedy: 

103 outer_args = get_args(self.type) 

104 inner_type = outer_args[0] 

105 inner_args = get_args(inner_type) 

106 if not inner_args: 

107 raise ComponentError( 

108 f"Variadic input '{self.name}' must have a type argument, e.g. Variadic[int]. " 

109 f"Got bare {inner_type!r} without a type argument." 

110 ) 

111 self.type = inner_args[0] 

112 

113 

114class InputSocketTypeDescriptor(TypedDict): 

115 """ 

116 Describes the type of `InputSocket`. 

117 """ 

118 

119 type: type | UnionType 

120 is_mandatory: bool 

121 

122 

123@dataclass 

124class OutputSocket: 

125 """ 

126 Represents an output of a `Component`. 

127 

128 :param name: 

129 The name of the output. 

130 :param type: 

131 The type of the output. 

132 :param receivers: 

133 The list of components that receive the output of this component. 

134 """ 

135 

136 name: str 

137 type: type 

138 receivers: list[str] = field(default_factory=list)