malepati/custom_template_working

1import mimetypes2from fastapi import HTTPException3import os, asyncio4from typing import List, Optional, Tuple5import pdfplumber6 7from constants.documents import (8    PDF_MIME_TYPES,9    POWERPOINT_TYPES,10    TEXT_MIME_TYPES,11    WORD_TYPES,12)13from services.docling_service import DoclingService14 15 16class DocumentsLoader:17 18    def __init__(self, file_paths: List[str]):19        self._file_paths = file_paths20 21        self.docling_service = DoclingService()22 23        self._documents: List[str] = []24        self._images: List[List[str]] = []25 26    @property27    def documents(self):28        return self._documents29 30    @property31    def images(self):32        return self._images33 34    async def load_documents(35        self,36        temp_dir: Optional[str] = None,37        load_text: bool = True,38        load_images: bool = False,39    ):40        """If load_images is True, temp_dir must be provided"""41 42        documents: List[str] = []43        images: List[str] = []44 45        for file_path in self._file_paths:46            if not os.path.exists(file_path):47                raise HTTPException(48                    status_code=404, detail=f"File {file_path} not found"49                )50 51            document = ""52            imgs = []53 54            mime_type = mimetypes.guess_type(file_path)[0]55            if mime_type in PDF_MIME_TYPES:56                document, imgs = await self.load_pdf(57                    file_path, load_text, load_images, temp_dir58                )59            elif mime_type in TEXT_MIME_TYPES:60                document = await self.load_text(file_path)61            elif mime_type in POWERPOINT_TYPES:62                document = self.load_powerpoint(file_path)63            elif mime_type in WORD_TYPES:64                document = self.load_msword(file_path)65 66            documents.append(document)67            images.append(imgs)68 69        self._documents = documents70        self._images = images71 72    async def load_pdf(73        self,74        file_path: str,75        load_text: bool,76        load_images: bool,77        temp_dir: Optional[str] = None,78    ) -> Tuple[str, List[str]]:79        image_paths = []80        document: str = ""81 82        if load_text:83            document = self.docling_service.parse_to_markdown(file_path)84 85        if load_images:86            image_paths = await self.get_page_images_from_pdf_async(file_path, temp_dir)87 88        return document, image_paths89 90    async def load_text(self, file_path: str) -> str:91        with open(file_path, "r") as file:92            return await asyncio.to_thread(file.read)93 94    def load_msword(self, file_path: str) -> str:95        return self.docling_service.parse_to_markdown(file_path)96 97    def load_powerpoint(self, file_path: str) -> str:98        return self.docling_service.parse_to_markdown(file_path)99 100    @classmethod101    def get_page_images_from_pdf(cls, file_path: str, temp_dir: str) -> List[str]:102        with pdfplumber.open(file_path) as pdf:103            images = []104            for page in pdf.pages:105                img = page.to_image(resolution=150)106                image_path = os.path.join(temp_dir, f"page_{page.page_number}.png")107                img.save(image_path)108                images.append(image_path)109            return images110 111    @classmethod112    async def get_page_images_from_pdf_async(cls, file_path: str, temp_dir: str):113        return await asyncio.to_thread(114            cls.get_page_images_from_pdf, file_path, temp_dir115        )116