Mistral OCR / Document AI - PDF Chunking

from mistralai.client import Mistral

api_key = "API_KEY"
client = Mistral(api_key)

import base64

def encode_pdf_to_base64(pdf_path: str) -> str:
    """Encode the contents of a PDF file to a base64 string.

    Reads the raw bytes of the file and converts them to an ASCII-safe
    base64 string that can be embedded in a JSON request body.
    """
    try:
        # Open in binary mode ('rb') — PDFs are binary, not text.
        with open(pdf_path, "rb") as pdf_file:
            pdf_data = pdf_file.read()
            return base64.b64encode(pdf_data).decode("utf-8")
    except FileNotFoundError:
        print(f"File not found: {pdf_path}")
        raise
    except IOError as e:
        print(f"Error reading file {pdf_path}: {e}")
        raise
    except Exception as e:
        print(f"Error encoding PDF to base64: {e}")
        raise

from PyPDF2 import PdfReader, PdfWriter
from PyPDF2.errors import PdfReadError

def count_pdf_pages(pdf_path: str) -> int:
    """Count the number of pages in a PDF file.

    Uses PyPDF2's PdfReader, which parses the PDF structure without
    rendering each page — so it's fast even for large documents.
    """
    try:
        with open(pdf_path, "rb") as pdf_file:
            pdf_reader = PdfReader(pdf_file)
            return len(pdf_reader.pages)
    except FileNotFoundError:
        print(f"File not found: {pdf_path}")
        raise
    except PdfReadError as e:
        print(f"Error reading PDF structure: {e}")
        raise
    except Exception as e:
        print(f"Error counting PDF pages: {e}")
        raise

import os

def get_pdf_size_in_mb(pdf_path: str) -> float:
    """Get the size of a PDF file in MB.

    os.path.getsize() returns bytes; we divide by 1024*1024 to convert
    to megabytes for a human-friendly comparison against MAX_SIZE_MB.
    """
    try:
        size_in_bytes = os.path.getsize(pdf_path)
        return size_in_bytes / (1024 * 1024)  # Convert to MB
    except FileNotFoundError:
        print(f"File not found: {pdf_path}")
        raise
    except OSError as e:
        print(f"Error getting file size: {e}")
        raise

def split_pdf_by_pages(pdf_path: str, max_pages: int) -> list:
    """Split a PDF into multiple PdfWriters with a maximum number of pages.

    Returns a list of PdfWriter objects (in-memory).
    """
    try:
        with open(pdf_path, "rb") as pdf_file:
            pdf_reader = PdfReader(pdf_file)
            pdf_writers = []

            # Step through pages in blocks of `max_pages`.
            for i in range(0, len(pdf_reader.pages), max_pages):
                pdf_writer = PdfWriter()
                # Add every page in the current block to this writer.
                for page_num in range(i, min(i + max_pages, len(pdf_reader.pages))):
                    pdf_writer.add_page(pdf_reader.pages[page_num])
                pdf_writers.append(pdf_writer)

        return pdf_writers
    except FileNotFoundError:
        print(f"File not found: {pdf_path}")
        raise
    except PdfReadError as e:
        print(f"Error reading PDF structure: {e}")
        raise
    except Exception as e:
        print(f"Error splitting PDF: {e}")
        raise

import time

def send_to_ocr(base64_pdf: str, retries: int = 10):
    """Send the base64 encoded PDF to Mistral OCR.

    Retries with exponential backoff.
    """
    for i in range(retries):
        try:
            ocr_response = client.ocr.process(
                model="mistral-ocr-latest",
                document={
                    "type": "document_url",
                    "document_url": f"data:application/pdf;base64,{base64_pdf}"
                },
                table_format="html", # default is None
                # extract_header=True, # default is False
                # extract_footer=True, # default is False
                include_image_base64=True
            )
            return ocr_response
        except Exception as e:
            time.sleep(min(4*2**i,32))
    print(e)
    raise

import math

def split_pdf_recursive(pdf_path: str, max_size_mb: float, max_pages: int,
                         output_dir: str = ".", prefix: str = "chunk",
                         _depth: int = 0) -> list:
    """Recursively split a PDF into on-disk chunks that fit both limits.

    Returns a list of file paths to the leaf chunks (each fitting both the
    size and page limits, or a single page that cannot be split further).
    """
    indent = "  " * _depth
    size_mb = get_pdf_size_in_mb(pdf_path)
    num_pages = count_pdf_pages(pdf_path)
    print(f"{indent}PDF: {os.path.basename(pdf_path)} | size: {size_mb:.2f} MB | pages: {num_pages}")

    # Base case: fits both limits — no split needed.
    if size_mb <= max_size_mb and num_pages <= max_pages:
        print(f"{indent}→ fits limits, no split needed")
        return [pdf_path]

    # Edge case: a single page already exceeds the size limit — cannot split further.
    if num_pages <= 1:
        print(f"{indent}⚠ single page exceeds size limit ({size_mb:.2f} MB); left as-is")
        return [pdf_path]

    # Recursive case: choose a chunk size that respects both constraints.
    chunks_for_pages = math.ceil(num_pages / max_pages)
    chunks_for_size = math.ceil(size_mb / max_size_mb)
    num_chunks = max(chunks_for_pages, chunks_for_size)
    pages_per_chunk = math.ceil(num_pages / num_chunks)

    print(f"{indent}→ exceeds limits; splitting into {num_chunks} chunk(s) of ~{pages_per_chunk} page(s)")
    pdf_writers = split_pdf_by_pages(pdf_path, pages_per_chunk)

    chunk_paths = []
    for i, pdf_writer in enumerate(pdf_writers):
        split_pdf_path = os.path.join(output_dir, f"{prefix}_{_depth}_{i}.pdf")
        try:
            with open(split_pdf_path, "wb") as split_pdf_file:
                pdf_writer.write(split_pdf_file)
            # Recurse: re-measures this chunk and splits again if still too large.
            leaf_paths = split_pdf_recursive(
                split_pdf_path, max_size_mb, max_pages,
                output_dir=output_dir, prefix=prefix, _depth=_depth + 1
            )
            chunk_paths.extend(leaf_paths)
        except Exception as e:
            print(f"{indent}Error splitting part {i + 1}: {e}")
            # Clean up a partially-written chunk file on failure.
            try:
                if os.path.exists(split_pdf_path):
                    os.remove(split_pdf_path)
            except OSError:
                pass

    return chunk_paths