Coverage for haystack/components/routers/file_type_router.py: 92%

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import mimetypes 

6import re 

7from collections import defaultdict 

8from pathlib import Path 

9from typing import Any 

10 

11from haystack import component, default_from_dict, default_to_dict, logging 

12from haystack.components.converters.utils import get_bytestream_from_source, normalize_metadata 

13from haystack.dataclasses import ByteStream 

14 

15from haystack.utils.misc import _guess_mime_type # ruff: isort: skip 

16 

17# We import CUSTOM_MIMETYPES here to prevent breaking change from moving to haystack.utils.misc 

18from haystack.utils.misc import CUSTOM_MIMETYPES # noqa: F401 

19 

20logger = logging.getLogger(__name__) 

21 

22 

23@component 

24class FileTypeRouter: 

25 """ 

26 Categorizes files or byte streams by their MIME types, helping in context-based routing. 

27 

28 FileTypeRouter supports both exact MIME type matching and regex patterns. 

29 

30 For file paths, MIME types come from extensions; byte streams use metadata. 

31 Each entry in `mime_types` is matched against a source's MIME type by exact equality first, 

32 falling back to regex `fullmatch` if equality misses. So `"image/svg+xml"` routes 

33 `image/svg+xml` streams correctly via the equality check (without `+` being interpreted as a 

34 regex quantifier), and patterns like `"audio/.*"` keep matching every audio subtype. 

35 

36 ### Usage example 

37 

38 ```python 

39 from haystack.components.routers import FileTypeRouter 

40 from pathlib import Path 

41 

42 # Exact MIME matching — `+`-containing IANA types like image/svg+xml work correctly 

43 router = FileTypeRouter(mime_types=["text/plain", "application/pdf", "image/svg+xml"]) 

44 

45 # Regex matching — catch every audio subtype 

46 router_with_regex = FileTypeRouter(mime_types=[r"audio/.*", r"text/plain"]) 

47 

48 sources = [Path("file.txt"), Path("document.pdf"), Path("song.mp3")] 

49 print(router.run(sources=sources)) 

50 print(router_with_regex.run(sources=sources)) 

51 

52 # Expected output: 

53 # {'text/plain': [ 

54 # PosixPath('file.txt')], 'application/pdf': [PosixPath('document.pdf')], 'unclassified': [PosixPath('song.mp3') 

55 # ]} 

56 # {'audio/.*': [ 

57 # PosixPath('song.mp3')], 'text/plain': [PosixPath('file.txt')], 'unclassified': [PosixPath('document.pdf') 

58 # ]} 

59 ``` 

60 """ 

61 

62 def __init__( 

63 self, mime_types: list[str], additional_mimetypes: dict[str, str] | None = None, raise_on_failure: bool = False 

64 ) -> None: 

65 """ 

66 Initialize the FileTypeRouter component. 

67 

68 :param mime_types: 

69 A list of MIME types or regex patterns to classify the input files or byte streams. 

70 (for example: `["text/plain", "audio/x-wav", "image/jpeg"]`). 

71 

72 :param additional_mimetypes: 

73 A dictionary containing the MIME type to add to the mimetypes package to prevent unsupported or non-native 

74 packages from being unclassified. 

75 (for example: `{"application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx"}`). 

76 

77 :param raise_on_failure: 

78 If True, raises FileNotFoundError when a file path doesn't exist. 

79 If False (default), only emits a warning when a file path doesn't exist. 

80 """ 

81 if not mime_types: 

82 raise ValueError("The list of mime types cannot be empty.") 

83 

84 if additional_mimetypes: 

85 for mime, ext in additional_mimetypes.items(): 

86 mimetypes.add_type(mime, ext) 

87 

88 self.mime_type_patterns = [] 

89 for mime_type in mime_types: 

90 try: 

91 pattern = re.compile(mime_type) 

92 except re.error as e: 

93 raise ValueError(f"Invalid MIME type or regex pattern '{mime_type}'.") from e 

94 self.mime_type_patterns.append(pattern) 

95 

96 # the actual output type is list[Union[Path, ByteStream]], 

97 # but this would cause PipelineConnectError with Converters 

98 component.set_output_types( 

99 self, 

100 unclassified=list[str | Path | ByteStream], 

101 failed=list[str | Path | ByteStream], 

102 **dict.fromkeys(mime_types, list[str | Path | ByteStream]), 

103 ) 

104 self.mime_types = mime_types 

105 self._additional_mimetypes = additional_mimetypes 

106 self._raise_on_failure = raise_on_failure 

107 

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

109 """ 

110 Serializes the component to a dictionary. 

111 

112 :returns: 

113 Dictionary with serialized data. 

114 """ 

115 return default_to_dict( 

116 self, 

117 mime_types=self.mime_types, 

118 additional_mimetypes=self._additional_mimetypes, 

119 raise_on_failure=self._raise_on_failure, 

120 ) 

121 

122 @classmethod 

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

124 """ 

125 Deserializes the component from a dictionary. 

126 

127 :param data: 

128 The dictionary to deserialize from. 

129 :returns: 

130 The deserialized component. 

131 """ 

132 return default_from_dict(cls, data) 

133 

134 def run( 

135 self, sources: list[str | Path | ByteStream], meta: dict[str, Any] | list[dict[str, Any]] | None = None 

136 ) -> dict[str, list[ByteStream | Path]]: 

137 """ 

138 Categorize files or byte streams according to their MIME types. 

139 

140 :param sources: 

141 A list of file paths or byte streams to categorize. 

142 

143 :param meta: 

144 Optional metadata to attach to the sources. 

145 When provided, the sources are internally converted to ByteStream objects and the metadata is added. 

146 This value can be a list of dictionaries or a single dictionary. 

147 If it's a single dictionary, its content is added to the metadata of all ByteStream objects. 

148 If it's a list, its length must match the number of sources, as they are zipped together. 

149 

150 :returns: A dictionary where the keys are MIME types and the values are lists of data sources. 

151 Two extra keys may be returned: `"unclassified"` when a source's MIME type doesn't match any pattern 

152 and `"failed"` when a source cannot be processed (for example, a file path that doesn't exist). 

153 :raises TypeError: If a source is not a Path, str, or ByteStream. 

154 """ 

155 

156 mime_types: defaultdict[str, list[Path | ByteStream]] = defaultdict(list) 

157 meta_list = normalize_metadata(meta=meta, sources_count=len(sources)) 

158 

159 for source, meta_dict in zip(sources, meta_list, strict=True): 

160 if isinstance(source, str): 

161 source = Path(source) 

162 

163 if isinstance(source, Path): 

164 if not source.exists(): 

165 if self._raise_on_failure: 

166 raise FileNotFoundError(f"File not found: {source}") 

167 logger.warning("File not found: {source}. Skipping it.", source=source) 

168 mime_types["failed"].append(source) 

169 continue 

170 

171 mime_type = _guess_mime_type(source) 

172 

173 elif isinstance(source, ByteStream): 

174 mime_type = source.mime_type 

175 else: 

176 raise TypeError(f"Unsupported data source type: {type(source).__name__}") 

177 

178 # If we have metadata, we convert the source to ByteStream and add the metadata 

179 if meta_dict: 

180 try: 

181 source = get_bytestream_from_source(source) 

182 except Exception as e: 

183 if self._raise_on_failure: 

184 raise e 

185 logger.warning("Could not read {source}. Skipping it. Error: {error}", source=source, error=e) 

186 mime_types["failed"].append(source) 

187 continue 

188 

189 source.meta.update(meta_dict) 

190 

191 matched = False 

192 if mime_type: 

193 # Try exact equality first so MIMEs containing regex metacharacters (e.g. the `+` in 

194 # `image/svg+xml`) match themselves before the regex fallback gets a chance to misread them. 

195 for bucket_key, pattern in zip(self.mime_types, self.mime_type_patterns, strict=True): 

196 if mime_type == bucket_key or pattern.fullmatch(mime_type): 

197 mime_types[bucket_key].append(source) 

198 matched = True 

199 break 

200 if not matched: 

201 mime_types["unclassified"].append(source) 

202 

203 return dict(mime_types)