Coverage for haystack/components/routers/file_type_router.py: 92%
78 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-10-09 14:34 +0000
1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
2#
3# SPDX-License-Identifier: Apache-2.0
5import mimetypes
6import re
7from collections import defaultdict
8from dataclasses import replace
9from pathlib import Path
10from typing import Any
12from haystack import component, default_from_dict, default_to_dict, logging
13from haystack.components.converters.utils import get_bytestream_from_source, normalize_metadata
14from haystack.dataclasses import ByteStream
16from haystack.utils.misc import _guess_mime_type # ruff: isort: skip
18# We import CUSTOM_MIMETYPES here to prevent breaking change from moving to haystack.utils.misc
19from haystack.utils.misc import CUSTOM_MIMETYPES # noqa: F401
21logger = logging.getLogger(__name__)
24@component
25class FileTypeRouter:
26 """
27 Categorizes files or byte streams by their MIME types, helping in context-based routing.
29 FileTypeRouter supports both exact MIME type matching and regex patterns.
31 For file paths, MIME types come from extensions; byte streams use metadata.
32 Each entry in `mime_types` is matched against a source's MIME type by exact equality first,
33 falling back to regex `fullmatch` if equality misses. So `"image/svg+xml"` routes
34 `image/svg+xml` streams correctly via the equality check (without `+` being interpreted as a
35 regex quantifier), and patterns like `"audio/.*"` keep matching every audio subtype.
37 ### Usage example
39 ```python
40 from haystack.components.routers import FileTypeRouter
41 from pathlib import Path
43 # Exact MIME matching — `+`-containing IANA types like image/svg+xml work correctly
44 router = FileTypeRouter(mime_types=["text/plain", "application/pdf", "image/svg+xml"])
46 # Regex matching — catch every audio subtype
47 router_with_regex = FileTypeRouter(mime_types=[r"audio/.*", r"text/plain"])
49 sources = [Path("file.txt"), Path("document.pdf"), Path("song.mp3")]
50 print(router.run(sources=sources))
51 print(router_with_regex.run(sources=sources))
53 # Expected output:
54 # {'text/plain': [
55 # PosixPath('file.txt')], 'application/pdf': [PosixPath('document.pdf')], 'unclassified': [PosixPath('song.mp3')
56 # ]}
57 # {'audio/.*': [
58 # PosixPath('song.mp3')], 'text/plain': [PosixPath('file.txt')], 'unclassified': [PosixPath('document.pdf')
59 # ]}
60 ```
61 """
63 def __init__(
64 self, mime_types: list[str], additional_mimetypes: dict[str, str] | None = None, raise_on_failure: bool = False
65 ) -> None:
66 """
67 Initialize the FileTypeRouter component.
69 :param mime_types:
70 A list of MIME types or regex patterns to classify the input files or byte streams.
71 (for example: `["text/plain", "audio/x-wav", "image/jpeg"]`).
72 `"unclassified"` and `"failed"` are reserved output names and cannot be used here.
74 :param additional_mimetypes:
75 A dictionary containing the MIME type to add to the mimetypes package to prevent unsupported or non-native
76 packages from being unclassified.
77 (for example: `{"application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx"}`).
79 :param raise_on_failure:
80 If True, raises FileNotFoundError when a file path doesn't exist.
81 If False (default), only emits a warning when a file path doesn't exist.
82 :raises ValueError:
83 If `mime_types` is empty, contains an invalid regex, or uses the reserved names
84 `"unclassified"` or `"failed"`.
85 """
86 if not mime_types:
87 raise ValueError("The list of mime types cannot be empty.")
89 reserved = {"unclassified", "failed"}
90 collisions = reserved.intersection(mime_types)
91 if collisions:
92 names = ", ".join(repr(name) for name in sorted(collisions))
93 raise ValueError(
94 f"MIME type(s) {names} are reserved output names in FileTypeRouter "
95 "('unclassified' for unmatched sources, 'failed' for unreadable files). "
96 "Rename the MIME type to something else."
97 )
99 if additional_mimetypes:
100 for mime, ext in additional_mimetypes.items():
101 mimetypes.add_type(mime, ext)
103 self.mime_type_patterns = []
104 for mime_type in mime_types:
105 try:
106 pattern = re.compile(mime_type)
107 except re.error as e:
108 raise ValueError(f"Invalid MIME type or regex pattern '{mime_type}'.") from e
109 self.mime_type_patterns.append(pattern)
111 # the actual output type is list[Union[Path, ByteStream]],
112 # but this would cause PipelineConnectError with Converters
113 component.set_output_types(
114 self,
115 unclassified=list[str | Path | ByteStream],
116 failed=list[str | Path | ByteStream],
117 **dict.fromkeys(mime_types, list[str | Path | ByteStream]),
118 )
119 self.mime_types = mime_types
120 self._additional_mimetypes = additional_mimetypes
121 self._raise_on_failure = raise_on_failure
123 def to_dict(self) -> dict[str, Any]:
124 """
125 Serializes the component to a dictionary.
127 :returns:
128 Dictionary with serialized data.
129 """
130 return default_to_dict(
131 self,
132 mime_types=self.mime_types,
133 additional_mimetypes=self._additional_mimetypes,
134 raise_on_failure=self._raise_on_failure,
135 )
137 @classmethod
138 def from_dict(cls, data: dict[str, Any]) -> "FileTypeRouter":
139 """
140 Deserializes the component from a dictionary.
142 :param data:
143 The dictionary to deserialize from.
144 :returns:
145 The deserialized component.
146 """
147 return default_from_dict(cls, data)
149 def run(
150 self, sources: list[str | Path | ByteStream], meta: dict[str, Any] | list[dict[str, Any]] | None = None
151 ) -> dict[str, list[ByteStream | Path]]:
152 """
153 Categorize files or byte streams according to their MIME types.
155 :param sources:
156 A list of file paths or byte streams to categorize.
158 :param meta:
159 Optional metadata to attach to the sources.
160 When provided, the sources are internally converted to ByteStream objects and the metadata is added.
161 This value can be a list of dictionaries or a single dictionary.
162 If it's a single dictionary, its content is added to the metadata of all ByteStream objects.
163 If it's a list, its length must match the number of sources, as they are zipped together.
165 :returns: A dictionary where the keys are MIME types and the values are lists of data sources.
166 Two extra keys may be returned: `"unclassified"` when a source's MIME type doesn't match any pattern
167 and `"failed"` when a source cannot be processed (for example, a file path that doesn't exist).
168 :raises TypeError: If a source is not a Path, str, or ByteStream.
169 """
171 mime_types: defaultdict[str, list[Path | ByteStream]] = defaultdict(list)
172 meta_list = normalize_metadata(meta=meta, sources_count=len(sources))
174 for source, meta_dict in zip(sources, meta_list, strict=True):
175 if isinstance(source, str):
176 source = Path(source)
178 if isinstance(source, Path):
179 if not source.exists():
180 if self._raise_on_failure:
181 raise FileNotFoundError(f"File not found: {source}")
182 logger.warning("File not found: {source}. Skipping it.", source=source)
183 mime_types["failed"].append(source)
184 continue
186 mime_type = _guess_mime_type(source)
188 elif isinstance(source, ByteStream):
189 mime_type = source.mime_type
190 else:
191 raise TypeError(f"Unsupported data source type: {type(source).__name__}")
193 # If we have metadata, we convert the source to ByteStream and add the metadata
194 if meta_dict:
195 try:
196 source = get_bytestream_from_source(source)
197 except Exception as e:
198 if self._raise_on_failure:
199 raise e
200 logger.warning("Could not read {source}. Skipping it. Error: {error}", source=source, error=e)
201 mime_types["failed"].append(source)
202 continue
204 # `get_bytestream_from_source` hands back a ByteStream source unchanged, so writing the
205 # metadata onto it would modify the object the caller still holds. Build a copy instead.
206 source = replace(source, meta={**source.meta, **meta_dict})
208 matched = False
209 if mime_type:
210 # Try exact equality first so MIMEs containing regex metacharacters (e.g. the `+` in
211 # `image/svg+xml`) match themselves before the regex fallback gets a chance to misread them.
212 for bucket_key, pattern in zip(self.mime_types, self.mime_type_patterns, strict=True):
213 if mime_type == bucket_key or pattern.fullmatch(mime_type):
214 mime_types[bucket_key].append(source)
215 matched = True
216 break
217 if not matched:
218 mime_types["unclassified"].append(source)
220 return dict(mime_types)