Coverage for haystack/components/routers/conditional

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>

3# SPDX-License-Identifier: Apache-2.0

5import ast

6import contextlib

7from collections.abc import Callable, Mapping, Sequence

8from typing import Any, TypedDict, get_args, get_origin

10from jinja2 import Environment, TemplateSyntaxError

11from jinja2.nativetypes import NativeEnvironment

12from typing_extensions import NotRequired

14from haystack import component, default_from_dict, default_to_dict, logging

15from haystack.core.errors import DeserializationError

16from haystack.core.serialization_security import _is_unsafe_deserialization

17from haystack.utils import deserialize_callable, deserialize_type, serialize_callable, serialize_type

18from haystack.utils.jinja2_extensions import _extract_template_variables_and_assignments

19from haystack.utils.jinja2_sandbox import HaystackSandboxedEnvironment

20from haystack.utils.type_serialization import _is_union_type

22logger = logging.getLogger(__name__)

25class NoRouteSelectedException(Exception):

26 """Exception raised when no route is selected in ConditionalRouter."""

29class RouteConditionException(Exception):

30 """Exception raised when there is an error parsing or evaluating the condition expression in ConditionalRouter."""

33class Route(TypedDict):

34 condition: str

35 output: str | list[str]

36 output_name: str | list[str]

37 output_type: type | list[type]

38 output_passthrough: NotRequired[bool]

41@component

42class ConditionalRouter:

43 """

44 Routes data based on specific conditions.

46 You define these conditions in a list of dictionaries called `routes`.

47 Each dictionary in this list represents a single route. Each route has these four elements:

48 - `condition`: A Jinja2 string expression that determines if the route is selected.

49 - `output`: A Jinja2 expression defining the route's output value.

50 - `output_type`: The type of the output data (for example, `str`, `list[int]`).

51 - `output_name`: The name you want to use to publish `output`. This name is used to connect

52 the router to other components in the pipeline.

54 An optional field `output_passthrough` can be set to `True` to treat `output` as a variable name

55 instead of a Jinja2 template, passing the variable value directly. This is useful for routing

56 complex non-basic types (dataclasses, Pydantic models, etc.) without Jinja2 processing.

58 ### Usage example

60 ```python

61 from haystack.components.routers import ConditionalRouter

63 routes = [

64 {

65 "condition": "{{streams|length > 2}}",

66 "output": "{{streams}}",

67 "output_name": "enough_streams",

68 "output_type": list[int],

69 },

70 {

71 "condition": "{{streams|length <= 2}}",

72 "output": "{{streams}}",

73 "output_name": "insufficient_streams",

74 "output_type": list[int],

75 },

76 ]

77 router = ConditionalRouter(routes)

78 # When 'streams' has more than 2 items, 'enough_streams' output will activate, emitting the list [1, 2, 3]

79 kwargs = {"streams": [1, 2, 3], "query": "Haystack"}

80 result = router.run(**kwargs)

81 assert result == {"enough_streams": [1, 2, 3]}

82 ```

84 In this example, we configure two routes. The first route sends the 'streams' value to 'enough_streams' if the

85 stream count exceeds two. The second route directs 'streams' to 'insufficient_streams' if there

86 are two or fewer streams.

88 In the pipeline setup, the Router connects to other components using the output names. For example,

89 'enough_streams' might connect to a component that processes streams, while

90 'insufficient_streams' might connect to a component that fetches more streams.

93 Here is a pipeline that uses `ConditionalRouter` and routes the fetched `ByteStreams` to

94 different components depending on the number of streams fetched:

96 ```python

97 from haystack import Pipeline

98 from haystack.dataclasses import ByteStream

99 from haystack.components.routers import ConditionalRouter

100

101 routes = [

102 {"condition": "{{count > 5}}",

103 "output": "Processing many items",

104 "output_name": "many_items",

105 "output_type": str,

106 },

107 {"condition": "{{count <= 5}}",

108 "output": "Processing few items",

109 "output_name": "few_items",

110 "output_type": str,

111 },

112 ]

113

114 pipe = Pipeline()

115 pipe.add_component("router", ConditionalRouter(routes))

116

117 # Run with count > 5

118 result = pipe.run({"router": {"count": 10}})

119 print(result)

120 # >> {'router': {'many_items': 'Processing many items'}}

121

122 # Run with count <= 5

123 result = pipe.run({"router": {"count": 3}})

124 print(result)

125 # >> {'router': {'few_items': 'Processing few items'}}

126 ```

127

128 ### Passthrough routing for non-basic types

129

130 Without `output_passthrough`, the router renders `output` as a Jinja2 template, which converts

131 the value to its string representation. Custom types cannot survive that round-trip:

132

133 ```python

134 # Without output_passthrough — the object is silently converted to a string

135 routes = [

136 {

137 "condition": "{{True}}",

138 "output": "{{query}}",

139 "output_name": "out",

140 "output_type": ParsedQuery,

141 }

142 ]

143 router = ConditionalRouter(routes)

144 result = router.run(query=ParsedQuery(text="hello", intent="search", entities=[]))

145 # result["out"] == "ParsedQuery(text='hello', intent='search', entities=[])"

146 # ^^^ str, not ParsedQuery — the object was destroyed

147 ```

148

149 Set `output_passthrough: True` to skip Jinja2 entirely and pass the value directly from kwargs:

150

151 ```python

152 from haystack.components.routers import ConditionalRouter

153 from dataclasses import dataclass, field

154

155 @dataclass

156 class ParsedQuery:

157 text: str

158 intent: str # "search" | "chat"

159 entities: list[str] = field(default_factory=list)

160

161 routes = [

162 {

163 "condition": "{{query.intent == 'search'}}",

164 "output": "query", # variable name, not a Jinja2 template

165 "output_name": "search_query",

166 "output_type": ParsedQuery,

167 "output_passthrough": True,

168 },

169 {

170 "condition": "{{query.intent == 'chat'}}",

171 "output": "query",

172 "output_name": "chat_query",

173 "output_type": ParsedQuery,

174 "output_passthrough": True,

175 },

176 ]

177

178 router = ConditionalRouter(routes)

179 query = ParsedQuery(text="What is Haystack?", intent="search", entities=["Haystack"])

180 result = router.run(query=query)

181

182 assert isinstance(result["search_query"], ParsedQuery) # type preserved

183 assert result["search_query"] is query # same object, no copying

184 ```

185 """

186

187 def __init__(

188 self,

189 routes: list[Route],

190 custom_filters: dict[str, Callable] | None = None,

191 unsafe: bool = False,

192 validate_output_type: bool = False,

193 optional_variables: list[str] | None = None,

194 ) -> None:

195 """

196 Initializes the `ConditionalRouter` with a list of routes detailing the conditions for routing.

197

198 :param routes: A list of dictionaries, each defining a route.

199 Each route has these four elements:

200 - `condition`: A Jinja2 string expression that determines if the route is selected.

201 - `output`: A Jinja2 expression defining the route's output value, or a plain variable name

202 if `output_passthrough` is `True`.

203 - `output_type`: The type of the output data (for example, `str`, `list[int]`).

204 - `output_name`: The name you want to use to publish `output`. This name is used to connect

205 the router to other components in the pipeline.

206 - `output_passthrough` (optional): If `True`, treats `output` as a plain variable name and

207 passes the value directly from the input kwargs, skipping all Jinja2 processing. Useful

208 for routing complex non-basic types without template transformation.

209 Note: if the variable named in `output` is also listed in `optional_variables`, a missing

210 value at runtime will route `None` downstream rather than raising a `ValueError`.

211 :param custom_filters: A dictionary of custom Jinja2 filters used in the condition expressions.

212 For example, passing `{"my_filter": my_filter_fcn}` where:

213 - `my_filter` is the name of the custom filter.

214 - `my_filter_fcn` is a callable that takes `my_var:str` and returns `my_var[:3]`.

215 `{{ my_var|my_filter }}` can then be used inside a route condition expression:

216 `"condition": "{{ my_var|my_filter == 'foo' }}"`.

217 :param unsafe:

218 Enable execution of arbitrary code in the Jinja template.

219 This should only be used if you trust the source of the template as it can be lead to remote code execution.

220 :param validate_output_type:

221 Enable validation of routes' output.

222 If a route output doesn't match the declared type a ValueError is raised running.

223 :param optional_variables:

224 A list of variable names that are optional in your route conditions and outputs.

225 If these variables are not provided at runtime, they will be set to `None`.

226 This allows you to write routes that can handle missing inputs gracefully without raising errors.

227

228 Example usage with a default fallback route in a Pipeline:

229 ```python

230 from haystack import Pipeline

231 from haystack.components.routers import ConditionalRouter

232

233 routes = [

234 {

235 "condition": '{{ path == "rag" }}',

236 "output": "{{ question }}",

237 "output_name": "rag_route",

238 "output_type": str

239 },

240 {

241 "condition": "{{ True }}", # fallback route

242 "output": "{{ question }}",

243 "output_name": "default_route",

244 "output_type": str

245 }

246 ]

247

248 router = ConditionalRouter(routes, optional_variables=["path"])

249 pipe = Pipeline()

250 pipe.add_component("router", router)

251

252 # When 'path' is provided in the pipeline:

253 result = pipe.run(data={"router": {"question": "What?", "path": "rag"}})

254 assert result["router"] == {"rag_route": "What?"}

255

256 # When 'path' is not provided, fallback route is taken:

257 result = pipe.run(data={"router": {"question": "What?"}})

258 assert result["router"] == {"default_route": "What?"}

259 ```

260

261 This pattern is particularly useful when:

262 - You want to provide default/fallback behavior when certain inputs are missing

263 - Some variables are only needed for specific routing conditions

264 - You're building flexible pipelines where not all inputs are guaranteed to be present

265 """

266 self.routes: list[Route] = routes

267 self.custom_filters = custom_filters or {}

268 self._unsafe = unsafe

269 self._validate_output_type = validate_output_type

270 self.optional_variables = optional_variables or []

271

272 # Create a Jinja environment to inspect variables in the condition templates

273 if self._unsafe:

274 msg = (

275 "Unsafe mode is enabled. This allows execution of arbitrary code in the Jinja template. "

276 "Use this only if you trust the source of the template."

277 )

278 logger.warning(msg)

279

280 self._env = NativeEnvironment() if self._unsafe else HaystackSandboxedEnvironment()

281 self._env.filters.update(self.custom_filters)

282

283 self._validate_routes(routes)

284 # Inspect the routes to determine input and output types.

285 input_types: set[str] = set() # let's just store the name, type will always be Any

286 output_types: dict[str, type | list[type]] = {}

287

288 for route in routes:

289 output_passthrough = route.get("output_passthrough", False)

290 outputs = route["output"] if isinstance(route["output"], list) else [route["output"]]

291

292 if output_passthrough:

293 # For passthrough routes, output values are plain variable names — treat them as inputs

294 route_input_names = self._extract_variables(self._env, [route["condition"]])

295 route_input_names.update(outputs)

296 else:

297 # For normal routes, extract variables from both condition and output templates

298 route_input_names = self._extract_variables(self._env, [route["condition"]] + outputs)

299

300 input_types.update(route_input_names)

301

302 # extract outputs

303 output_names = route["output_name"] if isinstance(route["output_name"], list) else [route["output_name"]]

304 output_types_list = (

305 route["output_type"] if isinstance(route["output_type"], list) else [route["output_type"]]

306 )

307

308 output_types.update(dict(zip(output_names, output_types_list, strict=True)))

309

310 # remove optional variables from mandatory input types

311 mandatory_input_types = input_types - set(self.optional_variables)

312

313 # warn about unused optional variables

314 unused_optional_vars = set(self.optional_variables) - input_types if self.optional_variables else None

315 if unused_optional_vars:

316 logger.warning(

317 "The following optional variables are specified but not used in any route: {unused_optional_vars}. "

318 "Check if there's a typo in variable names.",

319 unused_optional_vars=unused_optional_vars,

320 )

321

322 # add mandatory input types

323 component.set_input_types(self, **dict.fromkeys(mandatory_input_types, Any))

324

325 # now add optional input types

326 for optional_var_name in self.optional_variables:

327 component.set_input_type(self, name=optional_var_name, type=Any, default=None)

328

329 # set output types

330 component.set_output_types(self, **output_types) # type: ignore[arg-type]

331

332 def to_dict(self) -> dict[str, Any]:

333 """

334 Serializes the component to a dictionary.

335

336 :returns:

337 Dictionary with serialized data.

338 """

339 serialized_routes = []

340 for route in self.routes:

341 serialized_output_type = (

342 [serialize_type(t) for t in route["output_type"]]

343 if isinstance(route["output_type"], list)

344 else serialize_type(route["output_type"])

345 )

346 serialized_routes.append({**route, "output_type": serialized_output_type})

347 se_filters = {name: serialize_callable(filter_func) for name, filter_func in self.custom_filters.items()}

348 return default_to_dict(

349 self,

350 routes=serialized_routes,

351 custom_filters=se_filters,

352 unsafe=self._unsafe,

353 validate_output_type=self._validate_output_type,

354 optional_variables=self.optional_variables,

355 )

356

357 @classmethod

358 def from_dict(cls, data: dict[str, Any]) -> "ConditionalRouter":

359 """

360 Deserializes the component from a dictionary.

361

362 :param data:

363 The dictionary to deserialize from.

364 :returns:

365 The deserialized component.

366 """

367 # Copy so the caller's data stays serialized; nested routes and filters are copied below too

368 init_params = dict(data.get("init_parameters", {}))

369

370 # `unsafe=True` swaps the Jinja sandbox for a NativeEnvironment that executes arbitrary code.

371 # Honor it from serialized data only when the whole pipeline is being loaded in unsafe mode;

372 # otherwise a hostile pipeline could disable the sandbox on its own in default safe mode.

373 if init_params.get("unsafe") and not _is_unsafe_deserialization():

374 raise DeserializationError(

375 "Refusing to deserialize a ConditionalRouter with unsafe=True while loading in safe mode. "

376 "If you trust the source of this data, load it with Pipeline.load(..., unsafe=True)."

377 )

378

379 custom_filters = init_params.get("custom_filters")

380 if custom_filters and not _is_unsafe_deserialization():

381 raise DeserializationError(

382 "Refusing to deserialize a ConditionalRouter with custom filters while loading in safe mode. "

383 "Custom filters are arbitrary callables that can execute during pipeline loading. "

384 "If you trust the source of this data, load it with Pipeline.load(..., unsafe=True)."

385 )

386

387 routes = init_params.get("routes")

388 if routes is not None:

389 init_params["routes"] = routes = [dict(route) for route in routes]

390 for route in routes:

391 # output_type needs to be deserialized from a string to a type

392 if isinstance(route["output_type"], list):

393 route["output_type"] = [deserialize_type(t) for t in route["output_type"]]

394 else:

395 route["output_type"] = deserialize_type(route["output_type"])

396

397 if custom_filters:

398 init_params["custom_filters"] = {

399 name: deserialize_callable(filter_func) if filter_func else None

400 for name, filter_func in custom_filters.items()

401 }

402 return default_from_dict(cls, {**data, "init_parameters": init_params})

403

404 def run(self, **kwargs: Any) -> dict[str, Any]:

405 """

406 Executes the routing logic.

407

408 Executes the routing logic by evaluating the specified boolean condition expressions for each route in the

409 order they are listed. The method directs the flow of data to the output specified in the first route whose

410 `condition` is True.

411

412 :param kwargs: All variables used in the `condition` expressed in the routes. When the component is used in a

413 pipeline, these variables are passed from the previous component's output.

414

415 :returns: A dictionary where the key is the `output_name` of the selected route and the value is the `output`

416 of the selected route.

417

418 :raises NoRouteSelectedException:

419 If no `condition' in the routes is `True`.

420 :raises RouteConditionException:

421 If there is an error parsing or evaluating the `condition` expression in the routes.

422 :raises ValueError:

423 If type validation is enabled and the route output doesn't match the declared type, or if

424 `output_passthrough` is `True` and the variable named in `output` is not found in kwargs.

425 """

426 for route in self.routes:

427 try:

428 t = self._env.from_string(route["condition"])

429 rendered = t.render(**kwargs)

430 if not self._unsafe:

431 rendered = ast.literal_eval(rendered)

432 if not rendered:

433 continue

434

435 # Handle multiple outputs

436 outputs = route["output"] if isinstance(route["output"], list) else [route["output"]]

437 output_types = (

438 route["output_type"] if isinstance(route["output_type"], list) else [route["output_type"]]

439 )

440 output_names = (

441 route["output_name"] if isinstance(route["output_name"], list) else [route["output_name"]]

442 )

443 output_passthrough = route.get("output_passthrough", False)

444

445 result = {}

446 for output, output_type, output_name in zip(outputs, output_types, output_names, strict=True):

447 if output_passthrough:

448 # output is a plain variable name — retrieve directly from kwargs, no Jinja2 processing

449 if output not in kwargs:

450 raise ValueError( # noqa: TRY301

451 f"Variable '{output}' not found in inputs for passthrough route '{output_name}'. "

452 f"Ensure '{output}' is passed as an input to the router."

453 )

454 output_value = kwargs[output]

455 else:

456 # Standard Jinja2 template evaluation

457 t_output = self._env.from_string(output)

458 output_value = t_output.render(**kwargs)

459

460 # We suppress the exception in case the output is already a string, otherwise

461 # we try to evaluate it and would fail.

462 # This must be done cause the output could be different literal structures.

463 # This doesn't support any user types.

464 # When the declared output_type is str we skip literal evaluation so that a

465 # rendered string that happens to be a valid Python literal (e.g. "1,000" -> (1, 0),

466 # "42" -> 42, "None" -> None) is returned unchanged instead of being coerced to

467 # another type, which would violate the declared output_type.

468 with contextlib.suppress(Exception):

469 if not self._unsafe and output_type is not str:

470 output_value = ast.literal_eval(output_value)

471

472 # Validate output type if needed

473 if self._validate_output_type and not self._output_matches_type(output_value, output_type):

474 raise ValueError(f"Route '{output_name}' type doesn't match expected type") # noqa: TRY301

475

476 result[output_name] = output_value

477

478 return result

479

480 except Exception as e:

481 # If this was a type-validation failure or missing passthrough variable, let it propagate

482 if isinstance(e, ValueError):

483 raise

484 msg = f"Error evaluating condition for route '{route}': {e}"

485 raise RouteConditionException(msg) from e

486

487 raise NoRouteSelectedException(f"No route fired. Routes: {self.routes}")

488

489 def _validate_routes(self, routes: list[Route]) -> None:

490 """

491 Validates a list of routes.

492

493 :param routes: A list of routes.

494 """

495 for route in routes:

496 try:

497 keys = set(route.keys())

498 except AttributeError as e:

499 raise ValueError(f"Route must be a dictionary, got: {route}") from e

500

501 mandatory_fields = {"condition", "output", "output_type", "output_name"}

502 has_all_mandatory_fields = mandatory_fields.issubset(keys)

503 if not has_all_mandatory_fields:

504 raise ValueError(

505 f"Route must contain 'condition', 'output', 'output_type' and 'output_name' fields: {route}"

506 )

507

508 # Validate outputs are consistent

509 outputs = route["output"] if isinstance(route["output"], list) else [route["output"]]

510 output_types = route["output_type"] if isinstance(route["output_type"], list) else [route["output_type"]]

511 output_names = route["output_name"] if isinstance(route["output_name"], list) else [route["output_name"]]

512

513 # Check lengths match

514 if not len(outputs) == len(output_types) == len(output_names):

515 raise ValueError(f"Route output, output_type and output_name must have same length: {route}")

516

517 # Condition is always a Jinja2 template — validate it

518 if not self._validate_template(self._env, route["condition"]):

519 condition_value = route["condition"]

520 if not isinstance(condition_value, str):

521 raise ValueError(

522 f"Invalid template for condition: {condition_value!r} (type: {type(condition_value).__name__})."

523 f"Condition must be a string representing a valid Jinja2 template. "

524 f"For example, use {str(condition_value)!r} instead of {condition_value!r}."

525 )

526 raise ValueError(f"Invalid template for condition: {condition_value}")

527

528 # Only validate output as Jinja2 template when output_passthrough is False (default)

529 output_passthrough = route.get("output_passthrough", False)

530 if not output_passthrough:

531 for output in outputs:

532 if not self._validate_template(self._env, output):

533 if not isinstance(output, str):

534 raise ValueError(

535 f"Invalid template for output: {output!r} (type: {type(output).__name__}). "

536 f"Output must be a string representing a valid Jinja2 template. "

537 f"For example, use {str(output)!r} instead of {output!r}."

538 )

539 raise ValueError(f"Invalid template for output: {output}")

540

541 @staticmethod

542 def _extract_variables(env: Environment, templates: list[str]) -> set[str]:

543 """

544 Extracts all variables from a list of Jinja template strings.

545

546 :param env: A Jinja environment.

547 :param templates: A list of Jinja template strings.

548 :returns: A set of variable names.

549 """

550 variables = set()

551 for template in templates:

552 assigned_variables, template_variables = _extract_template_variables_and_assignments(

553 env=env, template=template

554 )

555 variables.update(template_variables - assigned_variables)

556 return variables

557

558 def _validate_template(self, env: Environment, template_text: str) -> bool:

559 """

560 Validates a template string by parsing it with Jinja.

561

562 :param env: A Jinja environment.

563 :param template_text: A Jinja template string.

564 :returns: `True` if the template is valid, `False` otherwise.

565 """

566 # Check if template_text is a string before attempting to parse

567 if not isinstance(template_text, str):

568 return False

569 try:

570 env.parse(template_text)

571 return True

572 except TemplateSyntaxError:

573 return False

574

575 def _output_matches_type(self, value: Any, expected_type: type) -> bool: # noqa: PLR0911

576 """

577 Checks whether `value` type matches the `expected_type`.

578 """

579 # Handle Any type

580 if expected_type is Any:

581 return True

582

583 # Get the origin type (List, Dict, etc) and type arguments

584 origin = get_origin(expected_type)

585 args = get_args(expected_type)

586

587 # Handle basic types (int, str, etc)

588 if origin is None:

589 return isinstance(value, expected_type)

590

591 # Handle Sequence types (List, Tuple, etc)

592 if isinstance(origin, type) and issubclass(origin, Sequence):

593 if isinstance(value, (str, bytes)):

594 return False

595 if not isinstance(value, Sequence):

596 return False

597 # Empty sequence is valid

598 if not value:

599 return True

600 # Check each element against the sequence's type parameter

601 return all(self._output_matches_type(item, args[0]) for item in value)

602

603 # Handle Mapping types (Dict, etc)

604 if isinstance(origin, type) and issubclass(origin, Mapping):

605 if not isinstance(value, Mapping):

606 return False

607 # Empty mapping is valid

608 if not value:

609 return True

610 key_type, value_type = args

611 # Check all keys and values match their respective types

612 return all(

613 self._output_matches_type(k, key_type) and self._output_matches_type(v, value_type)

614 for k, v in value.items()

615 )

616

617 # Handle Union types (including Optional and X | Y syntax)

618 if _is_union_type(origin):

619 return any(self._output_matches_type(value, arg) for arg in args)

620

621 return False

Coverage for haystack/components/routers/conditional_router.py: 96%

179 statements