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:
@silvanocerza: That's expected. 

Because of the changes to fix the security issue we had to make this breaking change. The `ConditionalRouter` can only return the following Python literal types now: strings, bytes, numbers, tuples, lists, dicts, sets, booleans, `None` and `Ellipsis`. 
Any other will be returned as a string.

I would suggest creating a custom Component to cover those cases where the `ConditionalRouter` can't be used as before.

Something like this should do the trick:

```python
from typing import List

from haystack import component
from haystack.dataclasses import ChatMessage

@component
class Decision:
    @component.output_types(short_streams_message=ChatMessage, long_streams_message=ChatMessage)
    def run(self, streams: List[str], message: ChatMessage):
        if len(streams) > 2:
            return {"long_streams_message": message}
        return {"short_streams_message": message}
```@gustavo-has-stone: I see. So the parameter `output_type` will not work anymore?

We did not expect to see a breaking change with a [patch version change](https://docs.haystack.deepset.ai/docs/breaking-change-policy). Should we worry about those updates too?@silvanocerza: `output_type` has never been necessary to convert the type, it's only used to set the correct output type for the Component. Even before this change you could put any type in there but it wouldn't guarantee that the Component will return that type when calling `run`. That's only necessary when calling `Pipeline.connect()` to validate the connection.

We decided to release this security change as a patch release even if it's a breaking change as it's a pretty critical, it could lead to remote code execution.@mathislucka: Perhaps we could add an option to the `ConditionalRouter` so that it can route inputs to the selected route's output without rendering a jinja2 template?

I also had this case, where I wanted to route a list of Documents depending on a reply generated from the LLM. I don't have to go through jinja to route these documents because I don't want to change anything about them.

So something like:
```
routes = [
{"condition": "{{...}}", output: "documents", output_name: "a_documents", "output_type": List[Document], "use_literal_output": True}
]
router = ConditionalRouter(routes=routes)

router.run(documents=[Document(content="hey")])

# returns the document passed in unchanged
```


https://github.com/deepset-ai/haystack/blob/e17d0c41926849975166e70ec46411f402c8098e/haystack/components/routers/conditional_router.py#L206

This part would need to be changed to something like:
```
if route.get("use_literal_output"):
    output = kwargs[route["output"]]
else:
    ...
```

Would that solve your use case too @gustavo-has-stone ?
@gustavo-has-stone: In our scenario, we were routing a variable of a custom type (similar to `ChatMessage`) based on the value of an attribute of this variable. I did not fully understand your example. Could you provide more details?@alexsniffin: Hi, I ran into this problem too when upgrading to 2.3.1. I use a custom `dataclass` type for my pipeline which isn't working as expected anymore. Is there a suggested fix for this without creating a custom router?@silvanocerza: @gustavo-has-stone @alexsniffin

We discussed a bit internally and decided to give the possibility to enable the previous behaviour by creating `ConditionalRouter` and `OutputAdapter` with `unsafe=True`.

I just created a PR to bring that change in and we'll release it on the upcoming `2.4.0` some time next week. For the time being, and if you consider it safe, you can still rely on the previous behaviour by using `2.3.0` instead of the latest `2.3.1`.@gustavo-has-stone: I don't think that using an unsafe version is a good solution for us. We will try to find an alternative solution for our use case@dmvieira: Hey @silvanocerza Idk if it`s a good solution because you are publishing a version that the only way to use previews behavior is forcing unsafe. It forces people to open security breaches when need a feature.

An approach that could help us is using Haystack `ChatMessage` as Router output. Do you know if is it possible?@silvanocerza: > I don't think that using an unsafe version is a good solution for us. We will try to find an alternative solution for our use case

It's unsafe only if the Jinja templates used are users inputs too. If you know the templates it's completely safe to use.

> Hey @silvanocerza Idk if it`s a good solution because you are publishing a version that the only way to use previews behavior is forcing unsafe. It forces people to open security breaches when need a feature.
> 
> An approach that could help us is using Haystack `ChatMessage` as Router output. Do you know if is it possible?

Setting `unsafe` to `True` doesn't automatically open you to security breaches. It enables a behaviour that is considered unsafe in certain circumstances, that is when `ConditionalRouter` or `OutputAdapter` use templates that can be considered users input. 

If you know what you're doing and know you're not using users inputs as Jinja template that is completely fine to use.@gustavo-has-stone: Let me see if I understand. If we try to route based on some controlled condition (e.g., message length), we will be fine. However, if we try to use the content of a message for conditional routing, we could be in trouble. Is that right?@silvanocerza: > Let me see if I understand. If we try to route based on some controlled condition (e.g., message length), we will be fine. However, if we try to use the content of a message for conditional routing, we could be in trouble. Is that right?

No. You are vulnerable if you let your users write their own templates from scratch.

A simple example:

```python
from typing import List
from haystack.components.routers import ConditionalRouter

user_input = input("Write a Jinja template:")

routes = [
    {
        "condition": user_input,
        "output": "{{streams}}",
        "output_name": "enough_streams",
        "output_type": List[int],
    },
    {
        "condition": "{{streams|length <= 2}}",
        "output": "{{streams}}",
        "output_name": "insufficient_streams",
        "output_type": List[int],
    },
]
router = ConditionalRouter(routes)
...
router.run(whatever_input)
```@dmvieira: Thank you @silvanocerza . It sounds like a plan! 😄 

Solution patch:
diff --git a/haystack/components/converters/output_adapter.py b/haystack/components/converters/output_adapter.py
index 64cac7996a..50ffd53919 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 @@ def __init__(self, template: str, output_type: TypeAlias, custom_filters: Option
             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 @@ def run(self, **kwargs):
             # 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 @@ def to_dict(self) -> Dict[str, Any]:
         """
         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 @@ def from_dict(cls, data: Dict[str, Any]) -> "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 a72bb4fc08..ccac555d3d 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 @@ def __init__(self, routes: List[Dict], custom_filters: Optional[Dict[str, Callab
             - `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 @@ def to_dict(self) -> Dict[str, Any]:
             # 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 @@ def run(self, **kwargs):
             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 @@ def run(self, **kwargs):
                     # 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 @@ def _validate_routes(self, routes: List[Dict]):
 
         :param routes: A list of routes.
         """
-        env = NativeEnvironment()
         for route in routes:
             try:
                 keys = set(route.keys())
@@ -239,10 +253,10 @@ def _validate_routes(self, routes: List[Dict]):
                     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 0000000000..d6242432e3
--- /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 25b5eb7e07..547ce433e6 100644
--- a/test/components/converters/test_output_adapter.py
+++ b/test/components/converters/test_output_adapter.py
@@ -8,6 +8,7 @@
 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 @@ def test_output_adapter_from_dict_custom_filters_none(self):
                     "template": "{{ documents[0].content}}",
                     "output_type": "str",
                     "custom_filters": None,
+                    "unsafe": False,
                 },
             }
         )
@@ -157,6 +159,7 @@ def test_output_adapter_from_dict_custom_filters_none(self):
         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 @@ def run(self):
         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 b324cf8067..47873e3d3f 100644
--- a/test/components/routers/test_conditional_router.py
+++ b/test/components/routers/test_conditional_router.py
@@ -307,3 +307,24 @@ def test_sede_with_custom_filter(self):
         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}