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"]`). 

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 or if both `mime_type_meta_field` and `file_path_meta_field` are 

83 not provided. 

84 """ 

85 if not mime_types: 

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

87 

88 if mime_type_meta_field is None and file_path_meta_field is None: 

89 raise ValueError( 

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

91 "types." 

92 ) 

93 self.mime_type_meta_field = mime_type_meta_field 

94 self.file_path_meta_field = file_path_meta_field 

95 

96 if additional_mimetypes: 

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

98 mimetypes.add_type(mime, ext) 

99 

100 self._mime_type_patterns = [] 

101 for mime_type in mime_types: 

102 try: 

103 pattern = re.compile(mime_type) 

104 except re.error as e: 

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

106 self._mime_type_patterns.append(pattern) 

107 

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

109 self.mime_types = mime_types 

110 self.additional_mimetypes = additional_mimetypes 

111 

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

113 """ 

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

115 

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

117 standard Python `mimetypes` module and custom mappings. 

118 

119 :param documents: 

120 A list of documents to be categorized. 

121 

122 :returns: 

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

124 """ 

125 mime_types = defaultdict(list) 

126 

127 for doc in documents: 

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

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

130 

131 if mime_type is None and file_path: 

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

133 mime_type = _guess_mime_type(Path(file_path)) 

134 

135 matched = False 

136 if mime_type: 

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

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

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

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

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

142 mime_types[bucket_key].append(doc) 

143 matched = True 

144 break 

145 if not matched: 

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

147 

148 return dict(mime_types)