Coverage for haystack/components/joiners/branch.py: 85%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5from typing import Any 

6 

7from haystack import component, default_from_dict, default_to_dict 

8from haystack.core.component.types import GreedyVariadic 

9from haystack.utils import deserialize_type, serialize_type 

10 

11 

12@component 

13class BranchJoiner: 

14 """ 

15 A component that merges multiple input branches of a pipeline into a single output stream. 

16 

17 `BranchJoiner` receives multiple inputs of the same data type and forwards the first received value 

18 to its output. This is useful for scenarios where multiple branches need to converge before proceeding. 

19 

20 ### Common Use Cases: 

21 - **Loop Handling:** `BranchJoiner` helps close loops in pipelines. For example, if a pipeline component validates 

22 or modifies incoming data and produces an error-handling branch, `BranchJoiner` can merge both branches and send 

23 (or resend in the case of a loop) the data to the component that evaluates errors. See "Usage example" below. 

24 

25 - **Decision-Based Merging:** `BranchJoiner` reconciles branches coming from Router components (such as 

26 `ConditionalRouter`, `TextLanguageRouter`). Suppose a `TextLanguageRouter` directs user queries to different 

27 Retrievers based on the detected language. Each Retriever processes its assigned query and passes the results 

28 to `BranchJoiner`, which consolidates them into a single output before passing them to the next component, such 

29 as a `PromptBuilder`. 

30 

31 ### Example Usage: 

32 ```python 

33 import json 

34 

35 from haystack import Pipeline 

36 from haystack.components.generators.chat import OpenAIChatGenerator 

37 from haystack.components.joiners import BranchJoiner 

38 from haystack.components.validators import JsonSchemaValidator 

39 from haystack.dataclasses import ChatMessage 

40 

41 # Define a schema for validation 

42 person_schema = { 

43 "type": "object", 

44 "properties": { 

45 "first_name": {"type": "string", "pattern": "^[A-Z][a-z]+$"}, 

46 "last_name": {"type": "string", "pattern": "^[A-Z][a-z]+$"}, 

47 "nationality": {"type": "string", "enum": ["Italian", "Portuguese", "American"]}, 

48 }, 

49 "required": ["first_name", "last_name", "nationality"] 

50 } 

51 

52 # Initialize a pipeline 

53 pipe = Pipeline() 

54 

55 # Add components to the pipeline 

56 pipe.add_component("joiner", BranchJoiner(list[ChatMessage])) 

57 pipe.add_component("generator", OpenAIChatGenerator(model="gpt-4.1-mini")) 

58 pipe.add_component("validator", JsonSchemaValidator(json_schema=person_schema)) 

59 

60 # And connect them 

61 pipe.connect("joiner", "generator") 

62 pipe.connect("generator.replies", "validator.messages") 

63 pipe.connect("validator.validation_error", "joiner") 

64 

65 result = pipe.run( 

66 data={ 

67 "generator": {"generation_kwargs": {"response_format": {"type": "json_object"}}}, 

68 "joiner": {"value": [ChatMessage.from_user("Create json from Peter Parker")]}} 

69 ) 

70 

71 print(json.loads(result["validator"]["validated"][0].text)) 

72 

73 

74 # >> {'first_name': 'Peter', 'last_name': 'Parker', 'nationality': 'American', 'name': 'Spider-Man', 'occupation': 

75 # >> 'Superhero', 'age': 23, 'location': 'New York City'} 

76 ``` 

77 

78 Note that `BranchJoiner` can manage only one data type at a time. In this case, `BranchJoiner` is created for 

79 passing `list[ChatMessage]`. This determines the type of data that `BranchJoiner` will receive from the upstream 

80 connected components and also the type of data that `BranchJoiner` will send through its output. 

81 

82 In the code example, `BranchJoiner` receives a looped back `list[ChatMessage]` from the `JsonSchemaValidator` and 

83 sends it down to the `OpenAIChatGenerator` for re-generation. We can have multiple loopback connections in the 

84 pipeline. In this instance, the downstream component is only one (the `OpenAIChatGenerator`), but the pipeline could 

85 have more than one downstream component. 

86 """ 

87 

88 def __init__(self, type_: type) -> None: 

89 """ 

90 Creates a `BranchJoiner` component. 

91 

92 :param type_: The expected data type of inputs and outputs. 

93 """ 

94 self.type_ = type_ 

95 component.set_input_types(self, value=GreedyVariadic[type_]) # type: ignore 

96 component.set_output_types(self, value=type_) 

97 

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

99 """ 

100 Serializes the component into a dictionary. 

101 

102 :returns: 

103 Dictionary with serialized data. 

104 """ 

105 return default_to_dict(self, type_=serialize_type(self.type_)) 

106 

107 @classmethod 

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

109 """ 

110 Deserializes a `BranchJoiner` instance from a dictionary. 

111 

112 :param data: The dictionary containing serialized component data. 

113 :returns: 

114 A deserialized `BranchJoiner` instance. 

115 """ 

116 data["init_parameters"]["type_"] = deserialize_type(data["init_parameters"]["type_"]) 

117 return default_from_dict(cls, data) 

118 

119 def run(self, **kwargs: Any) -> dict[str, Any]: 

120 """ 

121 Executes the `BranchJoiner`, selecting the first available input value and passing it downstream. 

122 

123 :param **kwargs: The input data. Must be of the type declared by `type_` during initialization. 

124 :returns: 

125 A dictionary with a single key `value`, containing the first input received. 

126 """ 

127 if (inputs_count := len(kwargs["value"])) != 1: 

128 raise ValueError(f"BranchJoiner expects only one input, but {inputs_count} were received.") 

129 return {"value": kwargs["value"][0]}