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

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

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5 

6from haystack import component 

7from haystack.dataclasses import Document 

8 

9 

10@component 

11class DocumentLengthRouter: 

12 """ 

13 Categorizes documents based on the length of the `content` field and routes them to the appropriate output. 

14 

15 A common use case for DocumentLengthRouter is handling documents obtained from PDFs that contain non-text 

16 content, such as scanned pages or images. This component can detect empty or low-content documents and route them to 

17 components that perform OCR, generate captions, or compute image embeddings. 

18 

19 ### Usage example 

20 

21 ```python 

22 from haystack.components.routers import DocumentLengthRouter 

23 from haystack.dataclasses import Document 

24 

25 docs = [ 

26 Document(content="Short"), 

27 Document(content="Long document "*20), 

28 ] 

29 

30 router = DocumentLengthRouter(threshold=10) 

31 

32 result = router.run(documents=docs) 

33 print(result) 

34 

35 # { 

36 # "short_documents": [Document(content="Short", ...)], 

37 # "long_documents": [Document(content="Long document ...", ...)], 

38 # } 

39 ``` 

40 """ 

41 

42 def __init__(self, *, threshold: int = 10) -> None: 

43 """ 

44 Initialize the DocumentLengthRouter component. 

45 

46 :param threshold: 

47 The threshold for the number of characters in the document `content` field. Documents where `content` is 

48 None or whose character count is less than or equal to the threshold will be routed to the `short_documents` 

49 output. Otherwise, they will be routed to the `long_documents` output. 

50 To route only documents with None content to `short_documents`, set the threshold to a negative number. 

51 """ 

52 self.threshold = threshold 

53 

54 @component.output_types(short_documents=list[Document], long_documents=list[Document]) 

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

56 """ 

57 Categorize input documents into groups based on the length of the `content` field. 

58 

59 :param documents: 

60 A list of documents to be categorized. 

61 

62 :returns: A dictionary with the following keys: 

63 - `short_documents`: A list of documents where `content` is None or the length of `content` is less than or 

64 equal to the threshold. 

65 - `long_documents`: A list of documents where the length of `content` is greater than the threshold. 

66 """ 

67 short_documents = [] 

68 long_documents = [] 

69 

70 for doc in documents: 

71 if doc.content is None or len(doc.content) <= self.threshold: 

72 short_documents.append(doc) 

73 else: 

74 long_documents.append(doc) 

75 

76 return {"short_documents": short_documents, "long_documents": long_documents}