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,27 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import sys
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from lazy_imports import LazyImporter
|
||||
|
||||
_import_structure = {
|
||||
"conditional_router": ["ConditionalRouter"],
|
||||
"document_length_router": ["DocumentLengthRouter"],
|
||||
"document_type_router": ["DocumentTypeRouter"],
|
||||
"file_type_router": ["FileTypeRouter"],
|
||||
"llm_messages_router": ["LLMMessagesRouter"],
|
||||
"metadata_router": ["MetadataRouter"],
|
||||
}
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from .conditional_router import ConditionalRouter as ConditionalRouter
|
||||
from .document_length_router import DocumentLengthRouter as DocumentLengthRouter
|
||||
from .document_type_router import DocumentTypeRouter as DocumentTypeRouter
|
||||
from .file_type_router import FileTypeRouter as FileTypeRouter
|
||||
from .llm_messages_router import LLMMessagesRouter as LLMMessagesRouter
|
||||
from .metadata_router import MetadataRouter as MetadataRouter
|
||||
else:
|
||||
sys.modules[__name__] = LazyImporter(name=__name__, module_file=__file__, import_structure=_import_structure)
|
||||
@@ -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
|
||||
@@ -0,0 +1,76 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
|
||||
from haystack import component
|
||||
from haystack.dataclasses import Document
|
||||
|
||||
|
||||
@component
|
||||
class DocumentLengthRouter:
|
||||
"""
|
||||
Categorizes documents based on the length of the `content` field and routes them to the appropriate output.
|
||||
|
||||
A common use case for DocumentLengthRouter is handling documents obtained from PDFs that contain non-text
|
||||
content, such as scanned pages or images. This component can detect empty or low-content documents and route them to
|
||||
components that perform OCR, generate captions, or compute image embeddings.
|
||||
|
||||
### Usage example
|
||||
|
||||
```python
|
||||
from haystack.components.routers import DocumentLengthRouter
|
||||
from haystack.dataclasses import Document
|
||||
|
||||
docs = [
|
||||
Document(content="Short"),
|
||||
Document(content="Long document "*20),
|
||||
]
|
||||
|
||||
router = DocumentLengthRouter(threshold=10)
|
||||
|
||||
result = router.run(documents=docs)
|
||||
print(result)
|
||||
|
||||
# {
|
||||
# "short_documents": [Document(content="Short", ...)],
|
||||
# "long_documents": [Document(content="Long document ...", ...)],
|
||||
# }
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, *, threshold: int = 10) -> None:
|
||||
"""
|
||||
Initialize the DocumentLengthRouter component.
|
||||
|
||||
:param threshold:
|
||||
The threshold for the number of characters in the document `content` field. Documents where `content` is
|
||||
None or whose character count is less than or equal to the threshold will be routed to the `short_documents`
|
||||
output. Otherwise, they will be routed to the `long_documents` output.
|
||||
To route only documents with None content to `short_documents`, set the threshold to a negative number.
|
||||
"""
|
||||
self.threshold = threshold
|
||||
|
||||
@component.output_types(short_documents=list[Document], long_documents=list[Document])
|
||||
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
|
||||
"""
|
||||
Categorize input documents into groups based on the length of the `content` field.
|
||||
|
||||
:param documents:
|
||||
A list of documents to be categorized.
|
||||
|
||||
:returns: A dictionary with the following keys:
|
||||
- `short_documents`: A list of documents where `content` is None or the length of `content` is less than or
|
||||
equal to the threshold.
|
||||
- `long_documents`: A list of documents where the length of `content` is greater than the threshold.
|
||||
"""
|
||||
short_documents = []
|
||||
long_documents = []
|
||||
|
||||
for doc in documents:
|
||||
if doc.content is None or len(doc.content) <= self.threshold:
|
||||
short_documents.append(doc)
|
||||
else:
|
||||
long_documents.append(doc)
|
||||
|
||||
return {"short_documents": short_documents, "long_documents": long_documents}
|
||||
@@ -0,0 +1,145 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import mimetypes
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
|
||||
from haystack import component
|
||||
from haystack.dataclasses import Document
|
||||
from haystack.utils.misc import _guess_mime_type
|
||||
|
||||
|
||||
@component
|
||||
class DocumentTypeRouter:
|
||||
"""
|
||||
Routes documents by their MIME types.
|
||||
|
||||
DocumentTypeRouter is used to dynamically route documents within a pipeline based on their MIME types.
|
||||
It supports exact MIME type matches and regex patterns.
|
||||
|
||||
MIME types can be extracted directly from document metadata or inferred from file paths using standard or
|
||||
user-supplied MIME type mappings.
|
||||
|
||||
### Usage example
|
||||
|
||||
```python
|
||||
from haystack.components.routers import DocumentTypeRouter
|
||||
from haystack.dataclasses import Document
|
||||
|
||||
docs = [
|
||||
Document(content="Example text", meta={"file_path": "example.txt"}),
|
||||
Document(content="Another document", meta={"mime_type": "application/pdf"}),
|
||||
Document(content="Unknown type")
|
||||
]
|
||||
|
||||
router = DocumentTypeRouter(
|
||||
mime_type_meta_field="mime_type",
|
||||
file_path_meta_field="file_path",
|
||||
mime_types=["text/plain", "application/pdf"]
|
||||
)
|
||||
|
||||
result = router.run(documents=docs)
|
||||
print(result)
|
||||
```
|
||||
|
||||
Expected output:
|
||||
```python
|
||||
{
|
||||
"text/plain": [Document(...)],
|
||||
"application/pdf": [Document(...)],
|
||||
"unclassified": [Document(...)]
|
||||
}
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
mime_types: list[str],
|
||||
mime_type_meta_field: str | None = None,
|
||||
file_path_meta_field: str | None = None,
|
||||
additional_mimetypes: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Initialize the DocumentTypeRouter component.
|
||||
|
||||
:param mime_types:
|
||||
A list of MIME types or regex patterns to classify the input documents.
|
||||
(for example: `["text/plain", "audio/x-wav", "image/jpeg"]`).
|
||||
:param mime_type_meta_field:
|
||||
Optional name of the metadata field that holds the MIME type.
|
||||
:param file_path_meta_field:
|
||||
Optional name of the metadata field that holds the file path. Used to infer the MIME type if
|
||||
`mime_type_meta_field` is not provided or missing in a document.
|
||||
:param additional_mimetypes:
|
||||
Optional dictionary mapping MIME types to file extensions to enhance or override the standard
|
||||
`mimetypes` module. Useful when working with uncommon or custom file types.
|
||||
For example: `{"application/vnd.custom-type": ".custom"}`.
|
||||
|
||||
:raises ValueError: If `mime_types` is empty or if both `mime_type_meta_field` and `file_path_meta_field` are
|
||||
not provided.
|
||||
"""
|
||||
if not mime_types:
|
||||
raise ValueError("The list of mime types cannot be empty.")
|
||||
|
||||
if mime_type_meta_field is None and file_path_meta_field is None:
|
||||
raise ValueError(
|
||||
"At least one of 'mime_type_meta_field' or 'file_path_meta_field' must be provided to determine MIME "
|
||||
"types."
|
||||
)
|
||||
self.mime_type_meta_field = mime_type_meta_field
|
||||
self.file_path_meta_field = file_path_meta_field
|
||||
|
||||
if additional_mimetypes:
|
||||
for mime, ext in additional_mimetypes.items():
|
||||
mimetypes.add_type(mime, ext)
|
||||
|
||||
self._mime_type_patterns = []
|
||||
for mime_type in mime_types:
|
||||
try:
|
||||
pattern = re.compile(mime_type)
|
||||
except re.error as e:
|
||||
raise ValueError(f"Invalid regex pattern '{mime_type}'.") from e
|
||||
self._mime_type_patterns.append(pattern)
|
||||
|
||||
component.set_output_types(self, unclassified=list[Document], **dict.fromkeys(mime_types, list[Document]))
|
||||
self.mime_types = mime_types
|
||||
self.additional_mimetypes = additional_mimetypes
|
||||
|
||||
def run(self, documents: list[Document]) -> dict[str, list[Document]]:
|
||||
"""
|
||||
Categorize input documents into groups based on their MIME type.
|
||||
|
||||
MIME types can either be directly available in document metadata or derived from file paths using the
|
||||
standard Python `mimetypes` module and custom mappings.
|
||||
|
||||
:param documents:
|
||||
A list of documents to be categorized.
|
||||
|
||||
:returns:
|
||||
A dictionary where the keys are MIME types (or `"unclassified"`) and the values are lists of documents.
|
||||
"""
|
||||
mime_types = defaultdict(list)
|
||||
|
||||
for doc in documents:
|
||||
mime_type = doc.meta.get(self.mime_type_meta_field) if self.mime_type_meta_field else None
|
||||
file_path = doc.meta.get(self.file_path_meta_field) if self.file_path_meta_field else None
|
||||
|
||||
if mime_type is None and file_path:
|
||||
# if mime_type is not provided, try to guess it from the file path
|
||||
mime_type = _guess_mime_type(Path(file_path))
|
||||
|
||||
matched = False
|
||||
if mime_type:
|
||||
for pattern in self._mime_type_patterns:
|
||||
if pattern.fullmatch(mime_type):
|
||||
mime_types[pattern.pattern].append(doc)
|
||||
matched = True
|
||||
break
|
||||
if not matched:
|
||||
mime_types["unclassified"].append(doc)
|
||||
|
||||
return dict(mime_types)
|
||||
@@ -0,0 +1,203 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import mimetypes
|
||||
import re
|
||||
from collections import defaultdict
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
from haystack import component, default_from_dict, default_to_dict, logging
|
||||
from haystack.components.converters.utils import get_bytestream_from_source, normalize_metadata
|
||||
from haystack.dataclasses import ByteStream
|
||||
|
||||
from haystack.utils.misc import _guess_mime_type # ruff: isort: skip
|
||||
|
||||
# We import CUSTOM_MIMETYPES here to prevent breaking change from moving to haystack.utils.misc
|
||||
from haystack.utils.misc import CUSTOM_MIMETYPES # noqa: F401
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@component
|
||||
class FileTypeRouter:
|
||||
"""
|
||||
Categorizes files or byte streams by their MIME types, helping in context-based routing.
|
||||
|
||||
FileTypeRouter supports both exact MIME type matching and regex patterns.
|
||||
|
||||
For file paths, MIME types come from extensions; byte streams use metadata.
|
||||
Each entry in `mime_types` is matched against a source's MIME type by exact equality first,
|
||||
falling back to regex `fullmatch` if equality misses. So `"image/svg+xml"` routes
|
||||
`image/svg+xml` streams correctly via the equality check (without `+` being interpreted as a
|
||||
regex quantifier), and patterns like `"audio/.*"` keep matching every audio subtype.
|
||||
|
||||
### Usage example
|
||||
|
||||
```python
|
||||
from haystack.components.routers import FileTypeRouter
|
||||
from pathlib import Path
|
||||
|
||||
# Exact MIME matching — `+`-containing IANA types like image/svg+xml work correctly
|
||||
router = FileTypeRouter(mime_types=["text/plain", "application/pdf", "image/svg+xml"])
|
||||
|
||||
# Regex matching — catch every audio subtype
|
||||
router_with_regex = FileTypeRouter(mime_types=[r"audio/.*", r"text/plain"])
|
||||
|
||||
sources = [Path("file.txt"), Path("document.pdf"), Path("song.mp3")]
|
||||
print(router.run(sources=sources))
|
||||
print(router_with_regex.run(sources=sources))
|
||||
|
||||
# Expected output:
|
||||
# {'text/plain': [
|
||||
# PosixPath('file.txt')], 'application/pdf': [PosixPath('document.pdf')], 'unclassified': [PosixPath('song.mp3')
|
||||
# ]}
|
||||
# {'audio/.*': [
|
||||
# PosixPath('song.mp3')], 'text/plain': [PosixPath('file.txt')], 'unclassified': [PosixPath('document.pdf')
|
||||
# ]}
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, mime_types: list[str], additional_mimetypes: dict[str, str] | None = None, raise_on_failure: bool = False
|
||||
) -> None:
|
||||
"""
|
||||
Initialize the FileTypeRouter component.
|
||||
|
||||
:param mime_types:
|
||||
A list of MIME types or regex patterns to classify the input files or byte streams.
|
||||
(for example: `["text/plain", "audio/x-wav", "image/jpeg"]`).
|
||||
|
||||
:param additional_mimetypes:
|
||||
A dictionary containing the MIME type to add to the mimetypes package to prevent unsupported or non-native
|
||||
packages from being unclassified.
|
||||
(for example: `{"application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx"}`).
|
||||
|
||||
:param raise_on_failure:
|
||||
If True, raises FileNotFoundError when a file path doesn't exist.
|
||||
If False (default), only emits a warning when a file path doesn't exist.
|
||||
"""
|
||||
if not mime_types:
|
||||
raise ValueError("The list of mime types cannot be empty.")
|
||||
|
||||
if additional_mimetypes:
|
||||
for mime, ext in additional_mimetypes.items():
|
||||
mimetypes.add_type(mime, ext)
|
||||
|
||||
self.mime_type_patterns = []
|
||||
for mime_type in mime_types:
|
||||
try:
|
||||
pattern = re.compile(mime_type)
|
||||
except re.error as e:
|
||||
raise ValueError(f"Invalid MIME type or regex pattern '{mime_type}'.") from e
|
||||
self.mime_type_patterns.append(pattern)
|
||||
|
||||
# the actual output type is list[Union[Path, ByteStream]],
|
||||
# but this would cause PipelineConnectError with Converters
|
||||
component.set_output_types(
|
||||
self,
|
||||
unclassified=list[str | Path | ByteStream],
|
||||
failed=list[str | Path | ByteStream],
|
||||
**dict.fromkeys(mime_types, list[str | Path | ByteStream]),
|
||||
)
|
||||
self.mime_types = mime_types
|
||||
self._additional_mimetypes = additional_mimetypes
|
||||
self._raise_on_failure = raise_on_failure
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""
|
||||
Serializes the component to a dictionary.
|
||||
|
||||
:returns:
|
||||
Dictionary with serialized data.
|
||||
"""
|
||||
return default_to_dict(
|
||||
self,
|
||||
mime_types=self.mime_types,
|
||||
additional_mimetypes=self._additional_mimetypes,
|
||||
raise_on_failure=self._raise_on_failure,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "FileTypeRouter":
|
||||
"""
|
||||
Deserializes the component from a dictionary.
|
||||
|
||||
:param data:
|
||||
The dictionary to deserialize from.
|
||||
:returns:
|
||||
The deserialized component.
|
||||
"""
|
||||
return default_from_dict(cls, data)
|
||||
|
||||
def run(
|
||||
self, sources: list[str | Path | ByteStream], meta: dict[str, Any] | list[dict[str, Any]] | None = None
|
||||
) -> dict[str, list[ByteStream | Path]]:
|
||||
"""
|
||||
Categorize files or byte streams according to their MIME types.
|
||||
|
||||
:param sources:
|
||||
A list of file paths or byte streams to categorize.
|
||||
|
||||
:param meta:
|
||||
Optional metadata to attach to the sources.
|
||||
When provided, the sources are internally converted to ByteStream objects and the metadata is added.
|
||||
This value can be a list of dictionaries or a single dictionary.
|
||||
If it's a single dictionary, its content is added to the metadata of all ByteStream objects.
|
||||
If it's a list, its length must match the number of sources, as they are zipped together.
|
||||
|
||||
:returns: A dictionary where the keys are MIME types and the values are lists of data sources.
|
||||
Two extra keys may be returned: `"unclassified"` when a source's MIME type doesn't match any pattern
|
||||
and `"failed"` when a source cannot be processed (for example, a file path that doesn't exist).
|
||||
:raises TypeError: If a source is not a Path, str, or ByteStream.
|
||||
"""
|
||||
|
||||
mime_types: defaultdict[str, list[Path | ByteStream]] = defaultdict(list)
|
||||
meta_list = normalize_metadata(meta=meta, sources_count=len(sources))
|
||||
|
||||
for source, meta_dict in zip(sources, meta_list, strict=True):
|
||||
if isinstance(source, str):
|
||||
source = Path(source)
|
||||
|
||||
if isinstance(source, Path):
|
||||
if not source.exists():
|
||||
if self._raise_on_failure:
|
||||
raise FileNotFoundError(f"File not found: {source}")
|
||||
logger.warning("File not found: {source}. Skipping it.", source=source)
|
||||
mime_types["failed"].append(source)
|
||||
continue
|
||||
|
||||
mime_type = _guess_mime_type(source)
|
||||
|
||||
elif isinstance(source, ByteStream):
|
||||
mime_type = source.mime_type
|
||||
else:
|
||||
raise TypeError(f"Unsupported data source type: {type(source).__name__}")
|
||||
|
||||
# If we have metadata, we convert the source to ByteStream and add the metadata
|
||||
if meta_dict:
|
||||
try:
|
||||
source = get_bytestream_from_source(source)
|
||||
except Exception as e:
|
||||
if self._raise_on_failure:
|
||||
raise e
|
||||
logger.warning("Could not read {source}. Skipping it. Error: {error}", source=source, error=e)
|
||||
mime_types["failed"].append(source)
|
||||
continue
|
||||
|
||||
source.meta.update(meta_dict)
|
||||
|
||||
matched = False
|
||||
if mime_type:
|
||||
# Try exact equality first so MIMEs containing regex metacharacters (e.g. the `+` in
|
||||
# `image/svg+xml`) match themselves before the regex fallback gets a chance to misread them.
|
||||
for bucket_key, pattern in zip(self.mime_types, self.mime_type_patterns, strict=True):
|
||||
if mime_type == bucket_key or pattern.fullmatch(mime_type):
|
||||
mime_types[bucket_key].append(source)
|
||||
matched = True
|
||||
break
|
||||
if not matched:
|
||||
mime_types["unclassified"].append(source)
|
||||
|
||||
return dict(mime_types)
|
||||
@@ -0,0 +1,233 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
import re
|
||||
from typing import Any
|
||||
|
||||
from haystack import component, default_from_dict, default_to_dict
|
||||
from haystack.components.generators.chat.types import ChatGenerator
|
||||
from haystack.core.serialization import component_to_dict
|
||||
from haystack.dataclasses import ChatMessage, ChatRole
|
||||
from haystack.utils import deserialize_chatgenerator_inplace
|
||||
from haystack.utils.async_utils import _execute_component_async
|
||||
|
||||
|
||||
@component
|
||||
class LLMMessagesRouter:
|
||||
"""
|
||||
Routes Chat Messages to different connections using a generative Language Model to perform classification.
|
||||
|
||||
This component can be used with general-purpose LLMs and with specialized LLMs for moderation like Llama Guard.
|
||||
|
||||
### Usage example
|
||||
|
||||
```python
|
||||
from haystack.components.generators.chat import HuggingFaceAPIChatGenerator
|
||||
from haystack.components.routers.llm_messages_router import LLMMessagesRouter
|
||||
from haystack.dataclasses import ChatMessage
|
||||
|
||||
# initialize a Chat Generator with a generative model for moderation
|
||||
chat_generator = HuggingFaceAPIChatGenerator(
|
||||
api_type="serverless_inference_api",
|
||||
api_params={"model": "openai/gpt-oss-safeguard-20b", "provider": "groq"},
|
||||
)
|
||||
|
||||
router = LLMMessagesRouter(chat_generator=chat_generator,
|
||||
output_names=["unsafe", "safe"],
|
||||
output_patterns=["unsafe", "safe"])
|
||||
|
||||
|
||||
print(router.run([ChatMessage.from_user("How to rob a bank?")]))
|
||||
|
||||
# {
|
||||
# 'chat_generator_text': 'unsafe\\nS2',
|
||||
# 'unsafe': [
|
||||
# ChatMessage(
|
||||
# _role=<ChatRole.USER: 'user'>,
|
||||
# _content=[TextContent(text='How to rob a bank?')],
|
||||
# _name=None,
|
||||
# _meta={}
|
||||
# )
|
||||
# ]
|
||||
# }
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
chat_generator: ChatGenerator,
|
||||
output_names: list[str],
|
||||
output_patterns: list[str],
|
||||
system_prompt: str | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Initialize the LLMMessagesRouter component.
|
||||
|
||||
:param chat_generator: A ChatGenerator instance which represents the LLM.
|
||||
:param output_names: A list of output connection names. These can be used to connect the router to other
|
||||
components.
|
||||
:param output_patterns: A list of regular expressions to be matched against the output of the LLM. Each pattern
|
||||
corresponds to an output name. Patterns are evaluated in order.
|
||||
When using moderation models, refer to the model card to understand the expected outputs.
|
||||
:param system_prompt: An optional system prompt to customize the behavior of the LLM.
|
||||
For moderation models, refer to the model card for supported customization options.
|
||||
|
||||
:raises ValueError: If output_names and output_patterns are not non-empty lists of the same length.
|
||||
"""
|
||||
if not output_names or not output_patterns or len(output_names) != len(output_patterns):
|
||||
raise ValueError("`output_names` and `output_patterns` must be non-empty lists of the same length")
|
||||
|
||||
self._chat_generator = chat_generator
|
||||
self._system_prompt = system_prompt
|
||||
self._output_names = output_names
|
||||
self._output_patterns = output_patterns
|
||||
|
||||
self._compiled_patterns = [re.compile(pattern) for pattern in output_patterns]
|
||||
|
||||
component.set_output_types(
|
||||
self, **{"chat_generator_text": str, **dict.fromkeys(output_names + ["unmatched"], list[ChatMessage])}
|
||||
)
|
||||
|
||||
def warm_up(self) -> None:
|
||||
"""Warm up the underlying chat generator."""
|
||||
if hasattr(self._chat_generator, "warm_up"):
|
||||
self._chat_generator.warm_up()
|
||||
|
||||
async def warm_up_async(self) -> None:
|
||||
"""Warm up the underlying chat generator on the serving event loop."""
|
||||
if hasattr(self._chat_generator, "warm_up_async"):
|
||||
await self._chat_generator.warm_up_async()
|
||||
elif hasattr(self._chat_generator, "warm_up"):
|
||||
self._chat_generator.warm_up()
|
||||
|
||||
def close(self) -> None:
|
||||
"""Release the underlying chat generator's resources."""
|
||||
if hasattr(self._chat_generator, "close"):
|
||||
self._chat_generator.close()
|
||||
|
||||
async def close_async(self) -> None:
|
||||
"""Release the underlying chat generator's async resources."""
|
||||
if hasattr(self._chat_generator, "close_async"):
|
||||
await self._chat_generator.close_async()
|
||||
elif hasattr(self._chat_generator, "close"):
|
||||
self._chat_generator.close()
|
||||
|
||||
def run(self, messages: list[ChatMessage]) -> dict[str, str | list[ChatMessage]]:
|
||||
"""
|
||||
Classify the messages based on LLM output and route them to the appropriate output connection.
|
||||
|
||||
:param messages: A list of ChatMessages to be routed. Only user and assistant messages are supported.
|
||||
|
||||
:returns: A dictionary with the following keys:
|
||||
- "chat_generator_text": The text output of the LLM, useful for debugging.
|
||||
- "output_names": Each contains the list of messages that matched the corresponding pattern.
|
||||
- "unmatched": The messages that did not match any of the output patterns.
|
||||
|
||||
:raises ValueError: If messages is an empty list or contains messages with unsupported roles.
|
||||
"""
|
||||
if not messages:
|
||||
raise ValueError("`messages` must be a non-empty list.")
|
||||
if not all(message.is_from(ChatRole.USER) or message.is_from(ChatRole.ASSISTANT) for message in messages):
|
||||
msg = (
|
||||
"`messages` must contain only user and assistant messages. To customize the behavior of the "
|
||||
"`chat_generator`, you can use the `system_prompt` parameter."
|
||||
)
|
||||
raise ValueError(msg)
|
||||
|
||||
self.warm_up()
|
||||
|
||||
messages_for_inference = []
|
||||
if self._system_prompt:
|
||||
messages_for_inference.append(ChatMessage.from_system(self._system_prompt))
|
||||
messages_for_inference.extend(messages)
|
||||
|
||||
chat_generator_text = self._chat_generator.run(messages=messages_for_inference)["replies"][0].text
|
||||
|
||||
output = {"chat_generator_text": chat_generator_text}
|
||||
|
||||
for output_name, pattern in zip(self._output_names, self._compiled_patterns, strict=True):
|
||||
if pattern.search(chat_generator_text):
|
||||
output[output_name] = messages
|
||||
break
|
||||
else:
|
||||
output["unmatched"] = messages
|
||||
|
||||
return output
|
||||
|
||||
async def run_async(self, messages: list[ChatMessage]) -> dict[str, str | list[ChatMessage]]:
|
||||
"""
|
||||
Asynchronously classify the messages based on LLM output and route them to the appropriate output connection.
|
||||
|
||||
This is the asynchronous version of the `run` method. It has the same parameters and return values
|
||||
but can be used with `await` in an async code. If the chat generator only implements a synchronous
|
||||
`run` method, it is executed in a thread to avoid blocking the event loop.
|
||||
|
||||
:param messages: A list of ChatMessages to be routed. Only user and assistant messages are supported.
|
||||
|
||||
:returns: A dictionary with the following keys:
|
||||
- "chat_generator_text": The text output of the LLM, useful for debugging.
|
||||
- "output_names": Each contains the list of messages that matched the corresponding pattern.
|
||||
- "unmatched": The messages that did not match any of the output patterns.
|
||||
|
||||
:raises ValueError: If messages is an empty list or contains messages with unsupported roles.
|
||||
"""
|
||||
if not messages:
|
||||
raise ValueError("`messages` must be a non-empty list.")
|
||||
if not all(message.is_from(ChatRole.USER) or message.is_from(ChatRole.ASSISTANT) for message in messages):
|
||||
msg = (
|
||||
"`messages` must contain only user and assistant messages. To customize the behavior of the "
|
||||
"`chat_generator`, you can use the `system_prompt` parameter."
|
||||
)
|
||||
raise ValueError(msg)
|
||||
|
||||
await self.warm_up_async()
|
||||
|
||||
messages_for_inference = []
|
||||
if self._system_prompt:
|
||||
messages_for_inference.append(ChatMessage.from_system(self._system_prompt))
|
||||
messages_for_inference.extend(messages)
|
||||
|
||||
generator_result = await _execute_component_async(self._chat_generator, messages=messages_for_inference)
|
||||
chat_generator_text = generator_result["replies"][0].text
|
||||
|
||||
output = {"chat_generator_text": chat_generator_text}
|
||||
|
||||
for output_name, pattern in zip(self._output_names, self._compiled_patterns, strict=True):
|
||||
if pattern.search(chat_generator_text):
|
||||
output[output_name] = messages
|
||||
break
|
||||
else:
|
||||
output["unmatched"] = messages
|
||||
|
||||
return output
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""
|
||||
Serialize this component to a dictionary.
|
||||
|
||||
:returns:
|
||||
The serialized component as a dictionary.
|
||||
"""
|
||||
return default_to_dict(
|
||||
self,
|
||||
chat_generator=component_to_dict(obj=self._chat_generator, name="chat_generator"),
|
||||
output_names=self._output_names,
|
||||
output_patterns=self._output_patterns,
|
||||
system_prompt=self._system_prompt,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "LLMMessagesRouter":
|
||||
"""
|
||||
Deserialize this component from a dictionary.
|
||||
|
||||
:param data:
|
||||
The dictionary representation of this component.
|
||||
:returns:
|
||||
The deserialized component instance.
|
||||
"""
|
||||
if data["init_parameters"].get("chat_generator"):
|
||||
deserialize_chatgenerator_inplace(data["init_parameters"], key="chat_generator")
|
||||
|
||||
return default_from_dict(cls, data)
|
||||
@@ -0,0 +1,163 @@
|
||||
# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
|
||||
#
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
|
||||
from typing import Any
|
||||
|
||||
from haystack import Document, component, default_from_dict, default_to_dict
|
||||
from haystack.dataclasses import ByteStream
|
||||
from haystack.utils import deserialize_type, serialize_type
|
||||
from haystack.utils.filters import document_matches_filter
|
||||
|
||||
|
||||
@component
|
||||
class MetadataRouter:
|
||||
"""
|
||||
Routes documents or byte streams to different connections based on their metadata fields.
|
||||
|
||||
Specify the routing rules in the `init` method.
|
||||
If a document or byte stream does not match any of the rules, it's routed to a connection named "unmatched".
|
||||
|
||||
|
||||
### Usage examples
|
||||
|
||||
**Routing Documents by metadata:**
|
||||
```python
|
||||
from haystack import Document
|
||||
from haystack.components.routers import MetadataRouter
|
||||
|
||||
docs = [Document(content="Paris is the capital of France.", meta={"language": "en"}),
|
||||
Document(content="Berlin ist die Haupststadt von Deutschland.", meta={"language": "de"})]
|
||||
|
||||
router = MetadataRouter(rules={"en": {"field": "meta.language", "operator": "==", "value": "en"}})
|
||||
|
||||
print(router.run(documents=docs))
|
||||
# {'en': [Document(id=..., content: 'Paris is the capital of France.', meta: {'language': 'en'})],
|
||||
# 'unmatched': [Document(id=..., content: 'Berlin ist die Haupststadt von Deutschland.', meta: {'language': 'de'})]}
|
||||
```
|
||||
|
||||
**Routing ByteStreams by metadata:**
|
||||
```python
|
||||
from haystack.dataclasses import ByteStream
|
||||
from haystack.components.routers import MetadataRouter
|
||||
|
||||
streams = [
|
||||
ByteStream.from_string("Hello world", meta={"language": "en"}),
|
||||
ByteStream.from_string("Bonjour le monde", meta={"language": "fr"})
|
||||
]
|
||||
|
||||
router = MetadataRouter(
|
||||
rules={"english": {"field": "meta.language", "operator": "==", "value": "en"}},
|
||||
output_type=list[ByteStream]
|
||||
)
|
||||
|
||||
result = router.run(documents=streams)
|
||||
# {'english': [ByteStream(...)], 'unmatched': [ByteStream(...)]}
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, rules: dict[str, dict], output_type: type = list[Document]) -> None:
|
||||
"""
|
||||
Initializes the MetadataRouter component.
|
||||
|
||||
:param rules: A dictionary defining how to route documents or byte streams to output connections based on their
|
||||
metadata. Keys are output connection names, and values are dictionaries of
|
||||
[filtering expressions](https://docs.haystack.deepset.ai/docs/metadata-filtering) in Haystack.
|
||||
For example:
|
||||
```python
|
||||
{
|
||||
"edge_1": {
|
||||
"operator": "AND",
|
||||
"conditions": [
|
||||
{"field": "meta.created_at", "operator": ">=", "value": "2023-01-01"},
|
||||
{"field": "meta.created_at", "operator": "<", "value": "2023-04-01"},
|
||||
],
|
||||
},
|
||||
"edge_2": {
|
||||
"operator": "AND",
|
||||
"conditions": [
|
||||
{"field": "meta.created_at", "operator": ">=", "value": "2023-04-01"},
|
||||
{"field": "meta.created_at", "operator": "<", "value": "2023-07-01"},
|
||||
],
|
||||
},
|
||||
"edge_3": {
|
||||
"operator": "AND",
|
||||
"conditions": [
|
||||
{"field": "meta.created_at", "operator": ">=", "value": "2023-07-01"},
|
||||
{"field": "meta.created_at", "operator": "<", "value": "2023-10-01"},
|
||||
],
|
||||
},
|
||||
"edge_4": {
|
||||
"operator": "AND",
|
||||
"conditions": [
|
||||
{"field": "meta.created_at", "operator": ">=", "value": "2023-10-01"},
|
||||
{"field": "meta.created_at", "operator": "<", "value": "2024-01-01"},
|
||||
],
|
||||
},
|
||||
}
|
||||
```
|
||||
:param output_type: The type of the output produced. Lists of Documents or ByteStreams can be specified.
|
||||
"""
|
||||
self.rules = rules
|
||||
self.output_type = output_type
|
||||
for rule in self.rules.values():
|
||||
if "operator" not in rule:
|
||||
raise ValueError(
|
||||
"Invalid filter syntax. See https://docs.haystack.deepset.ai/docs/metadata-filtering for details."
|
||||
)
|
||||
component.set_output_types(self, unmatched=self.output_type, **dict.fromkeys(rules, self.output_type))
|
||||
|
||||
def run(self, documents: list[Document] | list[ByteStream]) -> dict[str, list[Document] | list[ByteStream]]:
|
||||
"""
|
||||
Routes documents or byte streams to different connections based on their metadata fields.
|
||||
|
||||
If a document or byte stream does not match any of the rules, it's routed to a connection named "unmatched".
|
||||
|
||||
:param documents: A list of `Document` or `ByteStream` objects to be routed based on their metadata.
|
||||
|
||||
:returns: A dictionary where the keys are the names of the output connections (including `"unmatched"`)
|
||||
and the values are lists of `Document` or `ByteStream` objects that matched the corresponding rules.
|
||||
"""
|
||||
|
||||
unmatched: list[Document] | list[ByteStream] = []
|
||||
output: dict[str, list[Document] | list[ByteStream]] = {edge: [] for edge in self.rules}
|
||||
|
||||
for doc_or_bytestream in documents:
|
||||
current_obj_matched = False
|
||||
for edge, rule in self.rules.items():
|
||||
if document_matches_filter(filters=rule, document=doc_or_bytestream):
|
||||
# we need to ignore the arg-type here because the underlying
|
||||
# filter methods use type Union[Document, ByteStream]
|
||||
output[edge].append(doc_or_bytestream) # type: ignore[arg-type]
|
||||
current_obj_matched = True
|
||||
|
||||
if not current_obj_matched:
|
||||
unmatched.append(doc_or_bytestream) # type: ignore[arg-type]
|
||||
|
||||
output["unmatched"] = unmatched
|
||||
return output
|
||||
|
||||
def to_dict(self) -> dict[str, Any]:
|
||||
"""
|
||||
Serialize this component to a dictionary.
|
||||
|
||||
:returns:
|
||||
The serialized component as a dictionary.
|
||||
"""
|
||||
return default_to_dict(self, rules=self.rules, output_type=serialize_type(self.output_type))
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> "MetadataRouter":
|
||||
"""
|
||||
Deserialize this component from a dictionary.
|
||||
|
||||
:param data:
|
||||
The dictionary representation of this component.
|
||||
:returns:
|
||||
The deserialized component instance.
|
||||
"""
|
||||
init_params = data.get("init_parameters", {})
|
||||
if "output_type" in init_params:
|
||||
# Deserialize the output_type to its original type
|
||||
init_params["output_type"] = deserialize_type(init_params["output_type"])
|
||||
return default_from_dict(cls, data)
|
||||
Reference in New Issue
Block a user