document_loader.py

# document_loader.py
# ===================
# Loads pharma documents for the RAG system
# Supports PDF files and text documents
# Uses PyMuPDF for reliable PDF text extraction

import os
import fitz  # PyMuPDF
from typing import List, Dict

def load_pdf(file_path: str) -> List[Dict]:
    """
    Load a PDF file and extract text by page.
    
    Args:
        file_path: Path to the PDF file
    
    Returns:
        List of dictionaries with page text and metadata
    """
    
    print(f"Loading PDF: {file_path}")
    
    try:
        doc = fitz.open(file_path)
        pages = []
        
        for page_num in range(len(doc)):
            page = doc[page_num]
            text = page.get_text()
            
            # Skip empty pages
            if len(text.strip()) < 50:
                continue
            
            pages.append({
                "text": text,
                "page_number": page_num + 1,
                "source": os.path.basename(file_path),
                "file_path": file_path
            })
        
        doc.close()
        print(f"Extracted {len(pages)} pages from {os.path.basename(file_path)}")
        return pages
        
    except Exception as e:
        print(f"Error loading PDF: {e}")
        return []


def load_pdf_from_bytes(file_bytes: bytes, filename: str) -> List[Dict]:
    """
    Load a PDF from bytes (for Streamlit file uploads).
    
    Args:
        file_bytes: PDF file as bytes
        filename: Name of the file
    
    Returns:
        List of dictionaries with page text and metadata
    """
    
    print(f"Loading uploaded PDF: {filename}")
    
    try:
        doc = fitz.open(stream=file_bytes, filetype="pdf")
        pages = []
        
        for page_num in range(len(doc)):
            page = doc[page_num]
            text = page.get_text()
            
            if len(text.strip()) < 50:
                continue
            
            pages.append({
                "text": text,
                "page_number": page_num + 1,
                "source": filename,
                "file_path": filename
            })
        
        doc.close()
        print(f"Extracted {len(pages)} pages from {filename}")
        return pages
        
    except Exception as e:
        print(f"Error loading PDF bytes: {e}")
        return []


def chunk_pages(pages: List[Dict], 
                chunk_size: int = 1000, 
                chunk_overlap: int = 100) -> List[Dict]:
    """
    Split pages into smaller chunks for better RAG retrieval.
    
    Args:
        pages: List of page dictionaries
        chunk_size: Maximum characters per chunk
        chunk_overlap: Characters to overlap between chunks
    
    Returns:
        List of chunk dictionaries
    """
    
    chunks = []
    
    for page in pages:
        text = page["text"]
        
        # If page is short enough — keep as single chunk
        if len(text) <= chunk_size:
            chunks.append({
                "text": text,
                "page_number": page["page_number"],
                "source": page["source"],
                "chunk_id": f"{page['source']}_p{page['page_number']}_c0"
            })
            continue
        
        # Split long pages into overlapping chunks
        start = 0
        chunk_num = 0
        
        while start < len(text):
            end = start + chunk_size
            chunk_text = text[start:end]
            
            chunks.append({
                "text": chunk_text,
                "page_number": page["page_number"],
                "source": page["source"],
                "chunk_id": f"{page['source']}_p{page['page_number']}_c{chunk_num}"
            })
            
            start = end - chunk_overlap
            chunk_num += 1
    
    print(f"Created {len(chunks)} chunks from {len(pages)} pages")
    return chunks


def get_document_stats(pages: List[Dict]) -> Dict:
    """
    Get statistics about a loaded document.
    """
    
    if not pages:
        return {"pages": 0, "total_chars": 0, "avg_chars_per_page": 0}
    
    total_chars = sum(len(p["text"]) for p in pages)
    
    return {
        "pages": len(pages),
        "total_chars": total_chars,
        "avg_chars_per_page": total_chars // len(pages),
        "source": pages[0]["source"]
    }


# Quick test
if __name__ == "__main__":
    # Test with a sample text to verify the module works
    sample_pages = [
        {
            "text": "This is a sample pharma document. " * 50,
            "page_number": 1,
            "source": "test_document.pdf",
            "file_path": "test_document.pdf"
        }
    ]
    
    chunks = chunk_pages(sample_pages)
    stats = get_document_stats(sample_pages)
    
    print(f"Test passed!")
    print(f"Stats: {stats}")
    print(f"Chunks created: {len(chunks)}")