Coverage for haystack/components/converters/output_adapter.py: 99%

71 statements  

« prev     ^ index     » next       coverage.py v7.16.1, created at 2026-10-09 14:34 +0000

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import ast 

6import contextlib 

7from collections.abc import Callable 

8from typing import Any, TypeAlias 

9 

10import jinja2.runtime 

11from jinja2 import TemplateSyntaxError 

12from jinja2.nativetypes import NativeEnvironment 

13 

14from haystack import component, default_from_dict, default_to_dict, logging 

15from haystack.core.errors import DeserializationError 

16from haystack.core.serialization_security import _is_unsafe_deserialization 

17from haystack.utils import deserialize_callable, deserialize_type, serialize_callable, serialize_type 

18from haystack.utils.jinja2_extensions import _extract_template_variables_and_assignments 

19from haystack.utils.jinja2_sandbox import HaystackSandboxedEnvironment 

20 

21logger = logging.getLogger(__name__) 

22 

23 

24class OutputAdaptationException(Exception): 

25 """Exception raised when there is an error during output adaptation.""" 

26 

27 

28@component 

29class OutputAdapter: 

30 """ 

31 Adapts output of a Component using Jinja templates. 

32 

33 Usage example: 

34 ```python 

35 from haystack import Document 

36 from haystack.components.converters import OutputAdapter 

37 

38 adapter = OutputAdapter(template="{{ documents[0].content }}", output_type=str) 

39 documents = [Document(content="Test content")] 

40 result = adapter.run(documents=documents) 

41 

42 assert result["output"] == "Test content" 

43 ``` 

44 """ 

45 

46 def __init__( 

47 self, 

48 template: str, 

49 output_type: TypeAlias, 

50 custom_filters: dict[str, Callable] | None = None, 

51 unsafe: bool = False, 

52 ) -> None: 

53 """ 

54 Create an OutputAdapter component. 

55 

56 :param template: 

57 A Jinja template that defines how to adapt the input data. 

58 The variables in the template define the input of this instance. 

59 e.g. 

60 With this template: 

61 ``` 

62 {{ documents[0].content }} 

63 ``` 

64 The Component input will be `documents`. 

65 :param output_type: 

66 The type of output this instance will return. 

67 :param custom_filters: 

68 A dictionary of custom Jinja filters used in the template. 

69 :param unsafe: 

70 Enable execution of arbitrary code in the Jinja template. 

71 This should only be used if you trust the source of the template as it can be lead to remote code execution. 

72 """ 

73 self.custom_filters = {**(custom_filters or {})} 

74 input_types: set[str] = set() 

75 

76 self._unsafe = unsafe 

77 

78 if self._unsafe: 

79 msg = ( 

80 "Unsafe mode is enabled. This allows execution of arbitrary code in the Jinja template. " 

81 "Use this only if you trust the source of the template." 

82 ) 

83 logger.warning(msg) 

84 self._env = ( 

85 NativeEnvironment() 

86 if self._unsafe 

87 else HaystackSandboxedEnvironment(undefined=jinja2.runtime.StrictUndefined) 

88 ) 

89 

90 try: 

91 self._env.parse(template) # Validate template syntax 

92 self.template = template 

93 except TemplateSyntaxError as e: 

94 raise ValueError(f"Invalid Jinja template '{template}': {e}") from e 

95 

96 for name, filter_func in self.custom_filters.items(): 

97 self._env.filters[name] = filter_func 

98 

99 # b) extract variables in the template 

100 assigned_variables, template_variables = _extract_template_variables_and_assignments( 

101 env=self._env, template=self.template 

102 ) 

103 route_input_names = template_variables - assigned_variables 

104 input_types.update(route_input_names) 

105 

106 # the env is not needed, discarded automatically 

107 component.set_input_types(self, **dict.fromkeys(input_types, Any)) 

108 component.set_output_types(self, output=output_type) 

109 self.output_type = output_type 

110 

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

112 """ 

113 Renders the Jinja template with the provided inputs. 

114 

115 :param kwargs: 

116 Must contain all variables used in the `template` string. 

117 :returns: 

118 A dictionary with the following keys: 

119 - `output`: Rendered Jinja template. 

120 

121 :raises OutputAdaptationException: If template rendering fails. 

122 """ 

123 # check if kwargs are empty 

124 if not kwargs: 

125 raise ValueError("No input data provided for output adaptation") 

126 for name, filter_func in self.custom_filters.items(): 

127 self._env.filters[name] = filter_func 

128 adapted_outputs = {} 

129 try: 

130 adapted_output_template = self._env.from_string(self.template) 

131 output_result = adapted_output_template.render(**kwargs) 

132 if isinstance(output_result, jinja2.runtime.Undefined): 

133 raise OutputAdaptationException(f"Undefined variable in the template {self.template}; kwargs: {kwargs}") # noqa: TRY301 

134 

135 # We suppress the exception in case the output is already a string, otherwise 

136 # we try to evaluate it and would fail. 

137 # This must be done cause the output could be different literal structures. 

138 # This doesn't support any user types. 

139 # When the declared output_type is str we skip literal evaluation so that a 

140 # rendered string that happens to be a valid Python literal (e.g. "1,000" -> (1, 0), 

141 # "42" -> 42, "None" -> None) is returned unchanged instead of being coerced to 

142 # another type, which would violate the declared output_type. 

143 with contextlib.suppress(Exception): 

144 if not self._unsafe and self.output_type is not str: 

145 output_result = ast.literal_eval(output_result) 

146 

147 adapted_outputs["output"] = output_result 

148 except Exception as e: 

149 raise OutputAdaptationException(f"Error adapting {self.template} with {kwargs}: {e}") from e 

150 return adapted_outputs 

151 

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

153 """ 

154 Serializes the component to a dictionary. 

155 

156 :returns: 

157 Dictionary with serialized data. 

158 """ 

159 se_filters = {name: serialize_callable(filter_func) for name, filter_func in self.custom_filters.items()} 

160 return default_to_dict( 

161 self, 

162 template=self.template, 

163 output_type=serialize_type(self.output_type), 

164 custom_filters=se_filters, 

165 unsafe=self._unsafe, 

166 ) 

167 

168 @classmethod 

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

170 """ 

171 Deserializes the component from a dictionary. 

172 

173 :param data: 

174 The dictionary to deserialize from. 

175 :returns: 

176 The deserialized component. 

177 """ 

178 init_params = data.get("init_parameters", {}) 

179 

180 # `unsafe=True` swaps the Jinja sandbox for a NativeEnvironment that executes arbitrary code. 

181 # Honor it from serialized data only when the whole pipeline is being loaded in unsafe mode; 

182 # otherwise a hostile pipeline could disable the sandbox on its own in default safe mode. 

183 if init_params.get("unsafe") and not _is_unsafe_deserialization(): 

184 raise DeserializationError( 

185 "Refusing to deserialize an OutputAdapter with unsafe=True while loading in safe mode. " 

186 "If you trust the source of this data, load it with Pipeline.load(..., unsafe=True)." 

187 ) 

188 

189 custom_filters = init_params.get("custom_filters", {}) 

190 if custom_filters and not _is_unsafe_deserialization(): 

191 raise DeserializationError( 

192 "Refusing to deserialize an OutputAdapter with custom filters while loading in safe mode. " 

193 "Custom filters are arbitrary callables that can execute during pipeline loading. " 

194 "If you trust the source of this data, load it with Pipeline.load(..., unsafe=True)." 

195 ) 

196 

197 init_params["output_type"] = deserialize_type(init_params["output_type"]) 

198 

199 if custom_filters: 

200 init_params["custom_filters"] = { 

201 name: deserialize_callable(filter_func) if filter_func else None 

202 for name, filter_func in custom_filters.items() 

203 } 

204 return default_from_dict(cls, data)