pdf_processor.py

import pdfplumber
import hashlib
import os
import gc
from typing import List, Optional
from concurrent.futures import ThreadPoolExecutor
from tqdm.auto import tqdm
from langchain_core.documents import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
import re

def normalize_text(text: str) -> str:
    """Normalize text for consistent hashing"""
    return " ".join(text.strip().lower().split())

def content_hash(text: str, length: int = 16) -> str:
    """Generate consistent hash for text content"""
    return hashlib.sha256(normalize_text(text).encode()).hexdigest()[:length]

def process_page(file_path: str, page_num: int) -> Optional[Document]:
    """Process individual PDF page with progress tracking"""
    try:
        with pdfplumber.open(file_path) as pdf:
            if page_num >= len(pdf.pages):
                tqdm.write(f"⚠️ Page number {page_num} out of range in {file_path}")
                return None
                
            # Extract text from the page
            page = pdf.pages[page_num]
            text = page.extract_text()
            
            if not text:
                return None
                
            # Handle LaTeX-style sections and subsections
            text = re.sub(r"\\section\*{(.*?)}", r"\n\n## \1\n\n", text)  # Convert \section* to Markdown-style headings
            text = re.sub(r"\\subsection\*{(.*?)}", r"\n\n### \1\n\n", text)  # Convert \subsection* to subheadings
            
            # Remove excessive whitespace
            text = re.sub(r'\s+', ' ', text)  # Replace multiple spaces with a single space
            text = re.sub(r'\n{3,}', '\n\n', text)  # Replace multiple newlines with double newlines
            
            # Fix common formatting issues
            text = text.replace('\n• ', ' • ')  # Bullet points
            text = re.sub(r'\n(\d+\. )', r'\1', text)  # Numbered lists
            text = re.sub(r'-\n(\w)', r'\1', text)  # Hyphenated words split across lines
            
            # Normalize the final text
            text = text.strip()
            
            return Document(
                page_content=text,
                metadata={
                    "source": file_path,
                    "page": page_num + 1,
                    "content_hash": content_hash(text)
                }
            )
    except Exception as e:
        tqdm.write(f"⚠️ Error processing page {page_num+1} in {file_path}: {str(e)}")
        return None

def process_large_pdf(file_path: str, max_workers: int = 4) -> List[Document]:
    """Process PDF with parallel processing and progress tracking"""
    try:
        # Get the total number of pages first
        with pdfplumber.open(file_path) as pdf:
            total_pages = len(pdf.pages)
        
        # Use thread pool to process pages in parallel with limited workers
        results = []
        workers = min(max_workers, total_pages, os.cpu_count() or 4)
        
        with ThreadPoolExecutor(max_workers=workers) as executor:
            futures = [
                executor.submit(process_page, file_path, page_num)
                for page_num in range(total_pages)
            ]
            
            # Gather results with progress tracking
            results = list(tqdm(
                (future.result() for future in futures),
                total=len(futures),
                desc=f"📖 {os.path.basename(file_path)}",
                unit="page",
                leave=False
            ))
        
        # Force garbage collection
        gc.collect()
        
        # Filter out None results
        return [doc for doc in results if doc is not None]
    except Exception as e:
        tqdm.write(f"🚨 Failed to process {file_path}: {str(e)}")
        return []

def pdf_chunker(documents: List[Document], chunk_size: int = 800, overlap: int = 200) -> List[Document]:
    """Split documents into chunks with overlap"""
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,  # Increased for rulebook content
        chunk_overlap=overlap,
        separators=[
            "\n\n## ",  # Split at major headings (Markdown-style ##)
            "\n\n### ",  # Split at subheadings
            "\n\n",  # Paragraph breaks
            r"\.\s+(?=[A-Z])",  # Split at sentences followed by uppercase (proper nouns)
            "\n"
        ],
        length_function=len,
        is_separator_regex=True,
        keep_separator=True
    )
    
    chunks = []
    total_docs = len(documents)
    
    with tqdm(total=total_docs, desc="✂️ Chunking documents", unit="doc") as pbar:
        for i, doc in enumerate(documents):
            try:
                doc_chunks = text_splitter.split_documents([doc])
                
                for j, chunk in enumerate(doc_chunks):
                    chunk_text = chunk.page_content
                    chunk_hash = content_hash(chunk_text)
                    
                    chunk.metadata.update({
                        "chunk_id": f"{doc.metadata['content_hash']}-{chunk_hash}",
                        "chunk_number": j + 1,
                        "total_chunks": len(doc_chunks)
                    })
                
                chunks.extend(doc_chunks)
                pbar.update(1)
                
                # Garbage collect every 10 documents
                if i % 10 == 0:
                    gc.collect()
                    
            except Exception as e:
                tqdm.write(f"⚠️ Error chunking document {doc.metadata['source']}: {str(e)}")
                continue
    
    # Force garbage collection before deduplication
    gc.collect()
    
    return deduplicate_chunks(chunks)

def deduplicate_chunks(chunks: List[Document]) -> List[Document]:
    """Remove duplicate chunks with progress tracking"""
    seen_hashes = set()
    unique_chunks = []
    
    with tqdm(chunks, desc="🚿 Deduplicating", unit="chunk") as pbar:
        for chunk in pbar:
            chunk_hash = chunk.metadata["chunk_id"]
            if chunk_hash not in seen_hashes:
                seen_hashes.add(chunk_hash)
                unique_chunks.append(chunk)
    
    tqdm.write(f"🧹 Removed {len(chunks)-len(unique_chunks)} duplicates")
    
    return unique_chunks