19852408366

Committed 02 Dec 2025 08:42AM UTC coverage: 92.192% (+0.03%) from 92.16%

Build # 19852408366

Build Type

Pull #10093

github

Committed by

web-flow

Commit Message

Merge cfee3ac4b into 050c25cee

Pull Request Pull Request #10093: fix: Fix variable extraction in Jinja2 templates returning set variables

Run Details

14063 of 15254 relevant lines covered (92.19%)

0.92 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

98.55

haystack/components/converters/output_adapter.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import ast
import contextlib
from typing import Any, Callable, Optional

import jinja2.runtime
from jinja2 import Environment, TemplateSyntaxError, meta
from jinja2.nativetypes import NativeEnvironment
from jinja2.sandbox import SandboxedEnvironment
from typing_extensions import TypeAlias

from haystack import component, default_from_dict, default_to_dict, logging
from haystack.utils import deserialize_callable, deserialize_type, serialize_callable, serialize_type
from haystack.utils.jinja2_extensions import _collect_assigned_variables

logger = logging.getLogger(__name__)


class OutputAdaptationException(Exception):
    """Exception raised when there is an error during output adaptation."""


@component
class OutputAdapter:
    """
    Adapts output of a Component using Jinja templates.

    Usage example:
    ```python
    from haystack import Document
    from haystack.components.converters import OutputAdapter

    adapter = OutputAdapter(template="{{ documents[0].content }}", output_type=str)
    documents = [Document(content="Test content"]
    result = adapter.run(documents=documents)

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

    def __init__(
        self,
        template: str,
        output_type: TypeAlias,
        custom_filters: Optional[dict[str, Callable]] = None,
        unsafe: bool = False,
    ) -> None:
        """
        Create an OutputAdapter component.

        :param template:
            A Jinja template that defines how to adapt the input data.
            The variables in the template define the input of this instance.
            e.g.
            With this template:
            ```
            {{ documents[0].content }}
            ```
            The Component input will be `documents`.
        :param output_type:
            The type of output this instance will return.
        :param custom_filters:
            A dictionary of custom Jinja filters used in the template.
        :param unsafe:
            Enable execution of arbitrary code in the Jinja template.
            This should only be used if you trust the source of the template as it can be lead to remote code execution.
        """
        self.custom_filters = {**(custom_filters or {})}
        input_types: set[str] = set()

        self._unsafe = unsafe

        if self._unsafe:
            msg = (
                "Unsafe mode is enabled. This allows execution of arbitrary code in the Jinja template. "
                "Use this only if you trust the source of the template."
            )
            logger.warning(msg)
        self._env = (
            NativeEnvironment() if self._unsafe else SandboxedEnvironment(undefined=jinja2.runtime.StrictUndefined)
        )

        try:
            self._env.parse(template)  # Validate template syntax
            self.template = template
        except TemplateSyntaxError as e:
            raise ValueError(f"Invalid Jinja template '{template}': {e}") from e

        for name, filter_func in self.custom_filters.items():
            self._env.filters[name] = filter_func

        # b) extract variables in the template
        route_input_names = self._extract_variables(self._env)
        input_types.update(route_input_names)

        # the env is not needed, discarded automatically
        component.set_input_types(self, **dict.fromkeys(input_types, Any))
        component.set_output_types(self, **{"output": output_type})
        self.output_type = output_type

    def run(self, **kwargs):
        """
        Renders the Jinja template with the provided inputs.

        :param kwargs:
            Must contain all variables used in the `template` string.
        :returns:
            A dictionary with the following keys:
            - `output`: Rendered Jinja template.

        :raises OutputAdaptationException: If template rendering fails.
        """
        # check if kwargs are empty
        if not kwargs:
            raise ValueError("No input data provided for output adaptation")
        for name, filter_func in self.custom_filters.items():
            self._env.filters[name] = filter_func
        adapted_outputs = {}
        try:
            adapted_output_template = self._env.from_string(self.template)
            output_result = adapted_output_template.render(**kwargs)
            if isinstance(output_result, jinja2.runtime.Undefined):
                raise OutputAdaptationException(f"Undefined variable in the template {self.template}; kwargs: {kwargs}")

            # We suppress the exception in case the output is already a string, otherwise
            # we try to evaluate it and would fail.
            # This must be done cause the output could be different literal structures.
            # This doesn't support any user types.
            with contextlib.suppress(Exception):
                if not self._unsafe:
                    output_result = ast.literal_eval(output_result)

            adapted_outputs["output"] = output_result
        except Exception as e:
            raise OutputAdaptationException(f"Error adapting {self.template} with {kwargs}: {e}") from e
        return adapted_outputs

    def to_dict(self) -> dict[str, Any]:
        """
        Serializes the component to a dictionary.

        :returns:
            Dictionary with serialized data.
        """
        se_filters = {name: serialize_callable(filter_func) for name, filter_func in self.custom_filters.items()}
        return default_to_dict(
            self,
            template=self.template,
            output_type=serialize_type(self.output_type),
            custom_filters=se_filters,
            unsafe=self._unsafe,
        )

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "OutputAdapter":
        """
        Deserializes the component from a dictionary.

        :param data:
            The dictionary to deserialize from.
        :returns:
            The deserialized component.
        """
        init_params = data.get("init_parameters", {})
        init_params["output_type"] = deserialize_type(init_params["output_type"])

        custom_filters = init_params.get("custom_filters", {})
        if custom_filters:
            init_params["custom_filters"] = {
                name: deserialize_callable(filter_func) if filter_func else None
                for name, filter_func in custom_filters.items()
            }
        return default_from_dict(cls, data)

    def _extract_variables(self, env: Environment) -> set[str]:
        """
        Extracts all variables from a list of Jinja template strings.

        :param env: A Jinja environment.
        :returns: A set of variable names extracted from the template strings.
        """
        jinja2_ast = env.parse(self.template)
        template_variables = meta.find_undeclared_variables(jinja2_ast)
        assigned_variables = _collect_assigned_variables(jinja2_ast)
        return template_variables - assigned_variables

1	# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2	#
3	# SPDX-License-Identifier: Apache-2.0
4
5	import ast	1✔
6	import contextlib	1✔
7	from typing import Any, Callable, Optional	1✔
8
9	import jinja2.runtime	1✔
10	from jinja2 import Environment, TemplateSyntaxError, meta	1✔
11	from jinja2.nativetypes import NativeEnvironment	1✔
12	from jinja2.sandbox import SandboxedEnvironment	1✔
13	from typing_extensions import TypeAlias	1✔
14
15	from haystack import component, default_from_dict, default_to_dict, logging	1✔
16	from haystack.utils import deserialize_callable, deserialize_type, serialize_callable, serialize_type	1✔
17	from haystack.utils.jinja2_extensions import _collect_assigned_variables	1✔
18
19	logger = logging.getLogger(__name__)	1✔
20
21
22	class OutputAdaptationException(Exception):	1✔
23	"""Exception raised when there is an error during output adaptation."""
24
25
26	@component	1✔
27	class OutputAdapter:	1✔
28	"""
29	Adapts output of a Component using Jinja templates.
30
31	Usage example:
32	```python
33	from haystack import Document
34	from haystack.components.converters import OutputAdapter
35
36	adapter = OutputAdapter(template="{{ documents[0].content }}", output_type=str)
37	documents = [Document(content="Test content"]
38	result = adapter.run(documents=documents)
39
40	assert result["output"] == "Test content"
41	```
42	"""
43
44	def __init__(	1✔
45	self,
46	template: str,
47	output_type: TypeAlias,
48	custom_filters: Optional[dict[str, Callable]] = None,
49	unsafe: bool = False,
50	) -> None:
51	"""
52	Create an OutputAdapter component.
53
54	:param template:
55	A Jinja template that defines how to adapt the input data.
56	The variables in the template define the input of this instance.
57	e.g.
58	With this template:
59	```
60	{{ documents[0].content }}
61	```
62	The Component input will be `documents`.
63	:param output_type:
64	The type of output this instance will return.
65	:param custom_filters:
66	A dictionary of custom Jinja filters used in the template.
67	:param unsafe:
68	Enable execution of arbitrary code in the Jinja template.
69	This should only be used if you trust the source of the template as it can be lead to remote code execution.
70	"""
71	self.custom_filters = {**(custom_filters or {})}	1✔
72	input_types: set[str] = set()	1✔
73
74	self._unsafe = unsafe	1✔
75
76	if self._unsafe:	1✔
77	msg = (	1✔
78	"Unsafe mode is enabled. This allows execution of arbitrary code in the Jinja template. "
79	"Use this only if you trust the source of the template."
80	)
81	logger.warning(msg)	1✔
82	self._env = (	1✔
83	NativeEnvironment() if self._unsafe else SandboxedEnvironment(undefined=jinja2.runtime.StrictUndefined)
84	)
85
86	try:	1✔
87	self._env.parse(template) # Validate template syntax	1✔
88	self.template = template	1✔
89	except TemplateSyntaxError as e:	1✔
90	raise ValueError(f"Invalid Jinja template '{template}': {e}") from e	1✔
91
92	for name, filter_func in self.custom_filters.items():	1✔
93	self._env.filters[name] = filter_func	1✔
94
95	# b) extract variables in the template
96	route_input_names = self._extract_variables(self._env)	1✔
97	input_types.update(route_input_names)	1✔
98
99	# the env is not needed, discarded automatically
100	component.set_input_types(self, **dict.fromkeys(input_types, Any))	1✔
101	component.set_output_types(self, **{"output": output_type})	1✔
102	self.output_type = output_type	1✔
103
104	def run(self, **kwargs):	1✔
105	"""
106	Renders the Jinja template with the provided inputs.
107
108	:param kwargs:
109	Must contain all variables used in the `template` string.
110	:returns:
111	A dictionary with the following keys:
112	- `output`: Rendered Jinja template.
113
114	:raises OutputAdaptationException: If template rendering fails.
115	"""
116	# check if kwargs are empty
117	if not kwargs:	1✔
118	raise ValueError("No input data provided for output adaptation")	1✔
119	for name, filter_func in self.custom_filters.items():	1✔
120	self._env.filters[name] = filter_func	1✔
121	adapted_outputs = {}	1✔
122	try:	1✔
123	adapted_output_template = self._env.from_string(self.template)	1✔
124	output_result = adapted_output_template.render(**kwargs)	1✔
125	if isinstance(output_result, jinja2.runtime.Undefined):	1✔
126	raise OutputAdaptationException(f"Undefined variable in the template {self.template}; kwargs: {kwargs}")	×
127
128	# We suppress the exception in case the output is already a string, otherwise
129	# we try to evaluate it and would fail.
130	# This must be done cause the output could be different literal structures.
131	# This doesn't support any user types.
132	with contextlib.suppress(Exception):	1✔
133	if not self._unsafe:	1✔
134	output_result = ast.literal_eval(output_result)	1✔
135
136	adapted_outputs["output"] = output_result	1✔
137	except Exception as e:	1✔
138	raise OutputAdaptationException(f"Error adapting {self.template} with {kwargs}: {e}") from e	1✔
139	return adapted_outputs	1✔
140
141	def to_dict(self) -> dict[str, Any]:	1✔
142	"""
143	Serializes the component to a dictionary.
144
145	:returns:
146	Dictionary with serialized data.
147	"""
148	se_filters = {name: serialize_callable(filter_func) for name, filter_func in self.custom_filters.items()}	1✔
149	return default_to_dict(	1✔
150	self,
151	template=self.template,
152	output_type=serialize_type(self.output_type),
153	custom_filters=se_filters,
154	unsafe=self._unsafe,
155	)
156
157	@classmethod	1✔
158	def from_dict(cls, data: dict[str, Any]) -> "OutputAdapter":	1✔
159	"""
160	Deserializes the component from a dictionary.
161
162	:param data:
163	The dictionary to deserialize from.
164	:returns:
165	The deserialized component.
166	"""
167	init_params = data.get("init_parameters", {})	1✔
168	init_params["output_type"] = deserialize_type(init_params["output_type"])	1✔
169
170	custom_filters = init_params.get("custom_filters", {})	1✔
171	if custom_filters:	1✔
172	init_params["custom_filters"] = {	1✔
173	name: deserialize_callable(filter_func) if filter_func else None
174	for name, filter_func in custom_filters.items()
175	}
176	return default_from_dict(cls, data)	1✔
177
178	def _extract_variables(self, env: Environment) -> set[str]:	1✔
179	"""
180	Extracts all variables from a list of Jinja template strings.
181
182	:param env: A Jinja environment.
183	:returns: A set of variable names extracted from the template strings.
184	"""
185	jinja2_ast = env.parse(self.template)	1✔
186	template_variables = meta.find_undeclared_variables(jinja2_ast)	1✔
187	assigned_variables = _collect_assigned_variables(jinja2_ast)	1✔
188	return template_variables - assigned_variables	1✔

deepset-ai / haystack / 19852408366

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous