Add PDF chunking and --max-pages flag

mistral_ocr/cli.py

@@ -86,6 +86,12 @@
     default=False,
     help="Extract page footers (default: False). OCR 3+ only.",
 )
+@click.option(
+    "--max-pages",
+    type=click.IntRange(min=1),
+    default=None,
+    help="Maximum number of PDF pages to process (default: all pages)",
+)
 @click.option(
     "--workers",
     "-w",
@@ -128,6 +134,7 @@ def main(
     table_format: str | None,
     extract_headers: bool,
     extract_footers: bool,
+    max_pages: int | None,
     workers: int,
     reprocess: bool,
     dry_run: bool,
@@ -271,6 +278,11 @@ def main(
             and ctx.get_parameter_source("workers") != click.core.ParameterSource.DEFAULT
         ):
             config.max_workers = workers
+        if (
+            "max_pages" in ctx.params
+            and ctx.get_parameter_source("max_pages") != click.core.ParameterSource.DEFAULT
+        ):
+            config.max_pages = max_pages
 
         config.quiet = quiet
 

mistral_ocr/config.py

@@ -21,6 +21,7 @@ class Config:
     extract_footer: bool = False
     include_metadata: bool = True
     include_page_headings: bool = True
+    max_pages: int | None = None  # None = no limit
     max_workers: int = 1
     max_retries: int = 3
     retry_base_delay: float = 1.0
@@ -65,6 +66,16 @@ def from_env(cls, env_file: Path | None = None) -> "Config":
         if max_retries < 0:
             raise ValueError(f"MAX_RETRIES must be non-negative, got: {max_retries}")
 
+        max_pages_str = os.getenv("MAX_PAGES", "")
+        max_pages_val: int | None = None
+        if max_pages_str:
+            try:
+                max_pages_val = int(max_pages_str)
+            except ValueError as e:
+                raise ValueError(f"MAX_PAGES must be an integer, got: {max_pages_str!r}") from e
+            if max_pages_val <= 0:
+                raise ValueError(f"MAX_PAGES must be positive, got: {max_pages_val}")
+
         try:
             retry_base_delay = float(os.getenv("RETRY_BASE_DELAY", "1.0"))
         except ValueError as e:
@@ -83,6 +94,7 @@ def from_env(cls, env_file: Path | None = None) -> "Config":
             table_format=table_fmt,
             extract_header=os.getenv("EXTRACT_HEADER", "false").lower() == "true",
             extract_footer=os.getenv("EXTRACT_FOOTER", "false").lower() == "true",
+            max_pages=max_pages_val,
             max_workers=max(1, int(os.getenv("MAX_WORKERS", "1"))),
             max_retries=max_retries,
             retry_base_delay=retry_base_delay,

mistral_ocr/processor.py

@@ -2,6 +2,7 @@
 
 import logging
 import shutil
+import tempfile
 import threading
 import time
 from concurrent.futures import ThreadPoolExecutor, as_completed
@@ -18,18 +19,23 @@
     create_data_uri,
     determine_output_path,
     format_file_size,
+    get_pdf_page_count,
     get_supported_files,
     load_metadata,
     make_unique_basename,
     save_base64_image,
     save_metadata,
+    split_pdf,
 )
 
 logger = logging.getLogger(__name__)
 
 # Shared console instance — CLI sets .quiet on this directly
 console = Console()
 
+# Mistral API limit: max pages per single OCR request
+MAX_PAGES_PER_REQUEST = 1000
+
 
 class OCRProcessor:
     """OCR processor using Mistral AI API."""
@@ -88,25 +94,100 @@ def _call_with_retry(self, **ocr_kwargs: object) -> object:
         # Unreachable, but keeps mypy happy
         raise RuntimeError("Retry loop exited unexpectedly")
 
+    def _build_ocr_kwargs(self, document: dict) -> dict:
+        """Build common OCR API kwargs."""
+        ocr_kwargs: dict = {
+            "model": self.config.model,
+            "document": document,
+            "include_image_base64": self.config.include_images,
+        }
+        if self.config.table_format:
+            ocr_kwargs["table_format"] = self.config.table_format
+        if self.config.extract_header:
+            ocr_kwargs["extract_header"] = True
+        if self.config.extract_footer:
+            ocr_kwargs["extract_footer"] = True
+        return ocr_kwargs
+
+    def _upload_and_process(self, file_path: Path) -> object:
+        """Upload a file via Mistral files API and process with OCR."""
+        with open(file_path, "rb") as f:
+            uploaded = self.client.files.upload(
+                file={"file_name": file_path.name, "content": f},
+                purpose="ocr",
+            )
+        try:
+            document = {"type": "file", "file_id": uploaded.id}
+            return self._call_with_retry(**self._build_ocr_kwargs(document))
+        finally:
+            try:
+                self.client.files.delete(file_id=uploaded.id)
+            except Exception:
+                logger.debug("Failed to delete uploaded file %s", uploaded.id)
+
+    def _process_pdf(self, file_path: Path) -> object:
+        """Process a PDF, chunking if needed for the API page limit."""
+        page_count = get_pdf_page_count(file_path)
+        max_pages = self.config.max_pages
+        effective_pages = min(page_count, max_pages) if max_pages else page_count
+
+        if effective_pages <= MAX_PAGES_PER_REQUEST:
+            # Small enough — upload directly (may be truncated by max_pages via chunking)
+            if max_pages and page_count > max_pages:
+                logger.debug("Truncating %d-page PDF to %d pages", page_count, max_pages)
+                return self._process_pdf_chunked(file_path, page_count)
+            logger.debug("Uploading PDF directly (%d pages)", page_count)
+            return self._upload_and_process(file_path)
+
+        logger.debug(
+            "PDF has %d pages (processing %d), splitting into chunks of %d",
+            page_count,
+            effective_pages,
+            MAX_PAGES_PER_REQUEST,
+        )
+        return self._process_pdf_chunked(file_path, page_count)
+
+    def _process_pdf_chunked(self, file_path: Path, total_pages: int) -> object:
+        """Split a PDF into chunks, process each, and reassemble pages."""
+        from types import SimpleNamespace
+
+        max_pages = self.config.max_pages
+
+        with tempfile.TemporaryDirectory(prefix="mistral_ocr_") as tmp:
+            chunks = split_pdf(
+                file_path,
+                Path(tmp),
+                max_pages_per_chunk=MAX_PAGES_PER_REQUEST,
+                max_pages=max_pages,
+            )
+
+            all_pages = []
+            for chunk_path, start_page, chunk_count in chunks:
+                logger.debug(
+                    "Processing chunk: pages %d-%d", start_page + 1, start_page + chunk_count
+                )
+                response = self._upload_and_process(chunk_path)
+
+                # Reindex pages to their position in the original document.
+                # We pass through the original page objects to preserve all
+                # OCR 3 fields (tables, headers, footers, hyperlinks, dimensions).
+                for local_idx, page in enumerate(getattr(response, "pages", [])):
+                    page.index = start_page + local_idx
+                    all_pages.append(page)
+
+        result = SimpleNamespace(pages=all_pages)
+
+        # Add truncation note if we limited pages
+        if max_pages and total_pages > max_pages:
+            result.truncated = f"Processed {max_pages} of {total_pages} pages (--max-pages)"
+        return result
+
     def process_file(self, file_path: Path) -> dict | None:
         """Process a single file with OCR."""
         try:
-            # Validate file size
             file_size_mb = file_path.stat().st_size / (1024 * 1024)
             logger.debug("File size: %.2f MB", file_size_mb)
-            self.config.validate_file_size(file_path)
-
-            # Create data URI for the file
-            logger.debug("Creating data URI for %s file...", file_path.suffix)
-            data_uri = create_data_uri(file_path)
-
-            # Determine document type based on file extension
-            if file_path.suffix.lower() in DOCUMENT_EXTENSIONS:
-                document = {"type": "document_url", "document_url": data_uri}
-            else:
-                document = {"type": "image_url", "image_url": data_uri}
 
-            # Process with Mistral OCR
             if not hasattr(self.client, "ocr"):
                 raise AttributeError(
                     "OCR endpoint not available in Mistral client. "
@@ -116,20 +197,17 @@ def process_file(self, file_path: Path) -> dict | None:
 
             logger.debug("Sending to Mistral OCR API (model=%s)...", self.config.model)
 
-            # Build API kwargs (only pass OCR 3 params if set)
-            ocr_kwargs = {
-                "model": self.config.model,
-                "document": document,
-                "include_image_base64": self.config.include_images,
-            }
-            if self.config.table_format:
-                ocr_kwargs["table_format"] = self.config.table_format
-            if self.config.extract_header:
-                ocr_kwargs["extract_header"] = True
-            if self.config.extract_footer:
-                ocr_kwargs["extract_footer"] = True
-
-            response = self._call_with_retry(**ocr_kwargs)
+            if file_path.suffix.lower() == ".pdf":
+                response = self._process_pdf(file_path)
+            else:
+                # Images and other documents: validate size, use data URI
+                self.config.validate_file_size(file_path)
+                data_uri = create_data_uri(file_path)
+                if file_path.suffix.lower() in DOCUMENT_EXTENSIONS:
+                    document = {"type": "document_url", "document_url": data_uri}
+                else:
+                    document = {"type": "image_url", "image_url": data_uri}
+                response = self._call_with_retry(**self._build_ocr_kwargs(document))
 
             return {"file_path": file_path, "response": response, "success": True}
 
@@ -186,6 +264,10 @@ def save_results(
                     f"**Original:** [{base_name}{file_path.suffix}](./{base_name}{file_path.suffix})\n\n"
                 )
 
+            # Truncation note from PDF chunking
+            if hasattr(response, "truncated") and response.truncated:
+                markdown_content.append(f"**Note:** {response.truncated}\n\n")
+
             markdown_content.append("---\n\n")
 
         # Process each page

mistral_ocr/utils.py

@@ -90,6 +90,52 @@ def get_supported_files(
     return sorted(files)
 
 
+def get_pdf_page_count(file_path: Path) -> int:
+    """Return the number of pages in a PDF."""
+    from pypdf import PdfReader
+
+    return len(PdfReader(str(file_path)).pages)
+
+
+def split_pdf(
+    file_path: Path,
+    output_dir: Path,
+    *,
+    max_pages_per_chunk: int = 1000,
+    max_pages: int | None = None,
+) -> list[tuple[Path, int, int]]:
+    """Split a PDF into page-bounded chunks.
+
+    Returns a list of (chunk_path, start_page_index, page_count) tuples.
+    """
+    from pypdf import PdfReader, PdfWriter
+
+    reader = PdfReader(str(file_path))
+    total_pages = len(reader.pages)
+    pages_to_process = min(total_pages, max_pages) if max_pages else total_pages
+
+    output_dir.mkdir(parents=True, exist_ok=True)
+    chunks: list[tuple[Path, int, int]] = []
+    start = 0
+    chunk_idx = 0
+
+    while start < pages_to_process:
+        end = min(start + max_pages_per_chunk, pages_to_process)
+        writer = PdfWriter()
+        for i in range(start, end):
+            writer.add_page(reader.pages[i])
+
+        chunk_path = output_dir / f"{file_path.stem}_chunk{chunk_idx + 1}.pdf"
+        with open(chunk_path, "wb") as f:
+            writer.write(f)
+
+        chunks.append((chunk_path, start, end - start))
+        start = end
+        chunk_idx += 1
+
+    return chunks
+
+
 def determine_output_path(
     input_path: Path,
     output_path: Path | None = None,

tests/test_concurrent.py

@@ -35,8 +35,8 @@ def test_sequential_processes_all_files(self, tmp_path):
         proc = _make_processor(max_workers=1)
         input_dir = tmp_path / "input"
         input_dir.mkdir()
-        for name in ["a.pdf", "b.pdf", "c.pdf"]:
-            (input_dir / name).write_bytes(b"%PDF-1.4 test")
+        for name in ["a.png", "b.png", "c.png"]:
+            (input_dir / name).write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50)
 
         success, total = proc.process_directory(input_dir)
         assert total == 3
@@ -49,8 +49,8 @@ def test_concurrent_processes_all_files(self, tmp_path):
         proc = _make_processor(max_workers=3)
         input_dir = tmp_path / "input"
         input_dir.mkdir()
-        for name in ["a.pdf", "b.pdf", "c.pdf"]:
-            (input_dir / name).write_bytes(b"%PDF-1.4 test")
+        for name in ["a.png", "b.png", "c.png"]:
+            (input_dir / name).write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50)
 
         success, total = proc.process_directory(input_dir)
         assert total == 3
@@ -66,7 +66,7 @@ def test_concurrent_handles_failures(self, tmp_path):
         original_process = proc.process_file
 
         def selective_fail(file_path):
-            if file_path.name == "b.pdf":
+            if file_path.name == "b.png":
                 proc.errors.append({"file": str(file_path), "error": "test error"})
                 return None
             return original_process(file_path)
@@ -75,8 +75,8 @@ def selective_fail(file_path):
 
         input_dir = tmp_path / "input"
         input_dir.mkdir()
-        for name in ["a.pdf", "b.pdf", "c.pdf"]:
-            (input_dir / name).write_bytes(b"%PDF-1.4 test")
+        for name in ["a.png", "b.png", "c.png"]:
+            (input_dir / name).write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50)
 
         success, total = proc.process_directory(input_dir)
         assert total == 3
@@ -96,8 +96,8 @@ def track_thread(*args, **kwargs):
 
         input_dir = tmp_path / "input"
         input_dir.mkdir()
-        for name in ["a.pdf", "b.pdf", "c.pdf"]:
-            (input_dir / name).write_bytes(b"%PDF-1.4 test")
+        for name in ["a.png", "b.png", "c.png"]:
+            (input_dir / name).write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50)
 
         proc.process_directory(input_dir)
         # With 3 workers and 3 files, we should see more than 1 thread