Coverage for src/local_deep_research/document_loaders/bounded_pdf.py: 97%

1"""Bound PDF text extraction for byte-backed document loaders. 

2 

3RAG uploads and Zotero sync both call ``extract_text_from_bytes``. The 

4registry's ``PyPDFLoader.load()`` materializes text from every PDF page, so 

5those entry points need the same ceilings as the other extraction paths. 

6Pypdf still flattens the page tree before the first page is yielded, and one 

7``extract_text()`` call cannot be interrupted without process isolation. 

8""" 

9 

10from __future__ import annotations 

11 

12import io 

13import time 

14 

15from langchain_core.documents import Document 

16from loguru import logger 

17from pypdf import PdfReader 

18 

19from ..utilities import pdf_extraction_limits as limits 

20from ..utilities.pdf_stream_release import PypdfWalkReleaser 

21from ..utilities.resource_utils import safe_close 

22 

23 

24def load_bounded_pdf(content: bytes, source: str) -> list[Document]: 

25 """Return page documents, truncating at the shared extraction ceilings.""" 

26 deadline = time.thread_time() + limits.MAX_PDF_EXTRACTION_CPU_SECONDS 

27 documents: list[Document] = [] 

28 retained_chars = 0 

29 text_pages = 0 

30 truncated_by: str | None = None 

31 

32 with io.BytesIO(content) as stream: 

33 reader = PdfReader(stream) 

34 try: 

35 # Pypdf materializes its page tree here; the limit below bounds 

36 # text extraction, not the construction of that tree. 

37 total_pages = len(reader.pages) 

38 releaser = PypdfWalkReleaser(reader) 

39 for page_number in range( 

40 min(total_pages, limits.MAX_PDF_EXTRACTION_PAGES) 

41 ): 

42 if time.thread_time() >= deadline: 

43 truncated_by = "CPU-time" 

44 break 

45 separator_chars = 2 if text_pages else 0 

46 remaining = ( 

47 limits.MAX_PDF_EXTRACTED_CHARS 

48 - retained_chars 

49 - separator_chars 

50 ) 

51 if remaining <= 0: 

52 truncated_by = "character" 

53 break 

54 

55 page = reader.pages[page_number] 

56 releaser.before_page() 

57 try: 

58 page_text = (page.extract_text() or "").strip() 

59 finally: 

60 # The reader otherwise retains decoded streams and parsed 

61 # objects after each page has yielded its text. 

62 releaser.after_page() 

63 

64 if page_text: 64 ↛ 69line 64 didn't jump to line 69 because the condition on line 64 was always true

65 kept_text = page_text[:remaining] 

66 retained_chars += separator_chars + len(kept_text) 

67 text_pages += 1 

68 else: 

69 kept_text = "" 

70 

71 documents.append( 

72 Document( 

73 page_content=kept_text, 

74 metadata={ 

75 "source": source, 

76 "page": page_number, 

77 "total_pages": total_pages, 

78 }, 

79 ) 

80 ) 

81 if len(kept_text) < len(page_text): 

82 truncated_by = "character" 

83 break 

84 if ( 

85 retained_chars >= limits.MAX_PDF_EXTRACTED_CHARS 

86 and page_number + 1 < total_pages 

87 ): 

88 truncated_by = "character" 

89 break 

90 if ( 

91 truncated_by is None 

92 and total_pages > limits.MAX_PDF_EXTRACTION_PAGES 

93 ): 

94 truncated_by = "page" 

95 finally: 

96 safe_close(reader, "PDF byte-loader reader") 

97 

98 if truncated_by: 

99 for document in documents: 

100 document.metadata["truncated"] = True 

101 logger.warning( 

102 "PDF document-loader extraction stopped at {} ceiling " 

103 "after {} of {} pages", 

104 truncated_by, 

105 len(documents), 

106 total_pages, 

107 ) 

108 return documents