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 init_params = data.get("init_parameters", {})

368

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

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

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

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

373 raise DeserializationError(

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

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

376 )

377

378 custom_filters = init_params.get("custom_filters", {})

379 if custom_filters and not _is_unsafe_deserialization():

380 raise DeserializationError(

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

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

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

384 )

385

386 routes = init_params.get("routes")

387 for route in routes:

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

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

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

391 else:

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

393

394 # Since the custom_filters are typed as optional in the init signature, we catch the

395 # case where they are not present in the serialized data and set them to an empty dict.

396 if custom_filters is not None:

397 for name, filter_func in custom_filters.items():

398 init_params["custom_filters"][name] = deserialize_callable(filter_func) if filter_func else None

399 return default_from_dict(cls, data)

400

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

402 """

403 Executes the routing logic.

404

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

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

407 `condition` is True.

408

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

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

411

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

413 of the selected route.

414

415 :raises NoRouteSelectedException:

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

417 :raises RouteConditionException:

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

419 :raises ValueError:

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

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

422 """

423 for route in self.routes:

424 try:

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

426 rendered = t.render(**kwargs)

427 if not self._unsafe:

428 rendered = ast.literal_eval(rendered)

429 if not rendered:

430 continue

431

432 # Handle multiple outputs

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

434 output_types = (

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

436 )

437 output_names = (

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

439 )

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

441

442 result = {}

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

444 if output_passthrough:

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

446 if output not in kwargs:

447 raise ValueError( # noqa: TRY301

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

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

450 )

451 output_value = kwargs[output]

452 else:

453 # Standard Jinja2 template evaluation

454 t_output = self._env.from_string(output)

455 output_value = t_output.render(**kwargs)

456

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

458 # we try to evaluate it and would fail.

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

460 # This doesn't support any user types.

461 with contextlib.suppress(Exception):

462 if not self._unsafe:

463 output_value = ast.literal_eval(output_value)

464

465 # Validate output type if needed

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

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

468

469 result[output_name] = output_value

470

471 return result

472

473 except Exception as e:

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

475 if isinstance(e, ValueError):

476 raise

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

478 raise RouteConditionException(msg) from e

479

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

481

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

483 """

484 Validates a list of routes.

485

486 :param routes: A list of routes.

487 """

488 for route in routes:

489 try:

490 keys = set(route.keys())

491 except AttributeError as e:

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

493

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

495 has_all_mandatory_fields = mandatory_fields.issubset(keys)

496 if not has_all_mandatory_fields:

497 raise ValueError(

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

499 )

500

501 # Validate outputs are consistent

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

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

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

505

506 # Check lengths match

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

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

509

510 # Condition is always a Jinja2 template — validate it

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

512 condition_value = route["condition"]

513 if not isinstance(condition_value, str):

514 raise ValueError(

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

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

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

518 )

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

520

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

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

523 if not output_passthrough:

524 for output in outputs:

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

526 if not isinstance(output, str):

527 raise ValueError(

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

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

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

531 )

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

533

534 @staticmethod

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

536 """

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

538

539 :param env: A Jinja environment.

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

541 :returns: A set of variable names.

542 """

543 variables = set()

544 for template in templates:

545 assigned_variables, template_variables = _extract_template_variables_and_assignments(

546 env=env, template=template

547 )

548 variables.update(template_variables - assigned_variables)

549 return variables

550

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

552 """

553 Validates a template string by parsing it with Jinja.

554

555 :param env: A Jinja environment.

556 :param template_text: A Jinja template string.

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

558 """

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

560 if not isinstance(template_text, str):

561 return False

562 try:

563 env.parse(template_text)

564 return True

565 except TemplateSyntaxError:

566 return False

567

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

569 """

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

571 """

572 # Handle Any type

573 if expected_type is Any:

574 return True

575

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

577 origin = get_origin(expected_type)

578 args = get_args(expected_type)

579

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

581 if origin is None:

582 return isinstance(value, expected_type)

583

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

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

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

587 return False

588 if not isinstance(value, Sequence):

589 return False

590 # Empty sequence is valid

591 if not value:

592 return True

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

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

595

596 # Handle Mapping types (Dict, etc)

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

598 if not isinstance(value, Mapping):

599 return False

600 # Empty mapping is valid

601 if not value:

602 return True

603 key_type, value_type = args

604 # Check all keys and values match their respective types

605 return all(

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

607 for k, v in value.items()

608 )

609

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

611 if _is_union_type(origin):

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

613

614 return False

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

178 statements