Underground-Digital/Workflow-Engine

1"""Abstract interface for document loader implementations."""2 3from collections.abc import Iterator4from typing import Optional5 6from core.rag.extractor.blob.blob import Blob7from core.rag.extractor.extractor_base import BaseExtractor8from core.rag.models.document import Document9from extensions.ext_storage import storage10 11 12class PdfExtractor(BaseExtractor):13    """Load pdf files.14 15 16    Args:17        file_path: Path to the file to load.18    """19 20    def __init__(self, file_path: str, file_cache_key: Optional[str] = None):21        """Initialize with file path."""22        self._file_path = file_path23        self._file_cache_key = file_cache_key24 25    def extract(self) -> list[Document]:26        plaintext_file_key = ""27        plaintext_file_exists = False28        if self._file_cache_key:29            try:30                text = storage.load(self._file_cache_key).decode("utf-8")31                plaintext_file_exists = True32                return [Document(page_content=text)]33            except FileNotFoundError:34                pass35        documents = list(self.load())36        text_list = []37        for document in documents:38            text_list.append(document.page_content)39        text = "\n\n".join(text_list)40 41        # save plaintext file for caching42        if not plaintext_file_exists and plaintext_file_key:43            storage.save(plaintext_file_key, text.encode("utf-8"))44 45        return documents46 47    def load(48        self,49    ) -> Iterator[Document]:50        """Lazy load given path as pages."""51        blob = Blob.from_path(self._file_path)52        yield from self.parse(blob)53 54    def parse(self, blob: Blob) -> Iterator[Document]:55        """Lazily parse the blob."""56        import pypdfium257 58        with blob.as_bytes_io() as file_path:59            pdf_reader = pypdfium2.PdfDocument(file_path, autoclose=True)60            try:61                for page_number, page in enumerate(pdf_reader):62                    text_page = page.get_textpage()63                    content = text_page.get_text_range()64                    text_page.close()65                    page.close()66                    metadata = {"source": blob.source, "page": page_number}67                    yield Document(page_content=content, metadata=metadata)68            finally:69                pdf_reader.close()70