Coverage for open_webui/retrieval/loaders/pdf.py: 11%

1import datetime as dt 

2import io 

3import logging 

4from pathlib import Path 

5 

6from langchain_core.document_loaders import BaseLoader 

7from langchain_core.documents import Document 

8 

9log = logging.getLogger(__name__) 

10 

11 

12class PDFLoader(BaseLoader): 

13 def __init__(self, file_path, *, extract_images=False, mode='page'): 

14 if mode not in ('single', 'page'): 

15 raise ValueError("PDF mode must be 'single' or 'page'") 

16 self.file_path = str(Path(file_path).expanduser()) 

17 self.extract_images = extract_images 

18 self.mode = mode 

19 self.ocr = None 

20 

21 def lazy_load(self): 

22 from pypdf import PdfReader 

23 

24 with open(self.file_path, 'rb') as file: 

25 reader = PdfReader(file) 

26 metadata = {'producer': 'PyPDF', 'creator': 'PyPDF', 'creationdate': ''} 

27 for key, value in (reader.metadata or {}).items(): 

28 key = key.removeprefix('/').lower() 

29 value = value if type(value) in (str, int) else str(value) 

30 if key in ('creationdate', 'moddate') and isinstance(value, str): 

31 try: 

32 value = dt.datetime.strptime(value.replace("'", ''), 'D:%Y%m%d%H%M%S%z').isoformat() 

33 except ValueError: 

34 pass 

35 metadata[key] = ( 

36 value.strip() 

37 if isinstance(value, str) and key not in ('creationdate', 'moddate', 'page_count', 'file_path') 

38 else value 

39 ) 

40 metadata.update(source=self.file_path, total_pages=len(reader.pages)) 

41 labels = reader.page_labels if self.mode == 'page' else None 

42 texts = [] 

43 for index, page in enumerate(reader.pages): 

44 text = page.extract_text() 

45 if self.extract_images: 

46 image_text = self._extract_images(page) 

47 if image_text: 

48 text = self._merge_image_text(text, image_text) 

49 text = text.strip() 

50 if self.mode == 'page': 

51 yield Document(page_content=text, metadata={**metadata, 'page': index, 'page_label': labels[index]}) 

52 else: 

53 texts.append(text) 

54 if self.mode == 'single': 

55 yield Document(page_content='\n\f'.join(texts), metadata=metadata) 

56 

57 @staticmethod 

58 def _merge_image_text(text, image_text): 

59 # Insert before the final paragraphs/footer where possible, matching existing chunks. 

60 position, separator = len(text), '\n\n' 

61 for _ in range(2): 

62 for delimiter in ('\n\n\n', '\n\n'): 

63 found = text.rfind(delimiter, 0, position) 

64 if found >= 0: 

65 position, separator = found, delimiter 

66 break 

67 else: 

68 break 

69 return text[:position] + separator + image_text + text[position:] 

70 

71 def _extract_images(self, page): 

72 import numpy as np 

73 from PIL import Image, UnidentifiedImageError 

74 

75 if '/Resources' not in page or '/XObject' not in page['/Resources']: 

76 return '' 

77 texts = [] 

78 xobjects = page['/Resources']['/XObject'] 

79 for name in xobjects: 

80 try: 

81 stream = xobjects[name] 

82 if stream.get('/Subtype') != '/Image': 

83 continue 

84 try: 

85 # Encoded images, including CMYK JPEGs, can go straight to Pillow. 

86 image = Image.open(io.BytesIO(stream.get_data())) 

87 except UnidentifiedImageError: 

88 image = stream.decode_as_image() 

89 pixels = np.array(image.convert('RGB')) 

90 except Exception as e: 

91 log.warning('Skipping unreadable PDF image %s: %s', name, e) 

92 continue 

93 

94 if self.ocr is None: 

95 from rapidocr import RapidOCR 

96 

97 self.ocr = RapidOCR() 

98 result = self.ocr(pixels) 

99 if result and result.txts: 

100 texts.append('\n'.join(result.txts).strip()) 

101 return '\n\n' + '\n'.join(filter(None, texts)) + '\n\n' if any(texts) else ''