main
py 44 lines 1.78 KB
Raw
1 """PDF parser with PyMuPDF primary and Tesseract OCR fallback."""
2
3 import os
4
5 from helpers.print_style import PrintStyle
6 from plugins._document_query.helpers.fetch import FetchedDocument
7
8 from .base import BaseParser
9
10
11 class PdfParser(BaseParser):
12 mimetypes = ["application/pdf"]
13
14 def _parse_sync(self, document: FetchedDocument, config: dict) -> str:
15 with document.local_file() as file_path:
16 if not os.path.exists(file_path):
17 raise ValueError(f"Temporary file not found: {file_path}")
18 contents = self._parse_with_pymupdf(file_path)
19 if not contents:
20 if not config.get("pdf_ocr_fallback", True):
21 raise ValueError("PyMuPDF returned no content and OCR fallback is disabled")
22 contents = self._parse_with_ocr(file_path)
23 return contents
24
25 def _parse_with_pymupdf(self, file_path: str) -> str:
26 from langchain_community.document_loaders.pdf import PyMuPDFLoader
27 from langchain_community.document_loaders.parsers.images import TesseractBlobParser
28
29 try:
30 loader = PyMuPDFLoader(
31 file_path, mode="single", extract_tables="markdown",
32 extract_images=True, images_inner_format="text",
33 images_parser=TesseractBlobParser(), pages_delimiter="\n",
34 )
35 return "\n".join(e.page_content for e in loader.load())
36 except Exception as e:
37 PrintStyle.error(f"PyMuPDF parsing failed: {e}")
38 return ""
39
40 def _parse_with_ocr(self, file_path: str) -> str:
41 import pdf2image, pytesseract
42 PrintStyle.debug(f"FALLBACK: OCR for {file_path}")
43 pages = pdf2image.convert_from_path(file_path)
44 return "\n\n".join(pytesseract.image_to_string(p) for p in pages)