diff --git a/mistral_ocr/cli.py b/mistral_ocr/cli.py index be2ba08..76c101f 100644 --- a/mistral_ocr/cli.py +++ b/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 diff --git a/mistral_ocr/config.py b/mistral_ocr/config.py index d4e3cac..a95fa33 100644 --- a/mistral_ocr/config.py +++ b/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, diff --git a/mistral_ocr/processor.py b/mistral_ocr/processor.py index ad9fba3..e5a0340 100644 --- a/mistral_ocr/processor.py +++ b/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,11 +19,13 @@ 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__) @@ -30,6 +33,9 @@ # 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 diff --git a/mistral_ocr/utils.py b/mistral_ocr/utils.py index 8aada97..8cfd154 100644 --- a/mistral_ocr/utils.py +++ b/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, diff --git a/tests/test_concurrent.py b/tests/test_concurrent.py index c35dfab..fbed4de 100644 --- a/tests/test_concurrent.py +++ b/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 diff --git a/tests/test_pdf_chunking.py b/tests/test_pdf_chunking.py new file mode 100644 index 0000000..e027b54 --- /dev/null +++ b/tests/test_pdf_chunking.py @@ -0,0 +1,236 @@ +"""Tests for PDF chunking and upload-based processing.""" + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from mistral_ocr.config import Config +from mistral_ocr.processor import MAX_PAGES_PER_REQUEST, OCRProcessor +from mistral_ocr.utils import get_pdf_page_count, split_pdf + +# --------------------------------------------------------------------------- +# Utility tests (split_pdf, get_pdf_page_count) +# --------------------------------------------------------------------------- + + +def _make_pdf(path: Path, num_pages: int) -> Path: + """Create a minimal PDF with *num_pages* blank pages using pypdf.""" + from pypdf import PdfWriter + + writer = PdfWriter() + for _ in range(num_pages): + writer.add_blank_page(width=72, height=72) + with open(path, "wb") as f: + writer.write(f) + return path + + +class TestGetPdfPageCount: + def test_counts_pages(self, tmp_path: Path) -> None: + pdf = _make_pdf(tmp_path / "doc.pdf", 5) + assert get_pdf_page_count(pdf) == 5 + + def test_single_page(self, tmp_path: Path) -> None: + pdf = _make_pdf(tmp_path / "one.pdf", 1) + assert get_pdf_page_count(pdf) == 1 + + +class TestSplitPdf: + def test_no_split_needed(self, tmp_path: Path) -> None: + pdf = _make_pdf(tmp_path / "small.pdf", 3) + out = tmp_path / "chunks" + chunks = split_pdf(pdf, out, max_pages_per_chunk=10) + assert len(chunks) == 1 + chunk_path, start, count = chunks[0] + assert start == 0 + assert count == 3 + assert chunk_path.exists() + + def test_splits_into_multiple_chunks(self, tmp_path: Path) -> None: + pdf = _make_pdf(tmp_path / "big.pdf", 7) + out = tmp_path / "chunks" + chunks = split_pdf(pdf, out, max_pages_per_chunk=3) + assert len(chunks) == 3 + # chunk1: pages 0-2, chunk2: 3-5, chunk3: 6 + assert chunks[0] == (out / "big_chunk1.pdf", 0, 3) + assert chunks[1] == (out / "big_chunk2.pdf", 3, 3) + assert chunks[2] == (out / "big_chunk3.pdf", 6, 1) + # Verify each chunk has the right page count + assert get_pdf_page_count(chunks[0][0]) == 3 + assert get_pdf_page_count(chunks[1][0]) == 3 + assert get_pdf_page_count(chunks[2][0]) == 1 + + def test_max_pages_limits_output(self, tmp_path: Path) -> None: + pdf = _make_pdf(tmp_path / "many.pdf", 10) + out = tmp_path / "chunks" + chunks = split_pdf(pdf, out, max_pages_per_chunk=5, max_pages=7) + assert len(chunks) == 2 + assert chunks[0][2] == 5 # first chunk: 5 pages + assert chunks[1][2] == 2 # second chunk: 2 pages (7-5) + + def test_max_pages_none_processes_all(self, tmp_path: Path) -> None: + pdf = _make_pdf(tmp_path / "all.pdf", 4) + out = tmp_path / "chunks" + chunks = split_pdf(pdf, out, max_pages_per_chunk=100, max_pages=None) + assert len(chunks) == 1 + assert chunks[0][2] == 4 + + +# --------------------------------------------------------------------------- +# Config: MAX_PAGES env var +# --------------------------------------------------------------------------- + + +class TestConfigMaxPages: + def test_default_is_none(self) -> None: + config = Config(api_key="test") + assert config.max_pages is None + + @patch.dict("os.environ", {"MISTRAL_API_KEY": "key", "MAX_PAGES": "50"}) + def test_from_env(self) -> None: + config = Config.from_env() + assert config.max_pages == 50 + + @patch.dict("os.environ", {"MISTRAL_API_KEY": "key", "MAX_PAGES": "abc"}) + def test_invalid_raises(self) -> None: + with pytest.raises(ValueError, match="MAX_PAGES must be an integer"): + Config.from_env() + + @patch.dict("os.environ", {"MISTRAL_API_KEY": "key", "MAX_PAGES": "0"}) + def test_zero_raises(self) -> None: + with pytest.raises(ValueError, match="MAX_PAGES must be positive"): + Config.from_env() + + @patch.dict("os.environ", {"MISTRAL_API_KEY": "key", "MAX_PAGES": "-5"}) + def test_negative_raises(self) -> None: + with pytest.raises(ValueError, match="MAX_PAGES must be positive"): + Config.from_env() + + +# --------------------------------------------------------------------------- +# Processor: _process_pdf routing +# --------------------------------------------------------------------------- + + +def _make_processor(**overrides: object) -> OCRProcessor: + defaults = {"api_key": "test-key", "include_images": False, "save_original_images": False} + defaults.update(overrides) + config = Config(**defaults) + with patch("mistral_ocr.processor.Mistral"): + return OCRProcessor(config) + + +def _fake_page(index: int) -> SimpleNamespace: + return SimpleNamespace(index=index, markdown=f"Page {index + 1} text") + + +class TestProcessPdf: + """Test _process_pdf routing logic with mocked upload.""" + + @patch("mistral_ocr.processor.get_pdf_page_count", return_value=5) + def test_small_pdf_uploads_directly(self, mock_count: MagicMock) -> None: + proc = _make_processor() + proc._upload_and_process = MagicMock( + return_value=SimpleNamespace(pages=[_fake_page(i) for i in range(5)]) + ) + result = proc._process_pdf(Path("test.pdf")) + proc._upload_and_process.assert_called_once_with(Path("test.pdf")) + assert len(result.pages) == 5 + + @patch("mistral_ocr.processor.get_pdf_page_count", return_value=5) + def test_max_pages_triggers_chunking(self, mock_count: MagicMock) -> None: + proc = _make_processor(max_pages=3) + proc._process_pdf_chunked = MagicMock( + return_value=SimpleNamespace(pages=[_fake_page(i) for i in range(3)]) + ) + result = proc._process_pdf(Path("test.pdf")) + proc._process_pdf_chunked.assert_called_once_with(Path("test.pdf"), 5) + assert len(result.pages) == 3 + + @patch("mistral_ocr.processor.get_pdf_page_count", return_value=1500) + def test_large_pdf_triggers_chunking(self, mock_count: MagicMock) -> None: + proc = _make_processor() + proc._process_pdf_chunked = MagicMock( + return_value=SimpleNamespace(pages=[_fake_page(i) for i in range(1500)]) + ) + result = proc._process_pdf(Path("test.pdf")) + proc._process_pdf_chunked.assert_called_once_with(Path("test.pdf"), 1500) + assert len(result.pages) == 1500 + + +class TestProcessPdfChunked: + """Test _process_pdf_chunked with real split_pdf but mocked API calls.""" + + def test_reassembles_pages_with_correct_indices(self, tmp_path: Path) -> None: + pdf = _make_pdf(tmp_path / "doc.pdf", 5) + proc = _make_processor() + + def fake_upload(chunk_path: Path) -> SimpleNamespace: + from mistral_ocr.utils import get_pdf_page_count as gpc + + n = gpc(chunk_path) + return SimpleNamespace(pages=[_fake_page(i) for i in range(n)]) + + proc._upload_and_process = MagicMock(side_effect=fake_upload) + result = proc._process_pdf_chunked(pdf, total_pages=5) + # Should have 5 pages with indices 0-4 + assert len(result.pages) == 5 + assert [p.index for p in result.pages] == [0, 1, 2, 3, 4] + + def test_truncation_note_when_max_pages(self, tmp_path: Path) -> None: + pdf = _make_pdf(tmp_path / "doc.pdf", 10) + proc = _make_processor(max_pages=3) + + def fake_upload(chunk_path: Path) -> SimpleNamespace: + from mistral_ocr.utils import get_pdf_page_count as gpc + + n = gpc(chunk_path) + return SimpleNamespace(pages=[_fake_page(i) for i in range(n)]) + + proc._upload_and_process = MagicMock(side_effect=fake_upload) + result = proc._process_pdf_chunked(pdf, total_pages=10) + assert len(result.pages) == 3 + assert hasattr(result, "truncated") + assert "3 of 10" in result.truncated + + def test_no_truncation_note_when_all_pages(self, tmp_path: Path) -> None: + pdf = _make_pdf(tmp_path / "doc.pdf", 3) + proc = _make_processor() + + def fake_upload(chunk_path: Path) -> SimpleNamespace: + from mistral_ocr.utils import get_pdf_page_count as gpc + + n = gpc(chunk_path) + return SimpleNamespace(pages=[_fake_page(i) for i in range(n)]) + + proc._upload_and_process = MagicMock(side_effect=fake_upload) + result = proc._process_pdf_chunked(pdf, total_pages=3) + assert not hasattr(result, "truncated") + + +class TestBuildOcrKwargs: + def test_basic_kwargs(self) -> None: + proc = _make_processor() + doc = {"type": "file", "file_id": "abc"} + kwargs = proc._build_ocr_kwargs(doc) + assert kwargs["model"] == "mistral-ocr-latest" + assert kwargs["document"] == doc + assert kwargs["include_image_base64"] is False + + def test_with_table_format(self) -> None: + proc = _make_processor(table_format="html") + kwargs = proc._build_ocr_kwargs({"type": "file", "file_id": "x"}) + assert kwargs["table_format"] == "html" + + def test_with_headers_footers(self) -> None: + proc = _make_processor(extract_header=True, extract_footer=True) + kwargs = proc._build_ocr_kwargs({"type": "file", "file_id": "x"}) + assert kwargs["extract_header"] is True + assert kwargs["extract_footer"] is True + + +class TestMaxPagesPerRequest: + def test_constant_value(self) -> None: + assert MAX_PAGES_PER_REQUEST == 1000 diff --git a/tests/test_retry.py b/tests/test_retry.py index 989fe89..672c750 100644 --- a/tests/test_retry.py +++ b/tests/test_retry.py @@ -116,8 +116,8 @@ def test_forwards_all_kwargs(self, processor): class TestProcessFileRetry: def test_process_file_retries_transient_api_error(self, processor, tmp_path): """Integration: process_file retries on transient errors.""" - pdf = tmp_path / "doc.pdf" - pdf.write_bytes(b"%PDF-1.4 fake content") + img = tmp_path / "doc.png" + img.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50) mock_response = SimpleNamespace(pages=[]) processor.client.ocr.process.side_effect = [ @@ -125,7 +125,7 @@ def test_process_file_retries_transient_api_error(self, processor, tmp_path): mock_response, ] - result = processor.process_file(pdf) + result = processor.process_file(img) assert result is not None assert result["success"] is True assert processor.client.ocr.process.call_count == 2