Repository: deepset-ai/haystack

Issue:
ConditionalRouter output type is not working anymore
**Describe the bug**
We tried to update the Haystack version to 2.3.1 due to the [security update](https://github.com/deepset-ai/haystack/releases/tag/v2.3.1) and our code stopped working. We identified that the `ConditionalRouter` is returning a string instead of the type assigned in `output_type`

**Error message**
We are getting errors related to the fact that a string is now being returned. Here is an example of the error when running the provided code:

`AttributeError: 'str' object has no attribute 'content'`

**Expected behavior**
We expected that the `output_type` provided in the route would be returned

**Additional context**
The problem is related to version 2.3.1. Our provided example works on 2.3.0

**To Reproduce**
The following code works on haystack 2.3.0 but it does not work on haystack 2.3.1:

```
from haystack.components.routers import ConditionalRouter
from haystack.dataclasses import ChatMessage, ChatRole

routes = [
    {
        "condition": "{{streams|length > 2}}",
        "output": "{{message}}",
        "output_name": "long_streams_message",
        "output_type": ChatMessage,
    },
    {
        "condition": "{{streams|length <= 2}}",
        "output": "{{message}}",
        "output_name": "short_streams_message",
        "output_type": ChatMessage,
    },
]

message = ChatMessage(content='Hi', role=ChatRole('user'), name='', meta={})
print(message.content)

router = ConditionalRouter(routes)

kwargs = {"streams": [1, 2, 3], "message": message}
result = router.run(**kwargs)
print(result['long_streams_message'].content)
```


**FAQ Check**
- [x] Have you had a look at [our new FAQ page](https://docs.haystack.deepset.ai/docs/faq)?

**System:**
 - OS: Ubuntu 22.04 LTS
 - GPU/CPU: CPU
 - Haystack version (commit or version number): 2.3.1


Discussion:
coveralls: ## Pull Request Test Coverage Report for [Build 10304167842](https://coveralls.io/builds/69120385)


### Details

* **0** of **0**   changed or added relevant lines in **0** files are covered.
* **3** unchanged lines in **2** files lost coverage.
* Overall coverage increased (+**0.02%**) to **90.189%**

---


|  Files with Coverage Reduction | New Missed Lines | % |
| :-----|--------------|--: |
| [components/converters/output_adapter.py](https://coveralls.io/builds/69120385/source?filename=components%2Fconverters%2Foutput_adapter.py#L124) | 1 | 98.51% |
| [components/routers/conditional_router.py](https://coveralls.io/builds/69120385/source?filename=components%2Frouters%2Fconditional_router.py#L233) | 2 | 97.75% |
<!-- | **Total:** | **3** |  | -->

|  Totals | [![Coverage Status](https://coveralls.io/builds/69120385/badge)](https://coveralls.io/builds/69120385) |
| :-- | --: |
| Change from base [Build 10301663157](https://coveralls.io/builds/69116910): |  0.02% |
| Covered Lines: | 6931 |
| Relevant Lines: | 7685 |

---
##### 💛  - [Coveralls](https://coveralls.io)


Patch:
diff --git a/haystack/components/converters/output_adapter.py b/haystack/components/converters/output_adapter.py
index 64cac799..50ffd539 100644
--- a/haystack/components/converters/output_adapter.py
+++ b/haystack/components/converters/output_adapter.py
@@ -5,9 +5,11 @@
 import ast
 import contextlib
 from typing import Any, Callable, Dict, Optional, Set
+from warnings import warn
 
 import jinja2.runtime
-from jinja2 import TemplateSyntaxError, meta
+from jinja2 import Environment, TemplateSyntaxError, meta
+from jinja2.nativetypes import NativeEnvironment
 from jinja2.sandbox import SandboxedEnvironment
 from typing_extensions import TypeAlias
 
@@ -37,7 +39,13 @@ class OutputAdapter:
     ```
     """
 
-    def __init__(self, template: str, output_type: TypeAlias, custom_filters: Optional[Dict[str, Callable]] = None):
+    def __init__(
+        self,
+        template: str,
+        output_type: TypeAlias,
+        custom_filters: Optional[Dict[str, Callable]] = None,
+        unsafe: bool = False,
+    ):
         """
         Create an OutputAdapter component.
 
@@ -54,13 +62,25 @@ class OutputAdapter:
             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()
 
-        # Create a Jinja native environment, we need it to:
-        # a) add custom filters to the environment for filter compilation stage
-        self._env = SandboxedEnvironment(undefined=jinja2.runtime.StrictUndefined)
+        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."
+            )
+            warn(msg)
+        self._env = (
+            NativeEnvironment() if self._unsafe else SandboxedEnvironment(undefined=jinja2.runtime.StrictUndefined)
+        )
+
         try:
             self._env.parse(template)  # Validate template syntax
             self.template = template
@@ -108,7 +128,8 @@ class OutputAdapter:
             # This must be done cause the output could be different literal structures.
             # This doesn't support any user types.
             with contextlib.suppress(Exception):
-                output_result = ast.literal_eval(output_result)
+                if not self._unsafe:
+                    output_result = ast.literal_eval(output_result)
 
             adapted_outputs["output"] = output_result
         except Exception as e:
@@ -124,7 +145,11 @@ class OutputAdapter:
         """
         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
+            self,
+            template=self.template,
+            output_type=serialize_type(self.output_type),
+            custom_filters=se_filters,
+            unsafe=self._unsafe,
         )
 
     @classmethod
@@ -148,11 +173,11 @@ class OutputAdapter:
             }
         return default_from_dict(cls, data)
 
-    def _extract_variables(self, env: SandboxedEnvironment) -> Set[str]:
+    def _extract_variables(self, env: Environment) -> Set[str]:
         """
         Extracts all variables from a list of Jinja template strings.
 
-        :param env: A Jinja native environment.
+        :param env: A Jinja environment.
         :return: A set of variable names extracted from the template strings.
         """
         ast = env.parse(self.template)
diff --git a/haystack/components/routers/conditional_router.py b/haystack/components/routers/conditional_router.py
index a72bb4fc..ccac555d 100644
--- a/haystack/components/routers/conditional_router.py
+++ b/haystack/components/routers/conditional_router.py
@@ -5,6 +5,7 @@
 import ast
 import contextlib
 from typing import Any, Callable, Dict, List, Optional, Set
+from warnings import warn
 
 from jinja2 import Environment, TemplateSyntaxError, meta
 from jinja2.nativetypes import NativeEnvironment
@@ -106,7 +107,7 @@ class ConditionalRouter:
     ```
     """
 
-    def __init__(self, routes: List[Dict], custom_filters: Optional[Dict[str, Callable]] = None):
+    def __init__(self, routes: List[Dict], custom_filters: Optional[Dict[str, Callable]] = None, unsafe: bool = False):
         """
         Initializes the `ConditionalRouter` with a list of routes detailing the conditions for routing.
 
@@ -123,15 +124,26 @@ class ConditionalRouter:
             - `my_filter_fcn` is a callable that takes `my_var:str` and returns `my_var[:3]`.
               `{{ my_var|my_filter }}` can then be used inside a route condition expression:
                 `"condition": "{{ my_var|my_filter == 'foo' }}"`.
+        :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._validate_routes(routes)
         self.routes: List[dict] = routes
         self.custom_filters = custom_filters or {}
+        self._unsafe = unsafe
+
+        # Create a Jinja environment to inspect variables in the condition templates
+        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."
+            )
+            warn(msg)
 
-        # Create a Jinja native environment to inspect variables in the condition templates
-        self._env = SandboxedEnvironment()
+        self._env = NativeEnvironment() if self._unsafe else SandboxedEnvironment()
         self._env.filters.update(self.custom_filters)
 
+        self._validate_routes(routes)
         # Inspect the routes to determine input and output types.
         input_types: Set[str] = set()  # let's just store the name, type will always be Any
         output_types: Dict[str, str] = {}
@@ -158,7 +170,7 @@ class ConditionalRouter:
             # output_type needs to be serialized to a string
             route["output_type"] = serialize_type(route["output_type"])
         se_filters = {name: serialize_callable(filter_func) for name, filter_func in self.custom_filters.items()}
-        return default_to_dict(self, routes=self.routes, custom_filters=se_filters)
+        return default_to_dict(self, routes=self.routes, custom_filters=se_filters, unsafe=self._unsafe)
 
     @classmethod
     def from_dict(cls, data: Dict[str, Any]) -> "ConditionalRouter":
@@ -202,7 +214,9 @@ class ConditionalRouter:
             try:
                 t = self._env.from_string(route["condition"])
                 rendered = t.render(**kwargs)
-                if ast.literal_eval(rendered):
+                if not self._unsafe:
+                    rendered = ast.literal_eval(rendered)
+                if rendered:
                     # We now evaluate the `output` expression to determine the route output
                     t_output = self._env.from_string(route["output"])
                     output = t_output.render(**kwargs)
@@ -211,7 +225,8 @@ class ConditionalRouter:
                     # This must be done cause the output could be different literal structures.
                     # This doesn't support any user types.
                     with contextlib.suppress(Exception):
-                        output = ast.literal_eval(output)
+                        if not self._unsafe:
+                            output = ast.literal_eval(output)
                     # and return the output as a dictionary under the output_name key
                     return {route["output_name"]: output}
             except Exception as e:
@@ -225,7 +240,6 @@ class ConditionalRouter:
 
         :param routes: A list of routes.
         """
-        env = NativeEnvironment()
         for route in routes:
             try:
                 keys = set(route.keys())
@@ -239,10 +253,10 @@ class ConditionalRouter:
                     f"Route must contain 'condition', 'output', 'output_type' and 'output_name' fields: {route}"
                 )
             for field in ["condition", "output"]:
-                if not self._validate_template(env, route[field]):
+                if not self._validate_template(self._env, route[field]):
                     raise ValueError(f"Invalid template for field '{field}': {route[field]}")
 
-    def _extract_variables(self, env: SandboxedEnvironment, templates: List[str]) -> Set[str]:
+    def _extract_variables(self, env: Environment, templates: List[str]) -> Set[str]:
         """
         Extracts all variables from a list of Jinja template strings.
 
diff --git a/releasenotes/notes/unsafe-behaviour-e8b41d957113e0c3.yaml b/releasenotes/notes/unsafe-behaviour-e8b41d957113e0c3.yaml
new file mode 100644
index 00000000..d6242432
--- /dev/null
+++ b/releasenotes/notes/unsafe-behaviour-e8b41d957113e0c3.yaml
@@ -0,0 +1,8 @@
+---
+features:
+  - |
+    Add `unsafe` argument to enable behaviour that could lead to remote code execution in `ConditionalRouter` and `OutputAdapter`.
+    By default unsafe behaviour is not enabled, the user must set it explicitly to `True`.
+    This means that user types like `ChatMessage`, `Document`, and `Answer` can be used as output types when `unsafe` is `True`.
+    We recommend using `unsafe` behaviour only when the Jinja templates source is trusted.
+    For more info see the documentation for [`ConditionalRouter`](https://docs.haystack.deepset.ai/docs/conditionalrouter#unsafe-behaviour) and [`OutputAdapter`](https://docs.haystack.deepset.ai/docs/outputadapter#unsafe-behaviour)


Test patch:
diff --git a/test/components/converters/test_output_adapter.py b/test/components/converters/test_output_adapter.py
index 25b5eb7e..547ce433 100644
--- a/test/components/converters/test_output_adapter.py
+++ b/test/components/converters/test_output_adapter.py
@@ -8,6 +8,7 @@ import json
 import pytest
 
 from haystack import Pipeline, component
+from haystack.dataclasses import Document
 from haystack.components.converters import OutputAdapter
 from haystack.components.converters.output_adapter import OutputAdaptationException
 
@@ -150,6 +151,7 @@ class TestOutputAdapter:
                     "template": "{{ documents[0].content}}",
                     "output_type": "str",
                     "custom_filters": None,
+                    "unsafe": False,
                 },
             }
         )
@@ -157,6 +159,7 @@ class TestOutputAdapter:
         assert component.template == "{{ documents[0].content}}"
         assert component.output_type == str
         assert component.custom_filters == {}
+        assert not component._unsafe
 
     def test_output_adapter_in_pipeline(self):
         @component
@@ -179,3 +182,13 @@ class TestOutputAdapter:
         result = pipe.run(data={})
         assert result
         assert result["output_adapter"]["output"] == {"framework": "Haystack"}
+
+    def test_unsafe(self):
+        adapter = OutputAdapter(template="{{ documents[0] }}", output_type=Document, unsafe=True)
+        documents = [
+            Document(content="Test document"),
+            Document(content="Another test document"),
+            Document(content="Yet another test document"),
+        ]
+        res = adapter.run(documents=documents)
+        assert res["output"] == documents[0]
diff --git a/test/components/routers/test_conditional_router.py b/test/components/routers/test_conditional_router.py
index b324cf80..47873e3d 100644
--- a/test/components/routers/test_conditional_router.py
+++ b/test/components/routers/test_conditional_router.py
@@ -307,3 +307,24 @@ class TestRouter:
         assert deserialized_router.custom_filters == router.custom_filters
         assert deserialized_router.custom_filters["custom_filter_to_sede"]("123-456-789") == 123
         assert result == deserialized_router.run(**kwargs)
+
+    def test_unsafe(self):
+        routes = [
+            {
+                "condition": "{{streams|length < 2}}",
+                "output": "{{message}}",
+                "output_type": ChatMessage,
+                "output_name": "message",
+            },
+            {
+                "condition": "{{streams|length >= 2}}",
+                "output": "{{streams}}",
+                "output_type": List[int],
+                "output_name": "streams",
+            },
+        ]
+        router = ConditionalRouter(routes, unsafe=True)
+        streams = [1]
+        message = ChatMessage.from_user(content="This is a message")
+        res = router.run(streams=streams, message=message)
+        assert res == {"message": message}