chore: import upstream snapshot with attribution
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
CodeQL / Analyze (python) (push) Has been cancelled
Update Platform Components Table / update (push) Has been cancelled
Docker image release / Build base image (push) Has been cancelled
Sync docs with Docusaurus / sync (push) Has been cancelled
Tests / Check if changed (push) Has been cancelled
Tests / format (push) Has been cancelled
Tests / check-imports (push) Has been cancelled
Tests / Unit / macos-latest (push) Has been cancelled
Tests / Unit / ubuntu-latest (push) Has been cancelled
Tests / Unit / windows-latest (push) Has been cancelled
Tests / mypy (push) Has been cancelled
Tests / Integration / ubuntu-latest (push) Has been cancelled
Tests / Integration / macos-latest (push) Has been cancelled
Tests / Integration / windows-latest (push) Has been cancelled
Tests / notify-slack-on-failure (push) Has been cancelled
Tests / Mark tests as completed (push) Has been cancelled
This commit is contained in:
@@ -0,0 +1,595 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import ast
|
||||
import contextlib
|
||||
from collections.abc import Callable, Mapping, Sequence
|
||||
from typing import Any, TypedDict, get_args, get_origin
|
||||
|
||||
from jinja2 import Environment, TemplateSyntaxError
|
||||
from jinja2.nativetypes import NativeEnvironment
|
||||
from jinja2.sandbox import SandboxedEnvironment
|
||||
from typing_extensions import NotRequired
|
||||
|
||||
from haystack import component, default_from_dict, default_to_dict, logging
|
||||
from haystack.utils import deserialize_callable, deserialize_type, serialize_callable, serialize_type
|
||||
from haystack.utils.jinja2_extensions import _extract_template_variables_and_assignments
|
||||
from haystack.utils.type_serialization import _is_union_type
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class NoRouteSelectedException(Exception):
|
||||
"""Exception raised when no route is selected in ConditionalRouter."""
|
||||
|
||||
|
||||
class RouteConditionException(Exception):
|
||||
"""Exception raised when there is an error parsing or evaluating the condition expression in ConditionalRouter."""
|
||||
|
||||
|
||||
class Route(TypedDict):
|
||||
condition: str
|
||||
output: str | list[str]
|
||||
output_name: str | list[str]
|
||||
output_type: type | list[type]
|
||||
output_passthrough: NotRequired[bool]
|
||||
|
||||
|
||||
@component
|
||||
class ConditionalRouter:
|
||||
"""
|
||||
Routes data based on specific conditions.
|
||||
|
||||
You define these conditions in a list of dictionaries called `routes`.
|
||||
Each dictionary in this list represents a single route. Each route has these four elements:
|
||||
- `condition`: A Jinja2 string expression that determines if the route is selected.
|
||||
- `output`: A Jinja2 expression defining the route's output value.
|
||||
- `output_type`: The type of the output data (for example, `str`, `list[int]`).
|
||||
- `output_name`: The name you want to use to publish `output`. This name is used to connect
|
||||
the router to other components in the pipeline.
|
||||
|
||||
An optional field `output_passthrough` can be set to `True` to treat `output` as a variable name
|
||||
instead of a Jinja2 template, passing the variable value directly. This is useful for routing
|
||||
complex non-basic types (dataclasses, Pydantic models, etc.) without Jinja2 processing.
|
||||
|
||||
### Usage example
|
||||
|
||||
```python
|
||||
from haystack.components.routers import ConditionalRouter
|
||||
|
||||
routes = [
|
||||
{
|
||||
"condition": "{{streams|length > 2}}",
|
||||
"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)
|
||||
# When 'streams' has more than 2 items, 'enough_streams' output will activate, emitting the list [1, 2, 3]
|
||||
kwargs = {"streams": [1, 2, 3], "query": "Haystack"}
|
||||
result = router.run(**kwargs)
|
||||
assert result == {"enough_streams": [1, 2, 3]}
|
||||
```
|
||||
|
||||
In this example, we configure two routes. The first route sends the 'streams' value to 'enough_streams' if the
|
||||
stream count exceeds two. The second route directs 'streams' to 'insufficient_streams' if there
|
||||
are two or fewer streams.
|
||||
|
||||
In the pipeline setup, the Router connects to other components using the output names. For example,
|
||||
'enough_streams' might connect to a component that processes streams, while
|
||||
'insufficient_streams' might connect to a component that fetches more streams.
|
||||
|
||||
|
||||
Here is a pipeline that uses `ConditionalRouter` and routes the fetched `ByteStreams` to
|
||||
different components depending on the number of streams fetched:
|
||||
|
||||
```python
|
||||
from haystack import Pipeline
|
||||
from haystack.dataclasses import ByteStream
|
||||
from haystack.components.routers import ConditionalRouter
|
||||
|
||||
routes = [
|
||||
{"condition": "{{count > 5}}",
|
||||
"output": "Processing many items",
|
||||
"output_name": "many_items",
|
||||
"output_type": str,
|
||||
},
|
||||
{"condition": "{{count <= 5}}",
|
||||
"output": "Processing few items",
|
||||
"output_name": "few_items",
|
||||
"output_type": str,
|
||||
},
|
||||
]
|
||||
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("router", ConditionalRouter(routes))
|
||||
|
||||
# Run with count > 5
|
||||
result = pipe.run({"router": {"count": 10}})
|
||||
print(result)
|
||||
# >> {'router': {'many_items': 'Processing many items'}}
|
||||
|
||||
# Run with count <= 5
|
||||
result = pipe.run({"router": {"count": 3}})
|
||||
print(result)
|
||||
# >> {'router': {'few_items': 'Processing few items'}}
|
||||
```
|
||||
|
||||
### Passthrough routing for non-basic types
|
||||
|
||||
Without `output_passthrough`, the router renders `output` as a Jinja2 template, which converts
|
||||
the value to its string representation. Custom types cannot survive that round-trip:
|
||||
|
||||
```python
|
||||
# Without output_passthrough — the object is silently converted to a string
|
||||
routes = [
|
||||
{
|
||||
"condition": "{{True}}",
|
||||
"output": "{{query}}",
|
||||
"output_name": "out",
|
||||
"output_type": ParsedQuery,
|
||||
}
|
||||
]
|
||||
router = ConditionalRouter(routes)
|
||||
result = router.run(query=ParsedQuery(text="hello", intent="search", entities=[]))
|
||||
# result["out"] == "ParsedQuery(text='hello', intent='search', entities=[])"
|
||||
# ^^^ str, not ParsedQuery — the object was destroyed
|
||||
```
|
||||
|
||||
Set `output_passthrough: True` to skip Jinja2 entirely and pass the value directly from kwargs:
|
||||
|
||||
```python
|
||||
from haystack.components.routers import ConditionalRouter
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
@dataclass
|
||||
class ParsedQuery:
|
||||
text: str
|
||||
intent: str # "search" | "chat"
|
||||
entities: list[str] = field(default_factory=list)
|
||||
|
||||
routes = [
|
||||
{
|
||||
"condition": "{{query.intent == 'search'}}",
|
||||
"output": "query", # variable name, not a Jinja2 template
|
||||
"output_name": "search_query",
|
||||
"output_type": ParsedQuery,
|
||||
"output_passthrough": True,
|
||||
},
|
||||
{
|
||||
"condition": "{{query.intent == 'chat'}}",
|
||||
"output": "query",
|
||||
"output_name": "chat_query",
|
||||
"output_type": ParsedQuery,
|
||||
"output_passthrough": True,
|
||||
},
|
||||
]
|
||||
|
||||
router = ConditionalRouter(routes)
|
||||
query = ParsedQuery(text="What is Haystack?", intent="search", entities=["Haystack"])
|
||||
result = router.run(query=query)
|
||||
|
||||
assert isinstance(result["search_query"], ParsedQuery) # type preserved
|
||||
assert result["search_query"] is query # same object, no copying
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
routes: list[Route],
|
||||
custom_filters: dict[str, Callable] | None = None,
|
||||
unsafe: bool = False,
|
||||
validate_output_type: bool = False,
|
||||
optional_variables: list[str] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Initializes the `ConditionalRouter` with a list of routes detailing the conditions for routing.
|
||||
|
||||
:param routes: A list of dictionaries, each defining a route.
|
||||
Each route has these four elements:
|
||||
- `condition`: A Jinja2 string expression that determines if the route is selected.
|
||||
- `output`: A Jinja2 expression defining the route's output value, or a plain variable name
|
||||
if `output_passthrough` is `True`.
|
||||
- `output_type`: The type of the output data (for example, `str`, `list[int]`).
|
||||
- `output_name`: The name you want to use to publish `output`. This name is used to connect
|
||||
the router to other components in the pipeline.
|
||||
- `output_passthrough` (optional): If `True`, treats `output` as a plain variable name and
|
||||
passes the value directly from the input kwargs, skipping all Jinja2 processing. Useful
|
||||
for routing complex non-basic types without template transformation.
|
||||
Note: if the variable named in `output` is also listed in `optional_variables`, a missing
|
||||
value at runtime will route `None` downstream rather than raising a `ValueError`.
|
||||
:param custom_filters: A dictionary of custom Jinja2 filters used in the condition expressions.
|
||||
For example, passing `{"my_filter": my_filter_fcn}` where:
|
||||
- `my_filter` is the name of the custom filter.
|
||||
- `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.
|
||||
:param validate_output_type:
|
||||
Enable validation of routes' output.
|
||||
If a route output doesn't match the declared type a ValueError is raised running.
|
||||
:param optional_variables:
|
||||
A list of variable names that are optional in your route conditions and outputs.
|
||||
If these variables are not provided at runtime, they will be set to `None`.
|
||||
This allows you to write routes that can handle missing inputs gracefully without raising errors.
|
||||
|
||||
Example usage with a default fallback route in a Pipeline:
|
||||
```python
|
||||
from haystack import Pipeline
|
||||
from haystack.components.routers import ConditionalRouter
|
||||
|
||||
routes = [
|
||||
{
|
||||
"condition": '{{ path == "rag" }}',
|
||||
"output": "{{ question }}",
|
||||
"output_name": "rag_route",
|
||||
"output_type": str
|
||||
},
|
||||
{
|
||||
"condition": "{{ True }}", # fallback route
|
||||
"output": "{{ question }}",
|
||||
"output_name": "default_route",
|
||||
"output_type": str
|
||||
}
|
||||
]
|
||||
|
||||
router = ConditionalRouter(routes, optional_variables=["path"])
|
||||
pipe = Pipeline()
|
||||
pipe.add_component("router", router)
|
||||
|
||||
# When 'path' is provided in the pipeline:
|
||||
result = pipe.run(data={"router": {"question": "What?", "path": "rag"}})
|
||||
assert result["router"] == {"rag_route": "What?"}
|
||||
|
||||
# When 'path' is not provided, fallback route is taken:
|
||||
result = pipe.run(data={"router": {"question": "What?"}})
|
||||
assert result["router"] == {"default_route": "What?"}
|
||||
```
|
||||
|
||||
This pattern is particularly useful when:
|
||||
- You want to provide default/fallback behavior when certain inputs are missing
|
||||
- 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[Route] = routes
|
||||
self.custom_filters = custom_filters or {}
|
||||
self._unsafe = unsafe
|
||||
self._validate_output_type = validate_output_type
|
||||
self.optional_variables = optional_variables or []
|
||||
|
||||
# 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."
|
||||
)
|
||||
logger.warning(msg)
|
||||
|
||||
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, type | list[type]] = {}
|
||||
|
||||
for route in routes:
|
||||
output_passthrough = route.get("output_passthrough", False)
|
||||
outputs = route["output"] if isinstance(route["output"], list) else [route["output"]]
|
||||
|
||||
if output_passthrough:
|
||||
# For passthrough routes, output values are plain variable names — treat them as inputs
|
||||
route_input_names = self._extract_variables(self._env, [route["condition"]])
|
||||
route_input_names.update(outputs)
|
||||
else:
|
||||
# For normal routes, extract variables from both condition and output templates
|
||||
route_input_names = self._extract_variables(self._env, [route["condition"]] + outputs)
|
||||
|
||||
input_types.update(route_input_names)
|
||||
|
||||
# extract outputs
|
||||
output_names = route["output_name"] if isinstance(route["output_name"], list) else [route["output_name"]]
|
||||
output_types_list = (
|
||||
route["output_type"] if isinstance(route["output_type"], list) else [route["output_type"]]
|
||||
)
|
||||
|
||||
output_types.update(dict(zip(output_names, output_types_list, strict=True)))
|
||||
|
||||
# remove optional variables from mandatory input types
|
||||
mandatory_input_types = input_types - set(self.optional_variables)
|
||||
|
||||
# warn about unused optional variables
|
||||
unused_optional_vars = set(self.optional_variables) - input_types if self.optional_variables else None
|
||||
if unused_optional_vars:
|
||||
logger.warning(
|
||||
"The following optional variables are specified but not used in any route: {unused_optional_vars}. "
|
||||
"Check if there's a typo in variable names.",
|
||||
unused_optional_vars=unused_optional_vars,
|
||||
)
|
||||
|
||||
# add mandatory input types
|
||||
component.set_input_types(self, **dict.fromkeys(mandatory_input_types, Any))
|
||||
|
||||
# now add optional input types
|
||||
for optional_var_name in self.optional_variables:
|
||||
component.set_input_type(self, name=optional_var_name, type=Any, default=None)
|
||||
|
||||
# set output types
|
||||
component.set_output_types(self, **output_types) # type: ignore[arg-type]
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""
|
||||
Serializes the component to a dictionary.
|
||||
|
||||
:returns:
|
||||
Dictionary with serialized data.
|
||||
"""
|
||||
serialized_routes = []
|
||||
for route in self.routes:
|
||||
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,
|
||||
routes=serialized_routes,
|
||||
custom_filters=se_filters,
|
||||
unsafe=self._unsafe,
|
||||
validate_output_type=self._validate_output_type,
|
||||
optional_variables=self.optional_variables,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "ConditionalRouter":
|
||||
"""
|
||||
Deserializes the component from a dictionary.
|
||||
|
||||
:param data:
|
||||
The dictionary to deserialize from.
|
||||
:returns:
|
||||
The deserialized component.
|
||||
"""
|
||||
init_params = data.get("init_parameters", {})
|
||||
routes = init_params.get("routes")
|
||||
for route in routes:
|
||||
# output_type needs to be deserialized from a string to a 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.
|
||||
custom_filters = init_params.get("custom_filters", {})
|
||||
if custom_filters is not None:
|
||||
for name, filter_func in custom_filters.items():
|
||||
init_params["custom_filters"][name] = deserialize_callable(filter_func) if filter_func else None
|
||||
return default_from_dict(cls, data)
|
||||
|
||||
def run(self, **kwargs: Any) -> dict[str, Any]:
|
||||
"""
|
||||
Executes the routing logic.
|
||||
|
||||
Executes the routing logic by evaluating the specified boolean condition expressions for each route in the
|
||||
order they are listed. The method directs the flow of data to the output specified in the first route whose
|
||||
`condition` is True.
|
||||
|
||||
:param kwargs: All variables used in the `condition` expressed in the routes. When the component is used in a
|
||||
pipeline, these variables are passed from the previous component's output.
|
||||
|
||||
:returns: A dictionary where the key is the `output_name` of the selected route and the value is the `output`
|
||||
of the selected route.
|
||||
|
||||
:raises NoRouteSelectedException:
|
||||
If no `condition' in the routes is `True`.
|
||||
:raises RouteConditionException:
|
||||
If there is an error parsing or evaluating the `condition` expression in the routes.
|
||||
:raises ValueError:
|
||||
If type validation is enabled and the route output doesn't match the declared type, or if
|
||||
`output_passthrough` is `True` and the variable named in `output` is not found in kwargs.
|
||||
"""
|
||||
for route in self.routes:
|
||||
try:
|
||||
t = self._env.from_string(route["condition"])
|
||||
rendered = t.render(**kwargs)
|
||||
if not self._unsafe:
|
||||
rendered = ast.literal_eval(rendered)
|
||||
if not rendered:
|
||||
continue
|
||||
|
||||
# Handle multiple outputs
|
||||
outputs = route["output"] if isinstance(route["output"], list) else [route["output"]]
|
||||
output_types = (
|
||||
route["output_type"] if isinstance(route["output_type"], list) else [route["output_type"]]
|
||||
)
|
||||
output_names = (
|
||||
route["output_name"] if isinstance(route["output_name"], list) else [route["output_name"]]
|
||||
)
|
||||
output_passthrough = route.get("output_passthrough", False)
|
||||
|
||||
result = {}
|
||||
for output, output_type, output_name in zip(outputs, output_types, output_names, strict=True):
|
||||
if output_passthrough:
|
||||
# output is a plain variable name — retrieve directly from kwargs, no Jinja2 processing
|
||||
if output not in kwargs:
|
||||
raise ValueError( # noqa: TRY301
|
||||
f"Variable '{output}' not found in inputs for passthrough route '{output_name}'. "
|
||||
f"Ensure '{output}' is passed as an input to the router."
|
||||
)
|
||||
output_value = kwargs[output]
|
||||
else:
|
||||
# Standard Jinja2 template evaluation
|
||||
t_output = self._env.from_string(output)
|
||||
output_value = t_output.render(**kwargs)
|
||||
|
||||
# We suppress the exception in case the output is already a string, otherwise
|
||||
# we try to evaluate it and would fail.
|
||||
# This must be done cause the output could be different literal structures.
|
||||
# This doesn't support any user types.
|
||||
with contextlib.suppress(Exception):
|
||||
if not self._unsafe:
|
||||
output_value = ast.literal_eval(output_value)
|
||||
|
||||
# Validate output type if needed
|
||||
if self._validate_output_type and not self._output_matches_type(output_value, output_type):
|
||||
raise ValueError(f"Route '{output_name}' type doesn't match expected type") # noqa: TRY301
|
||||
|
||||
result[output_name] = output_value
|
||||
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
# If this was a type-validation failure or missing passthrough variable, let it propagate
|
||||
if isinstance(e, ValueError):
|
||||
raise
|
||||
msg = f"Error evaluating condition for route '{route}': {e}"
|
||||
raise RouteConditionException(msg) from e
|
||||
|
||||
raise NoRouteSelectedException(f"No route fired. Routes: {self.routes}")
|
||||
|
||||
def _validate_routes(self, routes: list[Route]) -> None:
|
||||
"""
|
||||
Validates a list of routes.
|
||||
|
||||
:param routes: A list of routes.
|
||||
"""
|
||||
for route in routes:
|
||||
try:
|
||||
keys = set(route.keys())
|
||||
except AttributeError as e:
|
||||
raise ValueError(f"Route must be a dictionary, got: {route}") from e
|
||||
|
||||
mandatory_fields = {"condition", "output", "output_type", "output_name"}
|
||||
has_all_mandatory_fields = mandatory_fields.issubset(keys)
|
||||
if not has_all_mandatory_fields:
|
||||
raise ValueError(
|
||||
f"Route must contain 'condition', 'output', 'output_type' and 'output_name' fields: {route}"
|
||||
)
|
||||
|
||||
# Validate outputs are consistent
|
||||
outputs = route["output"] if isinstance(route["output"], list) else [route["output"]]
|
||||
output_types = route["output_type"] if isinstance(route["output_type"], list) else [route["output_type"]]
|
||||
output_names = route["output_name"] if isinstance(route["output_name"], list) else [route["output_name"]]
|
||||
|
||||
# Check lengths match
|
||||
if not len(outputs) == len(output_types) == len(output_names):
|
||||
raise ValueError(f"Route output, output_type and output_name must have same length: {route}")
|
||||
|
||||
# Condition is always a Jinja2 template — validate it
|
||||
if not self._validate_template(self._env, route["condition"]):
|
||||
condition_value = route["condition"]
|
||||
if not isinstance(condition_value, str):
|
||||
raise ValueError(
|
||||
f"Invalid template for condition: {condition_value!r} (type: {type(condition_value).__name__})."
|
||||
f"Condition must be a string representing a valid Jinja2 template. "
|
||||
f"For example, use {str(condition_value)!r} instead of {condition_value!r}."
|
||||
)
|
||||
raise ValueError(f"Invalid template for condition: {condition_value}")
|
||||
|
||||
# Only validate output as Jinja2 template when output_passthrough is False (default)
|
||||
output_passthrough = route.get("output_passthrough", False)
|
||||
if not output_passthrough:
|
||||
for output in outputs:
|
||||
if not self._validate_template(self._env, output):
|
||||
if not isinstance(output, str):
|
||||
raise ValueError(
|
||||
f"Invalid template for output: {output!r} (type: {type(output).__name__}). "
|
||||
f"Output must be a string representing a valid Jinja2 template. "
|
||||
f"For example, use {str(output)!r} instead of {output!r}."
|
||||
)
|
||||
raise ValueError(f"Invalid template for output: {output}")
|
||||
|
||||
@staticmethod
|
||||
def _extract_variables(env: Environment, templates: list[str]) -> set[str]:
|
||||
"""
|
||||
Extracts all variables from a list of Jinja template strings.
|
||||
|
||||
:param env: A Jinja environment.
|
||||
:param templates: A list of Jinja template strings.
|
||||
:returns: A set of variable names.
|
||||
"""
|
||||
variables = set()
|
||||
for template in templates:
|
||||
assigned_variables, template_variables = _extract_template_variables_and_assignments(
|
||||
env=env, template=template
|
||||
)
|
||||
variables.update(template_variables - assigned_variables)
|
||||
return variables
|
||||
|
||||
def _validate_template(self, env: Environment, template_text: str) -> bool:
|
||||
"""
|
||||
Validates a template string by parsing it with Jinja.
|
||||
|
||||
:param env: A Jinja environment.
|
||||
:param template_text: A Jinja template string.
|
||||
:returns: `True` if the template is valid, `False` otherwise.
|
||||
"""
|
||||
# Check if template_text is a string before attempting to parse
|
||||
if not isinstance(template_text, str):
|
||||
return False
|
||||
try:
|
||||
env.parse(template_text)
|
||||
return True
|
||||
except TemplateSyntaxError:
|
||||
return False
|
||||
|
||||
def _output_matches_type(self, value: Any, expected_type: type) -> bool: # noqa: PLR0911
|
||||
"""
|
||||
Checks whether `value` type matches the `expected_type`.
|
||||
"""
|
||||
# Handle Any type
|
||||
if expected_type is Any:
|
||||
return True
|
||||
|
||||
# Get the origin type (List, Dict, etc) and type arguments
|
||||
origin = get_origin(expected_type)
|
||||
args = get_args(expected_type)
|
||||
|
||||
# Handle basic types (int, str, etc)
|
||||
if origin is None:
|
||||
return isinstance(value, expected_type)
|
||||
|
||||
# Handle Sequence types (List, Tuple, etc)
|
||||
if isinstance(origin, type) and issubclass(origin, Sequence):
|
||||
if isinstance(value, (str, bytes)):
|
||||
return False
|
||||
if not isinstance(value, Sequence):
|
||||
return False
|
||||
# Empty sequence is valid
|
||||
if not value:
|
||||
return True
|
||||
# Check each element against the sequence's type parameter
|
||||
return all(self._output_matches_type(item, args[0]) for item in value)
|
||||
|
||||
# Handle Mapping types (Dict, etc)
|
||||
if isinstance(origin, type) and issubclass(origin, Mapping):
|
||||
if not isinstance(value, Mapping):
|
||||
return False
|
||||
# Empty mapping is valid
|
||||
if not value:
|
||||
return True
|
||||
key_type, value_type = args
|
||||
# Check all keys and values match their respective types
|
||||
return all(
|
||||
self._output_matches_type(k, key_type) and self._output_matches_type(v, value_type)
|
||||
for k, v in value.items()
|
||||
)
|
||||
|
||||
# Handle Union types (including Optional and X | Y syntax)
|
||||
if _is_union_type(origin):
|
||||
return any(self._output_matches_type(value, arg) for arg in args)
|
||||
|
||||
return False
|
||||
Reference in New Issue
Block a user