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
5import itertools
6from collections.abc import Callable
7from enum import Enum
8from math import inf
9from typing import Any
11from haystack import component, default_from_dict, default_to_dict
12from haystack.core.component.types import Variadic
13from haystack.dataclasses.answer import ExtractedAnswer, GeneratedAnswer
15AnswerType = GeneratedAnswer | ExtractedAnswer
18class JoinMode(Enum):
19 """
20 Enum for AnswerJoiner join modes.
21 """
23 CONCATENATE = "concatenate"
25 def __str__(self) -> str:
26 return self.value
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
41@component
42class AnswerJoiner:
43 """
44 Merges multiple lists of `Answer` objects into a single list.
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.
50 ### Usage example
52 In this example, AnswerJoiner merges answers from two different Generators:
54 ```python
55 from haystack.components.builders import AnswerBuilder
56 from haystack.components.joiners import AnswerJoiner
58 from haystack.core.pipeline import Pipeline
60 from haystack.components.generators.chat import OpenAIChatGenerator
61 from haystack.dataclasses import ChatMessage
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)]
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())
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")
80 results = pipe.run(data={"llm_1": {"messages": messages},
81 "llm_2": {"messages": messages},
82 "aba": {"query": query},
83 "abb": {"query": query}})
84 ```
85 """
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.
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.
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
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.
122 :param answers:
123 Nested list of Answers to be merged.
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.
129 :returns:
130 A dictionary with the following keys:
131 - `answers`: Merged list of Answers
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)
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 )
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}
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.
159 :param answer_lists: List of lists of Answers to be flattened.
160 """
161 return list(itertools.chain.from_iterable(answer_lists))
163 def to_dict(self) -> dict[str, Any]:
164 """
165 Serializes the component to a dictionary.
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)
172 @classmethod
173 def from_dict(cls, data: dict[str, Any]) -> "AnswerJoiner":
174 """
175 Deserializes the component from a dictionary.
177 :param data:
178 The dictionary to deserialize from.
179 :returns:
180 The deserialized component.
181 """
182 return default_from_dict(cls, data)