Coverage for haystack/components/joiners/answer_joiner.py: 98%

54 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 

4 

5import itertools 

6from collections.abc import Callable 

7from enum import Enum 

8from math import inf 

9from typing import Any 

10 

11from haystack import component, default_from_dict, default_to_dict 

12from haystack.core.component.types import Variadic 

13from haystack.dataclasses.answer import ExtractedAnswer, GeneratedAnswer 

14 

15AnswerType = GeneratedAnswer | ExtractedAnswer 

16 

17 

18class JoinMode(Enum): 

19 """ 

20 Enum for AnswerJoiner join modes. 

21 """ 

22 

23 CONCATENATE = "concatenate" 

24 

25 def __str__(self) -> str: 

26 return self.value 

27 

28 @staticmethod 

29 def from_str(string: str) -> "JoinMode": 

30 """ 

31 Convert a string to a JoinMode enum. 

32 """ 

33 enum_map = {e.value: e for e in JoinMode} 

34 mode = enum_map.get(string) 

35 if mode is None: 

36 msg = f"Unknown join mode '{string}'. Supported modes in AnswerJoiner are: {list(enum_map.keys())}" 

37 raise ValueError(msg) 

38 return mode 

39 

40 

41@component 

42class AnswerJoiner: 

43 """ 

44 Merges multiple lists of `Answer` objects into a single list. 

45 

46 Use this component to combine answers from different Generators into a single list. 

47 Currently, the component supports only one join mode: `CONCATENATE`. 

48 This mode concatenates multiple lists of answers into a single list. 

49 

50 ### Usage example 

51 

52 In this example, AnswerJoiner merges answers from two different Generators: 

53 

54 ```python 

55 from haystack.components.builders import AnswerBuilder 

56 from haystack.components.joiners import AnswerJoiner 

57 

58 from haystack.core.pipeline import Pipeline 

59 

60 from haystack.components.generators.chat import OpenAIChatGenerator 

61 from haystack.dataclasses import ChatMessage 

62 

63 

64 query = "What's Natural Language Processing?" 

65 messages = [ChatMessage.from_system("You are a helpful, respectful and honest assistant. Be super concise."), 

66 ChatMessage.from_user(query)] 

67 

68 pipe = Pipeline() 

69 pipe.add_component("llm_1", OpenAIChatGenerator()) 

70 pipe.add_component("llm_2", OpenAIChatGenerator()) 

71 pipe.add_component("aba", AnswerBuilder()) 

72 pipe.add_component("abb", AnswerBuilder()) 

73 pipe.add_component("joiner", AnswerJoiner()) 

74 

75 pipe.connect("llm_1.replies", "aba") 

76 pipe.connect("llm_2.replies", "abb") 

77 pipe.connect("aba.answers", "joiner") 

78 pipe.connect("abb.answers", "joiner") 

79 

80 results = pipe.run(data={"llm_1": {"messages": messages}, 

81 "llm_2": {"messages": messages}, 

82 "aba": {"query": query}, 

83 "abb": {"query": query}}) 

84 ``` 

85 """ 

86 

87 def __init__( 

88 self, join_mode: str | JoinMode = JoinMode.CONCATENATE, top_k: int | None = None, sort_by_score: bool = False 

89 ) -> None: 

90 """ 

91 Creates an AnswerJoiner component. 

92 

93 :param join_mode: 

94 Specifies the join mode to use. Available modes: 

95 - `concatenate`: Concatenates multiple lists of Answers into a single list. 

96 :param top_k: 

97 The maximum number of Answers to return. Must be `None` or greater than 0. 

98 :param sort_by_score: 

99 If `True`, sorts the documents by score in descending order. 

100 If a document has no score, it is handled as if its score is -infinity. 

101 

102 :raises ValueError: 

103 If `top_k` is not `None` and is less than or equal to 0. 

104 """ 

105 if top_k is not None and top_k <= 0: 

106 raise ValueError("top_k must be greater than 0.") 

107 if isinstance(join_mode, str): 

108 join_mode = JoinMode.from_str(join_mode) 

109 join_mode_functions: dict[JoinMode, Callable[[list[list[AnswerType]]], list[AnswerType]]] = { 

110 JoinMode.CONCATENATE: self._concatenate 

111 } 

112 self.join_mode_function: Callable[[list[list[AnswerType]]], list[AnswerType]] = join_mode_functions[join_mode] 

113 self.join_mode = join_mode 

114 self.top_k = top_k 

115 self.sort_by_score = sort_by_score 

116 

117 @component.output_types(answers=list[AnswerType]) 

118 def run(self, answers: Variadic[list[AnswerType]], top_k: int | None = None) -> dict[str, Any]: 

119 """ 

120 Joins multiple lists of Answers into a single list depending on the `join_mode` parameter. 

121 

122 :param answers: 

123 Nested list of Answers to be merged. 

124 

125 :param top_k: 

126 The maximum number of Answers to return. Overrides the instance's `top_k` if provided. 

127 A value of 0 returns no answers. Must not be negative. 

128 

129 :returns: 

130 A dictionary with the following keys: 

131 - `answers`: Merged list of Answers 

132 

133 :raises ValueError: 

134 If `top_k` is negative. 

135 """ 

136 answers_list = list(answers) 

137 join_function = self.join_mode_function 

138 output_answers: list[AnswerType] = join_function(answers_list) 

139 

140 if self.sort_by_score: 

141 output_answers = sorted( 

142 output_answers, 

143 key=lambda answer: score if (score := getattr(answer, "score", None)) is not None else -inf, 

144 reverse=True, 

145 ) 

146 

147 if top_k is not None: 

148 if top_k < 0: 

149 raise ValueError("top_k must not be negative.") 

150 output_answers = output_answers[:top_k] 

151 elif self.top_k is not None: 

152 output_answers = output_answers[: self.top_k] 

153 return {"answers": output_answers} 

154 

155 def _concatenate(self, answer_lists: list[list[AnswerType]]) -> list[AnswerType]: 

156 """ 

157 Concatenate multiple lists of Answers, flattening them into a single list and sorting by score. 

158 

159 :param answer_lists: List of lists of Answers to be flattened. 

160 """ 

161 return list(itertools.chain.from_iterable(answer_lists)) 

162 

163 def to_dict(self) -> dict[str, Any]: 

164 """ 

165 Serializes the component to a dictionary. 

166 

167 :returns: 

168 Dictionary with serialized data. 

169 """ 

170 return default_to_dict(self, join_mode=str(self.join_mode), top_k=self.top_k, sort_by_score=self.sort_by_score) 

171 

172 @classmethod 

173 def from_dict(cls, data: dict[str, Any]) -> "AnswerJoiner": 

174 """ 

175 Deserializes the component from a dictionary. 

176 

177 :param data: 

178 The dictionary to deserialize from. 

179 :returns: 

180 The deserialized component. 

181 """ 

182 return default_from_dict(cls, data)