Extend Jinja Environment to allow Haystack data classes in ConditionalRouter

haystack/components/routers/conditional_router.py

@@ -8,14 +8,17 @@
 from warnings import warn
 
 from jinja2 import Environment, TemplateSyntaxError, meta
-from jinja2.nativetypes import NativeEnvironment
+from jinja2.nativetypes import NativeEnvironment, NativeTemplate, Template
 from jinja2.sandbox import SandboxedEnvironment
 
 from haystack import component, default_from_dict, default_to_dict, logging
+from haystack.dataclasses import Answer, ByteStream, ChatMessage, Document, SparseEmbedding, StreamingChunk
 from haystack.utils import deserialize_callable, deserialize_type, serialize_callable, serialize_type
 
 logger = logging.getLogger(__name__)
 
+haystack_dataclass_types = (ByteStream, ChatMessage, Document, Answer, SparseEmbedding, StreamingChunk)
+
 
 class NoRouteSelectedException(Exception):
     """Exception raised when no route is selected in ConditionalRouter."""
@@ -25,6 +28,41 @@ class RouteConditionException(Exception):
     """Exception raised when there is an error parsing or evaluating the condition expression in ConditionalRouter."""
 
 
+class NativeSandboxedTemplate(NativeTemplate, Template):
+    """
+    A template class that returns native Python objects and also respects the sandbox security checks.
+    """
+
+    pass
+
+
+class NativeSandboxedEnvironment(SandboxedEnvironment, NativeEnvironment):
+    """
+    An environment that combines sandbox restrictions with native rendering.
+    """
+
+    # We tell the environment to use our custom template class by default.
+
+    def from_string(self, source):
+        """
+        Override from_string to ensure the sandbox logic + native logic are used together.
+        """
+        template_class = NativeSandboxedTemplate
+
+        return SandboxedEnvironment.from_string(self, source, template_class=template_class)
+
+    def is_safe_attribute(self, obj, attr="", value=""):
+        """
+        Whitelist Haystack dataclasses so the sandbox won't block them.
+        """
+
+        if not isinstance(obj, haystack_dataclass_types):
+            return False
+
+        # Otherwise, fallback to the default sandbox behavior
+        return SandboxedEnvironment.is_safe_attribute(self, obj, attr, value)
+
+
 @component
 class ConditionalRouter:
     """
@@ -195,6 +233,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
             warn(msg)
 
         self._env = NativeEnvironment() if self._unsafe else SandboxedEnvironment()
+        self._custom_env = NativeSandboxedEnvironment()
         self._env.filters.update(self.custom_filters)
 
         self._validate_routes(routes)
@@ -309,15 +348,24 @@ def run(self, **kwargs):
                 if not rendered:
                     continue
                 # We now evaluate the `output` expression to determine the route output
-                t_output = self._env.from_string(route["output"])
+                t_output = self._custom_env.from_string(route["output"])
                 output = t_output.render(**kwargs)
+
+                # Check if output is a list/sequence and validate accordingly
+                if isinstance(output, (list, tuple)):
+                    if all(self._custom_env.is_safe_attribute(item) for item in output):
+                        pass
+                elif self._custom_env.is_safe_attribute(output):
+                    pass
+
                 # 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 = ast.literal_eval(output)
+                else:
+                    with contextlib.suppress(Exception):
+                        if not self._unsafe and isinstance(output, str):
+                            output = ast.literal_eval(output)
             except Exception as e:
                 msg = f"Error evaluating condition for route '{route}': {e}"
                 raise RouteConditionException(msg) from e