fix: Fix serialization and deserialization of ConditionalRouter with multiple outputs (#9490)

* Fix sede of ConditionalRouter with multiple outputs

* Add reno

Author: sjrl

haystack/components/routers/conditional_router.py

@@ -4,7 +4,7 @@
 
 import ast
 import contextlib
-from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Set, Union, get_args, get_origin
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Set, TypedDict, Union, get_args, get_origin
 
 from jinja2 import Environment, TemplateSyntaxError, meta
 from jinja2.nativetypes import NativeEnvironment
@@ -24,6 +24,13 @@ class RouteConditionException(Exception):
     """Exception raised when there is an error parsing or evaluating the condition expression in ConditionalRouter."""
 
 
+class Route(TypedDict):
+    condition: str
+    output: Union[str, List[str]]
+    output_name: Union[str, List[str]]
+    output_type: Union[type, List[type]]
+
+
 @component
 class ConditionalRouter:
     """
@@ -108,7 +115,7 @@ class ConditionalRouter:
 
     def __init__(  # pylint: disable=too-many-positional-arguments
         self,
-        routes: List[Dict],
+        routes: List[Route],
         custom_filters: Optional[Dict[str, Callable]] = None,
         unsafe: bool = False,
         validate_output_type: bool = False,
@@ -179,7 +186,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
             - Some variables are only needed for specific routing conditions
             - You're building flexible pipelines where not all inputs are guaranteed to be present
         """
-        self.routes: List[dict] = routes
+        self.routes: List[Route] = routes
         self.custom_filters = custom_filters or {}
         self._unsafe = unsafe
         self._validate_output_type = validate_output_type
@@ -199,7 +206,7 @@ def __init__(  # pylint: disable=too-many-positional-arguments
         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] = {}
+        output_types: Dict[str, Union[type, List[type]]] = {}
 
         for route in routes:
             # extract inputs
@@ -248,8 +255,12 @@ def to_dict(self) -> Dict[str, Any]:
         """
         serialized_routes = []
         for route in self.routes:
-            # output_type needs to be serialized to a string
-            serialized_routes.append({**route, "output_type": serialize_type(route["output_type"])})
+            serialized_output_type = (
+                [serialize_type(t) for t in route["output_type"]]
+                if isinstance(route["output_type"], list)
+                else serialize_type(route["output_type"])
+            )
+            serialized_routes.append({**route, "output_type": serialized_output_type})
         se_filters = {name: serialize_callable(filter_func) for name, filter_func in self.custom_filters.items()}
         return default_to_dict(
             self,
@@ -274,7 +285,10 @@ def from_dict(cls, data: Dict[str, Any]) -> "ConditionalRouter":
         routes = init_params.get("routes")
         for route in routes:
             # output_type needs to be deserialized from a string to a type
-            route["output_type"] = deserialize_type(route["output_type"])
+            if isinstance(route["output_type"], list):
+                route["output_type"] = [deserialize_type(t) for t in route["output_type"]]
+            else:
+                route["output_type"] = deserialize_type(route["output_type"])
 
         # Since the custom_filters are typed as optional in the init signature, we catch the
         # case where they are not present in the serialized data and set them to an empty dict.
@@ -355,7 +369,7 @@ def run(self, **kwargs):
 
         raise NoRouteSelectedException(f"No route fired. Routes: {self.routes}")
 
-    def _validate_routes(self, routes: List[Dict]):
+    def _validate_routes(self, routes: List[Route]):
         """
         Validates a list of routes.
 
@@ -401,8 +415,7 @@ def _extract_variables(self, env: Environment, templates: List[str]) -> Set[str]
         """
         variables = set()
         for template in templates:
-            ast = env.parse(template)
-            variables.update(meta.find_undeclared_variables(ast))
+            variables.update(meta.find_undeclared_variables(env.parse(template)))
         return variables
 
     def _validate_template(self, env: Environment, template_text: str):

releasenotes/notes/fix-sede-conditional-router-multiple-outputs-6fe2e58c4c25f80c.yaml

@@ -0,0 +1,4 @@
+---
+fixes:
+  - |
+    In ConditionalRouter fixed the to_dict and from_dict methods to properly handle the case when output_type is a List of types or a List of strings. This occurs when a user specifies a route in ConditionalRouter to have multiple outputs.

test/components/routers/test_conditional_router.py

@@ -615,3 +615,24 @@ def test_multiple_outputs_validation(self):
                     }
                 ]
             )
+
+    def test_sede_multiple_outputs(self):
+        routes = [
+            {
+                "condition": "{{phone_num|get_area_code == 123}}",
+                "output": ["{{phone_num}}", "{{phone_num|get_area_code}}"],
+                "output_name": ["phone_num", "area_code"],
+                "output_type": [str, int],
+            },
+            {
+                "condition": "{{phone_num|get_area_code != 123}}",
+                "output": ["{{phone_num}}", "{{phone_num|get_area_code}}"],
+                "output_name": ["phone_num", "area_code"],
+                "output_type": [str, int],
+            },
+        ]
+
+        router = ConditionalRouter(routes, custom_filters={"get_area_code": custom_filter_to_sede})
+        reloaded_router = ConditionalRouter.from_dict(router.to_dict())
+        assert reloaded_router.custom_filters == router.custom_filters
+        assert reloaded_router.routes == router.routes