Coverage for haystack/components/routers/document_type_router.py: 100%

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 

9 

10from haystack import component 

11from haystack.dataclasses import Document 

12from haystack.utils.misc import _guess_mime_type 

13 

14 

15@component 

16class DocumentTypeRouter: 

17 """ 

18 Routes documents by their MIME types. 

19 

20 DocumentTypeRouter is used to dynamically route documents within a pipeline based on their MIME types. 

21 It supports exact MIME type matches and regex patterns. 

22 

23 MIME types can be extracted directly from document metadata or inferred from file paths using standard or 

24 user-supplied MIME type mappings. 

25 

26 ### Usage example 

27 

28 ```python 

29 from haystack.components.routers import DocumentTypeRouter 

30 from haystack.dataclasses import Document 

31 

32 docs = [ 

33 Document(content="Example text", meta={"file_path": "example.txt"}), 

34 Document(content="Another document", meta={"mime_type": "application/pdf"}), 

35 Document(content="Unknown type") 

36 ] 

37 

38 router = DocumentTypeRouter( 

39 mime_type_meta_field="mime_type", 

40 file_path_meta_field="file_path", 

41 mime_types=["text/plain", "application/pdf"] 

42 ) 

43 

44 result = router.run(documents=docs) 

45 print(result) 

46 ``` 

47 

48 Expected output: 

49 ```python 

50 { 

51 "text/plain": [Document(...)], 

52 "application/pdf": [Document(...)], 

53 "unclassified": [Document(...)] 

54 } 

55 ``` 

56 """ 

57 

58 def __init__( 

59 self, 

60 *, 

61 mime_types: list[str], 

62 mime_type_meta_field: str | None = None, 

63 file_path_meta_field: str | None = None, 

64 additional_mimetypes: dict[str, str] | None = None, 

65 ) -> None: 

66 """ 

67 Initialize the DocumentTypeRouter component. 

68 

69 :param mime_types: 

70 A list of MIME types or regex patterns to classify the input documents. 

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

72 :param mime_type_meta_field: 

73 Optional name of the metadata field that holds the MIME type. 

74 :param file_path_meta_field: 

75 Optional name of the metadata field that holds the file path. Used to infer the MIME type if 

76 `mime_type_meta_field` is not provided or missing in a document. 

77 :param additional_mimetypes: 

78 Optional dictionary mapping MIME types to file extensions to enhance or override the standard 

79 `mimetypes` module. Useful when working with uncommon or custom file types. 

80 For example: `{"application/vnd.custom-type": ".custom"}`. 

81 

82 :raises ValueError: If `mime_types` is empty, uses the reserved name `"unclassified"`, or if both 

83 `mime_type_meta_field` and `file_path_meta_field` are not provided. 

84 """ 

85 if not mime_types: 

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

87 

88 if "unclassified" in mime_types: 

89 raise ValueError( 

90 "'unclassified' is a reserved output name in DocumentTypeRouter for documents " 

91 "whose MIME type does not match any rule. Rename the MIME type to something else." 

92 ) 

93 

94 if mime_type_meta_field is None and file_path_meta_field is None: 

95 raise ValueError( 

96 "At least one of 'mime_type_meta_field' or 'file_path_meta_field' must be provided to determine MIME " 

97 "types." 

98 ) 

99 self.mime_type_meta_field = mime_type_meta_field 

100 self.file_path_meta_field = file_path_meta_field 

101 

102 if additional_mimetypes: 

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

104 mimetypes.add_type(mime, ext) 

105 

106 self._mime_type_patterns = [] 

107 for mime_type in mime_types: 

108 try: 

109 pattern = re.compile(mime_type) 

110 except re.error as e: 

111 raise ValueError(f"Invalid regex pattern '{mime_type}'.") from e 

112 self._mime_type_patterns.append(pattern) 

113 

114 component.set_output_types(self, unclassified=list[Document], **dict.fromkeys(mime_types, list[Document])) 

115 self.mime_types = mime_types 

116 self.additional_mimetypes = additional_mimetypes 

117 

118 def run(self, documents: list[Document]) -> dict[str, list[Document]]: 

119 """ 

120 Categorize input documents into groups based on their MIME type. 

121 

122 MIME types can either be directly available in document metadata or derived from file paths using the 

123 standard Python `mimetypes` module and custom mappings. 

124 

125 :param documents: 

126 A list of documents to be categorized. 

127 

128 :returns: 

129 A dictionary where the keys are MIME types (or `"unclassified"`) and the values are lists of documents. 

130 """ 

131 mime_types = defaultdict(list) 

132 

133 for doc in documents: 

134 mime_type = doc.meta.get(self.mime_type_meta_field) if self.mime_type_meta_field else None 

135 file_path = doc.meta.get(self.file_path_meta_field) if self.file_path_meta_field else None 

136 

137 if mime_type is None and file_path: 

138 # if mime_type is not provided, try to guess it from the file path 

139 mime_type = _guess_mime_type(Path(file_path)) 

140 

141 matched = False 

142 if mime_type: 

143 for bucket_key, pattern in zip(self.mime_types, self._mime_type_patterns, strict=True): 

144 # Match an exact MIME type first, so literal types containing regex 

145 # metacharacters (e.g. the '+' in 'image/svg+xml') are not misread as 

146 # a regex; fall back to regex matching for patterns like 'audio/.*'. 

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

148 mime_types[bucket_key].append(doc) 

149 matched = True 

150 break 

151 if not matched: 

152 mime_types["unclassified"].append(doc) 

153 

154 return dict(mime_types)