mve-micro-docs/app/services/pdf_text.py

"""PDF text extraction service using PyMuPDF, pdfplumber, and OCR."""
import fitz  # PyMuPDF
import pdfplumber
import logging
from typing import Tuple, Optional
from io import BytesIO
from .ocr_service import extract_text_with_ocr, is_pdf_scanned, OCRError

logger = logging.getLogger(__name__)


class PDFExtractionError(Exception):
    """Custom exception for PDF extraction errors."""
    pass


def extract_text_with_pymupdf(pdf_bytes: bytes) -> Tuple[str, int]:
    """
    Extract text from PDF using PyMuPDF (fitz).
    
    Args:
        pdf_bytes: PDF file content as bytes
        
    Returns:
        Tuple of (extracted_text, page_count)
        
    Raises:
        PDFExtractionError: If extraction fails
    """
    try:
        doc = fitz.open(stream=pdf_bytes, filetype="pdf")
        
        if doc.is_encrypted:
            raise PDFExtractionError("PDF is encrypted and cannot be read")
        
        page_count = len(doc)
        text_parts = []
        
        for page in doc:
            text_parts.append(page.get_text())
        
        doc.close()
        full_text = "\n".join(text_parts)
        
        logger.info(f"Extracted {len(full_text)} characters using PyMuPDF from {page_count} pages")
        return full_text, page_count
        
    except Exception as e:
        logger.warning(f"PyMuPDF extraction failed: {str(e)}")
        raise PDFExtractionError(f"PyMuPDF extraction failed: {str(e)}")


def extract_text_with_pdfplumber(pdf_bytes: bytes) -> Tuple[str, int]:
    """
    Extract text from PDF using pdfplumber (fallback method).
    
    Args:
        pdf_bytes: PDF file content as bytes
        
    Returns:
        Tuple of (extracted_text, page_count)
        
    Raises:
        PDFExtractionError: If extraction fails
    """
    try:
        with pdfplumber.open(BytesIO(pdf_bytes)) as pdf:
            page_count = len(pdf.pages)
            text_parts = []
            
            for page in pdf.pages:
                page_text = page.extract_text()
                if page_text:
                    text_parts.append(page_text)
            
            full_text = "\n".join(text_parts)
            
            logger.info(f"Extracted {len(full_text)} characters using pdfplumber from {page_count} pages")
            return full_text, page_count
            
    except Exception as e:
        logger.error(f"pdfplumber extraction failed: {str(e)}")
        raise PDFExtractionError(f"pdfplumber extraction failed: {str(e)}")


def extract_text_from_pdf(pdf_bytes: bytes, enable_ocr: bool = True, ocr_lang: str = "spa") -> Tuple[str, int, str]:
    """
    Extract text from PDF using available methods.
    Tries PyMuPDF first, falls back to pdfplumber, and uses OCR for scanned PDFs.
    
    Args:
        pdf_bytes: PDF file content as bytes
        enable_ocr: Whether to use OCR for scanned PDFs (default: True)
        ocr_lang: Language for OCR - "spa" for Spanish, "eng" for English (default: "spa")
        
    Returns:
        Tuple of (extracted_text, page_count, method_used)
        method_used can be: "pymupdf", "pdfplumber", "ocr"
        
    Raises:
        PDFExtractionError: If all extraction methods fail
    """
    # Try PyMuPDF first
    try:
        text, pages = extract_text_with_pymupdf(pdf_bytes)
        
        # Check if the PDF is scanned (no text content)
        if enable_ocr and is_pdf_scanned(text):
            logger.info("PDF appears to be scanned, attempting OCR extraction")
            try:
                ocr_text, ocr_pages = extract_text_with_ocr(pdf_bytes, lang=ocr_lang)
                return ocr_text, ocr_pages, "ocr"
            except OCRError as ocr_e:
                logger.error(f"OCR failed: {str(ocr_e)}")
                # Return the minimal text we got, if any
                if text.strip():
                    return text, pages, "pymupdf"
                raise PDFExtractionError(f"PDF appears to be scanned and OCR failed: {str(ocr_e)}")
        
        return text, pages, "pymupdf"
    except PDFExtractionError as e:
        logger.warning(f"PyMuPDF failed, trying pdfplumber: {str(e)}")
    
    # Fallback to pdfplumber
    try:
        text, pages = extract_text_with_pdfplumber(pdf_bytes)
        
        # Check if the PDF is scanned
        if enable_ocr and is_pdf_scanned(text):
            logger.info("PDF appears to be scanned (pdfplumber), attempting OCR extraction")
            try:
                ocr_text, ocr_pages = extract_text_with_ocr(pdf_bytes, lang=ocr_lang)
                return ocr_text, ocr_pages, "ocr"
            except OCRError as ocr_e:
                logger.error(f"OCR failed: {str(ocr_e)}")
                if text.strip():
                    return text, pages, "pdfplumber"
                raise PDFExtractionError(f"PDF appears to be scanned and OCR failed: {str(ocr_e)}")
        
        return text, pages, "pdfplumber"
    except PDFExtractionError as e:
        logger.error(f"All extraction methods failed: {str(e)}")
        raise PDFExtractionError("Failed to extract text from PDF using all available methods")