nvtitan/graphRAG

1"""2PDF Ingestion & Preprocessing Module3Handles extraction of text, tables, code blocks, and images from PDFs4"""5import fitz  # PyMuPDF6import pdfplumber7import pytesseract8from PIL import Image9import io10import re11from typing import List, Dict, Any, Optional, Tuple12from loguru import logger13from models import Chunk, ChunkType, PDFMetadata14from config import settings15import uuid16 17 18class PDFProcessor:19    """20    Comprehensive PDF processor that extracts:21    - Page-level text with character ranges22    - Tables (structured)23    - Code blocks (detected heuristically)24    - Images (with OCR)25    """26 27    def __init__(self):28        self.code_patterns = [29            re.compile(r'```[\s\S]*?```'),  # Markdown code blocks30            re.compile(r'def\s+\w+\s*\('),   # Python functions31            re.compile(r'class\s+\w+\s*[:\(]'),  # Python/Java classes32            re.compile(r'function\s+\w+\s*\('),  # JavaScript functions33            re.compile(r'public\s+class\s+\w+'),  # Java classes34        ]35 36    def process_pdf(self, filepath: str, pdf_id: str) -> Tuple[List[Chunk], PDFMetadata]:37        """38        Main entry point: process entire PDF and return chunks + metadata39 40        Args:41            filepath: Path to PDF file42            pdf_id: Unique identifier for this PDF43 44        Returns:45            Tuple of (chunks list, metadata object)46        """47        logger.info(f"Processing PDF: {filepath}")48 49        chunks: List[Chunk] = []50 51        # Open with PyMuPDF for text and images52        pdf_doc = fitz.open(filepath)53        num_pages = len(pdf_doc)54 55        # Open with pdfplumber for tables56        with pdfplumber.open(filepath) as plumber_pdf:57            for page_num in range(num_pages):58                logger.debug(f"Processing page {page_num + 1}/{num_pages}")59 60                # Extract from PyMuPDF61                fitz_page = pdf_doc[page_num]62                page_chunks = self._process_page(63                    fitz_page=fitz_page,64                    plumber_page=plumber_pdf.pages[page_num],65                    page_num=page_num + 1,  # 1-indexed66                    pdf_id=pdf_id67                )68                chunks.extend(page_chunks)69 70        pdf_doc.close()71 72        # Create metadata73        import os74        file_size = os.path.getsize(filepath)75        metadata = PDFMetadata(76            pdf_id=pdf_id,77            filename=os.path.basename(filepath),78            filepath=filepath,79            num_pages=num_pages,80            file_size_bytes=file_size,81            num_chunks=len(chunks),82            processing_status="completed"83        )84 85        logger.info(f"Extracted {len(chunks)} chunks from {num_pages} pages")86        return chunks, metadata87 88    def _process_page(89        self,90        fitz_page,91        plumber_page,92        page_num: int,93        pdf_id: str94    ) -> List[Chunk]:95        """Process a single page and return all chunks"""96        chunks: List[Chunk] = []97 98        # 1. Extract raw text with character positions99        page_text = fitz_page.get_text("text")100 101        # 2. Extract tables102        table_chunks = self._extract_tables(plumber_page, page_num, pdf_id)103        chunks.extend(table_chunks)104 105        # 3. Extract code blocks106        code_chunks = self._extract_code_blocks(page_text, page_num, pdf_id)107        chunks.extend(code_chunks)108 109        # 4. Extract images and run OCR110        image_chunks = self._extract_images(fitz_page, page_num, pdf_id)111        chunks.extend(image_chunks)112 113        # 5. Extract remaining text as paragraphs114        # Remove table and code regions from text before creating paragraph chunks115        cleaned_text = self._remove_extracted_regions(116            page_text,117            [c.text for c in code_chunks]118        )119 120        if cleaned_text.strip():121            para_chunk = Chunk(122                chunk_id=str(uuid.uuid4()),123                pdf_id=pdf_id,124                page_number=page_num,125                char_range=(0, len(cleaned_text)),126                type=ChunkType.PARAGRAPH,127                text=cleaned_text,128                metadata={"source": "text_extraction"}129            )130            chunks.append(para_chunk)131 132        return chunks133 134    def _extract_tables(self, plumber_page, page_num: int, pdf_id: str) -> List[Chunk]:135        """Extract tables from page using pdfplumber"""136        chunks = []137        tables = plumber_page.extract_tables()138 139        for idx, table in enumerate(tables):140            if not table:141                continue142 143            # Convert table to structured JSON144            table_json = self._table_to_json(table)145 146            # Convert table to text representation147            table_text = self._table_to_text(table)148 149            chunk = Chunk(150                chunk_id=str(uuid.uuid4()),151                pdf_id=pdf_id,152                page_number=page_num,153                char_range=(0, len(table_text)),154                type=ChunkType.TABLE,155                text=table_text,156                table_json=table_json,157                metadata={"table_index": idx, "num_rows": len(table)}158            )159            chunks.append(chunk)160 161        logger.debug(f"Extracted {len(chunks)} tables from page {page_num}")162        return chunks163 164    def _table_to_json(self, table: List[List[str]]) -> Dict[str, Any]:165        """Convert table to structured JSON"""166        if not table or len(table) < 2:167            return {"headers": [], "rows": []}168 169        headers = table[0]170        rows = table[1:]171 172        return {173            "headers": headers,174            "rows": [175                {headers[i]: cell for i, cell in enumerate(row) if i < len(headers)}176                for row in rows177            ]178        }179 180    def _table_to_text(self, table: List[List[str]]) -> str:181        """Convert table to readable text"""182        return "\n".join([" | ".join([str(cell) for cell in row]) for row in table])183 184    def _extract_code_blocks(self, text: str, page_num: int, pdf_id: str) -> List[Chunk]:185        """Extract code blocks using heuristic patterns"""186        chunks = []187 188        # Look for code patterns189        for pattern in self.code_patterns:190            matches = pattern.finditer(text)191            for match in matches:192                code_text = match.group(0)193                if len(code_text) < 20:  # Skip very short matches194                    continue195 196                chunk = Chunk(197                    chunk_id=str(uuid.uuid4()),198                    pdf_id=pdf_id,199                    page_number=page_num,200                    char_range=(match.start(), match.end()),201                    type=ChunkType.CODE,202                    text=code_text,203                    metadata={204                        "pattern": pattern.pattern,205                        "detected_language": self._detect_language(code_text)206                    }207                )208                chunks.append(chunk)209 210        # Also detect monospace font regions (if PDF has font info)211        # This is more advanced and would require font analysis212 213        logger.debug(f"Extracted {len(chunks)} code blocks from page {page_num}")214        return chunks215 216    def _detect_language(self, code: str) -> str:217        """Heuristically detect programming language"""218        if 'def ' in code and ':' in code:219            return 'python'220        elif 'function' in code or 'const' in code or 'let' in code:221            return 'javascript'222        elif 'public class' in code or 'private' in code:223            return 'java'224        elif '#include' in code:225            return 'c++'226        else:227            return 'unknown'228 229    def _extract_images(self, fitz_page, page_num: int, pdf_id: str) -> List[Chunk]:230        """Extract images and run OCR"""231        chunks = []232        image_list = fitz_page.get_images()233 234        for img_index, img in enumerate(image_list):235            try:236                xref = img[0]237                base_image = fitz_page.parent.extract_image(xref)238                image_bytes = base_image["image"]239 240                # Convert to PIL Image241                image = Image.open(io.BytesIO(image_bytes))242 243                # Run OCR244                ocr_text = pytesseract.image_to_string(image)245 246                if ocr_text.strip():247                    image_id = f"{pdf_id}_p{page_num}_img{img_index}"248 249                    chunk = Chunk(250                        chunk_id=str(uuid.uuid4()),251                        pdf_id=pdf_id,252                        page_number=page_num,253                        char_range=(0, len(ocr_text)),254                        type=ChunkType.IMAGE_TEXT,255                        text=ocr_text,256                        image_id=image_id,257                        metadata={258                            "image_format": base_image["ext"],259                            "image_index": img_index260                        }261                    )262                    chunks.append(chunk)263            except Exception as e:264                logger.warning(f"Failed to extract image {img_index} on page {page_num}: {e}")265 266        logger.debug(f"Extracted {len(chunks)} images from page {page_num}")267        return chunks268 269    def _remove_extracted_regions(self, text: str, code_blocks: List[str]) -> str:270        """Remove already-extracted code blocks from text"""271        for code in code_blocks:272            text = text.replace(code, "")273        return text274 275    def chunk_text(self, chunks: List[Chunk]) -> List[Chunk]:276        """277        Further chunk large text blocks into smaller overlapping chunks278 279        Args:280            chunks: Initial chunks from PDF extraction281 282        Returns:283            Refined chunks with proper overlap284        """285        refined_chunks = []286 287        for chunk in chunks:288            # Skip non-text chunks (tables, images already chunked)289            if chunk.type in [ChunkType.TABLE, ChunkType.CODE]:290                refined_chunks.append(chunk)291                continue292 293            # Split long paragraphs into smaller chunks with overlap294            text = chunk.text295            chunk_size = settings.chunk_size296            overlap = settings.chunk_overlap297 298            if len(text) <= chunk_size:299                refined_chunks.append(chunk)300                continue301 302            # Create overlapping windows303            for i in range(0, len(text), chunk_size - overlap):304                chunk_text = text[i:i + chunk_size]305 306                if len(chunk_text) < settings.min_chunk_size:307                    continue308 309                new_chunk = Chunk(310                    chunk_id=str(uuid.uuid4()),311                    pdf_id=chunk.pdf_id,312                    page_number=chunk.page_number,313                    char_range=(i, i + len(chunk_text)),314                    type=chunk.type,315                    text=chunk_text,316                    metadata={317                        **chunk.metadata,318                        "parent_chunk_id": chunk.chunk_id,319                        "window_index": i // (chunk_size - overlap)320                    }321                )322                refined_chunks.append(new_chunk)323 324        logger.info(f"Refined {len(chunks)} chunks into {len(refined_chunks)} chunks")325        return refined_chunks326