diff --git a/.env.example b/.env.example index f9922b6..67b92ca 100644 --- a/.env.example +++ b/.env.example @@ -10,14 +10,15 @@ MISTRAL_MODEL=mistral-ocr-latest # Optional: Maximum file size in MB (default: 50) MAX_FILE_SIZE_MB=50 -# Optional: Maximum number of pages (default: 1000) -MAX_PAGES=1000 - -# Optional: Output format (default: markdown) -OUTPUT_FORMAT=markdown +# Optional: Maximum number of PDF pages to process. Blank = all pages (default). +# (Not a default cap — leave unset to process every page.) +MAX_PAGES= # Optional: Include extracted images (default: true) INCLUDE_IMAGES=true +# Optional: Number of concurrent workers for batch runs (default: 1) +MAX_WORKERS=1 + # Optional: Enable verbose output (default: false) -VERBOSE=false \ No newline at end of file +VERBOSE=false diff --git a/.gitignore b/.gitignore index 3efad86..838f003 100644 --- a/.gitignore +++ b/.gitignore @@ -78,6 +78,7 @@ cover/ # Virtual environments .env .venv +.venv* env/ venv/ ENV/ @@ -116,6 +117,7 @@ dmypy.json # OCR output directories mistral_ocr_output/ *_ocr_output/ +ocr/ output/ results/ extracted/ diff --git a/README.md b/README.md index 2266a4a..16498e4 100644 --- a/README.md +++ b/README.md @@ -46,8 +46,8 @@ export MISTRAL_API_KEY="your_key_here" # Process a single file mistral-ocr document.pdf -# Process a directory -mistral-ocr ./documents --output-path ./results +# Process a directory (default output root is ./documents/ocr/) +mistral-ocr ./documents -o ./results # Preview what would be processed (no API calls) mistral-ocr ./documents --dry-run @@ -65,49 +65,59 @@ Arguments: INPUT_PATH Path to input file or directory (required) Options: - -o, --output-path PATH Output directory (default: /mistral_ocr_output/) + -o, --output-dir PATH Output root (default: /ocr/). Never required. --api-key TEXT Mistral API key (or set MISTRAL_API_KEY env var) --model TEXT OCR model (default: mistral-ocr-latest) --env-file PATH Path to .env file - --include-images/--no-images Extract embedded images (default: True) - --save-originals/--no-save-originals Copy original files to output (default: True) - --metadata/--no-metadata Include markdown header block (default: True) - --page-headings/--no-page-headings Include "## Page N" headings (default: True) + --include-images/--no-images Extract embedded figures (default: True) - --table-format [markdown|html] Extract tables separately (OCR 3+) + --table-format [markdown|html] Request tables inline in a given format (OCR 3+) --extract-headers/--no-extract-headers Extract page headers (OCR 3+) --extract-footers/--no-extract-footers Extract page footers (OCR 3+) + --max-pages N Max PDF pages to process (default: all pages) -w, --workers N Concurrent workers for batch processing (default: 1) - --reprocess Reprocess already-processed files - --add-timestamp/--no-timestamp Timestamp output folder name (default: False) + --reprocess Re-OCR files already recorded completed (checksum-based) --dry-run List files without calling the API - -q, --quiet Suppress all output except errors + -q, --quiet Suppress output except the written .md paths (for scripting) -v, --verbose Enable verbose/debug output --log-file PATH Write logs to file --version Show version --help Show this message ``` +> Output writing is owned by the shared +> [`ocr-output-contract`](https://github.com/r-uben/ocr-output-contract) package, so +> mistral's output structure is byte-identical to every sibling OCR engine CLI. The +> markdown body is always clean (`## Page N` headers, no header block, no YAML +> frontmatter); all provenance lives in the JSON sidecars. The `--save-originals`, +> `--metadata`, `--page-headings` and `--add-timestamp` flags are deprecated no-ops kept +> only for invocation compatibility. + ## Output structure +Default output root is `/ocr/` (`-o` overrides verbatim; never required). +Each source document gets one aggregated folder, mirroring the input subtree so +same-basename inputs in different directories never collide: + ``` -mistral_ocr_output/ +ocr/ ├── document_name/ -│ ├── document_name.pdf # original copy (if --save-originals) -│ ├── document_name.md # OCR markdown -│ ├── figures/ # extracted images -│ │ ├── page1_img1.png -│ │ └── page2_img1.png -│ └── tables/ # extracted tables (if --table-format) -│ └── page1_table1.md -├── another_document/ +│ ├── document_name.md # all pages, joined under "## Page N" headers (clean body) +│ ├── figures/ # extracted embedded images (normalised to PNG) +│ │ ├── figure_1_page1.png +│ │ └── figure_2_page2.png +│ └── metadata.json # per-document sidecar: status/checksum/model/backend/... +├── sub/dir/another_document/ │ └── ... -└── metadata.json # processing stats, file list, errors +└── metadata.json # root index, keyed by input-relative path ``` -Use `--no-metadata` and `--no-page-headings` for cleaner markdown output without the header block and page separators. +Resume is content-aware: a file recorded `completed` is skipped only when its SHA-256 +checksum still matches, so editing a file in place forces a re-OCR. Failures are recorded +with `status="failed"`, and any file/page failure drives a nonzero exit (uniform across +single-file and batch runs). ## Configuration @@ -118,12 +128,10 @@ All CLI options can also be set via environment variables or a `.env` file: | `--api-key` | `MISTRAL_API_KEY` | (required) | | `--model` | `MISTRAL_MODEL` | `mistral-ocr-latest` | | `--include-images` | `INCLUDE_IMAGES` | `true` | -| `--save-originals` | `SAVE_ORIGINAL_IMAGES` | `true` | -| `--metadata` | `INCLUDE_METADATA` | `true` | -| `--page-headings` | `INCLUDE_PAGE_HEADINGS` | `true` | | `--table-format` | `TABLE_FORMAT` | (none) | | `--extract-headers` | `EXTRACT_HEADER` | `false` | | `--extract-footers` | `EXTRACT_FOOTER` | `false` | +| `--max-pages` | `MAX_PAGES` | (all pages) | | `--workers` | `MAX_WORKERS` | `1` | | `--verbose` | `VERBOSE` | `false` | | | `MAX_FILE_SIZE_MB` | `50` | diff --git a/mistral_ocr/cli.py b/mistral_ocr/cli.py index 76c101f..e26d13f 100644 --- a/mistral_ocr/cli.py +++ b/mistral_ocr/cli.py @@ -1,4 +1,17 @@ -"""Command-line interface for Mistral OCR.""" +"""Command-line interface for Mistral OCR. + +Output is written through the shared ``ocr-output-contract`` package, so +mistral's output structure is byte-identical to every sibling engine:: + + mistral-ocr [-o DIR] [--model M] [--include-images/--no-images] + [--table-format markdown|html] [--extract-headers] [--extract-footers] + [--max-pages N] [-w N] [--reprocess] [--dry-run] [-q] [-v] + mistral-ocr --version + +The default output root is ``/ocr/``; ``-o`` overrides verbatim but +is never required. The exit code is nonzero if any file failed (uniform across +single-file and batch). +""" import logging import os @@ -6,6 +19,7 @@ from pathlib import Path import click +from ocr_output_contract import resolve_output_root from rich.logging import RichHandler from . import __version__ @@ -20,11 +34,12 @@ @click.command() @click.argument("input_path", type=click.Path(path_type=Path), required=True) @click.option( - "--output-path", + "--output-dir", "-o", + "output_dir", type=click.Path(path_type=Path), required=False, - help="Path to output directory (default: /mistral_ocr_output/)", + help="Output root (default: /ocr/). Writes /.md per document.", ) @click.option( "--api-key", @@ -35,8 +50,8 @@ @click.option( "--model", type=str, - default="mistral-ocr-latest", - help="Mistral OCR model to use (default: mistral-ocr-latest)", + default=None, + help="Mistral OCR model to use (default: mistral-ocr-latest / $MISTRAL_MODEL)", ) @click.option( "--env-file", @@ -45,45 +60,23 @@ ) @click.option( "--include-images/--no-images", - default=True, - help="Include extracted images in output (default: True)", -) -@click.option( - "--save-originals/--no-save-originals", - default=True, - help="Save original input images alongside OCR results (default: True)", -) -@click.option( - "--metadata/--no-metadata", - "include_metadata", - default=True, - help="Include markdown metadata header block (default: True)", -) -@click.option( - "--page-headings/--no-page-headings", - "include_page_headings", - default=True, - help="Include markdown headings for each OCR page (default: True)", -) -@click.option( - "--add-timestamp/--no-timestamp", - default=False, - help="Add timestamp to output folder name (default: False)", + default=None, + help="Extract embedded figures as figures/figure__page

.png (default: True)", ) @click.option( "--table-format", type=click.Choice(["markdown", "html"], case_sensitive=False), default=None, - help="Extract tables in a separate format (markdown or html). OCR 3+ only.", + help="Render extracted tables into the page body as markdown or html. OCR 3+ only.", ) @click.option( "--extract-headers/--no-extract-headers", - default=False, + default=None, help="Extract page headers (default: False). OCR 3+ only.", ) @click.option( "--extract-footers/--no-extract-footers", - default=False, + default=None, help="Extract page footers (default: False). OCR 3+ only.", ) @click.option( @@ -96,14 +89,14 @@ "--workers", "-w", type=click.IntRange(min=1), - default=1, + default=None, help="Number of concurrent workers for batch processing (default: 1)", ) @click.option( "--reprocess", is_flag=True, default=False, - help="Reprocess files even if they already exist in metadata (default: False)", + help="Reprocess files even if recorded completed with a matching checksum.", ) @click.option( "--dry-run", @@ -111,7 +104,16 @@ default=False, help="List files that would be processed without calling the API", ) -@click.option("--quiet", "-q", is_flag=True, help="Suppress all output except errors") +# Deprecated no-ops kept for invocation compatibility. The canonical output +# contract owns the markdown body (always clean: ## Page N, no header block, no +# frontmatter, no original copy), so these flags no longer change anything. +@click.option("--save-originals/--no-save-originals", default=None, hidden=True) +@click.option("--metadata/--no-metadata", "include_metadata", default=None, hidden=True) +@click.option("--page-headings/--no-page-headings", "page_headings", default=None, hidden=True) +@click.option("--add-timestamp/--no-timestamp", default=None, hidden=True) +@click.option( + "--quiet", "-q", is_flag=True, help="Suppress output except file paths (for scripting)" +) @click.option("--verbose", "-v", is_flag=True, help="Enable verbose output") @click.option( "--log-file", @@ -122,61 +124,45 @@ @click.version_option(version=__version__, prog_name="mistral-ocr") def main( input_path: Path, - output_path: Path | None, + output_dir: Path | None, api_key: str | None, - model: str, + model: str | None, env_file: Path | None, - include_images: bool, - save_originals: bool, - include_metadata: bool, - include_page_headings: bool, - add_timestamp: bool, + include_images: bool | None, table_format: str | None, - extract_headers: bool, - extract_footers: bool, + extract_headers: bool | None, + extract_footers: bool | None, max_pages: int | None, - workers: int, + workers: int | None, reprocess: bool, dry_run: bool, + save_originals: bool | None, # deprecated no-op + include_metadata: bool | None, # deprecated no-op + page_headings: bool | None, # deprecated no-op + add_timestamp: bool | None, # deprecated no-op quiet: bool, verbose: bool, log_file: Path | None, ) -> None: - """ - Mistral OCR - Process documents using Mistral AI's OCR API. - - This tool processes PDF and image files using Mistral's powerful OCR capabilities, - extracting text, tables, equations, and images with high accuracy. + """Mistral OCR - Process documents using Mistral AI's OCR API. + \b Examples: - - # Process a single PDF file mistral-ocr document.pdf - - # Process all files in a directory - mistral-ocr ./documents --output-path ./results - - # Use a specific .env file + mistral-ocr ./documents -o ./results mistral-ocr doc.pdf --env-file .env.production """ try: - # Resolve input path relative to original working directory if not input_path.is_absolute(): input_path = Path(ORIGINAL_CWD) / input_path - - # Check if input path exists if not input_path.exists(): raise ValueError(f"Input path does not exist: {input_path}") + if output_dir and not output_dir.is_absolute(): + output_dir = Path(ORIGINAL_CWD) / output_dir - # Resolve output path if provided - if output_path and not output_path.is_absolute(): - output_path = Path(ORIGINAL_CWD) / output_path - - # Quiet mode: suppress non-error output (uses processor's shared console) if quiet: console.quiet = True - # Configure logging — use verbose flag now, may upgrade later from config log_level = logging.DEBUG if verbose else logging.WARNING handlers: list[logging.Handler] = [] if not quiet: @@ -186,122 +172,72 @@ def main( file_handler.setFormatter( logging.Formatter("%(asctime)s %(levelname)s %(name)s: %(message)s") ) - file_handler.setLevel(logging.DEBUG) + # Track effective verbosity rather than pinning to DEBUG, so a plain + # --log-file does not capture httpx/SDK DEBUG bodies. + file_handler.setLevel(log_level) handlers.append(file_handler) logging.basicConfig(level=log_level, handlers=handlers, force=True) - # Print header - console.print("\n[bold blue]🔍 Mistral OCR[/bold blue]") - console.print("[dim]Powered by Mistral AI's OCR API[/dim]\n") + if not quiet: + console.print("\n[bold blue]Mistral OCR[/bold blue]") + console.print("[dim]Powered by Mistral AI's OCR API[/dim]\n") - # Dry-run: list files that would be processed, then exit (no API key needed) + # Dry-run: list files that would be processed (no API key needed). if dry_run: - if input_path.is_file(): - size = format_file_size(input_path.stat().st_size) - console.print(f" {input_path.name} ({size})") - console.print("\n[dim]1 file would be processed (dry run)[/dim]") - elif input_path.is_dir(): - files = get_supported_files(input_path) - if not files: - console.print("[yellow]No supported files found.[/yellow]") - else: - for f in files: - size = format_file_size(f.stat().st_size) - console.print(f" {f.relative_to(input_path)} ({size})") - console.print(f"\n[dim]{len(files)} file(s) would be processed (dry run)[/dim]") + _dry_run(input_path, output_dir) return - # Load configuration (requires API key — after dry-run check) - - # If API key is provided via CLI, set it before loading config - # (must happen before load_dotenv, which won't override existing vars) + # API key provided via CLI: set it before load_dotenv (won't override). if api_key: os.environ["MISTRAL_API_KEY"] = api_key - # Create config from environment config = Config.from_env(env_file) - # If config has VERBOSE=true but CLI didn't pass --verbose, upgrade log level if config.verbose and not verbose: logging.getLogger().setLevel(logging.DEBUG) - # Only override config with CLI options that were explicitly passed + # Only override config with CLI options that were explicitly passed, so + # env-var / .env precedence is respected rather than clobbered by defaults. ctx = click.get_current_context() - if ( - "model" in ctx.params - and ctx.get_parameter_source("model") != click.core.ParameterSource.DEFAULT - ): + + def _set(name: str) -> bool: + return ctx.get_parameter_source(name) != click.core.ParameterSource.DEFAULT + + if model is not None: config.model = model - if ( - "include_images" in ctx.params - and ctx.get_parameter_source("include_images") != click.core.ParameterSource.DEFAULT - ): + if include_images is not None: config.include_images = include_images - if ( - "save_originals" in ctx.params - and ctx.get_parameter_source("save_originals") != click.core.ParameterSource.DEFAULT - ): - config.save_original_images = save_originals - if ( - "include_metadata" in ctx.params - and ctx.get_parameter_source("include_metadata") != click.core.ParameterSource.DEFAULT - ): - config.include_metadata = include_metadata - if ( - "include_page_headings" in ctx.params - and ctx.get_parameter_source("include_page_headings") - != click.core.ParameterSource.DEFAULT - ): - config.include_page_headings = include_page_headings - if ( - "verbose" in ctx.params - and ctx.get_parameter_source("verbose") != click.core.ParameterSource.DEFAULT - ): - config.verbose = verbose - if ( - "table_format" in ctx.params - and ctx.get_parameter_source("table_format") != click.core.ParameterSource.DEFAULT - ): + if table_format is not None: config.table_format = table_format - if ( - "extract_headers" in ctx.params - and ctx.get_parameter_source("extract_headers") != click.core.ParameterSource.DEFAULT - ): + if extract_headers is not None: config.extract_header = extract_headers - if ( - "extract_footers" in ctx.params - and ctx.get_parameter_source("extract_footers") != click.core.ParameterSource.DEFAULT - ): + if extract_footers is not None: config.extract_footer = extract_footers - if ( - "workers" in ctx.params - and ctx.get_parameter_source("workers") != click.core.ParameterSource.DEFAULT - ): + if workers is not None: config.max_workers = workers - if ( - "max_pages" in ctx.params - and ctx.get_parameter_source("max_pages") != click.core.ParameterSource.DEFAULT - ): + if max_pages is not None: config.max_pages = max_pages - + if _set("verbose"): + config.verbose = verbose config.quiet = quiet - # Create processor processor = OCRProcessor(config) + outcome = processor.process(input_path, output_dir, reprocess=reprocess) - # Process input - processor.process(input_path, output_path, add_timestamp=add_timestamp, reprocess=reprocess) - - # Print summary - if processor.errors: - if verbose: - console.print("\n[yellow]⚠ Errors encountered:[/yellow]") - for error in processor.errors: - console.print(f" [red]• {error['file']}: {error['error']}[/red]") - console.print("\n[bold yellow]⚠ Processing complete with errors.[/bold yellow]\n") - sys.exit(1) + if quiet: + # Scripting contract: emit one output .md path per line on stdout. + for path in outcome.outputs: + click.echo(path) + elif outcome.has_failures: + console.print( + f"\n[bold yellow]Processing complete with {outcome.failed} failure(s).[/bold yellow]\n" + ) + else: + console.print("\n[bold green]Processing complete![/bold green]\n") - console.print("\n[bold green]✨ Processing complete![/bold green]\n") + # Uniform exit policy (canon SYS-02): nonzero if any file failed. + if outcome.exit_code != 0: + sys.exit(outcome.exit_code) except ValueError as e: console.print(f"\n[red]Error: {e}[/red]\n") @@ -315,5 +251,29 @@ def main( sys.exit(1) +def _dry_run(input_path: Path, output_dir: Path | None) -> None: + """List files that would be processed without calling the API. + + Discovery mirrors the real run exactly: it resolves the same output root and + delegates to the contract's discovery so the dry-run never under- or + over-reports relative to what would actually be OCR'd (and so it surfaces the + output-root exclusion rather than diverging from it). + """ + output_root = resolve_output_root(input_path, output_dir) + if input_path.is_file(): + size = format_file_size(input_path.stat().st_size) + console.print(f" {input_path.name} ({size})") + console.print("\n[dim]1 file would be processed (dry run)[/dim]") + else: + files = get_supported_files(input_path, output_root) + if not files: + console.print("[yellow]No supported files found.[/yellow]") + return + for f in files: + size = format_file_size(f.stat().st_size) + console.print(f" {f.relative_to(input_path)} ({size})") + console.print(f"\n[dim]{len(files)} file(s) would be processed (dry run)[/dim]") + + if __name__ == "__main__": main() diff --git a/mistral_ocr/config.py b/mistral_ocr/config.py index a95fa33..198888f 100644 --- a/mistral_ocr/config.py +++ b/mistral_ocr/config.py @@ -9,23 +9,26 @@ @dataclass class Config: - """Configuration for Mistral OCR.""" + """Configuration for Mistral OCR. + + The markdown body is owned by the canonical output contract: it is always + clean (``## Page N`` headers, no ``# OCR Results`` block, no YAML + frontmatter) and never copies the original input. The old + ``save_original_images`` / ``include_metadata`` / ``include_page_headings`` + knobs are therefore gone; provenance lives only in the JSON sidecars. + """ api_key: str model: str = "mistral-ocr-latest" max_file_size_mb: int = 50 include_images: bool = True - save_original_images: bool = True table_format: str | None = None # None, "markdown", or "html" extract_header: bool = False 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 - dry_run: bool = False quiet: bool = False verbose: bool = False @@ -90,7 +93,6 @@ def from_env(cls, env_file: Path | None = None) -> "Config": model=os.getenv("MISTRAL_MODEL", "mistral-ocr-latest"), max_file_size_mb=max_file_size_mb, include_images=os.getenv("INCLUDE_IMAGES", "true").lower() == "true", - save_original_images=os.getenv("SAVE_ORIGINAL_IMAGES", "true").lower() == "true", table_format=table_fmt, extract_header=os.getenv("EXTRACT_HEADER", "false").lower() == "true", extract_footer=os.getenv("EXTRACT_FOOTER", "false").lower() == "true", @@ -98,8 +100,6 @@ def from_env(cls, env_file: Path | None = None) -> "Config": max_workers=max(1, int(os.getenv("MAX_WORKERS", "1"))), max_retries=max_retries, retry_base_delay=retry_base_delay, - include_metadata=os.getenv("INCLUDE_METADATA", "true").lower() == "true", - include_page_headings=os.getenv("INCLUDE_PAGE_HEADINGS", "true").lower() == "true", verbose=os.getenv("VERBOSE", "false").lower() == "true", ) diff --git a/mistral_ocr/processor.py b/mistral_ocr/processor.py index e5a0340..8fb68fe 100644 --- a/mistral_ocr/processor.py +++ b/mistral_ocr/processor.py @@ -1,15 +1,74 @@ -"""Core OCR processing module using Mistral AI.""" - +"""Core OCR processing using the Mistral OCR API, emitting canonical output. + +Mistral is a *cloud document-OCR API*: a file is uploaded once and the API +returns a per-page result list (each page carrying ``.markdown`` plus optional +OCR 3 fields — page dimensions, header/footer, tables, hyperlinks, embedded +images). This module owns *how OCR happens* (the API round-trip, large-PDF +chunking + global page re-indexing, and folding the per-page OCR 3 extras into +each page's markdown). The shared ``ocr-output-contract`` package owns *where +the bytes go and what the metadata looks like*, so mistral's output is +byte-structure-identical to every sibling engine. + +The per-page text list this module produces is fed straight into the contract's +:func:`assemble_pages`, so every page of a document lands in ONE +``///.md`` under ``## Page N`` headers (the API's +whole-doc response is never dumped as a single blob). Embedded images are +normalised to PNG and named ``figures/figure__page

.png`` with resolving +links. Provenance lives only in the dual-level ``metadata.json`` sidecars — the +markdown body is clean (no ``# OCR Results`` header, no YAML frontmatter). + +Audit fixes baked in here: + +* **HIGH (idempotency)** — resume is now content-aware via the package's + :meth:`RootIndex.is_completed` (status==completed AND SHA256 match), replacing + the old resolved-path-equality skip that served stale output on in-place edits. +* **Metadata** — the non-standard ``files_processed``/``errors`` list schema is + gone wholesale; the dual-level ``DocMetadata`` + ``RootIndex`` carries + ``version``/``checksum``/``model``/``backend`` and records failures + (``status=failed``). +* **Retry** — the catch-all that retried every ``SDKError`` (incl. permanent + 4xx) now excludes non-429 4xx, so auth/validation errors fail fast. +* **MAX_FILE_SIZE_MB** — now enforced for PDFs too, not only the non-PDF branch. +* **Exit policy** — a :class:`RunOutcome` drives a uniform nonzero exit on any + failure, across single-file and batch. +""" + +from __future__ import annotations + +import io import logging -import shutil +import re import tempfile import threading import time from concurrent.futures import ThreadPoolExecutor, as_completed +from dataclasses import dataclass, field from pathlib import Path +from typing import Any from mistralai import Mistral -from mistralai import models as mistral_models +from ocr_output_contract import ( + FIGURES_DIRNAME, + DocMetadata, + RootIndex, + RunOutcome, + Status, + assemble_pages, + doc_dir_for, + failure_checksum, + figure_filename, + figure_markdown_link, + figures_dir_for, + markdown_path_for, + relative_key, + resolve_output_root, + run_fingerprint, + safe_checksum, + sha256_checksum, + utc_timestamp, + write_doc_metadata, +) +from PIL import Image from rich.console import Console from rich.progress import BarColumn, Progress, SpinnerColumn, TextColumn, TimeRemainingColumn @@ -17,14 +76,10 @@ from .utils import ( DOCUMENT_EXTENSIONS, create_data_uri, - determine_output_path, + decode_base64_image, format_file_size, get_pdf_page_count, get_supported_files, - load_metadata, - make_unique_basename, - save_base64_image, - save_metadata, split_pdf, ) @@ -33,12 +88,148 @@ # Shared console instance — CLI sets .quiet on this directly console = Console() +# Dedicated stderr console for per-file failures. It is NEVER muted by --quiet +# (the CLI only sets console.quiet), so SYS-02's "failures still emitted to +# stderr under --quiet" holds while stdout stays a clean one-path-per-line list +# of successful outputs for scripting. +err_console = Console(stderr=True) + +#: Backend identifier recorded in metadata. Mistral is a cloud OCR API. +BACKEND = "mistral-api" + # Mistral API limit: max pages per single OCR request MAX_PAGES_PER_REQUEST = 1000 +#: HTTP statuses that are transient and worth retrying. Everything else in the +#: 4xx range (401 auth, 400/422 validation, ...) is permanent and must fail fast. +_RETRYABLE_STATUSES = {429, 500, 502, 503, 504} + +#: Matches a markdown inline image link ``![alt](target)``, capturing the alt +#: text and the raw target separately. Used to rewrite/strip the API's inline +#: image placeholders so the produced body never carries a dangling local link. +_INLINE_IMAGE_RE = re.compile(r"!\[([^\]]*)\]\(([^)]+)\)") + + +def _is_external_target(target: str) -> bool: + """True for a target the conformance harness does not resolve on disk.""" + t = target.strip().lower() + return t.startswith(("http://", "https://", "data:", "mailto:", "//", "#")) or t.startswith( + " tuple[str, bool]: + """Rewrite an inline ``![..](image_id)`` placeholder to a canonical link. + + Replaces every inline image whose raw target equals ``image_id`` (the API's + placeholder reference, e.g. ``img-0.jpeg``) with a link to the canonical + figure file ``./figures/``, preserving the original alt text. Returns + ``(new_markdown, replaced)`` where ``replaced`` is True iff at least one + placeholder matched. Matching is on the exact target token so unrelated + image links are untouched. + """ + if not image_id: + return markdown, False + replaced = False + + def _sub(match: re.Match[str]) -> str: + nonlocal replaced + alt, raw = match.group(1), match.group(2).strip() + if raw == image_id: + replaced = True + label = alt or f"Figure (page {target})" + return ( + f"![{alt}](./{FIGURES_DIRNAME}/{target})" + if alt + else f"![{label}](./{FIGURES_DIRNAME}/{target})" + ) + return match.group(0) + + return _INLINE_IMAGE_RE.sub(_sub, markdown), replaced + + +def _strip_unresolved_image_links(markdown: str) -> str: + """Drop any inline image link whose LOCAL target was never written. + + A figure that failed to save (Pillow could not decode/encode it) leaves the + API's ``![..](img-N.jpeg)`` placeholder unrewritten. The conformance harness + resolves every local inline link on disk, so such a dangling placeholder + would FAIL conformance. We strip it (degrading to the alt text, or nothing) + rather than ship a broken link. External targets (http/data/anchors) are + left untouched. + """ + + def _sub(match: re.Match[str]) -> str: + alt, raw = match.group(1), match.group(2).strip() + # Already-canonical figure links and external targets are kept verbatim. + if raw.startswith(f"./{FIGURES_DIRNAME}/") or raw.startswith(f"{FIGURES_DIRNAME}/"): + return match.group(0) + if _is_external_target(raw): + return match.group(0) + # A bare local target (e.g. img-0.jpeg) that survived rewriting points at + # a figure that was never written: drop the link, keep any alt text. + return alt + + return _INLINE_IMAGE_RE.sub(_sub, markdown) + + +@dataclass(frozen=True) +class _SavedFigure: + """A figure that was written to disk, with the API id that referenced it. + + ``image_id`` is the API's original image reference (e.g. ``img-0.jpeg``) used + to rewrite the matching inline placeholder; ``figure_number`` / ``page_number`` + give the canonical ``figure__page

.png`` identity. + """ + + image_id: str | None + figure_number: int + page_number: int + + +@dataclass +class OCRResult: + """Result of OCR'ing one source document. + + ``pages`` holds the per-page markdown in order (already including any folded + OCR 3 extras). ``page_images`` maps a 1-indexed page number to the list of + ``(image_id, raw_bytes)`` pairs extracted from that page. The ``image_id`` is + the API's reference for the image (e.g. ``img-0.jpeg``); it is what the page + markdown's inline ``![img-0.jpeg](img-0.jpeg)`` placeholder points at, so we + carry it through to rewrite that placeholder to the canonical figure filename + (otherwise the body keeps a dangling local link that fails conformance). + ``status`` maps to the contract enum: a document that could not be + opened/processed at all is ``FAILED``; one whose API call returned zero pages + is also ``FAILED``. + """ + + file_path: Path + pages: list[str] + page_images: dict[int, list[tuple[str | None, bytes]]] = field(default_factory=dict) + processing_time: float = 0.0 + error: str | None = None + #: Non-fatal note (e.g. ``--max-pages`` truncation), recorded in metadata. + note: str | None = None + + @property + def page_count(self) -> int: + return len(self.pages) + + @property + def status(self) -> Status: + """``completed`` when the document yielded >=1 page; else ``failed``. + + The Mistral API does not surface per-page success/failure (it returns a + page list or raises), so there is no ``partial`` state for a single doc: + either the call produced pages (completed) or it did not (failed). + """ + if self.pages and not self.error: + return Status.COMPLETED + return Status.FAILED + class OCRProcessor: - """OCR processor using Mistral AI API.""" + """OCR processor using the Mistral AI API, routed through the contract.""" def __init__(self, config: Config): """Initialize the OCR processor.""" @@ -48,27 +239,37 @@ def __init__(self, config: Config): except (ValueError, TypeError, RuntimeError) as e: console.print(f"[red]Failed to initialize Mistral client: {e}[/red]") raise - self.errors: list[dict] = [] - self.processed_files: list[dict] = [] self._lock = threading.Lock() + # ------------------------------------------------------------------ + # API round-trip (mistral-owned: how OCR happens) + # ------------------------------------------------------------------ + @staticmethod def _is_retryable(error: Exception) -> bool: - """Check if an error is transient and worth retrying.""" + """Check if an error is transient and worth retrying. + + Fix (audit MEDIUM): a permanent 4xx ``SDKError`` (401 auth, 400/422 + validation) is NOT retried. Only the explicit transient statuses, the + SDK's typed rate-limit/server errors, and network-level errors retry. + An ``SDKError`` whose status we can read is gated on that status; one + whose status we cannot read is treated as non-retryable (fail fast) + rather than blindly retried. + """ # SDK-typed rate limit / server errors for exc_name in ("RateLimitError", "InternalServerError", "ServiceUnavailableError"): if type(error).__name__ == exc_name: return True - # httpx-level HTTP status errors - if hasattr(error, "response"): - status = getattr(error.response, "status_code", 0) - if status in (429, 500, 502, 503, 504): - return True - # Network-level transient errors - if isinstance(error, (TimeoutError, ConnectionError, OSError)): - return True - # Catch SDK errors that wrap HTTP status codes - return hasattr(mistral_models, "SDKError") and isinstance(error, mistral_models.SDKError) + # httpx-level / SDK HTTP status errors: only the transient set retries. + status = getattr(getattr(error, "response", None), "status_code", None) + if status is None: + status = getattr(error, "status_code", None) + if isinstance(status, int): + return status in _RETRYABLE_STATUSES + # Network-level transient errors retry; an SDKError with no readable + # status is permanent (fail fast), NOT a blanket retry — that catch-all + # was the old bug that retried permanent 4xx. + return isinstance(error, (TimeoutError, ConnectionError, OSError)) def _call_with_retry(self, **ocr_kwargs: object) -> object: """Call ocr.process with exponential backoff on transient errors.""" @@ -110,7 +311,11 @@ def _build_ocr_kwargs(self, document: dict) -> dict: return ocr_kwargs def _upload_and_process(self, file_path: Path) -> object: - """Upload a file via Mistral files API and process with OCR.""" + """Upload a file via Mistral files API and process with OCR. + + The uploaded file is always deleted in a ``finally`` block so transient + OCR files do not accumulate in the user's Mistral account/quota. + """ with open(file_path, "rb") as f: uploaded = self.client.files.upload( file={"file_name": file_path.name, "content": f}, @@ -126,13 +331,18 @@ def _upload_and_process(self, file_path: Path) -> object: 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.""" + """Process a PDF, chunking if needed for the API page limit. + + Fix (audit LOW): the configured ``MAX_FILE_SIZE_MB`` is enforced here for + PDFs too (the old code only validated the non-PDF branch). + """ + self.config.validate_file_size(file_path) + 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) @@ -148,7 +358,12 @@ def _process_pdf(self, file_path: Path) -> object: 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.""" + """Split a PDF into chunks, process each, and reassemble pages. + + Pages are re-indexed to their position in the original document and the + original page objects are passed through to preserve all OCR 3 fields + (tables, headers, footers, hyperlinks, dimensions). + """ from types import SimpleNamespace max_pages = self.config.max_pages @@ -167,433 +382,693 @@ def _process_pdf_chunked(self, file_path: Path, total_pages: int) -> object: "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: - file_size_mb = file_path.stat().st_size / (1024 * 1024) - logger.debug("File size: %.2f MB", file_size_mb) - - if not hasattr(self.client, "ocr"): - raise AttributeError( - "OCR endpoint not available in Mistral client. " - "Please ensure you have the latest mistralai package " - "and OCR access enabled for your API key." - ) + def _call_api(self, file_path: Path) -> object: + """Run the Mistral OCR API on one file (PDF chunked, others via data URI).""" + if not hasattr(self.client, "ocr"): + raise AttributeError( + "OCR endpoint not available in Mistral client. " + "Please ensure you have the latest mistralai package " + "and OCR access enabled for your API key." + ) + logger.debug("Sending to Mistral OCR API (model=%s)...", self.config.model) - logger.debug("Sending to Mistral OCR API (model=%s)...", self.config.model) + if file_path.suffix.lower() == ".pdf": + return self._process_pdf(file_path) - if file_path.suffix.lower() == ".pdf": - response = self._process_pdf(file_path) + # 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} + return self._call_with_retry(**self._build_ocr_kwargs(document)) + + # ------------------------------------------------------------------ + # Page parsing: API page list -> per-page markdown + figure bytes + # ------------------------------------------------------------------ + + def _render_page_markdown(self, page: Any) -> str: + """Fold one API page's OCR 3 extras into a single markdown body. + + The contract supplies the ``## Page N`` header at assembly time, so this + produces only the page *body*: optional page dimensions, header/footer + quotes, the OCR text, the structured tables, and hyperlinks. The API's + inline image placeholders (``![img-0.jpeg](img-0.jpeg)``) are kept here + verbatim and later rewritten to canonical figure links / stripped by + :meth:`_render_page_with_figures` once figures are saved, so the produced + body never carries a dangling local image link. + + Tables: when ``--table-format=markdown|html`` is set, Mistral OCR returns + tables in the structured ``page.tables`` field, NOT inline in + ``page.markdown`` (the markdown may carry a placeholder referencing the + table ``id``). Each structured table is rendered into the body so the + ``--table-format`` flag is not a silent no-op: a placeholder is replaced + in place when present, otherwise the table is appended. + """ + parts: list[str] = [] + + dims = getattr(page, "dimensions", None) + if dims: + w = getattr(dims, "width", None) + h = getattr(dims, "height", None) + if w and h: + parts.append(f"*Page size: {w} x {h}*") + + header = getattr(page, "header", None) + if header: + parts.append(f"> **Header:** {header}") + + markdown = getattr(page, "markdown", None) or "" + markdown, appended_tables = self._render_tables(markdown, getattr(page, "tables", None)) + if markdown: + parts.append(markdown) + if appended_tables: + parts.append(appended_tables) + + hyperlinks = getattr(page, "hyperlinks", None) + if hyperlinks: + lines = ["**Hyperlinks:**"] + for link in hyperlinks: + text = getattr(link, "text", "") or "" + url = getattr(link, "url", "") or getattr(link, "href", "") or "" + if url: + lines.append(f"- [{text or url}]({url})") + if len(lines) > 1: + parts.append("\n".join(lines)) + + footer = getattr(page, "footer", None) + if footer: + parts.append(f"> **Footer:** {footer}") + + return "\n\n".join(parts).strip() + + @staticmethod + def _render_one_table(table: Any) -> str: + """Render one structured OCR-3 table object into markdown body text. + + ``content`` already carries the table in the requested ``--table-format`` + (a markdown table or an HTML ````); HTML is passed through verbatim + (markdown renderers accept inline HTML), markdown likewise. Empty content + renders to an empty string. + """ + content: str = str(getattr(table, "content", "") or "") + return content.strip() + + @classmethod + def _render_tables(cls, markdown: str, tables: Any) -> tuple[str, str]: + """Fold structured ``page.tables`` into the page body. + + Returns ``(markdown, appended)`` where ``markdown`` has any in-body table + placeholder (a markdown link/text referencing the table ``id``) replaced + by the rendered table, and ``appended`` is the concatenation of tables + that had no placeholder to anchor them (so a populated ``tables`` field is + never silently discarded). When ``--table-format`` is NOT set the API + inlines tables in ``page.markdown`` and ``tables`` is empty, so this is a + no-op and the existing inline content is preserved. + """ + if not tables: + return markdown, "" + + appended: list[str] = [] + for table in tables: + rendered = cls._render_one_table(table) + if not rendered: + continue + table_id = getattr(table, "id", None) + placeholder = cls._table_placeholder(markdown, table_id) + if placeholder is not None: + markdown = markdown.replace(placeholder, rendered) 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)) + appended.append(rendered) + return markdown, "\n\n".join(appended) - return {"file_path": file_path, "response": response, "success": True} + @staticmethod + def _table_placeholder(markdown: str, table_id: Any) -> str | None: + """Return the exact placeholder substring in ``markdown`` for a table id. - except Exception as e: - error_msg = f"Error processing {file_path.name}: {str(e)}" - console.print(f"[red]{error_msg}[/red]") - logger.debug("Traceback for %s", file_path.name, exc_info=True) - self.errors.append({"file": str(file_path.resolve()), "error": str(e)}) + Mistral may embed a table placeholder such as ``[tbl-0.html](tbl-0.html)`` + or a bare ``tbl-0.html`` token referencing the table ``id``. Returns the + matched substring to replace, or ``None`` if no placeholder is present. + """ + if not table_id or not markdown: return None + tid = str(table_id) + linked = f"[{tid}]({tid})" + if linked in markdown: + return linked + if tid in markdown: + return tid + return None - def save_results( - self, - result: dict, - output_dir: Path, - is_single_file: bool = False, - base_dir: Path | None = None, - ) -> None: - """Save OCR results in a per-document folder structure. - - Output layout: - output_dir/ - ├── doc_name/ - │ ├── doc_name.pdf # original copy - │ ├── doc_name.md # OCR markdown - │ ├── figures/ # extracted images - │ └── tables/ # extracted tables - └── metadata.json + @staticmethod + def _page_image_bytes(page: Any) -> list[tuple[str | None, bytes]]: + """Decode the embedded base64 images on one API page into raw bytes. + + Returns ``(image_id, raw_bytes)`` pairs, preserving the API's image id + (e.g. ``img-0.jpeg``) so the page markdown's inline placeholder pointing + at that id can later be rewritten to the canonical figure filename. """ - file_path = result["file_path"] - response = result["response"] + out: list[tuple[str | None, bytes]] = [] + for image in getattr(page, "images", None) or []: + b64 = getattr(image, "image_base64", None) or getattr(image, "base64", None) + if not b64: + continue + image_id = getattr(image, "id", None) + try: + out.append( + (str(image_id) if image_id is not None else None, decode_base64_image(b64)) + ) + except Exception as exc: + logger.warning("Failed to decode embedded image: %s", exc) + return out + + def _parse_response(self, file_path: Path, response: Any, start: float) -> OCRResult: + """Turn an API response into an OCRResult (per-page text + figure bytes).""" + pages_obj = getattr(response, "pages", None) or [] + pages: list[str] = [] + page_images: dict[int, list[bytes]] = {} + for page in pages_obj: + page_no = getattr(page, "index", len(pages)) + 1 + pages.append(self._render_page_markdown(page)) + if self.config.include_images: + imgs = self._page_image_bytes(page) + if imgs: + page_images[page_no] = imgs + note = getattr(response, "truncated", None) if hasattr(response, "truncated") else None + return OCRResult( + file_path=file_path, + pages=pages, + page_images=page_images, + processing_time=time.time() - start, + note=note, + ) - # Per-document folder - base_name = make_unique_basename(file_path, base_dir=base_dir) - doc_dir = output_dir / base_name - doc_dir.mkdir(parents=True, exist_ok=True) + def process_file(self, file_path: Path) -> OCRResult: + """OCR a single file into an OCRResult (no output written).""" + start = time.time() + try: + response = self._call_api(file_path) + result = self._parse_response(file_path, response, start) + if not result.pages: + result.error = "Empty response from Mistral OCR (no pages returned)" + return result + except Exception as e: + logger.error("Error processing %s: %s", file_path.name, e) + logger.debug("Traceback for %s", file_path.name, exc_info=True) + return OCRResult( + file_path=file_path, + pages=[], + processing_time=time.time() - start, + error=str(e), + ) - # Copy original file into the document folder - if self.config.save_original_images: - original_copy = doc_dir / f"{base_name}{file_path.suffix}" - shutil.copy2(file_path, original_copy) - logger.debug("Saved original to %s", original_copy) + # ------------------------------------------------------------------ + # Output writing (all routed through the ocr-output-contract package) + # ------------------------------------------------------------------ - markdown_content = [] + def save_results(self, result: OCRResult, output_root: Path, rel_key: str) -> Path: + """Write the aggregated markdown + figures for one document. - # File header (optional metadata block) - if self.config.include_metadata: - markdown_content.append("# OCR Results\n\n") - markdown_content.append(f"**Original File:** {file_path.name}\n") - markdown_content.append(f"**Full Path:** `{file_path}`\n") - markdown_content.append(f"**Processed:** {time.strftime('%Y-%m-%d %H:%M:%S')}\n\n") + Layout is owned entirely by the contract: + ``///.md`` plus a ``figures/`` folder. + Figures are normalised to PNG, named ``figure__page

.png``, and + linked from the page that produced them. The body is clean (no + ``# OCR Results`` header, no frontmatter). + """ + doc_dir = doc_dir_for(output_root, rel_key) + doc_dir.mkdir(parents=True, exist_ok=True) + markdown_path = markdown_path_for(doc_dir, rel_key) + + if result.pages: + page_figures = self._save_figures(result, doc_dir) + pages = [ + self._render_page_with_figures(text, page_figures.get(idx, [])) + for idx, text in enumerate(result.pages, start=1) + ] + body = assemble_pages(pages) + if result.note: + body = f"> **Note:** {result.note}\n\n" + body + else: + body = "*[OCR Failed]*\n" - if self.config.save_original_images: - markdown_content.append( - f"**Original:** [{base_name}{file_path.suffix}](./{base_name}{file_path.suffix})\n\n" - ) + markdown_path.write_text(body, encoding="utf-8") + if self.config.verbose: + console.print(f"[green]Saved:[/green] {markdown_path}") + return markdown_path - # 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 - if hasattr(response, "pages"): - for page in response.pages: - if self.config.include_page_headings: - markdown_content.append(f"## Page {page.index + 1}\n\n") - - # Page dimensions (OCR 3) - if hasattr(page, "dimensions") and page.dimensions: - dims = page.dimensions - w = getattr(dims, "width", None) - h = getattr(dims, "height", None) - if w and h: - markdown_content.append(f"*Page size: {w} x {h}*\n\n") - - # Header (OCR 3) - if hasattr(page, "header") and page.header: - markdown_content.append(f"> **Header:** {page.header}\n\n") - - # Extracted text - if hasattr(page, "markdown"): - markdown_content.append(page.markdown) - markdown_content.append("\n\n") - - # Tables (OCR 3) - if hasattr(page, "tables") and page.tables: - tables_dir = doc_dir / "tables" - tables_dir.mkdir(parents=True, exist_ok=True) - ext = "html" if self.config.table_format == "html" else "md" - for tidx, table in enumerate(page.tables): - table_content = ( - getattr(table, "content", None) - or getattr(table, "markdown", None) - or str(table) - ) - table_filename = f"page{page.index + 1}_table{tidx + 1}.{ext}" - table_path = tables_dir / table_filename - with open(table_path, "w", encoding="utf-8") as tf: - tf.write(table_content) - markdown_content.append( - f"[Table {tidx + 1}](./tables/{table_filename})\n\n" + @staticmethod + def _render_page_with_figures(text: str, figures: list[_SavedFigure]) -> str: + """Resolve a page's figures into its markdown body. + + For each saved figure, the API's inline placeholder ``![..](image_id)`` + is rewritten IN PLACE to the canonical ``./figures/figure__page

.png`` + link, so no dangling local image link survives (the conformance harness + resolves every inline link on disk; an unrewritten ``img-0.jpeg`` link + would point at a file that was never written and FAIL conformance). A + figure whose id is NOT referenced inline (the engine extracted an image + the model did not place in the markdown) is appended at the end so it is + still surfaced with a resolving link rather than silently dropped. + + Finally, any inline image link still pointing at a bare LOCAL target + (e.g. a placeholder whose figure failed to save, or one the API emitted + with no extractable image bytes) is stripped, so the body can never carry + a dangling local image link regardless of what the API returned. + """ + for fig in figures: + link = figure_markdown_link(fig.figure_number, fig.page_number) + target = figure_filename(fig.figure_number, fig.page_number) + replaced = False + if fig.image_id: + text, replaced = _rewrite_inline_image_link(text, fig.image_id, target) + if not replaced: + text = f"{text}\n\n{link}" if text else link + return _strip_unresolved_image_links(text) + + def _save_figures(self, result: OCRResult, doc_dir: Path) -> dict[int, list[_SavedFigure]]: + """Persist extracted images as PNG; return per-page saved-figure records. + + Figures are numbered globally (``figure_1``, ``figure_2``, ...) across the + whole document and tagged with their source page, matching the canonical + ``figure__page

.png`` naming. Each returned record carries the API's + original ``image_id`` so the page renderer can rewrite the matching inline + placeholder. A figure that fails to save is simply not recorded, so its + inline placeholder is left for the renderer to strip (never a dangling + local link in the produced body). + """ + if not (result.page_images and self.config.include_images): + return {} + + figures_dir = figures_dir_for(doc_dir) + figures_dir.mkdir(parents=True, exist_ok=True) + saved: dict[int, list[_SavedFigure]] = {} + figure_counter = 0 + for page_no in sorted(result.page_images): + for image_id, raw in result.page_images[page_no]: + figure_counter += 1 + filename = figure_filename(figure_counter, page_no) + img_path = figures_dir / filename + try: + image = Image.open(io.BytesIO(raw)) + if image.mode not in ("RGB", "RGBA"): + image = image.convert("RGB") + image.save(img_path, format="PNG") + saved.setdefault(page_no, []).append( + _SavedFigure( + image_id=image_id, + figure_number=figure_counter, + page_number=page_no, ) + ) + except Exception as exc: + logger.warning( + "Failed to save figure %d (page %d): %s", figure_counter, page_no, exc + ) + figure_counter -= 1 + return saved + + def _run_fingerprint(self) -> str: + """Fingerprint of the run config that affects *what output is produced*. + + Beyond model + backend, mistral's OCR-3 toggles change the produced + markdown, so they are passed to the contract's ``run_fingerprint`` via the + ``extra`` dict (the v0.1.2 channel for engine-specific output-affecting + flags) rather than being string-mangled into ``task``: the resolved table + format, header / footer extraction, and embedded-image extraction. A + re-run under a different ``--table-format`` / ``--extract-headers`` / + ``--extract-footers`` / ``--no-images`` / ``--model`` therefore + reprocesses instead of silently reusing a cached result keyed only on the + input checksum. ``extra`` carries the RESOLVED effective values (including + falsy ones), so ``True``/``False`` stay distinct and key order is + irrelevant. + """ + extra = { + "table_format": self.config.table_format or None, + "extract_header": bool(self.config.extract_header), + "extract_footer": bool(self.config.extract_footer), + "include_images": bool(self.config.include_images), + } + fingerprint: str = run_fingerprint(model=self.config.model, backend=BACKEND, extra=extra) + return fingerprint + + def _build_doc_metadata( + self, result: OCRResult, markdown_path: Path, output_root: Path + ) -> DocMetadata: + """Assemble the per-document metadata record from a result.""" + status = result.status + error = result.error if status is not Status.COMPLETED else None + return DocMetadata( + status=status, + checksum=sha256_checksum(result.file_path), + model=self.config.model, + backend=BACKEND, + processing_time=result.processing_time, + timestamp=utc_timestamp(), + output_path=str(markdown_path.relative_to(output_root)), + pages=result.page_count, + error=error, + fingerprint=self._run_fingerprint(), + ) - # Hyperlinks (OCR 3) - if hasattr(page, "hyperlinks") and page.hyperlinks: - markdown_content.append("**Hyperlinks:**\n") - for link in page.hyperlinks: - text = getattr(link, "text", "") or "" - url = getattr(link, "url", "") or getattr(link, "href", "") or "" - if url: - markdown_content.append(f"- [{text or url}]({url})\n") - markdown_content.append("\n") - - # Figures - if self.config.include_images and hasattr(page, "images") and page.images: - figures_dir = doc_dir / "figures" - figures_dir.mkdir(parents=True, exist_ok=True) - - for idx, image in enumerate(page.images): - b64_data = getattr(image, "image_base64", None) or getattr( - image, "base64", None - ) - if b64_data: - img_id = getattr(image, "id", None) or f"img{idx + 1}" - img_ext = Path(img_id).suffix if "." in str(img_id) else ".png" - image_filename = f"page{page.index + 1}_img{idx + 1}{img_ext}" - image_path = figures_dir / image_filename - save_base64_image(b64_data, image_path) - markdown_content.append( - f"![Image {idx + 1}](./figures/{image_filename})\n\n" - ) - - # Footer (OCR 3) - if hasattr(page, "footer") and page.footer: - markdown_content.append(f"> **Footer:** {page.footer}\n\n") - - # Write markdown file - markdown_path = doc_dir / f"{base_name}.md" - with open(markdown_path, "w", encoding="utf-8") as f: - f.write("".join(markdown_content)) - - logger.debug("Saved results to %s", markdown_path) - - def process_directory( + def _persist( + self, result: OCRResult, output_root: Path, rel_key: str, index: RootIndex + ) -> tuple[DocMetadata, Path]: + """Write markdown, figures, and BOTH metadata levels for one document. + + Always writes output (markdown + per-doc + root metadata) regardless of + OCR success, so failures are recorded with ``status=failed`` per the canon. + + Persistence is also failure-isolated: if a save / metadata-write / index + record raises (e.g. a disk-full / permission / OSError), the exception is + NOT allowed to escape and abort the whole batch. Instead it is folded into + a ``status=failed`` record (best-effort persisted, both metadata levels), + and a failed :class:`DocMetadata` is returned so the caller marks this one + document failed and the batch continues — uniform per-file failure + accounting that covers I/O failures during persistence, not just OCR. + """ + markdown_path = markdown_path_for(doc_dir_for(output_root, rel_key), rel_key) + try: + markdown_path = self.save_results(result, output_root, rel_key) + meta = self._build_doc_metadata(result, markdown_path, output_root) + doc_dir = doc_dir_for(output_root, rel_key) + write_doc_metadata(doc_dir, rel_key, meta) + with self._lock: + index.record(rel_key, meta) + return meta, markdown_path + except Exception as exc: + logger.error("Failed to persist output for %s: %s", rel_key, exc) + logger.debug("Persistence traceback for %s", rel_key, exc_info=True) + failed_meta = self._build_failed_persist_metadata( + result, markdown_path, output_root, exc + ) + self._best_effort_record_failure(output_root, rel_key, index, failed_meta) + return failed_meta, markdown_path + + def _build_failed_persist_metadata( + self, result: OCRResult, markdown_path: Path, output_root: Path, exc: Exception + ) -> DocMetadata: + """Build a ``status=failed`` record for a persistence I/O failure.""" + try: + output_rel = str(markdown_path.relative_to(output_root)) + except ValueError: + output_rel = str(markdown_path) + # failure_checksum: the real digest if the input is still readable, else + # the canonical ``sha256:`` UNREADABLE_CHECKSUM sentinel (v0.1.3). The + # schema requires a ``sha256:`` checksum even on a failure record, and the + # sentinel can never equal a real digest (unlike the old bare ``"sha256:"``, + # which had no digest and could spuriously match other failure records). + checksum = failure_checksum(result.file_path) + return DocMetadata( + status=Status.FAILED, + checksum=checksum, + model=self.config.model, + backend=BACKEND, + processing_time=result.processing_time, + timestamp=utc_timestamp(), + output_path=output_rel, + pages=result.page_count, + error=f"persistence failed: {exc}", + fingerprint=self._run_fingerprint(), + ) + + def _best_effort_record_failure( + self, output_root: Path, rel_key: str, index: RootIndex, meta: DocMetadata + ) -> None: + """Try to persist a failed record without ever re-raising. + + The original persistence already failed, so these writes may also fail; + each is isolated so the batch still continues and the in-memory + :class:`RunOutcome` (driven by the returned meta) remains authoritative + even when nothing reaches disk. + """ + try: + write_doc_metadata(doc_dir_for(output_root, rel_key), rel_key, meta) + except Exception: + logger.debug("Best-effort per-doc failure write also failed for %s", rel_key) + try: + with self._lock: + index.record(rel_key, meta) + except Exception: + logger.debug("Best-effort root-index failure write also failed for %s", rel_key) + + def _record_unreadable_input( + self, file_path: Path, output_root: Path, rel_key: str, index: RootIndex + ) -> DocMetadata: + """Record a ``status=failed`` entry for an input unreadable at pre-check. + + Discovery can yield a file that is unreadable by the time the idempotency + pre-check hashes it (permission denied, deleted/replaced mid-run, a broken + symlink that passed discovery). Using the contract's ``safe_checksum`` the + ``OSError`` no longer propagates and aborts the whole run (the SYS-02 "one + bad file aborts the batch" failure mode); instead it is folded into a + durable per-file failure and the batch CONTINUES. + """ + meta = DocMetadata( + status=Status.FAILED, + # failure_checksum yields a valid ``sha256:`` value -- the real digest + # if the file is now readable, else the canonical UNREADABLE_CHECKSUM + # sentinel -- never the old bare ``"sha256:"`` (no digest). The schema + # requires a ``sha256:`` checksum even on a failure record, and the + # sentinel can never equal a real digest, so a later readable run + # reprocesses instead of treating it as already done. + checksum=failure_checksum(file_path), + model=self.config.model, + backend=BACKEND, + processing_time=0.0, + timestamp=utc_timestamp(), + output_path=str(markdown_path_for(doc_dir_for(output_root, rel_key), rel_key)), + pages=0, + error=f"input unreadable at idempotency pre-check: {file_path}", + fingerprint=self._run_fingerprint(), + ) + self._best_effort_record_failure(output_root, rel_key, index, meta) + return meta + + # ------------------------------------------------------------------ + # Orchestration + # ------------------------------------------------------------------ + + def process( self, - input_dir: Path, - output_dir: Path | None = None, - add_timestamp: bool = False, + input_path: Path, + output_path: Path | None = None, reprocess: bool = False, - ) -> tuple[int, int]: - """Process all supported files in a directory.""" - output_path = determine_output_path(input_dir, output_dir, add_timestamp=add_timestamp) - - # Exclude the output directory from file discovery - files = get_supported_files( - input_dir, - exclude_paths=[output_path.resolve()], - ) + ) -> RunOutcome: + """Process an input path (file or directory). Returns a RunOutcome. + Output goes to ``resolve_output_root(input_path, output_path)`` — default + ``/ocr/``; ``-o`` overrides; never required. The returned + :class:`RunOutcome` carries the uniform exit policy: nonzero if any file + failed, across both single-file and batch runs. + """ + if input_path.is_file(): + return self._process_single_file(input_path, output_path, reprocess) + if input_path.is_dir(): + return self._process_directory(input_path, output_path, reprocess) + raise ValueError(f"Input path does not exist: {input_path}") + + def _process_single_file( + self, file_path: Path, output_path: Path | None, reprocess: bool + ) -> RunOutcome: + """Process a single file. Scan root is the file's parent (rel key = name).""" + outcome = RunOutcome() + output_root = resolve_output_root(file_path, output_path) + output_root.mkdir(parents=True, exist_ok=True) + rel_key = relative_key(file_path, file_path.parent) + index = RootIndex(output_root) + + # SYS-02: an input unreadable at the pre-check (permission denied, deleted + # mid-run, broken symlink) must be recorded status=failed and NOT abort — + # safe_checksum returns None instead of raising OSError. + checksum = safe_checksum(file_path) + if checksum is None: + meta = self._record_unreadable_input(file_path, output_root, rel_key, index) + outcome.add(meta.status, detail=rel_key) + err_console.print(f"FAILED {rel_key}: {meta.error}") + return outcome + + # Fix (audit HIGH): content-aware skip — status==completed AND checksum + # match AND output still on disk AND a matching run fingerprint (model / + # OCR-3 toggles) — so an in-place edit, a deleted output, or a config + # change all force reprocessing instead of serving stale output. + if not reprocess and index.is_completed( + rel_key, checksum, fingerprint=self._run_fingerprint() + ): + console.print(f"[yellow]Already processed:[/yellow] {file_path.name}") + console.print("[dim]Use --reprocess to force reprocessing[/dim]") + # Emit the existing .md path so the quiet scripting contract still + # reports a path for skipped-but-valid docs (one .md path per line). + doc_dir = doc_dir_for(output_root, rel_key) + outcome.add(Status.COMPLETED, output_path=str(markdown_path_for(doc_dir, rel_key))) + return outcome + + console.print(f"[blue]Processing:[/blue] {file_path}") + console.print(f"[blue]Output:[/blue] {output_root}\n") + + with self._progress() as progress: + if progress is not None: + progress.add_task(f"OCR {file_path.name}", total=None) + result = self.process_file(file_path) + + meta, markdown_path = self._persist(result, output_root, rel_key, index) + # Only successful outputs land in outcome.outputs (the quiet stdout + # scripting list). A failed doc's placeholder .md path is NOT echoed to + # stdout as if it succeeded; the failure goes to stderr instead. + if meta.status is Status.COMPLETED: + outcome.add(Status.COMPLETED, output_path=str(markdown_path)) + console.print("\n[green]Success[/green]") + console.print(f"[dim]Time: {result.processing_time:.2f}s[/dim]") + else: + outcome.add(meta.status, detail=rel_key) + console.print(f"\n[red]Failed:[/red] {meta.error}") + err_console.print(f"FAILED {rel_key}: {meta.error}") + return outcome + + def _process_directory( + self, dir_path: Path, output_path: Path | None, reprocess: bool + ) -> RunOutcome: + """Process all supported files in a directory, keyed on input-rel paths.""" + outcome = RunOutcome() + output_root = resolve_output_root(dir_path, output_path) + + # Discovery (contract iter_input_files) excludes the resolved output root + # so prior outputs are never re-ingested; it does NOT skip arbitrary dirs + # merely named 'ocr'. + files = get_supported_files(dir_path, output_root) if not files: console.print("[yellow]No supported files found in the directory.[/yellow]") - return 0, 0 - - # Load existing metadata to check for already processed files - existing_metadata = load_metadata(output_path) - existing_files_set = { - str(Path(item["file"]).resolve()) for item in existing_metadata["files_processed"] - } - - # Filter files based on reprocess flag - files_to_process = [] - skipped_files = [] - for file_path in files: - if str(file_path.resolve()) in existing_files_set and not reprocess: - skipped_files.append(file_path) - logger.debug("Skipping already processed: %s", file_path.name) + return outcome + + output_root.mkdir(parents=True, exist_ok=True) + index = RootIndex(output_root) + + files_to_process: list[tuple[Path, str]] = [] + for f in files: + rel_key = relative_key(f, dir_path) + # SYS-02: an input unreadable at the pre-check is recorded failed and + # the batch CONTINUES — safe_checksum returns None instead of raising + # an OSError that would abort the whole directory run. + checksum = safe_checksum(f) + if checksum is None: + meta = self._record_unreadable_input(f, output_root, rel_key, index) + outcome.add(meta.status, detail=rel_key) + err_console.print(f"FAILED {rel_key}: {meta.error}") + continue + if not reprocess and index.is_completed( + rel_key, checksum, fingerprint=self._run_fingerprint() + ): + if self.config.verbose: + console.print(f"[dim]Skipping: {rel_key}[/dim]") + # Emit the existing .md path so quiet scripting still reports a + # path for skipped-but-valid docs. + doc_dir = doc_dir_for(output_root, rel_key) + outcome.add(Status.COMPLETED, output_path=str(markdown_path_for(doc_dir, rel_key))) else: - files_to_process.append(file_path) - - if skipped_files: - console.print( - f"[yellow]Skipping {len(skipped_files)} already processed file(s)[/yellow]" - ) - console.print("[dim]Use --verbose to see which files were skipped[/dim]") + files_to_process.append((f, rel_key)) if not files_to_process: - console.print( - "[green]All files already processed. Use --reprocess to force reprocessing.[/green]" - ) - return 0, 0 + if outcome.has_failures: + console.print("[yellow]No files left to process (some inputs failed).[/yellow]") + else: + console.print("[green]All files already processed.[/green]") + console.print("[dim]Use --reprocess to force reprocessing.[/dim]") + return outcome workers = self.config.max_workers console.print(f"[blue]Processing {len(files_to_process)} file(s)...[/blue]") if workers > 1: console.print(f"[blue]Using {workers} concurrent workers[/blue]") - console.print(f"[blue]Output directory: {output_path}[/blue]\n") + console.print(f"[blue]Output:[/blue] {output_root}\n") start_time = time.time() - success_count = 0 - # Capture prior-session time once to avoid overcounting on incremental flushes - base_processing_time = existing_metadata.get("processing_time_seconds", 0) - - with Progress( - SpinnerColumn(), - TextColumn("[progress.description]{task.description}"), - BarColumn(), - TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), - TimeRemainingColumn(), - console=console, - ) as progress: - task = progress.add_task("Processing files...", total=len(files_to_process)) + with self._progress(total=len(files_to_process)) as progress: + task = ( + progress.add_task("Processing files...", total=len(files_to_process)) + if progress is not None + else None + ) if workers <= 1: - # Sequential processing (original behaviour) - for file_path in files_to_process: - file_size = format_file_size(file_path.stat().st_size) - progress.update( - task, description=f"Processing {file_path.name} ({file_size})..." - ) - success_count += self._process_and_save( - file_path, output_path, input_dir, base_processing_time, start_time - ) - progress.update(task, advance=1) + for file_path, rel_key in files_to_process: + if progress is not None and task is not None: + size = format_file_size(file_path.stat().st_size) + progress.update(task, description=f"Processing {rel_key} ({size})...") + result = self.process_file(file_path) + self._record(result, output_root, rel_key, index, outcome) + if progress is not None and task is not None: + progress.update(task, advance=1) else: - # Concurrent processing with ThreadPoolExecutor(max_workers=workers) as executor: futures = { - executor.submit(self.process_file, fp): fp for fp in files_to_process + executor.submit(self.process_file, fp): (fp, rk) + for fp, rk in files_to_process } for future in as_completed(futures): - file_path = futures[future] + _fp, rel_key = futures[future] result = future.result() - if result: - with self._lock: - try: - self.save_results( - result, - output_path, - is_single_file=False, - base_dir=input_dir, - ) - success_count += 1 - base_name = make_unique_basename(file_path, base_dir=input_dir) - self.processed_files.append( - { - "file": str(file_path.resolve()), - "size": file_path.stat().st_size, - "output": str( - output_path / base_name / f"{base_name}.md" - ), - } - ) - except (OSError, ValueError) as e: - console.print( - f"[red]Error saving results for {file_path.name}: {e}[/red]" - ) - self.errors.append( - { - "file": str(file_path.resolve()), - "error": f"Save failed: {e}", - } - ) - # Flush metadata under lock - processing_time = time.time() - start_time - save_metadata( - output_path, - self.processed_files, - processing_time, - self.errors, - base_processing_time=base_processing_time, - ) - progress.update(task, advance=1) - - return success_count, len(files_to_process) - - def _process_and_save( + self._record(result, output_root, rel_key, index, outcome) + if progress is not None and task is not None: + progress.update(task, advance=1) + + total = outcome.completed + outcome.failed + outcome.partial + console.print(f"\n[green]Completed:[/green] {outcome.completed}/{total} files") + if outcome.has_failures: + console.print(f"[red]Failures:[/red] {outcome.failed} failed") + console.print(f"[dim]Total time: {time.time() - start_time:.2f}s[/dim]") + return outcome + + def _record( self, - file_path: Path, - output_path: Path, - base_dir: Path, - base_processing_time: float, - start_time: float, - ) -> int: - """Process a single file and save results. Returns 1 on success, 0 on failure.""" - result = self.process_file(file_path) - if result: - try: - self.save_results(result, output_path, is_single_file=False, base_dir=base_dir) - base_name = make_unique_basename(file_path, base_dir=base_dir) - self.processed_files.append( - { - "file": str(file_path.resolve()), - "size": file_path.stat().st_size, - "output": str(output_path / base_name / f"{base_name}.md"), - } - ) - except (OSError, ValueError) as e: - console.print(f"[red]Error saving results for {file_path.name}: {e}[/red]") - self.errors.append({"file": str(file_path.resolve()), "error": f"Save failed: {e}"}) - return 0 - - # Flush metadata incrementally - processing_time = time.time() - start_time - save_metadata( - output_path, - self.processed_files, - processing_time, - self.errors, - base_processing_time=base_processing_time, - ) - return 1 - return 0 - - def process( - self, - input_path: Path, - output_path: Path | None = None, - add_timestamp: bool = False, - reprocess: bool = False, + result: OCRResult, + output_root: Path, + rel_key: str, + index: RootIndex, + outcome: RunOutcome, ) -> None: - """Process input path (file or directory).""" - if input_path.is_file(): - # Process single file - output_dir = determine_output_path(input_path, output_path, add_timestamp=add_timestamp) - - # Check if file already processed - existing_metadata = load_metadata(output_dir) - existing_files_set = { - str(Path(item["file"]).resolve()) for item in existing_metadata["files_processed"] - } - - if str(input_path.resolve()) in existing_files_set and not reprocess: - base_name = make_unique_basename(input_path) - output_file = output_dir / base_name / f"{base_name}.md" - console.print(f"[yellow]File already processed: {input_path.name}[/yellow]") - console.print(f"[dim]Output exists at: {output_file}[/dim]") - console.print("[dim]Use --reprocess to force reprocessing.[/dim]") - return - - console.print(f"[blue]Processing file: {input_path}[/blue]") - console.print(f"[blue]Output directory: {output_dir}[/blue]\n") - - start_time = time.time() - result = self.process_file(input_path) + """Persist one result and fold its status into the run outcome. - if result: - try: - self.save_results(result, output_dir, is_single_file=True) - base_name = make_unique_basename(input_path) - self.processed_files.append( - { - "file": str(input_path.resolve()), - "size": input_path.stat().st_size, - "output": str(output_dir / base_name / f"{base_name}.md"), - } - ) - except (OSError, ValueError) as e: - console.print(f"[red]Error saving results for {input_path.name}: {e}[/red]") - self.errors.append( - {"file": str(input_path.resolve()), "error": f"Save failed: {e}"} - ) - - # Save metadata - processing_time = time.time() - start_time - save_metadata(output_dir, self.processed_files, processing_time, self.errors) + Only successful docs contribute their .md path to ``outcome.outputs`` (the + quiet stdout scripting list); failures are routed to stderr instead, so a + scripting consumer never mistakes a failed placeholder for a success. + """ + meta, markdown_path = self._persist(result, output_root, rel_key, index) + if meta.status is Status.COMPLETED: + outcome.add(Status.COMPLETED, output_path=str(markdown_path)) + console.print(f" [green]OK[/green] {rel_key} ({result.processing_time:.1f}s)") + else: + outcome.add(meta.status, detail=rel_key) + console.print(f" [red]FAILED[/red] {rel_key}: {meta.error}") + err_console.print(f"FAILED {rel_key}: {meta.error}") + + def _progress(self, total: int | None = None) -> Progress | _NullProgress: + """A rich Progress, or a no-op context when quiet.""" + if self.config.quiet: + return _NullProgress() + return Progress( + SpinnerColumn(), + TextColumn("[progress.description]{task.description}"), + BarColumn(), + TextColumn("[progress.percentage]{task.percentage:>3.0f}%"), + TimeRemainingColumn(), + console=console, + transient=True, + ) - if self.errors: - console.print("\n[red]✗ Failed to save results[/red]") - else: - console.print("\n[green]✓ Successfully processed 1 file[/green]") - console.print(f"[dim]Processing time: {processing_time:.2f} seconds[/dim]") - else: - console.print("\n[red]✗ Failed to process file[/red]") - elif input_path.is_dir(): - # Process directory - success_count, total_count = self.process_directory( - input_path, output_path, add_timestamp, reprocess - ) +class _NullProgress: + """A context manager standing in for rich.Progress under --quiet.""" - console.print( - f"\n[green]✓ Successfully processed {success_count}/{total_count} files[/green]" - ) - if self.errors: - console.print(f"[red]✗ {len(self.errors)} file(s) failed[/red]") + def __enter__(self) -> None: + return None - else: - raise ValueError(f"Input path does not exist: {input_path}") + def __exit__(self, *exc: object) -> bool: + return False diff --git a/mistral_ocr/utils.py b/mistral_ocr/utils.py index 8cfd154..84b2b3e 100644 --- a/mistral_ocr/utils.py +++ b/mistral_ocr/utils.py @@ -1,12 +1,23 @@ -"""Utility functions for Mistral OCR.""" +"""Utility functions for Mistral OCR. + +Output-shape concerns (output-root resolution, input-relative keying, the +per-document layout, page assembly, metadata, figure naming and the exit-code +policy) are owned by the shared ``ocr-output-contract`` package and live in +:mod:`mistral_ocr.processor`. What remains here is mistral-specific I/O: MIME +sniffing / data-URI construction for the API, base64 image decoding, PDF page +counting + splitting for the API's per-request page cap, and supported-file +discovery for batch runs. +""" + +from __future__ import annotations import base64 -import json import mimetypes -import time from pathlib import Path -# Canonical extension sets — used by processor.py and get_supported_files() +from ocr_output_contract import iter_input_files + +# Canonical extension sets — used by processor.py and discovery. DOCUMENT_EXTENSIONS = {".pdf", ".docx", ".pptx"} IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".webp", ".gif", ".bmp", ".tiff", ".avif"} SUPPORTED_EXTENSIONS = DOCUMENT_EXTENSIONS | IMAGE_EXTENSIONS @@ -48,46 +59,25 @@ def create_data_uri(file_path: Path) -> str: return f"data:{mime_type};base64,{base64_data}" -def save_base64_image(base64_string: str, output_path: Path) -> None: - """Save a base64 encoded image to file.""" - image_data = base64.b64decode(base64_string) - output_path.parent.mkdir(parents=True, exist_ok=True) - with open(output_path, "wb") as f: - f.write(image_data) +def decode_base64_image(base64_string: str) -> bytes: + """Decode a base64 (optionally data-URI-prefixed) image into raw bytes.""" + if "," in base64_string and base64_string.lstrip().startswith("data:"): + base64_string = base64_string.split(",", 1)[1] + return base64.b64decode(base64_string) -def get_supported_files( - directory: Path, - exclude_dirs: list[str] | None = None, - exclude_paths: list[Path] | None = None, -) -> list[Path]: - """Get all supported files from a directory, excluding output directories. +def get_supported_files(directory: Path, output_root: Path) -> list[Path]: + """Get all supported input files under ``directory``, excluding outputs. - Args: - directory: Root directory to search. - exclude_dirs: Directory *names* to skip (matched against each path component). - exclude_paths: Resolved absolute paths to skip (any file underneath is excluded). + Discovery is delegated to the contract's :func:`iter_input_files`, which + recurses ``directory`` and prunes everything at or under the RESOLVED + ``output_root`` (so the engine never re-ingests its own ``.md``/figure + outputs on a rerun). It targets the *real* output directory by resolved path, + not any path component that merely happens to be named ``ocr`` — fixing the + "files under any directory named 'ocr' silently skipped" bug that was acutely + fatal under this user's own ``.../toolkits/ocr/...`` tree. """ - supported_extensions = SUPPORTED_EXTENSIONS - if exclude_dirs is None: - exclude_dirs = ["mistral_ocr_output"] - files = [] - - exclude_set = set(exclude_dirs) - resolved_excludes = [p.resolve() for p in (exclude_paths or [])] - for file_path in directory.rglob("*"): - if file_path.is_file() and file_path.suffix.lower() in supported_extensions: - resolved = file_path.resolve() - # Skip files inside excluded absolute paths - if any(ep in resolved.parents for ep in resolved_excludes): - continue - # Skip files inside excluded directory names (check parent dirs only) - rel_parts = file_path.relative_to(directory).parts[:-1] - if any(part in exclude_set for part in rel_parts): - continue - files.append(file_path) - - return sorted(files) + return list(iter_input_files(directory, output_root, suffixes=SUPPORTED_EXTENSIONS)) def get_pdf_page_count(file_path: Path) -> int: @@ -136,158 +126,11 @@ def split_pdf( return chunks -def determine_output_path( - input_path: Path, - output_path: Path | None = None, - default_folder_name: str = "mistral_ocr_output", - add_timestamp: bool = False, -) -> Path: - """Determine the output path for OCR results.""" - if output_path: - if output_path.exists() and not output_path.is_dir(): - raise ValueError(f"Output path exists and is not a directory: {output_path}") - output_path.mkdir(parents=True, exist_ok=True) - return output_path - - if input_path.is_file(): - parent_dir = input_path.parent - else: - parent_dir = input_path - - # Add timestamp if requested - if add_timestamp: - timestamp = time.strftime("%Y%m%d_%H%M%S") - folder_name = f"{default_folder_name}_{timestamp}" - else: - folder_name = default_folder_name - - output_dir = parent_dir / folder_name - output_dir.mkdir(parents=True, exist_ok=True) - return output_dir - - -def _empty_metadata() -> dict: - return { - "files_processed": [], - "total_files": 0, - "processing_time_seconds": 0, - "errors": [], - "error_count": 0, - } - - -def load_metadata(output_dir: Path) -> dict: - """Load existing metadata from JSON file. - - Returns empty metadata on missing or corrupt files. - """ - metadata_path = output_dir / "metadata.json" - if metadata_path.exists(): - try: - with open(metadata_path) as f: - return json.load(f) - except (json.JSONDecodeError, KeyError): - return _empty_metadata() - return _empty_metadata() - - -def save_metadata( - output_dir: Path, - files_processed: list[dict], - processing_time: float, - errors: list[dict], - base_processing_time: float | None = None, -) -> None: - """Save processing metadata to JSON file (append/update mode). - - Args: - processing_time: Elapsed time for the *current* session. - base_processing_time: Accumulated time from *prior* sessions. If None, - loaded from existing metadata (use this on the first call). Pass - the returned value on subsequent calls within the same session to - avoid overcounting. - """ - # Load existing metadata - existing_metadata = load_metadata(output_dir) - - if base_processing_time is None: - base_processing_time = existing_metadata.get("processing_time_seconds", 0) - - # Create a dict of existing files for quick lookup - existing_files = {item["file"]: item for item in existing_metadata["files_processed"]} - - # Update with new files (overwrite if exists, add if new) - for new_file in files_processed: - new_file["last_processed"] = time.strftime("%Y-%m-%d %H:%M:%S") - existing_files[new_file["file"]] = new_file - - # Merge errors: keep historical errors, append new ones (deduplicate by file) - existing_errors = {e["file"]: e for e in existing_metadata.get("errors", [])} - for err in errors: - err["last_seen"] = time.strftime("%Y-%m-%d %H:%M:%S") - existing_errors[err["file"]] = err - all_errors = list(existing_errors.values()) - - # Update metadata — total time = prior sessions + current session elapsed - metadata = { - "files_processed": list(existing_files.values()), - "total_files": len(existing_files), - "processing_time_seconds": base_processing_time + processing_time, - "errors": all_errors, - "error_count": len(all_errors), - "last_updated": time.strftime("%Y-%m-%d %H:%M:%S"), - } - - metadata_path = output_dir / "metadata.json" - tmp_path = metadata_path.with_suffix(".json.tmp") - with open(tmp_path, "w") as f: - json.dump(metadata, f, indent=2, default=str) - tmp_path.replace(metadata_path) - - def format_file_size(size_bytes: int) -> str: """Format file size in human-readable format.""" + size: float = float(size_bytes) for unit in ["B", "KB", "MB", "GB"]: - if size_bytes < 1024.0: - return f"{size_bytes:.2f} {unit}" - size_bytes /= 1024.0 - return f"{size_bytes:.2f} TB" - - -def make_unique_basename(file_path: Path, base_dir: Path | None = None) -> str: - """Create a unique base name for output files. - - When base_dir is provided (directory mode), includes the relative path - to disambiguate files with the same stem in different subdirectories. - E.g., subdir/report.pdf -> subdir__report - """ - stem = file_path.stem - if base_dir is not None: - try: - rel = file_path.parent.relative_to(base_dir) - if rel != Path("."): - prefix = str(rel).replace("/", "__").replace("\\", "__") - stem = f"{prefix}__{stem}" - except ValueError: - pass - return sanitize_filename(stem, max_length=200) - - -def sanitize_filename(filename: str, max_length: int | None = None) -> str: - """Sanitize filename by removing or replacing invalid characters.""" - invalid_chars = '<>:"/\\|?*' - for char in invalid_chars: - filename = filename.replace(char, "_") - - # Only truncate if max_length is specified - if max_length is not None: - # Truncate long filenames but keep extension - if len(filename) > max_length and "." in filename: - name, ext = filename.rsplit(".", 1) - if len(name) > max_length - len(ext) - 1: - name = name[: max_length - len(ext) - 4] + "..." - filename = f"{name}.{ext}" - elif len(filename) > max_length: - filename = filename[: max_length - 3] + "..." - - return filename + if size < 1024.0: + return f"{size:.2f} {unit}" + size /= 1024.0 + return f"{size:.2f} TB" diff --git a/pyproject.toml b/pyproject.toml index 4613a9e..6bc2e34 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -24,6 +24,7 @@ classifiers = [ "Topic :: Scientific/Engineering :: Artificial Intelligence", ] dependencies = [ + "ocr-output-contract @ git+https://github.com/r-uben/ocr-output-contract@v0.1.3", "mistralai>=1.0.0,<2.0.0", "python-dotenv>=1.0.0", "click>=8.1.7", @@ -55,6 +56,10 @@ build-backend = "hatchling.build" [tool.hatch.build.targets.wheel] packages = ["mistral_ocr"] +# Required so the git+https direct reference to ocr-output-contract is allowed. +[tool.hatch.metadata] +allow-direct-references = true + [tool.ruff] line-length = 100 diff --git a/tests/test_cli.py b/tests/test_cli.py index 86c69d1..c741592 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -4,6 +4,7 @@ import pytest from click.testing import CliRunner +from ocr_output_contract import RunOutcome, Status from mistral_ocr.cli import main @@ -86,14 +87,16 @@ def test_quiet_suppresses_output(self, runner, tmp_path, monkeypatch): with patch("mistral_ocr.cli.OCRProcessor") as mock_cls: mock_proc = mock_cls.return_value - mock_proc.errors = [] - mock_proc.process.return_value = None + outcome = RunOutcome() + outcome.add(Status.COMPLETED, output_path=str(tmp_path / "out.md")) + mock_proc.process.return_value = outcome result = runner.invoke(main, [str(pdf), "--quiet"]) assert result.exit_code == 0 - # Quiet mode: no banner, no completion message + # Quiet mode: no banner, no completion message; only the output path. assert "Mistral OCR" not in result.output assert "Processing complete" not in result.output + assert str(tmp_path / "out.md") in result.output def test_quiet_propagates_to_processor_console(self, runner, tmp_path, monkeypatch): """Quiet flag must set .quiet on the shared processor console.""" @@ -103,8 +106,7 @@ def test_quiet_propagates_to_processor_console(self, runner, tmp_path, monkeypat with patch("mistral_ocr.cli.OCRProcessor") as mock_cls: mock_proc = mock_cls.return_value - mock_proc.errors = [] - mock_proc.process.return_value = None + mock_proc.process.return_value = RunOutcome() runner.invoke(main, [str(pdf), "--quiet"]) @@ -115,3 +117,43 @@ def test_quiet_propagates_to_processor_console(self, runner, tmp_path, monkeypat # Reset for other tests proc_console.quiet = False + + def test_quiet_failure_not_on_stdout_goes_to_stderr(self, runner, tmp_path, monkeypatch): + """Under --quiet a failed doc must NOT pollute the stdout path list. + + The stdout scripting contract is success .md paths only; per-file + failures go to stderr. This drives the real processor (client mocked to + fail) so the failed placeholder path is genuinely produced, then asserts + it is absent from stdout and the failure is on stderr. + """ + _make_env(monkeypatch) + img = tmp_path / "broken.png" + img.write_bytes(b"\x89PNG\r\n\x1a\n") + + from types import SimpleNamespace + + def _fail_client(*_a, **_k): + client = SimpleNamespace() + client.ocr = SimpleNamespace( + process=lambda **_kw: (_ for _ in ()).throw(RuntimeError("API down")) + ) + client.files = SimpleNamespace( + upload=lambda **_kw: SimpleNamespace(id="u1"), + delete=lambda **_kw: None, + ) + return client + + from mistral_ocr.processor import console as proc_console + + with patch("mistral_ocr.processor.Mistral", _fail_client): + result = runner.invoke( + main, [str(img), "--quiet", "-o", str(tmp_path / "out")], catch_exceptions=False + ) + + proc_console.quiet = False # reset module-level state for other tests + + assert result.exit_code != 0 + # The failed placeholder .md path must NOT appear on stdout. + assert "broken.md" not in result.stdout + # The failure IS surfaced on stderr (SYS-02), even under --quiet. + assert "broken.png" in result.stderr diff --git a/tests/test_cli_overrides.py b/tests/test_cli_overrides.py index 779fbde..de927dc 100644 --- a/tests/test_cli_overrides.py +++ b/tests/test_cli_overrides.py @@ -1,10 +1,12 @@ """Tests for CLI config overrides and process() orchestration.""" +import threading from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest from click.testing import CliRunner +from ocr_output_contract import RunOutcome, doc_dir_for, markdown_path_for from mistral_ocr.cli import main from mistral_ocr.config import Config @@ -16,11 +18,14 @@ def runner(): return CliRunner() -def _mock_processor(): - """Return a mock OCRProcessor that does nothing.""" +def _mock_processor(outcome: RunOutcome | None = None): + """Return a mock OCRProcessor whose process() yields a real RunOutcome. + + The CLI inspects the returned RunOutcome (outputs / has_failures / + exit_code), so the mock must hand back a genuine outcome, not ``None``. + """ mock = MagicMock(spec=OCRProcessor) - mock.errors = [] - mock.process.return_value = None + mock.process.return_value = outcome if outcome is not None else RunOutcome() return mock @@ -109,27 +114,12 @@ def test_max_pages_override(self, runner, tmp_path, monkeypatch): config = mock_cls.call_args[0][0] assert config.max_pages == 100 - def test_no_metadata_override(self, runner, tmp_path, monkeypatch): - monkeypatch.setenv("MISTRAL_API_KEY", "key") - pdf = tmp_path / "doc.pdf" - pdf.write_bytes(b"%PDF") - - with patch("mistral_ocr.cli.OCRProcessor") as mock_cls: - mock_cls.return_value = _mock_processor() - runner.invoke(main, [str(pdf), "--no-metadata"]) - config = mock_cls.call_args[0][0] - assert config.include_metadata is False - - def test_no_page_headings_override(self, runner, tmp_path, monkeypatch): - monkeypatch.setenv("MISTRAL_API_KEY", "key") - pdf = tmp_path / "doc.pdf" - pdf.write_bytes(b"%PDF") - - with patch("mistral_ocr.cli.OCRProcessor") as mock_cls: - mock_cls.return_value = _mock_processor() - runner.invoke(main, [str(pdf), "--no-page-headings"]) - config = mock_cls.call_args[0][0] - assert config.include_page_headings is False + # NOTE: --no-metadata / --no-page-headings are now deprecated no-ops. The + # canonical output contract owns the markdown body (always clean: ## Page N, + # no header block, no frontmatter) and always writes the dual-level JSON + # sidecars, so there is no config field for them to override. The flags are + # kept hidden for invocation compatibility only; their override tests are + # gone deliberately. def test_defaults_not_overridden(self, runner, tmp_path, monkeypatch): """When no CLI flags are passed, config keeps env/default values.""" @@ -200,18 +190,20 @@ def test_verbose_flag(self, runner, tmp_path, monkeypatch): class TestProcessOrchestration: - """Test the process() method routing and skip logic.""" + """Test the process() method routing, skip logic and RunOutcome contract. + + The processor returns a :class:`RunOutcome` (the contract's exit-code + accumulator) and writes through the shared package, so these assert against + the canonical ``/ocr/`` output root and the RunOutcome tallies + rather than the removed ``processed_files`` / ``errors`` lists. + """ def _make_real_processor(self, **overrides): - defaults = {"api_key": "test", "save_original_images": False} + defaults = {"api_key": "test", "include_images": False} defaults.update(overrides) proc = OCRProcessor.__new__(OCRProcessor) proc.config = Config(**defaults) proc.client = MagicMock() - proc.errors = [] - proc.processed_files = [] - import threading - proc._lock = threading.Lock() return proc @@ -223,12 +215,13 @@ def test_single_file_success(self, tmp_path): response = SimpleNamespace(pages=[SimpleNamespace(index=0, markdown="Hello", images=[])]) proc.client.ocr.process.return_value = response - proc.process(img) - assert len(proc.processed_files) == 1 - assert len(proc.errors) == 0 - # Check output was created - out_dir = tmp_path / "mistral_ocr_output" - assert (out_dir / "doc" / "doc.md").exists() + outcome = proc.process(img) + assert outcome.completed == 1 + assert outcome.exit_code == 0 + # Output goes to the canonical default root /ocr/. + out_dir = tmp_path / "ocr" + assert markdown_path_for(doc_dir_for(out_dir, "doc.png"), "doc.png").exists() + assert (out_dir / "metadata.json").exists() def test_single_file_skip_already_processed(self, tmp_path): proc = self._make_real_processor() @@ -240,11 +233,32 @@ def test_single_file_skip_already_processed(self, tmp_path): proc.client.ocr.process.return_value = response proc.process(img) - # Second process — should skip + # Second process — content unchanged → checksum match → skip (no API call). proc.client.ocr.process.reset_mock() proc.process(img) proc.client.ocr.process.assert_not_called() + def test_in_place_edit_forces_reprocess(self, tmp_path): + """Audit HIGH fix: a same-path content change must NOT be skipped. + + The old resolved-path-equality skip served stale output on in-place + edits. The contract's checksum-based ``RootIndex.is_completed`` now + invalidates the cache when the bytes change. + """ + proc = self._make_real_processor() + img = tmp_path / "doc.png" + img.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50) + + response = SimpleNamespace(pages=[SimpleNamespace(index=0, markdown="v1", images=[])]) + proc.client.ocr.process.return_value = response + proc.process(img) + + # Replace the file's bytes at the SAME path → checksum differs → re-OCR. + img.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\xff" * 80) + proc.client.ocr.process.reset_mock() + proc.process(img) + proc.client.ocr.process.assert_called_once() + def test_single_file_reprocess(self, tmp_path): proc = self._make_real_processor() img = tmp_path / "doc.png" @@ -254,7 +268,7 @@ def test_single_file_reprocess(self, tmp_path): proc.client.ocr.process.return_value = response proc.process(img) - # Reprocess with flag + # Reprocess with flag → forces a second API call even with matching checksum. proc.process(img, reprocess=True) assert proc.client.ocr.process.call_count == 2 @@ -264,9 +278,15 @@ def test_single_file_failure(self, tmp_path): img.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50) proc.client.ocr.process.side_effect = RuntimeError("API down") - proc.process(img) - assert len(proc.errors) == 1 - assert len(proc.processed_files) == 0 + outcome = proc.process(img) + assert outcome.failed == 1 + assert outcome.completed == 0 + assert outcome.exit_code != 0 + # The failure is recorded (status=failed), never silently dropped. + import json + + entry = json.loads((tmp_path / "ocr" / "metadata.json").read_text())["files"]["doc.png"] + assert entry["status"] == "failed" def test_nonexistent_path_raises(self, tmp_path): proc = self._make_real_processor() @@ -282,5 +302,11 @@ def test_directory_processing(self, tmp_path): response = SimpleNamespace(pages=[SimpleNamespace(index=0, markdown="text", images=[])]) proc.client.ocr.process.return_value = response - proc.process(input_dir) - assert len(proc.processed_files) == 2 + outcome = proc.process(input_dir) + assert outcome.completed == 2 + assert outcome.exit_code == 0 + # Directory default root is /ocr/. PNG inputs disambiguate the + # doc folder as _png per the v0.1.1 same-stem/different-ext fix. + out_root = input_dir / "ocr" + assert markdown_path_for(doc_dir_for(out_root, "a.png"), "a.png").exists() + assert markdown_path_for(doc_dir_for(out_root, "b.png"), "b.png").exists() diff --git a/tests/test_concurrent.py b/tests/test_concurrent.py index fbed4de..f6a858a 100644 --- a/tests/test_concurrent.py +++ b/tests/test_concurrent.py @@ -1,4 +1,4 @@ -"""Tests for concurrent file processing.""" +"""Tests for concurrent file processing via the RunOutcome contract.""" import threading from types import SimpleNamespace @@ -19,75 +19,71 @@ def _make_processor(max_workers=1): proc.config = Config( api_key="test", max_workers=max_workers, - save_original_images=False, + include_images=False, ) proc.client = MagicMock() proc.client.ocr.process.return_value = _make_ocr_response() - proc.errors = [] - proc.processed_files = [] proc._lock = threading.Lock() return proc class TestConcurrentProcessing: def test_sequential_processes_all_files(self, tmp_path): - """Workers=1 processes files sequentially.""" + """Workers=1 processes files sequentially; RunOutcome tallies them all.""" proc = _make_processor(max_workers=1) input_dir = tmp_path / "input" input_dir.mkdir() 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 - assert success == 3 - assert len(proc.processed_files) == 3 - assert len(proc.errors) == 0 + outcome = proc.process(input_dir) + assert outcome.completed == 3 + assert outcome.failed == 0 + assert outcome.exit_code == 0 + assert len(outcome.outputs) == 3 def test_concurrent_processes_all_files(self, tmp_path): - """Workers>1 processes files concurrently and gets same results.""" + """Workers>1 processes files concurrently and gets the same results.""" proc = _make_processor(max_workers=3) input_dir = tmp_path / "input" input_dir.mkdir() 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 - assert success == 3 - assert len(proc.processed_files) == 3 - assert len(proc.errors) == 0 + outcome = proc.process(input_dir) + assert outcome.completed == 3 + assert outcome.failed == 0 + assert outcome.exit_code == 0 def test_concurrent_handles_failures(self, tmp_path): - """Concurrent mode handles individual file failures gracefully.""" - proc = _make_processor(max_workers=2) + """One failing file → recorded failed, others complete, nonzero exit.""" + from mistral_ocr.processor import OCRResult - # Make process_file fail for one specific file - original_process = proc.process_file + proc = _make_processor(max_workers=2) + original = proc.process_file - def selective_fail(file_path): + def patched(file_path): if file_path.name == "b.png": - proc.errors.append({"file": str(file_path), "error": "test error"}) - return None - return original_process(file_path) + return OCRResult(file_path=file_path, pages=[], error="test error") + return original(file_path) - proc.process_file = selective_fail + proc.process_file = patched input_dir = tmp_path / "input" input_dir.mkdir() 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 - assert success == 2 - assert len(proc.errors) == 1 + outcome = proc.process(input_dir) + assert outcome.completed == 2 + assert outcome.failed == 1 + assert outcome.exit_code != 0 def test_concurrent_uses_multiple_threads(self, tmp_path): - """Verify concurrent mode actually uses threads.""" + """Concurrent mode dispatches OCR calls off the main thread.""" seen_threads = set() - def track_thread(*args, **kwargs): + def track_thread(**kwargs): seen_threads.add(threading.current_thread().name) return _make_ocr_response() @@ -99,10 +95,10 @@ def track_thread(*args, **kwargs): 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 - # (not guaranteed with timing, so just check it completes) - assert len(proc.processed_files) == 3 + outcome = proc.process(input_dir) + assert outcome.completed == 3 + # Worker threads (not MainThread) handled the OCR calls. + assert any(t != "MainThread" for t in seen_threads) class TestWorkersConfig: diff --git a/tests/test_config.py b/tests/test_config.py index fdb63cd..95f9954 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -12,7 +12,6 @@ def test_config_defaults(): assert config.model == "mistral-ocr-latest" assert config.max_file_size_mb == 50 assert config.include_images is True - assert config.save_original_images is True assert config.table_format is None assert config.extract_header is False assert config.extract_footer is False diff --git a/tests/test_markdown_output.py b/tests/test_markdown_output.py deleted file mode 100644 index e1cf005..0000000 --- a/tests/test_markdown_output.py +++ /dev/null @@ -1,94 +0,0 @@ -"""Tests for markdown output flags (--metadata, --page-headings).""" - -from types import SimpleNamespace - -from mistral_ocr.config import Config -from mistral_ocr.processor import OCRProcessor - - -def _make_processor(**config_kwargs): - """Create an OCRProcessor without hitting the Mistral API.""" - config_kwargs.setdefault("save_original_images", False) - proc = OCRProcessor.__new__(OCRProcessor) - proc.config = Config(api_key="test", **config_kwargs) - proc.errors = [] - proc.processed_files = [] - return proc - - -def _fake_response(*page_texts): - return SimpleNamespace( - pages=[ - SimpleNamespace(index=i, markdown=text, images=[]) for i, text in enumerate(page_texts) - ] - ) - - -class TestIncludeMetadata: - def test_metadata_included_by_default(self, tmp_path): - proc = _make_processor() - output = tmp_path / "out" - output.mkdir() - proc.save_results( - {"file_path": tmp_path / "doc.pdf", "response": _fake_response("Hello")}, - output, - ) - md = (output / "doc" / "doc.md").read_text() - assert "# OCR Results" in md - assert "**Original File:**" in md - - def test_metadata_omitted(self, tmp_path): - proc = _make_processor(include_metadata=False) - output = tmp_path / "out" - output.mkdir() - proc.save_results( - {"file_path": tmp_path / "doc.pdf", "response": _fake_response("Hello")}, - output, - ) - md = (output / "doc" / "doc.md").read_text() - assert "# OCR Results" not in md - assert "**Original File:**" not in md - assert "Hello" in md - - -class TestIncludePageHeadings: - def test_page_headings_included_by_default(self, tmp_path): - proc = _make_processor() - output = tmp_path / "out" - output.mkdir() - proc.save_results( - {"file_path": tmp_path / "doc.pdf", "response": _fake_response("A", "B")}, - output, - ) - md = (output / "doc" / "doc.md").read_text() - assert "## Page 1" in md - assert "## Page 2" in md - - def test_page_headings_omitted(self, tmp_path): - proc = _make_processor(include_page_headings=False) - output = tmp_path / "out" - output.mkdir() - proc.save_results( - {"file_path": tmp_path / "doc.pdf", "response": _fake_response("A", "B")}, - output, - ) - md = (output / "doc" / "doc.md").read_text() - assert "## Page 1" not in md - assert "## Page 2" not in md - assert "A" in md - assert "B" in md - - -class TestBothDisabled: - def test_raw_text_only(self, tmp_path): - proc = _make_processor(include_metadata=False, include_page_headings=False) - output = tmp_path / "out" - output.mkdir() - proc.save_results( - {"file_path": tmp_path / "doc.pdf", "response": _fake_response("Content here")}, - output, - ) - md = (output / "doc" / "doc.md").read_text() - assert "# OCR Results" not in md - assert "## Page" not in md - assert "Content here" in md diff --git a/tests/test_output_contract.py b/tests/test_output_contract.py new file mode 100644 index 0000000..0a2fce9 --- /dev/null +++ b/tests/test_output_contract.py @@ -0,0 +1,451 @@ +"""Engine-level conformance: mistral's REAL output vs the shared contract. + +The contract *primitives* (path/key computation, page assembly, metadata +writers, the exit-code policy) are unit-tested inside the ``ocr-output-contract`` +package itself, so they are NOT re-tested here. + +What stays here is the engine-side proof: run mistral's actual processor with a +mocked Mistral API (no key, no network) over real inputs and assert the produced +output tree conforms to the family-wide contract via the package's reusable +:func:`ocr_output_contract.conformance.assert_conforms` harness — including the +cases that motivated the canon rollout: + +* multi-page document → ONE ``/.md`` with ``## Page 1`` + ``## Page 2`` + (the API's per-page list is fed into ``assemble_pages``, never dumped whole); +* a failure → recorded ``status="failed"`` with a nonzero exit (no silent drop), + replacing mistral's old list-based ``errors`` schema; +* a nested batch with same-basename inputs → both survive under input-relative + keys (the old basename/path keying lost one); +* embedded images → ``figures/figure__page

.png`` with resolving links. +""" + +from __future__ import annotations + +import json +import threading +from pathlib import Path +from types import SimpleNamespace + +from ocr_output_contract import ( + METADATA_FILENAME, + UNREADABLE_CHECKSUM, + doc_dir_for, + markdown_path_for, +) +from ocr_output_contract.conformance import ExpectedDoc, assert_conforms + +from mistral_ocr.config import Config +from mistral_ocr.processor import OCRProcessor + +# A valid 1x1 PNG so Pillow can open + re-encode an embedded figure. +_PNG_1x1 = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06" + b"\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00" + b"\x01\r\n-\xb4\x00\x00\x00\x00IEND\xaeB`\x82" +) + + +def _page(index: int, markdown: str, images: list | None = None) -> SimpleNamespace: + return SimpleNamespace(index=index, markdown=markdown, images=images or []) + + +def _b64_image(image_id: str = "fig.png") -> SimpleNamespace: + import base64 + + return SimpleNamespace(image_base64=base64.b64encode(_PNG_1x1).decode(), id=image_id) + + +class FakeMistralProcessor: + """Build an OCRProcessor whose ocr.process call is a deterministic stub. + + No Mistral client is constructed and no key is needed: the client attribute + is a stub whose ``ocr.process`` returns whatever ``responder`` yields. + """ + + @staticmethod + def make(responder, **config_kwargs) -> OCRProcessor: + proc = OCRProcessor.__new__(OCRProcessor) + proc.config = Config(api_key="test", **config_kwargs) + proc._lock = threading.Lock() + + class _Client: + class ocr: # noqa: N801 - mirrors the SDK attribute path + @staticmethod + def process(**kwargs): + return responder(kwargs) + + class files: # noqa: N801 + @staticmethod + def upload(**kwargs): + return SimpleNamespace(id="upload-1") + + @staticmethod + def delete(**kwargs): + return None + + proc.client = _Client() + return proc + + +def test_multipage_conforms_no_whole_doc_dump(tmp_path: Path) -> None: + """A multi-page response → one /.md with both pages present.""" + img = tmp_path / "sample.png" # image input → single API call, multi-page resp + img.write_bytes(_PNG_1x1) + + def responder(_kwargs): + return SimpleNamespace(pages=[_page(0, "PAGE-1-CONTENT"), _page(1, "PAGE-2-CONTENT")]) + + proc = FakeMistralProcessor.make(responder, include_images=False) + outcome = proc.process(img, output_path=tmp_path / "out") + assert outcome.exit_code == 0 + + assert_conforms( + tmp_path / "out", + [ExpectedDoc(rel_key="sample.png", pages=2, status="completed")], + require_failures_nonzero_exit=outcome.exit_code != 0, + ) + + doc_dir = doc_dir_for(tmp_path / "out", "sample.png") + body = markdown_path_for(doc_dir, "sample.png").read_text() + assert "## Page 1" in body and "## Page 2" in body + # Both pages survived — not just the first (no whole-doc dump / data loss). + assert "PAGE-1-CONTENT" in body and "PAGE-2-CONTENT" in body + + +def test_failure_recorded_and_nonzero_exit(tmp_path: Path) -> None: + """An API failure → status=failed recorded, nonzero exit (no silent drop).""" + img = tmp_path / "broken.png" + img.write_bytes(_PNG_1x1) + + def responder(_kwargs): + raise RuntimeError("API down") + + proc = FakeMistralProcessor.make(responder, include_images=False) + outcome = proc.process(img, output_path=tmp_path / "out") + assert outcome.exit_code != 0 + + assert_conforms( + tmp_path / "out", + [ExpectedDoc(rel_key="broken.png", status="failed")], + require_failures_nonzero_exit=True, + ) + doc_dir = doc_dir_for(tmp_path / "out", "broken.png") + doc_meta = json.loads((doc_dir / METADATA_FILENAME).read_text()) + assert doc_meta["status"] == "failed" + assert "error" in doc_meta + + +def test_empty_response_is_failure(tmp_path: Path) -> None: + """An API call returning zero pages is a FAILURE, not a 0-byte success.""" + img = tmp_path / "blank.png" + img.write_bytes(_PNG_1x1) + + proc = FakeMistralProcessor.make(lambda _k: SimpleNamespace(pages=[]), include_images=False) + outcome = proc.process(img, output_path=tmp_path / "out") + assert outcome.exit_code != 0 + + assert_conforms( + tmp_path / "out", + [ExpectedDoc(rel_key="blank.png", status="failed")], + require_failures_nonzero_exit=True, + ) + + +def test_nested_batch_no_basename_collision(tmp_path: Path) -> None: + """Two same-basename inputs in different subdirs both survive (rel-key).""" + root = tmp_path / "in" + (root / "a").mkdir(parents=True) + (root / "b").mkdir(parents=True) + (root / "a" / "intro.png").write_bytes(_PNG_1x1) + (root / "b" / "intro.png").write_bytes(_PNG_1x1) + + proc = FakeMistralProcessor.make( + lambda _k: SimpleNamespace(pages=[_page(0, "content")]), include_images=False + ) + outcome = proc.process(root, output_path=tmp_path / "out") + assert outcome.exit_code == 0 + + assert_conforms( + tmp_path / "out", + [ + ExpectedDoc(rel_key="a/intro.png", pages=1, status="completed"), + ExpectedDoc(rel_key="b/intro.png", pages=1, status="completed"), + ], + ) + + +def test_embedded_figures_conform(tmp_path: Path) -> None: + """Embedded images → figures/figure__page

.png with resolving links.""" + img = tmp_path / "withfig.png" + img.write_bytes(_PNG_1x1) + + def responder(_kwargs): + return SimpleNamespace(pages=[_page(0, "see figure", images=[_b64_image()])]) + + proc = FakeMistralProcessor.make(responder, include_images=True) + outcome = proc.process(img, output_path=tmp_path / "out") + assert outcome.exit_code == 0 + + assert_conforms( + tmp_path / "out", + [ExpectedDoc(rel_key="withfig.png", pages=1, status="completed", figures=[(1, 1)])], + ) + doc_dir = doc_dir_for(tmp_path / "out", "withfig.png") + body = markdown_path_for(doc_dir, "withfig.png").read_text() + assert "figures/figure_1_page1.png" in body + + +def test_embedded_image_placeholders_resolve_on_disk(tmp_path: Path) -> None: + """Byte-faithful real Mistral output → every inline image link resolves. + + Round-2 HIGH blocker: real Mistral ``page.markdown`` carries inline + placeholders like ``![img-0.jpeg](img-0.jpeg)`` whose target is the API's + ``page.images[i].id``. The old code kept ``page.markdown`` verbatim AND + appended a separate ``./figures/...`` link, leaving the original ``img-0.jpeg`` + as a DANGLING local link that the v0.1.2 conformance harness resolves on disk + and rejects ("dangling inline image link"). This fixture reproduces that real + shape (markdown placeholder + matching image id); the processor must rewrite + the placeholder in place to the canonical figure file so the produced body + has NO dangling link and assert_conforms passes. + """ + img = tmp_path / "paper.png" + img.write_bytes(_PNG_1x1) + + def responder(_kwargs): + # The REAL Mistral shape: an inline ![img-0.jpeg](img-0.jpeg) placeholder + # whose target equals page.images[0].id. (The old happy-path fixture had + # no placeholder, which is exactly why it could not catch this bug.) + return SimpleNamespace( + pages=[ + _page( + 0, + "Intro text.\n\n![img-0.jpeg](img-0.jpeg)\n\nTrailing text.", + images=[_b64_image("img-0.jpeg")], + ) + ] + ) + + proc = FakeMistralProcessor.make(responder, include_images=True) + outcome = proc.process(img, output_path=tmp_path / "out") + assert outcome.exit_code == 0 + + # The conformance harness resolves EVERY inline ![..](target) on disk; this + # would have FAILED before the placeholder-rewrite fix. + assert_conforms( + tmp_path / "out", + [ExpectedDoc(rel_key="paper.png", pages=1, status="completed", figures=[(1, 1)])], + ) + + body = markdown_path_for(doc_dir_for(tmp_path / "out", "paper.png"), "paper.png").read_text() + # The placeholder was rewritten IN PLACE to the canonical figure file... + assert "![img-0.jpeg](./figures/figure_1_page1.png)" in body + # ...and the dangling local link is GONE (no duplicate appended link either). + assert "(img-0.jpeg)" not in body + assert body.count("figure_1_page1.png") == 1 + assert (doc_dir_for(tmp_path / "out", "paper.png") / "figures" / "figure_1_page1.png").exists() + + +def test_unwritten_placeholder_is_stripped_not_dangling(tmp_path: Path) -> None: + """An inline placeholder with no extractable image is stripped, not dangling. + + If the API emits an ``![missing.jpeg](missing.jpeg)`` placeholder but no + decodable image bytes for it, there is no figure to point at. The body must + not ship a dangling local link; the placeholder is stripped (alt text kept) + so conformance still passes. + """ + img = tmp_path / "ghost.png" + img.write_bytes(_PNG_1x1) + + def responder(_kwargs): + # Placeholder present in markdown but NO images list → nothing to resolve. + return SimpleNamespace( + pages=[_page(0, "Body.\n\n![missing.jpeg](missing.jpeg)\n\nEnd.", images=[])] + ) + + proc = FakeMistralProcessor.make(responder, include_images=True) + outcome = proc.process(img, output_path=tmp_path / "out") + assert outcome.exit_code == 0 + + assert_conforms( + tmp_path / "out", + [ExpectedDoc(rel_key="ghost.png", pages=1, status="completed")], + ) + body = markdown_path_for(doc_dir_for(tmp_path / "out", "ghost.png"), "ghost.png").read_text() + assert "(missing.jpeg)" not in body + + +def test_unreadable_input_recorded_failed_batch_continues(tmp_path: Path) -> None: + """SYS-02: an input unreadable at the pre-check is failed, batch continues. + + safe_checksum returns None for an unreadable file instead of raising an + OSError that would abort the whole directory run with zero output. The bad + file is recorded status=failed (nonzero exit) and the good files still + process and conform. + """ + import os + import stat + + root = tmp_path / "in" + root.mkdir() + good = root / "good.png" + bad = root / "bad.png" + good.write_bytes(_PNG_1x1) + bad.write_bytes(_PNG_1x1) + # Make the input unreadable so safe_checksum cannot hash it. + os.chmod(bad, 0) + # chmod(0) does not deny reads when running as root (e.g. some CI); capture + # the effective readability so the sentinel assertion only fires when the + # input is genuinely unreadable. + bad_was_unreadable = not os.access(bad, os.R_OK) + + try: + proc = FakeMistralProcessor.make( + lambda _k: SimpleNamespace(pages=[_page(0, "content")]), include_images=False + ) + outcome = proc.process(root, output_path=tmp_path / "out") + + # The batch did NOT abort: good file completed, bad file failed. + assert outcome.completed == 1 + assert outcome.failed == 1 + assert outcome.exit_code != 0 + assert "bad.png" in outcome.failures + + # The unreadable input is durably recorded status=failed in the root index. + root_index = json.loads((tmp_path / "out" / "metadata.json").read_text()) + assert root_index["files"]["bad.png"]["status"] == "failed" + assert root_index["files"]["good.png"]["status"] == "completed" + # v0.1.3: the failure record must carry a VALID ``sha256:`` checksum + # (failure_checksum), never the old bare ``"sha256:"`` (no digest). When + # the input is genuinely unreadable it is the UNREADABLE_CHECKSUM sentinel. + bad_checksum = root_index["files"]["bad.png"]["checksum"] + assert bad_checksum.startswith("sha256:") + assert bad_checksum != "sha256:" + if bad_was_unreadable: + assert bad_checksum == UNREADABLE_CHECKSUM + finally: + os.chmod(bad, stat.S_IRUSR | stat.S_IWUSR) + + +def test_fingerprint_keys_extract_flags_via_extra(tmp_path: Path) -> None: + """A re-run toggling an OCR-3 extract flag reprocesses (extra-keyed fp). + + The output-affecting flags (table_format, extract_header/footer, + include_images) are passed to run_fingerprint via the v0.1.2 ``extra`` dict, + so flipping --extract-headers on a re-run changes the fingerprint and forces + reprocessing rather than serving a stale result keyed only on the checksum. + """ + img = tmp_path / "fp2.png" + img.write_bytes(_PNG_1x1) + calls = {"n": 0} + + def responder(_kwargs): + calls["n"] += 1 + return SimpleNamespace(pages=[_page(0, "content")]) + + out = tmp_path / "out" + FakeMistralProcessor.make(responder, include_images=False).process(img, output_path=out) + assert calls["n"] == 1 + + # Same input + checksum, but extract_header now on → different fingerprint. + proc2 = FakeMistralProcessor.make(responder, include_images=False, extract_header=True) + proc2.process(img, output_path=out) + assert calls["n"] == 2 # reprocessed, not silently reused + + +def test_checksum_resume_skip_emits_output_path(tmp_path: Path) -> None: + """A checksum-resumed (skipped-but-valid) doc still reports its .md path. + + The quiet scripting contract is one .md path per line on stdout, *including* + docs skipped on resume. The skip branch previously called add(COMPLETED) + without an output_path, so a rerun emitted nothing. This asserts the path is + reported AND that the API is not called again on the second run. + """ + img = tmp_path / "resume.png" + img.write_bytes(_PNG_1x1) + + calls = {"n": 0} + + def responder(_kwargs): + calls["n"] += 1 + return SimpleNamespace(pages=[_page(0, "content")]) + + out = tmp_path / "out" + proc1 = FakeMistralProcessor.make(responder, include_images=False) + first = proc1.process(img, output_path=out) + assert first.exit_code == 0 + assert calls["n"] == 1 + expected_md = str(markdown_path_for(doc_dir_for(out, "resume.png"), "resume.png")) + assert first.outputs == [expected_md] + + # Second run: a fresh processor (fresh RootIndex load from disk) resumes. + proc2 = FakeMistralProcessor.make(responder, include_images=False) + second = proc2.process(img, output_path=out) + assert second.exit_code == 0 + # The API was NOT called again (checksum + fingerprint + on-disk skip). + assert calls["n"] == 1 + # ...yet the skipped-but-valid doc STILL reports its .md path for scripting. + assert second.outputs == [expected_md] + + +def test_changed_fingerprint_reprocesses(tmp_path: Path) -> None: + """A re-run under a different run config (table_format) reprocesses, not skips.""" + img = tmp_path / "fp.png" + img.write_bytes(_PNG_1x1) + + calls = {"n": 0} + + def responder(_kwargs): + calls["n"] += 1 + return SimpleNamespace(pages=[_page(0, "content")]) + + out = tmp_path / "out" + FakeMistralProcessor.make(responder, include_images=False).process(img, output_path=out) + assert calls["n"] == 1 + + # Same input + checksum, but a different OCR-3 config → different fingerprint. + proc2 = FakeMistralProcessor.make(responder, include_images=False, table_format="markdown") + proc2.process(img, output_path=out) + assert calls["n"] == 2 # reprocessed, not silently reused + + +def test_persistence_failure_does_not_abort_batch(tmp_path: Path) -> None: + """A save/metadata I/O failure on one doc is recorded failed; batch continues. + + Persistence failures previously escaped to the CLI and aborted the whole + batch with no status=failed record. They must instead be folded into a + failed RunOutcome entry while the remaining files still process. + """ + root = tmp_path / "in" + root.mkdir() + for name in ("a.png", "boom.png", "c.png"): + (root / name).write_bytes(_PNG_1x1) + + proc = FakeMistralProcessor.make( + lambda _k: SimpleNamespace(pages=[_page(0, "content")]), include_images=False + ) + + # Make save_results raise for exactly one document; the others persist fine. + real_save = proc.save_results + + def flaky_save(result, output_root, rel_key): + if Path(rel_key).name == "boom.png": + raise OSError("disk full") + return real_save(result, output_root, rel_key) + + proc.save_results = flaky_save + + outcome = proc.process(root, output_path=tmp_path / "out") + + # Batch did NOT abort: the two good files completed, the bad one is failed. + assert outcome.completed == 2 + assert outcome.failed == 1 + assert outcome.exit_code != 0 + assert "boom.png" in outcome.failures + + # The failure is durably recorded (status=failed) in the root index, and the + # failed placeholder path is NOT echoed to the quiet stdout list. + root_index = json.loads((tmp_path / "out" / "metadata.json").read_text()) + assert root_index["files"]["boom.png"]["status"] == "failed" + failed_md = str(markdown_path_for(doc_dir_for(tmp_path / "out", "boom.png"), "boom.png")) + assert failed_md not in outcome.outputs + assert len(outcome.outputs) == 2 diff --git a/tests/test_pdf_chunking.py b/tests/test_pdf_chunking.py index e027b54..a21ad97 100644 --- a/tests/test_pdf_chunking.py +++ b/tests/test_pdf_chunking.py @@ -115,7 +115,7 @@ def test_negative_raises(self) -> None: def _make_processor(**overrides: object) -> OCRProcessor: - defaults = {"api_key": "test-key", "include_images": False, "save_original_images": False} + defaults = {"api_key": "test-key", "include_images": False} defaults.update(overrides) config = Config(**defaults) with patch("mistral_ocr.processor.Mistral"): @@ -127,38 +127,55 @@ def _fake_page(index: int) -> SimpleNamespace: class TestProcessPdf: - """Test _process_pdf routing logic with mocked upload.""" + """Test _process_pdf routing logic with mocked upload. + + The PDF must exist on disk: _process_pdf now validates file size first + (audit LOW fix — MAX_FILE_SIZE_MB is enforced for PDFs too, not only the + non-PDF branch). get_pdf_page_count is mocked, so the file's real page + count is irrelevant; only its existence + size matter. + """ @patch("mistral_ocr.processor.get_pdf_page_count", return_value=5) - def test_small_pdf_uploads_directly(self, mock_count: MagicMock) -> None: + def test_small_pdf_uploads_directly(self, mock_count: MagicMock, tmp_path: Path) -> None: + pdf = _make_pdf(tmp_path / "test.pdf", 1) 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")) + result = proc._process_pdf(pdf) + proc._upload_and_process.assert_called_once_with(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: + def test_max_pages_triggers_chunking(self, mock_count: MagicMock, tmp_path: Path) -> None: + pdf = _make_pdf(tmp_path / "test.pdf", 1) 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) + result = proc._process_pdf(pdf) + proc._process_pdf_chunked.assert_called_once_with(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: + def test_large_pdf_triggers_chunking(self, mock_count: MagicMock, tmp_path: Path) -> None: + pdf = _make_pdf(tmp_path / "test.pdf", 1) 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) + result = proc._process_pdf(pdf) + proc._process_pdf_chunked.assert_called_once_with(pdf, 1500) assert len(result.pages) == 1500 + @patch("mistral_ocr.processor.get_pdf_page_count", return_value=1) + def test_oversized_pdf_rejected(self, mock_count: MagicMock, tmp_path: Path) -> None: + """Audit LOW fix: MAX_FILE_SIZE_MB is enforced for the PDF branch too.""" + pdf = _make_pdf(tmp_path / "big.pdf", 1) + proc = _make_processor(max_file_size_mb=0.000001) # any real file exceeds this + with pytest.raises(ValueError, match="exceeds maximum allowed size"): + proc._process_pdf(pdf) + class TestProcessPdfChunked: """Test _process_pdf_chunked with real split_pdf but mocked API calls.""" diff --git a/tests/test_retry.py b/tests/test_retry.py index 672c750..5ad9937 100644 --- a/tests/test_retry.py +++ b/tests/test_retry.py @@ -119,13 +119,29 @@ def test_process_file_retries_transient_api_error(self, processor, tmp_path): img = tmp_path / "doc.png" img.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50) - mock_response = SimpleNamespace(pages=[]) + mock_response = SimpleNamespace( + pages=[SimpleNamespace(index=0, markdown="text", images=[])] + ) processor.client.ocr.process.side_effect = [ ConnectionError("transient"), mock_response, ] result = processor.process_file(img) - assert result is not None - assert result["success"] is True + assert result.status.value == "completed" + assert result.page_count == 1 assert processor.client.ocr.process.call_count == 2 + + def test_permanent_sdk_error_not_retried(self, processor, tmp_path): + """Audit fix: a permanent 4xx-style SDKError is not retried.""" + img = tmp_path / "doc.png" + img.write_bytes(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50) + + err = Exception("unauthorized") + err.status_code = 401 + processor.client.ocr.process.side_effect = err + + result = processor.process_file(img) + assert result.status.value == "failed" + # 401 is permanent -> exactly one attempt, no retries. + assert processor.client.ocr.process.call_count == 1 diff --git a/tests/test_save_results.py b/tests/test_save_results.py index d4cb5fa..6a2de75 100644 --- a/tests/test_save_results.py +++ b/tests/test_save_results.py @@ -1,280 +1,248 @@ -"""Tests for save_results rendering: OCR 3 features, images, truncation, originals.""" +"""Tests for mistral-specific page rendering and figure extraction. + +The output *shape* (layout, page assembly, metadata, figure naming) is owned by +``ocr-output-contract`` and is exercised end-to-end in +``tests/test_output_contract.py`` via the package's conformance harness. What +stays here is the mistral-owned glue: how one API page's OCR 3 extras +(dimensions, header/footer, hyperlinks) fold into a single page body, and how +embedded base64 images become canonically-named PNG figures with resolving +links. + +These call the real :class:`OCRProcessor` methods directly with a stub config +(no Mistral client), so they need no API key. +""" -import base64 from types import SimpleNamespace from mistral_ocr.config import Config -from mistral_ocr.processor import OCRProcessor +from mistral_ocr.processor import OCRProcessor, OCRResult -def _make_processor(**config_kwargs): +def _proc(**config_kwargs): + """An OCRProcessor with a stub config and no live Mistral client.""" proc = OCRProcessor.__new__(OCRProcessor) - config_kwargs.setdefault("save_original_images", False) proc.config = Config(api_key="test", **config_kwargs) - proc.errors = [] - proc.processed_files = [] return proc def _page(index=0, markdown="text", **kwargs): - return SimpleNamespace(index=index, markdown=markdown, images=[], **kwargs) - - -def _result(tmp_path, response, name="doc.pdf"): - fp = tmp_path / name - fp.write_bytes(b"%PDF-1.4 fake") - return {"file_path": fp, "response": response} + kwargs.setdefault("images", []) + return SimpleNamespace(index=index, markdown=markdown, **kwargs) # --------------------------------------------------------------------------- -# Original file copy +# Page body rendering (_render_page_markdown): OCR 3 extras folded inline # --------------------------------------------------------------------------- -class TestSaveOriginals: - def test_copies_original_when_enabled(self, tmp_path): - proc = _make_processor(save_original_images=True) - out = tmp_path / "out" - out.mkdir() - res = _result(tmp_path, SimpleNamespace(pages=[_page()])) - proc.save_results(res, out) - assert (out / "doc" / "doc.pdf").exists() - - def test_skips_copy_when_disabled(self, tmp_path): - proc = _make_processor(save_original_images=False) - out = tmp_path / "out" - out.mkdir() - res = _result(tmp_path, SimpleNamespace(pages=[_page()])) - proc.save_results(res, out) - assert not (out / "doc" / "doc.pdf").exists() - - def test_original_link_in_metadata(self, tmp_path): - proc = _make_processor(save_original_images=True, include_metadata=True) - out = tmp_path / "out" - out.mkdir() - res = _result(tmp_path, SimpleNamespace(pages=[_page()])) - proc.save_results(res, out) - md = (out / "doc" / "doc.md").read_text() - assert "**Original:**" in md - assert "doc.pdf" in md +class TestPageDimensions: + def test_dimensions_rendered(self): + page = _page(dimensions=SimpleNamespace(width=612, height=792)) + body = _proc()._render_page_markdown(page) + assert "612 x 792" in body + def test_no_dimensions_when_absent(self): + body = _proc()._render_page_markdown(_page()) + assert "Page size" not in body -# --------------------------------------------------------------------------- -# Truncation note -# --------------------------------------------------------------------------- +class TestHeaderFooter: + def test_header_rendered(self): + body = _proc()._render_page_markdown(_page(header="Chapter 1")) + assert "> **Header:** Chapter 1" in body -class TestTruncationNote: - def test_truncation_note_rendered(self, tmp_path): - proc = _make_processor(include_metadata=True) - out = tmp_path / "out" - out.mkdir() - response = SimpleNamespace( - pages=[_page()], truncated="Processed 50 of 200 pages (--max-pages)" - ) - res = _result(tmp_path, response) - proc.save_results(res, out) - md = (out / "doc" / "doc.md").read_text() - assert "**Note:** Processed 50 of 200" in md + def test_footer_rendered(self): + body = _proc()._render_page_markdown(_page(footer="Page 1 of 10")) + assert "> **Footer:** Page 1 of 10" in body - def test_no_truncation_note_when_absent(self, tmp_path): - proc = _make_processor(include_metadata=True) - out = tmp_path / "out" - out.mkdir() - res = _result(tmp_path, SimpleNamespace(pages=[_page()])) - proc.save_results(res, out) - md = (out / "doc" / "doc.md").read_text() - assert "**Note:**" not in md + def test_no_header_when_empty(self): + body = _proc()._render_page_markdown(_page(header="")) + assert "**Header:**" not in body -# --------------------------------------------------------------------------- -# Page dimensions (OCR 3) -# --------------------------------------------------------------------------- +class TestHyperlinks: + def test_hyperlinks_rendered(self): + link = SimpleNamespace(text="Example", url="https://example.com") + body = _proc()._render_page_markdown(_page(hyperlinks=[link])) + assert "**Hyperlinks:**" in body + assert "[Example](https://example.com)" in body + def test_hyperlink_without_text(self): + link = SimpleNamespace(text="", url="https://example.com") + body = _proc()._render_page_markdown(_page(hyperlinks=[link])) + assert "[https://example.com](https://example.com)" in body -class TestPageDimensions: - def test_dimensions_rendered(self, tmp_path): - proc = _make_processor(include_page_headings=True) - out = tmp_path / "out" - out.mkdir() - dims = SimpleNamespace(width=612, height=792) - page = _page(dimensions=dims) - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - md = (out / "doc" / "doc.md").read_text() - assert "612 x 792" in md - - def test_no_dimensions_when_absent(self, tmp_path): - proc = _make_processor() - out = tmp_path / "out" - out.mkdir() - res = _result(tmp_path, SimpleNamespace(pages=[_page()])) - proc.save_results(res, out) - md = (out / "doc" / "doc.md").read_text() - assert "Page size" not in md + def test_hyperlink_with_href_fallback(self): + link = SimpleNamespace(text="Link", href="https://test.com") # no 'url' attr + body = _proc()._render_page_markdown(_page(hyperlinks=[link])) + assert "[Link](https://test.com)" in body -# --------------------------------------------------------------------------- -# Headers and footers (OCR 3) -# --------------------------------------------------------------------------- +class TestPageBodyIsCleanBody: + def test_body_has_no_page_header(self): + """The page body must NOT carry its own ## Page header. + The contract's assemble_pages adds the canonical ``## Page N`` header, + so a doubled header would violate the marker count. The renderer emits + only the page body. + """ + body = _proc()._render_page_markdown(_page(markdown="hello")) + assert "## Page" not in body + assert body.strip() == "hello" -class TestHeaderFooter: - def test_header_rendered(self, tmp_path): - proc = _make_processor() - out = tmp_path / "out" - out.mkdir() - page = _page(header="Chapter 1") - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - md = (out / "doc" / "doc.md").read_text() - assert "> **Header:** Chapter 1" in md - - def test_footer_rendered(self, tmp_path): - proc = _make_processor() - out = tmp_path / "out" - out.mkdir() - page = _page(footer="Page 1 of 10") - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - md = (out / "doc" / "doc.md").read_text() - assert "> **Footer:** Page 1 of 10" in md - - def test_no_header_when_empty(self, tmp_path): - proc = _make_processor() - out = tmp_path / "out" - out.mkdir() - page = _page(header="") - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - md = (out / "doc" / "doc.md").read_text() - assert "**Header:**" not in md +def _table(tid="tbl-0", content="| A | B |\n|---|---|\n| 1 | 2 |", fmt="markdown"): + return SimpleNamespace(id=tid, content=content, format_=fmt) -# --------------------------------------------------------------------------- -# Tables (OCR 3) -# --------------------------------------------------------------------------- +class TestStructuredTables: + """Regression guard: page.tables must NOT be silently discarded. -class TestTables: - def test_table_saved_as_markdown(self, tmp_path): - proc = _make_processor(table_format="markdown") - out = tmp_path / "out" - out.mkdir() - table = SimpleNamespace(content="| A | B |\n|---|---|\n| 1 | 2 |") - page = _page(tables=[table]) - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - table_path = out / "doc" / "tables" / "page1_table1.md" - assert table_path.exists() - assert "| A | B |" in table_path.read_text() - md = (out / "doc" / "doc.md").read_text() - assert "[Table 1](./tables/page1_table1.md)" in md - - def test_table_saved_as_html(self, tmp_path): - proc = _make_processor(table_format="html") - out = tmp_path / "out" - out.mkdir() - table = SimpleNamespace(content="

1
") - page = _page(tables=[table]) - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - table_path = out / "doc" / "tables" / "page1_table1.html" - assert table_path.exists() - assert "" in table_path.read_text() - - def test_table_falls_back_to_markdown_attr(self, tmp_path): - proc = _make_processor(table_format="markdown") - out = tmp_path / "out" - out.mkdir() - table = SimpleNamespace(markdown="| X |") # no 'content' attr - page = _page(tables=[table]) - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - assert "| X |" in (out / "doc" / "tables" / "page1_table1.md").read_text() + With --table-format=markdown|html Mistral returns tables in the structured + page.tables field (NOT inline in page.markdown), so a renderer that ignores + page.tables makes --table-format a silent no-op and drops OCR-3 table + extraction. These assert the structured tables land in the rendered body. + """ + def test_table_appended_when_no_placeholder(self): + page = _page(markdown="Body text", tables=[_table()]) + body = _proc(table_format="markdown")._render_page_markdown(page) + assert "Body text" in body + # The table content survived (not dropped). + assert "| A | B |" in body and "| 1 | 2 |" in body -# --------------------------------------------------------------------------- -# Hyperlinks (OCR 3) -# --------------------------------------------------------------------------- + def test_table_replaces_inline_placeholder(self): + """A markdown placeholder referencing the table id is replaced in place.""" + page = _page( + markdown="See table below:\n\n[tbl-0](tbl-0)\n\nrest", + tables=[_table(tid="tbl-0", content="| X |\n|---|\n| 9 |")], + ) + body = _proc(table_format="markdown")._render_page_markdown(page) + assert "[tbl-0](tbl-0)" not in body # placeholder consumed + assert "| X |" in body and "| 9 |" in body + assert "rest" in body + + def test_html_table_passed_through(self): + page = _page( + markdown="Body", + tables=[_table(content="
v
", fmt="html")], + ) + body = _proc(table_format="html")._render_page_markdown(page) + assert "
v
" in body + def test_multiple_tables_all_preserved(self): + page = _page( + markdown="Body", + tables=[_table(tid="t1", content="TABLE-ONE"), _table(tid="t2", content="TABLE-TWO")], + ) + body = _proc(table_format="markdown")._render_page_markdown(page) + assert "TABLE-ONE" in body and "TABLE-TWO" in body -class TestHyperlinks: - def test_hyperlinks_rendered(self, tmp_path): - proc = _make_processor() - out = tmp_path / "out" - out.mkdir() - link = SimpleNamespace(text="Example", url="https://example.com") - page = _page(hyperlinks=[link]) - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - md = (out / "doc" / "doc.md").read_text() - assert "**Hyperlinks:**" in md - assert "[Example](https://example.com)" in md - - def test_hyperlink_without_text(self, tmp_path): - proc = _make_processor() - out = tmp_path / "out" - out.mkdir() - link = SimpleNamespace(text="", url="https://example.com") - page = _page(hyperlinks=[link]) - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - md = (out / "doc" / "doc.md").read_text() - assert "[https://example.com](https://example.com)" in md - - def test_hyperlink_with_href_fallback(self, tmp_path): - proc = _make_processor() - out = tmp_path / "out" - out.mkdir() - link = SimpleNamespace(text="Link", href="https://test.com") # no 'url' attr - page = _page(hyperlinks=[link]) - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - md = (out / "doc" / "doc.md").read_text() - assert "[Link](https://test.com)" in md + def test_no_tables_is_noop(self): + """When --table-format is unset the API inlines tables; render is a no-op.""" + page = _page(markdown="inline | table | here") + body = _proc()._render_page_markdown(page) + assert body.strip() == "inline | table | here" + + def test_empty_tables_list_is_noop(self): + page = _page(markdown="just text", tables=[]) + body = _proc(table_format="markdown")._render_page_markdown(page) + assert body.strip() == "just text" # --------------------------------------------------------------------------- -# Figures / images +# Figure extraction (_save_figures): canonical figure__page

.png naming # --------------------------------------------------------------------------- +# A valid 1x1 PNG so Pillow can open + re-encode it. +_PNG_1x1 = ( + b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x06" + b"\x00\x00\x00\x1f\x15\xc4\x89\x00\x00\x00\nIDATx\x9cc\x00\x01\x00\x00\x05\x00" + b"\x01\r\n-\xb4\x00\x00\x00\x00IEND\xaeB`\x82" +) + + +def _result_with_image(tmp_path, page_no=1, image_bytes=_PNG_1x1, image_id="img-0.jpeg"): + fp = tmp_path / "doc.png" + fp.write_bytes(_PNG_1x1) + return OCRResult( + file_path=fp, + pages=["text"], + # page_images now carries (image_id, raw_bytes) pairs so the inline + # placeholder ![..](image_id) can be rewritten to the canonical name. + page_images={page_no: [(image_id, image_bytes)]}, + ) + class TestFigures: - def test_image_saved_and_linked(self, tmp_path): - proc = _make_processor(include_images=True) - out = tmp_path / "out" - out.mkdir() - # 1x1 red PNG as base64 - b64 = base64.b64encode(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50).decode() - img = SimpleNamespace(image_base64=b64, id="fig1.png") - page = SimpleNamespace(index=0, markdown="text", images=[img]) - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - assert (out / "doc" / "figures" / "page1_img1.png").exists() - md = (out / "doc" / "doc.md").read_text() - assert "![Image 1](./figures/page1_img1.png)" in md + def test_image_saved_with_canonical_name_and_link(self, tmp_path): + proc = _proc(include_images=True) + doc_dir = tmp_path / "doc" + doc_dir.mkdir() + result = _result_with_image(tmp_path, page_no=1) + + saved = proc._save_figures(result, doc_dir) + + # Canonical naming: figure__page

.png (NOT page1_img1.png). + assert (doc_dir / "figures" / "figure_1_page1.png").exists() + assert not (doc_dir / "figures" / "page1_img1.png").exists() + # Returns a saved-figure record carrying the API image id + canonical + # number/page for the page that produced it. + rec = saved[1][0] + assert rec.image_id == "img-0.jpeg" + assert rec.figure_number == 1 + assert rec.page_number == 1 + + def test_figure_numbering_is_global_with_source_page(self, tmp_path): + """Figures are numbered globally across the doc, tagged by source page.""" + proc = _proc(include_images=True) + doc_dir = tmp_path / "doc" + doc_dir.mkdir() + result = OCRResult( + file_path=tmp_path / "doc.png", + pages=["a", "b"], + page_images={ + 1: [("img-0.jpeg", _PNG_1x1)], + 2: [("img-1.jpeg", _PNG_1x1), ("img-2.jpeg", _PNG_1x1)], + }, + ) + (tmp_path / "doc.png").write_bytes(_PNG_1x1) + + proc._save_figures(result, doc_dir) + figs = sorted(p.name for p in (doc_dir / "figures").iterdir()) + assert figs == [ + "figure_1_page1.png", + "figure_2_page2.png", + "figure_3_page2.png", + ] def test_no_figures_when_images_disabled(self, tmp_path): - proc = _make_processor(include_images=False) - out = tmp_path / "out" - out.mkdir() - b64 = base64.b64encode(b"\x89PNG\r\n\x1a\n" + b"\x00" * 50).decode() - img = SimpleNamespace(image_base64=b64, id="fig1.png") - page = SimpleNamespace(index=0, markdown="text", images=[img]) - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - assert not (out / "doc" / "figures").exists() - md = (out / "doc" / "doc.md").read_text() - assert "![Image" not in md - - def test_image_default_extension(self, tmp_path): - proc = _make_processor(include_images=True) - out = tmp_path / "out" - out.mkdir() - b64 = base64.b64encode(b"\x89PNG" + b"\x00" * 50).decode() - img = SimpleNamespace(image_base64=b64, id="no_ext") # no extension in id - page = SimpleNamespace(index=0, markdown="text", images=[img]) - res = _result(tmp_path, SimpleNamespace(pages=[page])) - proc.save_results(res, out) - assert (out / "doc" / "figures" / "page1_img1.png").exists() + proc = _proc(include_images=False) + doc_dir = tmp_path / "doc" + doc_dir.mkdir() + result = _result_with_image(tmp_path, page_no=1) + + links = proc._save_figures(result, doc_dir) + assert links == {} + assert not (doc_dir / "figures").exists() + + +# --------------------------------------------------------------------------- +# Truncation note (folded into the body by save_results, recorded in metadata) +# --------------------------------------------------------------------------- + + +class TestTruncationNote: + def test_note_prepended_to_body(self, tmp_path): + proc = _proc(include_images=False, verbose=False) + fp = tmp_path / "doc.png" + fp.write_bytes(_PNG_1x1) + result = OCRResult( + file_path=fp, + pages=["page one"], + note="Processed 50 of 200 pages (--max-pages)", + ) + md_path = proc.save_results(result, tmp_path / "out", "doc.png") + body = md_path.read_text() + assert "> **Note:** Processed 50 of 200 pages" in body + assert "## Page 1" in body # contract still owns page headers diff --git a/tests/test_utils.py b/tests/test_utils.py index ebdf276..2faa4fc 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,24 +1,34 @@ -"""Tests for utility functions.""" +"""Tests for utility functions (mistral-specific I/O only). + +Output-shape helpers (output-root resolution, keying, layout, metadata, figure +naming) now live in the ``ocr-output-contract`` package and are exercised via +``test_output_contract.py``, not here. +""" import base64 import pytest from mistral_ocr.utils import ( + DOCUMENT_EXTENSIONS, + IMAGE_EXTENSIONS, + SUPPORTED_EXTENSIONS, create_data_uri, - determine_output_path, + decode_base64_image, encode_file_to_base64, format_file_size, get_mime_type, get_supported_files, - load_metadata, - make_unique_basename, - sanitize_filename, - save_base64_image, - save_metadata, ) +class TestExtensionSets: + def test_supported_is_union(self): + assert SUPPORTED_EXTENSIONS == DOCUMENT_EXTENSIONS | IMAGE_EXTENSIONS + assert ".pdf" in DOCUMENT_EXTENSIONS + assert ".png" in IMAGE_EXTENSIONS + + class TestEncodeFileToBase64: def test_encodes_file(self, tmp_path): f = tmp_path / "test.txt" @@ -59,19 +69,15 @@ def test_creates_uri(self, tmp_path): assert uri.startswith("data:image/png;base64,") -class TestSaveBase64Image: - def test_saves_image(self, tmp_path): - data = base64.b64encode(b"fake image data").decode() - out = tmp_path / "subdir" / "img.png" - save_base64_image(data, out) - assert out.read_bytes() == b"fake image data" +class TestDecodeBase64Image: + def test_decodes_raw(self): + raw = base64.b64encode(b"pixels").decode() + assert decode_base64_image(raw) == b"pixels" - def test_strips_data_uri_prefix(self, tmp_path): - """save_base64_image should handle raw base64 (no prefix).""" + def test_strips_data_uri_prefix(self): raw = base64.b64encode(b"pixels").decode() - out = tmp_path / "img.png" - save_base64_image(raw, out) - assert out.read_bytes() == b"pixels" + uri = f"data:image/png;base64,{raw}" + assert decode_base64_image(uri) == b"pixels" class TestGetSupportedFiles: @@ -79,107 +85,52 @@ def test_finds_supported_files(self, tmp_path): (tmp_path / "doc.pdf").write_bytes(b"") (tmp_path / "img.png").write_bytes(b"") (tmp_path / "notes.txt").write_bytes(b"") - files = get_supported_files(tmp_path) + files = get_supported_files(tmp_path, tmp_path / "ocr") names = {f.name for f in files} assert names == {"doc.pdf", "img.png"} - def test_excludes_output_dir(self, tmp_path): - out = tmp_path / "mistral_ocr_output" + def test_excludes_resolved_output_root(self, tmp_path): + """The RESOLVED output root (its prior outputs) is excluded on a rerun.""" + out = tmp_path / "ocr" out.mkdir() (out / "result.pdf").write_bytes(b"") (tmp_path / "input.pdf").write_bytes(b"") - files = get_supported_files(tmp_path) + files = get_supported_files(tmp_path, out) assert len(files) == 1 assert files[0].name == "input.pdf" - def test_excludes_absolute_paths(self, tmp_path): + def test_does_not_exclude_arbitrary_dir_named_ocr(self, tmp_path): + """A non-output dir literally named 'ocr' is NOT skipped (the bug fix). + + Under the user's own .../toolkits/ocr/... tree the old name-based + exclusion silently dropped every input. Only the resolved output root is + excluded now, so an input under an 'ocr'-named ancestor is processed. + """ + ocr_dir = tmp_path / "ocr" + ocr_dir.mkdir() + (ocr_dir / "report.pdf").write_bytes(b"") + # Output root is elsewhere, so the 'ocr' input dir must survive. + files = get_supported_files(tmp_path, tmp_path / "results") + assert {f.name for f in files} == {"report.pdf"} + + def test_excludes_custom_output_root_inside_input(self, tmp_path): + """A custom -o dir nested in the input tree is excluded by resolved path.""" sub = tmp_path / "output" sub.mkdir() (sub / "file.pdf").write_bytes(b"") (tmp_path / "input.pdf").write_bytes(b"") - files = get_supported_files(tmp_path, exclude_paths=[sub]) - assert len(files) == 1 + files = get_supported_files(tmp_path, sub) + assert {f.name for f in files} == {"input.pdf"} def test_recursive(self, tmp_path): sub = tmp_path / "nested" sub.mkdir() (sub / "deep.jpg").write_bytes(b"") - files = get_supported_files(tmp_path) + files = get_supported_files(tmp_path, tmp_path / "ocr") assert len(files) == 1 assert files[0].name == "deep.jpg" -class TestDetermineOutputPath: - def test_default_output(self, tmp_path): - f = tmp_path / "doc.pdf" - f.write_bytes(b"") - result = determine_output_path(f) - assert result.name == "mistral_ocr_output" - assert result.parent == tmp_path - - def test_custom_output(self, tmp_path): - f = tmp_path / "doc.pdf" - f.write_bytes(b"") - custom = tmp_path / "custom_out" - result = determine_output_path(f, output_path=custom) - assert result == custom - assert custom.exists() - - def test_creates_nested_output(self, tmp_path): - f = tmp_path / "doc.pdf" - f.write_bytes(b"") - nested = tmp_path / "a" / "b" / "c" - result = determine_output_path(f, output_path=nested) - assert result == nested - assert nested.is_dir() - - def test_output_path_is_file_raises(self, tmp_path): - f = tmp_path / "doc.pdf" - f.write_bytes(b"") - existing_file = tmp_path / "block" - existing_file.write_bytes(b"") - with pytest.raises(ValueError, match="not a directory"): - determine_output_path(f, output_path=existing_file) - - def test_timestamp(self, tmp_path): - f = tmp_path / "doc.pdf" - f.write_bytes(b"") - result = determine_output_path(f, add_timestamp=True) - assert "mistral_ocr_output_" in result.name - - -class TestMetadata: - def test_load_missing(self, tmp_path): - meta = load_metadata(tmp_path) - assert meta["files_processed"] == [] - assert meta["total_files"] == 0 - - def test_save_and_load(self, tmp_path): - save_metadata(tmp_path, [{"file": "/a.pdf", "size": 100}], 1.5, []) - meta = load_metadata(tmp_path) - assert meta["total_files"] == 1 - assert meta["files_processed"][0]["file"] == "/a.pdf" - assert meta["processing_time_seconds"] == 1.5 - - def test_incremental_save(self, tmp_path): - save_metadata(tmp_path, [{"file": "/a.pdf", "size": 100}], 1.0, []) - save_metadata( - tmp_path, - [{"file": "/b.pdf", "size": 200}], - 2.0, - [], - base_processing_time=1.0, - ) - meta = load_metadata(tmp_path) - assert meta["total_files"] == 2 - assert meta["processing_time_seconds"] == 3.0 - - def test_corrupt_metadata(self, tmp_path): - (tmp_path / "metadata.json").write_text("not json{{{") - meta = load_metadata(tmp_path) - assert meta["files_processed"] == [] - - class TestFormatFileSize: @pytest.mark.parametrize( "size,expected", @@ -192,29 +143,3 @@ class TestFormatFileSize: ) def test_formatting(self, size, expected): assert format_file_size(size) == expected - - -class TestMakeUniqueBasename: - def test_simple(self, tmp_path): - f = tmp_path / "report.pdf" - assert make_unique_basename(f) == "report" - - def test_with_base_dir(self, tmp_path): - sub = tmp_path / "sub" - sub.mkdir() - f = sub / "report.pdf" - result = make_unique_basename(f, base_dir=tmp_path) - assert result == "sub__report" - - -class TestSanitizeFilename: - def test_removes_invalid_chars(self): - assert sanitize_filename("file<>:name") == "file___name" - - def test_truncates(self): - result = sanitize_filename("a" * 300, max_length=50) - assert len(result) <= 50 - - def test_no_truncation_by_default(self): - long = "a" * 300 - assert sanitize_filename(long) == long