File size: 4,605 Bytes
7ff7119 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | """PDF loader 3-szintű fallback-kel.
Réteg 1: PyMuPDF (fitz) — natív szövegkinyerés. Ha az oldal text >= SCANNED_THRESHOLD
karakter, ennyi elég, oldal "digitális".
Réteg 2: Tesseract OCR — ha az oldal natív szöveg < SCANNED_THRESHOLD, az oldalt
renderelt képpé alakítjuk és átadjuk az `ocr_image_bytes`-nak. Ha az
eredmény >= SCANNED_THRESHOLD, az oldalt is_scanned=False-ra állítjuk
és az OCR szöveget használjuk.
Réteg 3: Vision fallback — ha az OCR is < SCANNED_THRESHOLD, az oldal `is_scanned=True`
marad és a `image_bytes` mezőben tartjuk a renderelt képet. A downstream
extract subgraph `vision_extract_node` közvetlenül a képből (LLM vision)
nyer ki strukturált adatot.
Mindezt szinkronban csináljuk (PyMuPDF blocking C-binding), de a subgraph-ban
`asyncio.to_thread()` wrapper-rel hívjuk.
"""
from __future__ import annotations
from io import BytesIO
from graph.states.pipeline_state import IngestedDocument, PageContent
from ingest.ocr import SCANNED_THRESHOLD, ocr_image_bytes, tesseract_available
from ingest.tables import extract_tables_markdown
# Render DPI a vision-fallback-hez (200 DPI elég jó minőség Claude vision-nek)
RENDER_DPI = 200
def load_pdf(file_name: str, file_bytes: bytes) -> IngestedDocument:
"""Egy PDF betöltése IngestedDocument-té.
Args:
file_name: a fájl neve (a metadata-hoz)
file_bytes: a PDF bináris tartalma
Raises:
Az alsóbb rétegek hibái fel vannak fogva (try/except), de ha a PyMuPDF
kifejezetten nem tudja megnyitni a fájlt, akkor RuntimeError-ral fail-fast.
"""
import fitz # PyMuPDF
try:
pdf_doc = fitz.open(stream=file_bytes, filetype="pdf")
except Exception as e:
raise RuntimeError(f"Nem sikerult megnyitni a PDF-et: {file_name}: {e}") from e
try:
pages: list[PageContent] = []
any_scanned = False
ocr_enabled = tesseract_available()
for page_idx, page in enumerate(pdf_doc, start=1):
# 1. réteg: PyMuPDF natív
native_text = (page.get_text() or "").strip()
if len(native_text) >= SCANNED_THRESHOLD:
# Digitális oldal — natív szöveg elég
pages.append(PageContent(
page_number=page_idx,
text=native_text,
is_scanned=False,
image_bytes=None,
))
continue
# 2-3. réteg: oldalt képpé renderelünk
try:
pix = page.get_pixmap(dpi=RENDER_DPI)
image_bytes = pix.tobytes("png")
except Exception:
# Render fail — natív szöveggel megyünk tovább, mégha kevés is
pages.append(PageContent(
page_number=page_idx,
text=native_text,
is_scanned=True, # gyenge minőségű
image_bytes=None,
))
any_scanned = True
continue
# 2. réteg: Tesseract OCR (ha telepítve van)
ocr_text = ocr_image_bytes(image_bytes) if ocr_enabled else ""
if len(ocr_text) >= SCANNED_THRESHOLD:
# OCR sikerült — a natív szöveg helyett ezt használjuk
pages.append(PageContent(
page_number=page_idx,
text=ocr_text,
is_scanned=False,
image_bytes=image_bytes, # vision-extract opcionálisan használhatja
))
continue
# 3. réteg: vision fallback — a downstream extract LLM-vision-nel nyer ki
pages.append(PageContent(
page_number=page_idx,
text=native_text or ocr_text, # ami van (gyengébb), full_text-be megy RAG-hoz
is_scanned=True,
image_bytes=image_bytes,
))
any_scanned = True
# Aggregált full_text RAG-hoz
full_text = "\n\n".join(p.text for p in pages if p.text)
# Táblázatok kinyerése pdfplumber-rel (ha van)
tables_md, table_count = extract_tables_markdown(file_bytes)
return IngestedDocument(
file_name=file_name,
file_type="pdf",
pages=pages,
full_text=full_text,
tables_markdown=tables_md,
table_count=table_count,
is_scanned=any_scanned,
)
finally:
pdf_doc.close()
|