From 9c8bf7f4a000c87f31acabfab987af32cf9f91ac Mon Sep 17 00:00:00 2001 From: Daniel Vejsada Date: Fri, 23 Oct 2026 19:09:32 +0200 Subject: [PATCH 1/4] chore(deps): bump libraries; drop dependencies the fork no longer uses Bump every pin to its latest release within the current major line, and move the langchain family together (langchain-core 1.7.0 is required by the new langchain, langchain-openai, langchain-aws, langchain-google-genai and langchain-huggingface). boto3 rises for langchain-aws 1.9; aiofiles 25.1 unblocks current unstructured-client. tenacity stays at 9.1.4: google-genai requires <9.2. Held on purpose (breaking majors): pandas 3, SQLAlchemy 2.1, Pillow 13, chardet 7, sentence-transformers 6, mistralai 3, pgvector 0.5, asyncpg 0.32. Removed, since PDFs/images are OCR'd remotely via Mistral: - rapidocr-onnxruntime, opencv-python-headless (no imports) and the libgl1/libglib2.0-0 apt packages they needed in the full Dockerfile - networkx (no imports; only an optional langchain-mongodb extra) - pypdf from runtime deps (only a test fixture uses it; moved to test_requirements.txt, and unstructured-client still pulls it in) - mongomock (unreferenced) Co-Authored-By: Claude Opus 5.5 --- Dockerfile | 2 -- requirements.lite.txt | 46 +++++++++++++++++------------------ requirements.txt | 56 ++++++++++++++++++++----------------------- test_requirements.txt | 8 +++---- 4 files changed, 52 insertions(+), 60 deletions(-) diff --git a/Dockerfile b/Dockerfile index 8d7a836e..31c0b8c3 100644 --- a/Dockerfile +++ b/Dockerfile @@ -9,8 +9,6 @@ RUN apt-get update \ && apt-get install -y --no-install-recommends \ pandoc \ netcat-openbsd \ - libgl1 \ - libglib2.0-0 \ && rm -rf /var/lib/apt/lists/* COPY requirements.txt . diff --git a/requirements.lite.txt b/requirements.lite.txt index d617ab27..2555d97e 100644 --- a/requirements.lite.txt +++ b/requirements.lite.txt @@ -1,40 +1,38 @@ -langchain==1.3.10 +langchain==1.4.7 langchain-community==0.4.2 -langchain-openai==1.3.2 -langchain-core==1.4.8 -langchain-google-genai==4.2.5 -sqlalchemy==2.0.41 -python-dotenv==1.2.2 -fastapi==0.143.1 -starlette==1.7.0 # >=1.3.1 for CVE-2025-62727, CVE-2026-48818, CVE-2026-54283; fastapi no longer caps it -psycopg2-binary==2.9.9 +langchain-openai==1.7.2 +langchain-core==1.7.0 +langchain-google-genai==4.4.2 +sqlalchemy==2.0.55 +python-dotenv==1.2.4 +fastapi==0.161.1 +starlette==1.9.0 # >=1.3.1 for CVE-2025-62727, CVE-2026-48818, CVE-2026-54283; fastapi no longer caps it +psycopg2-binary==2.9.13 pgvector==0.2.5 -uvicorn==0.28.0 -pypdf==6.19.0 -unstructured==0.27.8 -# unstructured >=0.24 lazily downloads this spaCy model from GitHub on first use; install it at build time instead (same pin/hash unstructured 0.27.8 expects). +uvicorn==0.54.0 +unstructured==0.27.34 +# unstructured >=0.24 lazily downloads this spaCy model from GitHub on first use; install it at build time instead (same pin/hash unstructured 0.27.34 expects). en_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl#sha256=1932429db727d4bff3deed6b34cfc05df17794f4a52eeb26cf8928f7c1a0fb85 -markdown==3.8.2 -networkx==3.2.1 -pandas==2.2.1 +markdown==3.11.2 +pandas==2.3.3 openpyxl==3.1.5 docx2txt==0.9 -python-docx==1.1.2 # .docx header/footer extraction (pandoc drops them) -pypandoc==1.15 +python-docx==1.2.0 # .docx header/footer extraction (pandoc drops them) +pypandoc==1.17 python-oxmsg==0.0.2 # .msg (Outlook) extraction; also a transitive of unstructured Pillow==12.3.0 # normalize standalone images (incl. multi-page TIFF) for OCR PyJWT==2.15.1 asyncpg==0.29.0 python-multipart==0.0.32 -aiofiles==24.1.0 -pymongo==4.18.2 -langchain-mongodb==0.11.0 -cryptography==50.0.1 +aiofiles==25.1.0 +pymongo==4.19.0 +langchain-mongodb==0.12.1 +cryptography==50.0.2 python-magic==0.4.27 python-pptx==1.0.2 xlrd==2.0.2 -langchain-aws==1.6.0 -boto3==1.43.48 +langchain-aws==1.9.0 +boto3==1.43.118 chardet==5.2.0 langchain-ollama==1.1.0 tenacity==9.1.4 diff --git a/requirements.txt b/requirements.txt index 0d2eab37..047fe5b0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,47 +1,43 @@ -langchain==1.3.10 +langchain==1.4.7 langchain-community==0.4.2 -langchain-openai==1.3.2 -langchain-core==1.4.8 -langchain-aws==1.6.0 -langchain-text-splitters==1.1.2 -boto3==1.43.48 -sqlalchemy==2.0.41 -python-dotenv==1.2.2 -fastapi==0.143.1 -starlette==1.7.0 # >=1.3.1 for CVE-2025-62727, CVE-2026-48818, CVE-2026-54283; fastapi no longer caps it -psycopg2-binary==2.9.9 +langchain-openai==1.7.2 +langchain-core==1.7.0 +langchain-aws==1.9.0 +langchain-text-splitters==1.1.3 +boto3==1.43.118 +sqlalchemy==2.0.55 +python-dotenv==1.2.4 +fastapi==0.161.1 +starlette==1.9.0 # >=1.3.1 for CVE-2025-62727, CVE-2026-48818, CVE-2026-54283; fastapi no longer caps it +psycopg2-binary==2.9.13 pgvector==0.2.5 -uvicorn==0.28.0 -pypdf==6.19.0 -unstructured==0.27.8 -# unstructured >=0.24 lazily downloads this spaCy model from GitHub on first use; install it at build time instead (same pin/hash unstructured 0.27.8 expects). +uvicorn==0.54.0 +unstructured==0.27.34 +# unstructured >=0.24 lazily downloads this spaCy model from GitHub on first use; install it at build time instead (same pin/hash unstructured 0.27.34 expects). en_core_web_sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.8.0/en_core_web_sm-3.8.0-py3-none-any.whl#sha256=1932429db727d4bff3deed6b34cfc05df17794f4a52eeb26cf8928f7c1a0fb85 -markdown==3.8.2 -networkx==3.2.1 -pandas==2.2.1 +markdown==3.11.2 +pandas==2.3.3 openpyxl==3.1.5 docx2txt==0.9 -python-docx==1.1.2 # .docx header/footer extraction (pandoc drops them) -pypandoc==1.15 +python-docx==1.2.0 # .docx header/footer extraction (pandoc drops them) +pypandoc==1.17 python-oxmsg==0.0.2 # .msg (Outlook) extraction; also a transitive of unstructured Pillow==12.3.0 # normalize standalone images (incl. multi-page TIFF) for OCR PyJWT==2.15.1 asyncpg==0.29.0 python-multipart==0.0.32 -sentence_transformers==3.1.1 -aiofiles==24.1.0 -rapidocr-onnxruntime==1.4.4 -opencv-python-headless==4.9.0.80 -pymongo==4.18.2 -langchain-mongodb==0.11.0 +sentence_transformers==3.4.1 +aiofiles==25.1.0 +pymongo==4.19.0 +langchain-mongodb==0.12.1 langchain-ollama==1.1.0 -langchain-huggingface==1.2.2 -langchain-google-genai==4.2.5 -cryptography==50.0.1 +langchain-huggingface==1.2.3 +langchain-google-genai==4.4.2 +cryptography==50.0.2 python-magic==0.4.27 python-pptx==1.0.2 xlrd==2.0.2 -pydantic==2.13.4 +pydantic==2.14.1 chardet==5.2.0 tenacity==9.1.4 msoffcrypto-tool==6.0.0 diff --git a/test_requirements.txt b/test_requirements.txt index df657952..86722af7 100644 --- a/test_requirements.txt +++ b/test_requirements.txt @@ -1,5 +1,5 @@ -pytest==9.0.3 -pytest-asyncio==1.3.0 -testcontainers[postgres]==4.14.2 -mongomock==4.3.0 +pytest==9.1.1 +pytest-asyncio==1.4.0 +testcontainers[postgres]==4.15.0 +pypdf==6.21.2 # builds the text-layer PDF fixture in tests/utils/test_lazy_load.py httpx==0.28.1 From 8de8fd20a6f64f8f00d27bc4312f36ebcd61f58b Mon Sep 17 00:00:00 2001 From: Daniel Vejsada Date: Fri, 23 Oct 2026 19:09:32 +0200 Subject: [PATCH 2/4] feat(ocr): OCR images embedded in PDFs and inline their text Mistral OCR leaves only a placeholder (![img-0.jpeg](img-0.jpeg)) for a raster image embedded in a PDF page, such as a photo or scan of a page, without reading the text inside it. The PDF request now asks for the extracted images (include_image_base64, image_min_size, image_limit); each image is OCR'd in its own call on a small thread pool, and the recognized text replaces its placeholder on the page it belongs to. A failed image call logs a warning and keeps the placeholder; the main PDF call still raises. With the feature off, the request is unchanged. New settings: MISTRAL_OCR_EMBEDDED_IMAGES (default True), MISTRAL_OCR_IMAGE_MIN_SIZE (100), MISTRAL_OCR_MAX_EMBEDDED_IMAGES (50, <=0 for no cap), MISTRAL_OCR_IMAGE_CONCURRENCY (4). Each embedded image is billed as one extra OCR page. Co-Authored-By: Claude Opus 5.5 --- README.md | 6 +- app/config.py | 19 ++ app/utils/document_loader.py | 193 ++++++++++++-- tests/utils/test_document_loader.py | 389 ++++++++++++++++++++++++++++ 4 files changed, 580 insertions(+), 27 deletions(-) diff --git a/README.md b/README.md index 0d81a103..f838c85a 100644 --- a/README.md +++ b/README.md @@ -214,8 +214,12 @@ The following environment variables are required to run the application: - `MISTRAL_SERVER_URL`: (Optional) Base URL of the Mistral API. Leave unset for the SDK default (`https://api.mistral.ai`); set `https://api.eu.mistral.ai` to keep OCR traffic on Mistral's [EU regional endpoint](https://docs.mistral.ai/inference/regional-inference). Regional endpoints don't offer the Files API, which is fine here: documents and images are sent inline as base64. Confirm with Mistral that your OCR model is served on the regional endpoint before relying on it. - Note: LibreChat's own `mistral_ocr` strategy uploads through the Files API, so on an EU-only deployment route PDFs and images to this service's `/text` endpoint instead. - For EU embeddings as well, point `RAG_OPENAI_BASEURL` at an EU OpenAI-compatible endpoint (e.g. `https://eu.openrouter.ai/api/v1`) with a non-OpenAI model such as `mistralai/mistral-embed-2312` (1024 dimensions), and set `RAG_CHECK_EMBEDDING_CTX_LENGTH=False`: with the check on, inputs are sent as OpenAI tiktoken token IDs, which non-OpenAI models don't accept. Switching embedding models requires re-indexing existing files (see `EMBEDDINGS_DIMENSIONS`). -- Standalone image uploads (`.png`, `.jpg/.jpeg`, `.gif`, `.bmp`, `.tif/.tiff`, `.webp`) — including multi-page TIFF — are run through **Mistral OCR**, yielding one result per page. Requires `MISTRAL_API_KEY` (same key used for PDF OCR). This is separate from PDF handling; images embedded inside PDFs are not extracted. +- Standalone image uploads (`.png`, `.jpg/.jpeg`, `.gif`, `.bmp`, `.tif/.tiff`, `.webp`) — including multi-page TIFF — are run through **Mistral OCR**, yielding one result per page. Requires `MISTRAL_API_KEY` (same key used for PDF OCR). This is separate from PDF handling; images embedded inside PDFs are covered by `MISTRAL_OCR_EMBEDDED_IMAGES` below. - `IMAGE_OCR_MAX_PAGES`: (Optional) Max number of frames/pages of a single image upload (e.g. animated GIF or large multi-page TIFF) sent to OCR; frames beyond the cap are skipped with a warning. Default value is `100`. +- `MISTRAL_OCR_EMBEDDED_IMAGES`: (Optional) Boolean. Mistral OCR returns only a placeholder (`![img-0.jpeg](img-0.jpeg)`) for a raster image embedded in a PDF page, such as a photo or scan of a page, and does not read the text inside it. When "True" (default), those images are extracted, OCR'd in a second Mistral call each, and the recognized text replaces the placeholder. **Each embedded image is billed as an extra OCR page**; set "False" to OCR only the PDF itself. If an image's OCR call fails, its placeholder is kept and the document still loads. +- `MISTRAL_OCR_IMAGE_MIN_SIZE`: (Optional) Minimum image size in pixels for an embedded image to be extracted, so logos and icons are skipped. Default value is `100`. +- `MISTRAL_OCR_MAX_EMBEDDED_IMAGES`: (Optional) Max number of embedded images extracted and OCR'd per document; extraction stops at the cap (a warning is logged when it is reached), so further images may be left as placeholders. Set to `0` or below for no cap. Default value is `50`. +- `MISTRAL_OCR_IMAGE_CONCURRENCY`: (Optional) Number of embedded images OCR'd in parallel (minimum `1`). Default value is `4`. - `DEBUG_RAG_API`: (Optional) Set to "True" to show more verbose logging output in the server console, and to enable postgresql database routes - `DEBUG_PGVECTOR_QUERIES`: (Optional) Set to "True" to enable detailed PostgreSQL query logging for pgvector operations. Useful for debugging performance issues with vector database queries. - `CONSOLE_JSON`: (Optional) Set to "True" to log as json for Cloud Logging aggregations diff --git a/app/config.py b/app/config.py index 2a34257c..e56be63b 100644 --- a/app/config.py +++ b/app/config.py @@ -224,6 +224,25 @@ async def dispatch(self, request, call_next): # Override the Mistral API base URL, e.g. https://api.eu.mistral.ai for EU-only # deployments. Empty keeps the SDK default (https://api.mistral.ai). MISTRAL_SERVER_URL = get_env_variable("MISTRAL_SERVER_URL", "").rstrip("/") +# Mistral OCR returns only a placeholder (![img-0.jpeg](img-0.jpeg)) for a raster +# image embedded in a PDF page, without transcribing the text inside it. When +# enabled, the extracted images are requested and OCR'd separately, and the +# recognized text replaces the placeholder. Every image costs one extra OCR page. +MISTRAL_OCR_EMBEDDED_IMAGES = ( + get_env_variable("MISTRAL_OCR_EMBEDDED_IMAGES", "True").lower() == "true" +) +# Smallest image side (pixels) Mistral extracts, so logos/icons aren't OCR'd. +MISTRAL_OCR_IMAGE_MIN_SIZE = max( + 0, int(get_env_variable("MISTRAL_OCR_IMAGE_MIN_SIZE", "100")) +) +# Max embedded images extracted (and OCR'd) per document; <= 0 means no cap. +MISTRAL_OCR_MAX_EMBEDDED_IMAGES = int( + get_env_variable("MISTRAL_OCR_MAX_EMBEDDED_IMAGES", "50") +) +# Number of embedded images OCR'd in parallel (minimum 1). +MISTRAL_OCR_IMAGE_CONCURRENCY = max( + 1, int(get_env_variable("MISTRAL_OCR_IMAGE_CONCURRENCY", "4")) +) # DOCX handling for the /text endpoint: when enabled, .docx files are converted # via pandoc so tracked changes and comments are preserved (the embedding path diff --git a/app/utils/document_loader.py b/app/utils/document_loader.py index 44ab09db..1d24124a 100644 --- a/app/utils/document_loader.py +++ b/app/utils/document_loader.py @@ -5,6 +5,7 @@ import re import tempfile +from concurrent.futures import ThreadPoolExecutor from typing import Iterator, List, Optional import chardet @@ -17,6 +18,10 @@ MISTRAL_API_KEY, MISTRAL_OCR_MODEL, MISTRAL_SERVER_URL, + MISTRAL_OCR_EMBEDDED_IMAGES, + MISTRAL_OCR_IMAGE_MIN_SIZE, + MISTRAL_OCR_MAX_EMBEDDED_IMAGES, + MISTRAL_OCR_IMAGE_CONCURRENCY, DOCX_TEXT_USE_PANDOC, DOCX_TEXT_TRACK_CHANGES, DOCX_TEXT_INCLUDE_HEADERS_FOOTERS, @@ -328,45 +333,177 @@ def _mistral_ocr_client(): return Mistral(api_key=MISTRAL_API_KEY) -def _ocr_document(client, document_payload: dict, source: str, start_page: int = 1) -> List[Document]: +def _field(obj, name: str): + """Read ``name`` from an SDK model or a dict-like response object.""" + return obj.get(name) if isinstance(obj, dict) else getattr(obj, name, None) + + +def _ocr_pages(client, document_payload: dict, include_image_base64: bool = False, **request_kwargs) -> list: + """Run one Mistral OCR request and return the raw response pages (possibly empty). + + Errors propagate unlogged; each caller decides how severe a failure is. + """ + ocr_response = client.ocr.process( + model=MISTRAL_OCR_MODEL, + document=document_payload, + include_image_base64=include_image_base64, + **request_kwargs, + ) + pages = _field(ocr_response, "pages") + return list(pages) if pages else [] + + +def _page_markdowns(pages: list) -> List[str]: + """Markdown of each response page ("" when a page has none).""" + return [_field(page, "markdown") or "" for page in pages] + + +# Markdown image syntax (``![alt](target)``) — what OCR leaves for pictures and +# tables it did not transcribe. +_MARKDOWN_IMAGE_RE = re.compile(r"!\[[^\]]*\]\([^)]*\)") + +_IMAGE_MIME_BY_EXTENSION = {"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg"} + + +def _image_data_uri(image_id: str, image_base64: str) -> str: + """Return ``image_base64`` as a data URI, deriving the mime from the image id if bare.""" + if image_base64.startswith("data:"): + return image_base64 + extension = image_id.rsplit(".", 1)[-1].lower() if "." in image_id else "" + mime = _IMAGE_MIME_BY_EXTENSION.get(extension, "image/jpeg") + return f"data:{mime};base64,{image_base64}" + + +def _ocr_embedded_image(client, image_id: str, image_base64: str, source: str) -> str: + """OCR one embedded image and return its text, or "" if there is none or the call fails. + + A failure only costs this image's text (its placeholder is kept), never the document. + """ + data_uri = _image_data_uri(image_id, image_base64) + try: + pages = _ocr_pages(client, {"type": "image_url", "image_url": data_uri}) + except Exception as e: + logger.warning( + "OCR of embedded image %s in %s failed; keeping placeholder: %s", + image_id, + source, + e, + ) + return "" + + text = "\n\n".join( + _MARKDOWN_IMAGE_RE.sub("", md).strip() for md in _page_markdowns(pages) + ) + return text.strip() + + +def _replace_placeholder(markdown: str, image_id: str, text: str) -> str: + """Swap ``![id](id)`` for ``text`` as its own paragraph; append it if there is no placeholder. + + Newlines adjacent to the placeholder are folded into the paragraph break, so + the text is separated by exactly one blank line; the rest of the page is untouched. + """ + placeholder = f"![{image_id}]({image_id})" + if placeholder in markdown: + pattern = re.compile(rf"\n*{re.escape(placeholder)}\n*") + + def paragraph(match: "re.Match[str]") -> str: + before = "\n\n" if match.start() > 0 else "" + after = "\n\n" if match.end() < len(markdown) else "" + return f"{before}{text}{after}" + + return pattern.sub(paragraph, markdown) + return f"{markdown.rstrip()}\n\n{text}" if markdown.strip() else text + + +def _inline_embedded_images(client, pages: list, markdowns: List[str], source: str) -> List[str]: + """Return per-page markdown with embedded images' OCR text in place of their placeholders. + + Image ids repeat across pages (``img-0.jpeg`` on each), so images are matched + to their own page. Images are OCR'd concurrently but applied in page order, so + the result doesn't depend on which call finishes first. + """ + jobs = [] # (page position, image id, image base64) + for position, page in enumerate(pages): + for image in _field(page, "images") or []: + image_id = _field(image, "id") + image_base64 = _field(image, "image_base64") + if isinstance(image_id, str) and isinstance(image_base64, str) and image_base64: + jobs.append((position, image_id, image_base64)) + + # The request's image_limit makes Mistral stop extracting at the cap, so + # reaching it means more images may exist that we never received; the slice + # is a safety net should more arrive anyway. + if 0 < MISTRAL_OCR_MAX_EMBEDDED_IMAGES <= len(jobs): + logger.warning( + "%s reached MISTRAL_OCR_MAX_EMBEDDED_IMAGES=%d; some images may not have been extracted", + source, + MISTRAL_OCR_MAX_EMBEDDED_IMAGES, + ) + jobs = jobs[:MISTRAL_OCR_MAX_EMBEDDED_IMAGES] + if not jobs: + return markdowns + + with ThreadPoolExecutor(max_workers=MISTRAL_OCR_IMAGE_CONCURRENCY) as executor: + texts = list( + executor.map( + lambda job: _ocr_embedded_image(client, job[1], job[2], source), jobs + ) + ) + + result = list(markdowns) + for (position, image_id, _), text in zip(jobs, texts): + if text: + result[position] = _replace_placeholder(result[position], image_id, text) + return result + + +def _ocr_document( + client, + document_payload: dict, + source: str, + start_page: int = 1, + embedded_images: bool = False, +) -> List[Document]: """Run one Mistral OCR request and map its pages to Documents. Pages are numbered sequentially (1-based) from ``start_page`` in response order, so callers don't depend on the provider's own (0-based) page index. Returns a single empty placeholder when the response has no pages. + + With ``embedded_images``, the images Mistral extracts from the document are + OCR'd separately and their text is inlined where the page markdown only has + an image placeholder. """ + request_kwargs = {} + if embedded_images: + request_kwargs = { + "include_image_base64": True, + "image_min_size": MISTRAL_OCR_IMAGE_MIN_SIZE, + } + if MISTRAL_OCR_MAX_EMBEDDED_IMAGES > 0: + request_kwargs["image_limit"] = MISTRAL_OCR_MAX_EMBEDDED_IMAGES + try: - ocr_response = client.ocr.process( - model=MISTRAL_OCR_MODEL, - document=document_payload, - include_image_base64=False, - ) + pages = _ocr_pages(client, document_payload, **request_kwargs) except Exception as e: logger.error(f"Mistral OCR API call failed: {e}") raise - - pages = getattr(ocr_response, "pages", None) - # Some clients return dict-like response; handle both - if pages is None and isinstance(ocr_response, dict): - pages = ocr_response.get("pages") - if not pages: # Return an empty single document to avoid downstream crashes return [Document(page_content="", metadata={"source": source, "page": start_page})] - documents: List[Document] = [] - for offset, page in enumerate(pages): - markdown = ( - page.get("markdown") if isinstance(page, dict) else getattr(page, "markdown", None) - ) - documents.append( - Document( - page_content=markdown or "", - metadata={"source": source, "page": start_page + offset}, - ) - ) + markdowns = _page_markdowns(pages) + if embedded_images: + markdowns = _inline_embedded_images(client, pages, markdowns, source) - return documents + return [ + Document( + page_content=markdown, + metadata={"source": source, "page": start_page + offset}, + ) + for offset, markdown in enumerate(markdowns) + ] class SafePyPDFLoader: @@ -378,7 +515,9 @@ class SafePyPDFLoader: - metadata.source: original filepath - metadata.page: 1-based page index - Images embedded in the PDF are not extracted. + Mistral OCR leaves only a placeholder for a raster image embedded in a page + (e.g. a photo of a document). With ``MISTRAL_OCR_EMBEDDED_IMAGES`` enabled, + those images are OCR'd separately and their text replaces the placeholder. """ def __init__(self, filepath: str): @@ -397,7 +536,9 @@ def load(self) -> List[Document]: "type": "document_url", "document_url": f"data:application/pdf;base64,{self._encode_pdf_b64()}", } - return _ocr_document(client, payload, self.filepath) + return _ocr_document( + client, payload, self.filepath, embedded_images=MISTRAL_OCR_EMBEDDED_IMAGES + ) def lazy_load(self) -> Iterator[Document]: """Yield OCR-extracted pages one at a time. diff --git a/tests/utils/test_document_loader.py b/tests/utils/test_document_loader.py index b1b81a21..dca582db 100644 --- a/tests/utils/test_document_loader.py +++ b/tests/utils/test_document_loader.py @@ -1,5 +1,7 @@ import os from collections.abc import Iterator +from contextlib import ExitStack +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest @@ -242,6 +244,393 @@ def test_mistral_client_uses_configured_server_url(): client.ocr.process.assert_called_once() +# --- Embedded-image OCR inside PDFs ----------------------------------------- + + +def _ocr_page(markdown, images=()): + return SimpleNamespace(markdown=markdown, images=list(images)) + + +def _ocr_image(image_id, image_base64): + return SimpleNamespace(id=image_id, image_base64=image_base64) + + +def _load_pdf_with_images(pages, image_results=None, **settings): + """Load a PDF through a fake Mistral client that serves per-image OCR results. + + ``image_results`` maps an image's base64 payload to the text its OCR call + returns (a str, a list of page texts, or an Exception to raise). ``settings`` + override ``document_loader`` module settings (defaults are pinned so tests + don't depend on the environment). Returns ``(documents, client)``. + """ + from app.utils import document_loader + from app.utils.document_loader import SafePyPDFLoader + + settings = { + "MISTRAL_OCR_EMBEDDED_IMAGES": True, + "MISTRAL_OCR_IMAGE_MIN_SIZE": 100, + "MISTRAL_OCR_MAX_EMBEDDED_IMAGES": 50, + "MISTRAL_OCR_IMAGE_CONCURRENCY": 4, + **settings, + } + results = image_results or {} + + def process(**kwargs): + document = kwargs["document"] + if document["type"] == "document_url": + return pages if isinstance(pages, dict) else SimpleNamespace(pages=pages) + outcome = results[document["image_url"].split(",", 1)[1]] + if isinstance(outcome, Exception): + raise outcome + texts = outcome if isinstance(outcome, list) else [outcome] + return SimpleNamespace(pages=[SimpleNamespace(markdown=t) for t in texts]) + + fake_module, client = _make_mistral_module(pages=[]) + client.ocr.process.side_effect = process + + with ExitStack() as stack: + stack.enter_context( + patch.dict("sys.modules", {"mistralai": fake_module, "mistralai.client": fake_module}) + ) + stack.enter_context(patch.object(document_loader, "MISTRAL_API_KEY", "test-key")) + stack.enter_context( + patch.object(SafePyPDFLoader, "_encode_pdf_b64", return_value="b64data") + ) + for name, value in settings.items(): + stack.enter_context(patch.object(document_loader, name, value)) + documents = SafePyPDFLoader("scan.pdf").load() + return documents, client + + +def _image_calls(client): + return [ + c.kwargs + for c in client.ocr.process.call_args_list + if c.kwargs["document"]["type"] == "image_url" + ] + + +def test_pdf_embedded_image_text_replaces_placeholder(): + """The image's OCR text takes the placeholder's place; page text and metadata stay.""" + pages = [ + _ocr_page( + "Before\n\n![img-0.jpeg](img-0.jpeg)\n\nAfter", + [_ocr_image("img-0.jpeg", "AAA")], + ) + ] + + docs, client = _load_pdf_with_images(pages, {"AAA": " Text from the scan "}) + + assert len(docs) == 1 + assert docs[0].page_content == "Before\n\nText from the scan\n\nAfter" + assert docs[0].metadata == {"source": "scan.pdf", "page": 1} + assert "![img-0.jpeg]" not in docs[0].page_content + assert client.ocr.process.call_count == 2 + + +def test_pdf_embedded_images_request_kwargs_when_enabled(): + """The main request asks for image bytes with the configured size and limit.""" + docs, client = _load_pdf_with_images( + [_ocr_page("text")], + MISTRAL_OCR_EMBEDDED_IMAGES=True, + MISTRAL_OCR_IMAGE_MIN_SIZE=123, + MISTRAL_OCR_MAX_EMBEDDED_IMAGES=7, + ) + + client.ocr.process.assert_called_once() + kwargs = client.ocr.process.call_args.kwargs + assert kwargs["document"]["type"] == "document_url" + assert kwargs["include_image_base64"] is True + assert kwargs["image_min_size"] == 123 + assert kwargs["image_limit"] == 7 + + +def test_pdf_embedded_images_no_image_limit_when_uncapped(): + _, client = _load_pdf_with_images( + [_ocr_page("text")], + MISTRAL_OCR_EMBEDDED_IMAGES=True, + MISTRAL_OCR_MAX_EMBEDDED_IMAGES=0, + ) + + kwargs = client.ocr.process.call_args.kwargs + assert kwargs["include_image_base64"] is True + assert "image_limit" not in kwargs + + +def test_pdf_embedded_images_disabled_keeps_original_request(): + """With the feature off the request is as before and no image is OCR'd.""" + pages = [_ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "AAA")])] + + docs, client = _load_pdf_with_images( + pages, {"AAA": "never read"}, MISTRAL_OCR_EMBEDDED_IMAGES=False + ) + + client.ocr.process.assert_called_once() + kwargs = client.ocr.process.call_args.kwargs + assert kwargs["include_image_base64"] is False + assert "image_min_size" not in kwargs + assert "image_limit" not in kwargs + assert docs[0].page_content == "![img-0.jpeg](img-0.jpeg)" + + +@pytest.mark.parametrize("image_text", ["![img-0.jpeg](img-0.jpeg)", " ", ""]) +def test_pdf_embedded_image_without_text_keeps_placeholder(image_text): + """An image that yields only placeholders/whitespace leaves the original placeholder.""" + pages = [_ocr_page("A ![img-0.jpeg](img-0.jpeg) B", [_ocr_image("img-0.jpeg", "AAA")])] + + docs, _ = _load_pdf_with_images(pages, {"AAA": image_text}) + + assert docs[0].page_content == "A ![img-0.jpeg](img-0.jpeg) B" + + +def test_pdf_embedded_image_nested_placeholders_are_stripped(): + pages = [_ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "AAA")])] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "Real text ![img-0.jpeg](img-0.jpeg)"}) + + assert docs[0].page_content.strip() == "Real text" + + +def test_pdf_embedded_image_multiple_ocr_pages_joined_with_blank_line(): + pages = [_ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "AAA")])] + + docs, _ = _load_pdf_with_images(pages, {"AAA": ["first", "second"]}) + + assert docs[0].page_content.strip() == "first\n\nsecond" + + +def test_pdf_embedded_image_failure_keeps_placeholder_and_other_images(caplog): + """One failed image OCR call doesn't fail the document or the other images.""" + pages = [ + _ocr_page( + "![img-0.jpeg](img-0.jpeg) | ![img-1.jpeg](img-1.jpeg)", + [_ocr_image("img-0.jpeg", "SECRETB64"), _ocr_image("img-1.jpeg", "BBB")], + ) + ] + + with caplog.at_level("WARNING"): + docs, _ = _load_pdf_with_images( + pages, {"SECRETB64": RuntimeError("boom"), "BBB": "second image"} + ) + + # A non-fatal image failure is a warning, never an error. + assert [r.levelname for r in caplog.records if "img-0.jpeg" in r.getMessage()] == [ + "WARNING" + ] + assert not [r for r in caplog.records if r.levelname == "ERROR"] + content = docs[0].page_content + assert "![img-0.jpeg](img-0.jpeg)" in content + assert "second image" in content + assert "img-0.jpeg" in caplog.text and "scan.pdf" in caplog.text + assert "SECRETB64" not in caplog.text + + +def test_pdf_main_ocr_failure_still_raises_and_logs_error(caplog): + from app.utils import document_loader + from app.utils.document_loader import SafePyPDFLoader + + fake_module, _ = _make_mistral_module(raise_exc=RuntimeError("OCR boom")) + with patch.dict( + "sys.modules", {"mistralai": fake_module, "mistralai.client": fake_module} + ), patch.object(document_loader, "MISTRAL_API_KEY", "test-key"), patch.object( + document_loader, "MISTRAL_OCR_EMBEDDED_IMAGES", True + ), patch.object( + SafePyPDFLoader, "_encode_pdf_b64", return_value="b64data" + ), caplog.at_level("ERROR"), pytest.raises(RuntimeError, match="OCR boom"): + SafePyPDFLoader("scan.pdf").load() + + assert [r.levelname for r in caplog.records if "OCR boom" in r.getMessage()] == ["ERROR"] + + +def test_pdf_embedded_image_base64_without_prefix_gets_data_uri(): + pages = [ + _ocr_page( + "![a.png](a.png) ![b.jpg](b.jpg) ![c.avif](c.avif) ![d.jpeg](d.jpeg)", + [ + _ocr_image("a.png", "AAA"), + _ocr_image("b.jpg", "BBB"), + _ocr_image("c.avif", "CCC"), + _ocr_image("d.jpeg", "data:image/webp;base64,DDD"), + ], + ) + ] + + _, client = _load_pdf_with_images( + pages, {"AAA": "a", "BBB": "b", "CCC": "c", "DDD": "d"} + ) + + sent = sorted(c["document"]["image_url"] for c in _image_calls(client)) + assert sent == [ + "data:image/jpeg;base64,BBB", + "data:image/jpeg;base64,CCC", + "data:image/png;base64,AAA", + "data:image/webp;base64,DDD", + ] + + +def test_pdf_embedded_image_cap_limits_image_ocr_calls(caplog): + pages = [ + _ocr_page( + "![img-0.jpeg](img-0.jpeg) ![img-1.jpeg](img-1.jpeg)", + [_ocr_image("img-0.jpeg", "AAA"), _ocr_image("img-1.jpeg", "BBB")], + ) + ] + + with caplog.at_level("WARNING"): + docs, client = _load_pdf_with_images( + pages, + {"AAA": "first", "BBB": "second"}, + MISTRAL_OCR_MAX_EMBEDDED_IMAGES=1, + ) + + assert len(_image_calls(client)) == 1 + assert "first" in docs[0].page_content + assert "![img-1.jpeg](img-1.jpeg)" in docs[0].page_content + assert "MISTRAL_OCR_MAX_EMBEDDED_IMAGES" in caplog.text + + +def test_pdf_embedded_images_reaching_cap_warns_that_more_may_exist(caplog): + pages = [_ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "AAA")])] + + with caplog.at_level("WARNING"): + docs, client = _load_pdf_with_images( + pages, {"AAA": "text"}, MISTRAL_OCR_MAX_EMBEDDED_IMAGES=1 + ) + + assert len(_image_calls(client)) == 1 + assert "may not have been extracted" in caplog.text + + +def test_pdf_embedded_image_only_normalizes_newlines_around_placeholder(): + pages = [ + _ocr_page( + "![img-0.jpeg](img-0.jpeg)\n\nkept\n\n\n\nspacing", + [_ocr_image("img-0.jpeg", "AAA")], + ) + ] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "text"}) + + assert docs[0].page_content == "text\n\nkept\n\n\n\nspacing" + + +def test_pdf_embedded_image_ids_are_scoped_per_page(): + """Every page may have its own ``img-0.jpeg``; each gets its own text.""" + pages = [ + _ocr_page("P1 ![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "AAA")]), + _ocr_page("P2 ![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "BBB")]), + ] + + docs, _ = _load_pdf_with_images( + pages, {"AAA": "text one", "BBB": "text two"}, MISTRAL_OCR_IMAGE_CONCURRENCY=2 + ) + + assert "text one" in docs[0].page_content and "text two" not in docs[0].page_content + assert "text two" in docs[1].page_content and "text one" not in docs[1].page_content + assert [d.metadata["page"] for d in docs] == [1, 2] + + +def test_pdf_embedded_image_text_lands_at_own_placeholder_when_calls_finish_out_of_order(): + """The first image's OCR call only finishes after the second one has.""" + import threading + + from app.utils import document_loader + from app.utils.document_loader import SafePyPDFLoader + + second_done = threading.Event() + + def process(**kwargs): + document = kwargs["document"] + if document["type"] == "document_url": + return SimpleNamespace( + pages=[ + _ocr_page( + "![img-0.jpeg](img-0.jpeg)\n\n![img-1.jpeg](img-1.jpeg)", + [_ocr_image("img-0.jpeg", "AAA"), _ocr_image("img-1.jpeg", "BBB")], + ) + ] + ) + if document["image_url"].endswith("AAA"): + assert second_done.wait(timeout=10), "second image OCR never completed" + text = "first text" + else: + text = "second text" + second_done.set() + return SimpleNamespace(pages=[SimpleNamespace(markdown=text)]) + + fake_module, client = _make_mistral_module(pages=[]) + client.ocr.process.side_effect = process + + with patch.dict( + "sys.modules", {"mistralai": fake_module, "mistralai.client": fake_module} + ), patch.object(document_loader, "MISTRAL_API_KEY", "test-key"), patch.object( + document_loader, "MISTRAL_OCR_EMBEDDED_IMAGES", True + ), patch.object( + document_loader, "MISTRAL_OCR_IMAGE_CONCURRENCY", 2 + ), patch.object( + SafePyPDFLoader, "_encode_pdf_b64", return_value="b64data" + ): + docs = SafePyPDFLoader("scan.pdf").load() + + assert docs[0].page_content == "first text\n\nsecond text" + + +def test_pdf_embedded_image_without_placeholder_is_appended(): + pages = [_ocr_page("Page body", [_ocr_image("img-0.jpeg", "AAA")])] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "hidden text"}) + + assert docs[0].page_content == "Page body\n\nhidden text" + + +def test_pdf_embedded_images_support_dict_shaped_response(): + response = { + "pages": [ + { + "index": 0, + "markdown": "Top ![img-0.jpeg](img-0.jpeg)", + "images": [{"id": "img-0.jpeg", "image_base64": "AAA"}], + } + ] + } + + docs, _ = _load_pdf_with_images(response, {"AAA": "from dict"}) + + assert docs[0].page_content.split() == ["Top", "from", "dict"] + + +def test_pdf_embedded_images_skip_images_without_bytes(): + pages = [ + _ocr_page( + "![img-0.jpeg](img-0.jpeg)", + [_ocr_image("img-0.jpeg", None), _ocr_image("img-1.jpeg", "")], + ) + ] + + docs, client = _load_pdf_with_images(pages) + + client.ocr.process.assert_called_once() + assert docs[0].page_content == "![img-0.jpeg](img-0.jpeg)" + + +def test_image_ocr_loader_does_not_request_embedded_images(tmp_path): + from app.utils import document_loader + from app.utils.document_loader import ImageOCRLoader + + p = tmp_path / "scan.png" + _make_png(str(p)) + fake_module, client = _make_mistral_module(pages=[MagicMock(index=0, markdown="t")]) + + with patch.dict( + "sys.modules", {"mistralai": fake_module, "mistralai.client": fake_module} + ), patch.object(document_loader, "MISTRAL_API_KEY", "test-key"): + ImageOCRLoader(str(p)).load() + + kwargs = client.ocr.process.call_args.kwargs + assert kwargs["include_image_base64"] is False + assert "image_min_size" not in kwargs + + MARKDOWN_SAMPLE = ( "# Heading\n\n" "**bold** and *italic* text with a [link](https://example.com).\n\n" From d6683abaec87907ccc8fad6ee7c816e153e52b5e Mon Sep 17 00:00:00 2001 From: Daniel Vejsada Date: Fri, 23 Oct 2026 19:54:38 +0200 Subject: [PATCH 3/4] fix(ocr): address review of embedded-image OCR - OCR each distinct embedded image once (sha256 of its bytes) and reuse the text, and skip images below MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION (default 0.1) of their page, so logos/letterheads repeated on every page don't cost an extra billed call each or pollute every page. - Retry OCR calls on 429/5xx and connection errors with backoff via the SDK's RetryConfig, so rate limits no longer silently drop image text. - Match placeholders on the link target with any alt text. - MISTRAL_OCR_MAX_EMBEDDED_IMAGES is now at least 1, so the response size is always bounded; request cap+1 and warn only when exceeded. - Insert text inline (newlines collapsed, `|` escaped in table rows) when the placeholder shares its line, instead of breaking tables. - Use mimetypes for data-URI mime types. - Pin MISTRAL_OCR_EMBEDDED_IMAGES=False in the older PDF tests. Co-Authored-By: Claude Opus 5.5 --- README.md | 3 +- app/config.py | 13 +- app/utils/document_loader.py | 167 ++++++++++---- tests/utils/test_document_loader.py | 327 ++++++++++++++++++++++++++-- 4 files changed, 441 insertions(+), 69 deletions(-) diff --git a/README.md b/README.md index f838c85a..7ca6ec74 100644 --- a/README.md +++ b/README.md @@ -218,7 +218,8 @@ The following environment variables are required to run the application: - `IMAGE_OCR_MAX_PAGES`: (Optional) Max number of frames/pages of a single image upload (e.g. animated GIF or large multi-page TIFF) sent to OCR; frames beyond the cap are skipped with a warning. Default value is `100`. - `MISTRAL_OCR_EMBEDDED_IMAGES`: (Optional) Boolean. Mistral OCR returns only a placeholder (`![img-0.jpeg](img-0.jpeg)`) for a raster image embedded in a PDF page, such as a photo or scan of a page, and does not read the text inside it. When "True" (default), those images are extracted, OCR'd in a second Mistral call each, and the recognized text replaces the placeholder. **Each embedded image is billed as an extra OCR page**; set "False" to OCR only the PDF itself. If an image's OCR call fails, its placeholder is kept and the document still loads. - `MISTRAL_OCR_IMAGE_MIN_SIZE`: (Optional) Minimum image size in pixels for an embedded image to be extracted, so logos and icons are skipped. Default value is `100`. -- `MISTRAL_OCR_MAX_EMBEDDED_IMAGES`: (Optional) Max number of embedded images extracted and OCR'd per document; extraction stops at the cap (a warning is logged when it is reached), so further images may be left as placeholders. Set to `0` or below for no cap. Default value is `50`. +- `MISTRAL_OCR_MAX_EMBEDDED_IMAGES`: (Optional) Max number of embedded images extracted and OCR'd per document (minimum `1`); a warning is logged when the document has more, and the extra images are left as placeholders. Default value is `50`. +- `MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION`: (Optional) Embedded images covering less than this fraction of their page (`0` to `1`) are not OCR'd and keep their placeholder, so letterheads and logos repeated on every page cost nothing. Set to `0` to OCR every extracted image. Identical images are always OCR'd only once per document. Default value is `0.1`. - `MISTRAL_OCR_IMAGE_CONCURRENCY`: (Optional) Number of embedded images OCR'd in parallel (minimum `1`). Default value is `4`. - `DEBUG_RAG_API`: (Optional) Set to "True" to show more verbose logging output in the server console, and to enable postgresql database routes - `DEBUG_PGVECTOR_QUERIES`: (Optional) Set to "True" to enable detailed PostgreSQL query logging for pgvector operations. Useful for debugging performance issues with vector database queries. diff --git a/app/config.py b/app/config.py index e56be63b..8edad8c3 100644 --- a/app/config.py +++ b/app/config.py @@ -235,9 +235,16 @@ async def dispatch(self, request, call_next): MISTRAL_OCR_IMAGE_MIN_SIZE = max( 0, int(get_env_variable("MISTRAL_OCR_IMAGE_MIN_SIZE", "100")) ) -# Max embedded images extracted (and OCR'd) per document; <= 0 means no cap. -MISTRAL_OCR_MAX_EMBEDDED_IMAGES = int( - get_env_variable("MISTRAL_OCR_MAX_EMBEDDED_IMAGES", "50") +# Max embedded images extracted (and OCR'd) per document (minimum 1), which also +# bounds the size of the OCR response. +MISTRAL_OCR_MAX_EMBEDDED_IMAGES = max( + 1, int(get_env_variable("MISTRAL_OCR_MAX_EMBEDDED_IMAGES", "50")) +) +# Embedded images covering less than this fraction of their page (0-1) are not +# OCR'd, so letterheads and logos repeated on every page cost nothing. 0 disables. +MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION = min( + 1.0, + max(0.0, float(get_env_variable("MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION", "0.1"))), ) # Number of embedded images OCR'd in parallel (minimum 1). MISTRAL_OCR_IMAGE_CONCURRENCY = max( diff --git a/app/utils/document_loader.py b/app/utils/document_loader.py index 1d24124a..ee697354 100644 --- a/app/utils/document_loader.py +++ b/app/utils/document_loader.py @@ -2,6 +2,8 @@ import os import codecs +import hashlib +import mimetypes import re import tempfile @@ -20,6 +22,7 @@ MISTRAL_SERVER_URL, MISTRAL_OCR_EMBEDDED_IMAGES, MISTRAL_OCR_IMAGE_MIN_SIZE, + MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION, MISTRAL_OCR_MAX_EMBEDDED_IMAGES, MISTRAL_OCR_IMAGE_CONCURRENCY, DOCX_TEXT_USE_PANDOC, @@ -338,15 +341,35 @@ def _field(obj, name: str): return obj.get(name) if isinstance(obj, dict) else getattr(obj, name, None) +def _ocr_retries(): + """Retry policy for every OCR call. + + The SDK retries HTTP 429/500/502/503/504 (its per-operation default for + ``ocr.process``) and connection errors, with exponential backoff that honors + a ``Retry-After`` header. + """ + from mistralai.client import utils + + return utils.RetryConfig( + "backoff", + utils.BackoffStrategy( + initial_interval=500, max_interval=10_000, exponent=1.5, max_elapsed_time=60_000 + ), + retry_connection_errors=True, + ) + + def _ocr_pages(client, document_payload: dict, include_image_base64: bool = False, **request_kwargs) -> list: """Run one Mistral OCR request and return the raw response pages (possibly empty). - Errors propagate unlogged; each caller decides how severe a failure is. + Transient failures are retried (see :func:`_ocr_retries`); other errors + propagate unlogged, as each caller decides how severe a failure is. """ ocr_response = client.ocr.process( model=MISTRAL_OCR_MODEL, document=document_payload, include_image_base64=include_image_base64, + retries=_ocr_retries(), **request_kwargs, ) pages = _field(ocr_response, "pages") @@ -362,15 +385,14 @@ def _page_markdowns(pages: list) -> List[str]: # tables it did not transcribe. _MARKDOWN_IMAGE_RE = re.compile(r"!\[[^\]]*\]\([^)]*\)") -_IMAGE_MIME_BY_EXTENSION = {"png": "image/png", "jpg": "image/jpeg", "jpeg": "image/jpeg"} - def _image_data_uri(image_id: str, image_base64: str) -> str: """Return ``image_base64`` as a data URI, deriving the mime from the image id if bare.""" if image_base64.startswith("data:"): return image_base64 - extension = image_id.rsplit(".", 1)[-1].lower() if "." in image_id else "" - mime = _IMAGE_MIME_BY_EXTENSION.get(extension, "image/jpeg") + mime = mimetypes.guess_type(image_id)[0] + if not mime or not mime.startswith("image/"): + mime = "image/jpeg" return f"data:{mime};base64,{image_base64}" @@ -398,63 +420,124 @@ def _ocr_embedded_image(client, image_id: str, image_base64: str, source: str) - def _replace_placeholder(markdown: str, image_id: str, text: str) -> str: - """Swap ``![id](id)`` for ``text`` as its own paragraph; append it if there is no placeholder. + """Swap the image's placeholder (``![alt](id)``) for ``text``; append it if there is none. - Newlines adjacent to the placeholder are folded into the paragraph break, so - the text is separated by exactly one blank line; the rest of the page is untouched. + A placeholder alone on its line becomes its own paragraph: adjacent newlines + are folded into one blank line of separation and the rest of the page is + untouched. Elsewhere (a table cell, list item or sentence) the text goes + inline with its newlines collapsed, and ``|`` escaped inside a table row. """ - placeholder = f"![{image_id}]({image_id})" - if placeholder in markdown: - pattern = re.compile(rf"\n*{re.escape(placeholder)}\n*") - - def paragraph(match: "re.Match[str]") -> str: - before = "\n\n" if match.start() > 0 else "" - after = "\n\n" if match.end() < len(markdown) else "" - return f"{before}{text}{after}" - - return pattern.sub(paragraph, markdown) - return f"{markdown.rstrip()}\n\n{text}" if markdown.strip() else text + pattern = re.compile(rf"!\[[^\]]*\]\({re.escape(image_id)}\)") + matches = list(pattern.finditer(markdown)) + if not matches: + return f"{markdown.rstrip()}\n\n{text}" if markdown.strip() else text + + inline_text = re.sub(r"\s*\n\s*", " ", text) + for match in reversed(matches): # later first, so earlier offsets stay valid + start, end = match.span() + line_start = markdown.rfind("\n", 0, start) + 1 + line_end = markdown.find("\n", end) + line_end = len(markdown) if line_end == -1 else line_end + if markdown[line_start:start].strip() or markdown[end:line_end].strip(): + is_table_row = markdown[line_start:line_end].lstrip().startswith("|") + replacement = inline_text.replace("|", "\\|") if is_table_row else inline_text + markdown = markdown[:start] + replacement + markdown[end:] + continue + while line_start > 0 and markdown[line_start - 1] == "\n": + line_start -= 1 + while line_end < len(markdown) and markdown[line_end] == "\n": + line_end += 1 + lead = "\n\n" if line_start > 0 else "" + trail = "\n\n" if line_end < len(markdown) else "" + markdown = markdown[:line_start] + lead + text + trail + markdown[line_end:] + return markdown + + +def _is_number(value) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) + + +def _below_min_page_fraction(image, page) -> bool: + """Whether ``image`` covers less of ``page`` than MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION. + + The bounding box and the page dimensions are both taken to be in the pixel + space of the rendered page (UNVERIFIED against the live API). Images are kept + whenever the fraction can't be computed (missing coordinates or dimensions). + """ + if MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION <= 0: + return False + dimensions = _field(page, "dimensions") + page_width, page_height = _field(dimensions, "width"), _field(dimensions, "height") + corners = [ + _field(image, name) + for name in ("top_left_x", "top_left_y", "bottom_right_x", "bottom_right_y") + ] + if not all(_is_number(v) for v in [page_width, page_height, *corners]): + return False + left, top, right, bottom = corners + if page_width <= 0 or page_height <= 0 or right <= left or bottom <= top: + return False + page_area = page_width * page_height + return (right - left) * (bottom - top) < MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION * page_area -def _inline_embedded_images(client, pages: list, markdowns: List[str], source: str) -> List[str]: - """Return per-page markdown with embedded images' OCR text in place of their placeholders. +def _collect_embedded_images(pages: list, source: str) -> list: + """Embedded images worth OCR'ing as ``(page position, image id, image base64)``. - Image ids repeat across pages (``img-0.jpeg`` on each), so images are matched - to their own page. Images are OCR'd concurrently but applied in page order, so - the result doesn't depend on which call finishes first. + Mistral is asked for one image more than MISTRAL_OCR_MAX_EMBEDDED_IMAGES, so + receiving more than the cap means the document has images beyond it; those + are dropped with a warning. Images too small relative to their page are skipped. """ - jobs = [] # (page position, image id, image base64) + found = [] # (page position, page, image, image id, image base64) for position, page in enumerate(pages): for image in _field(page, "images") or []: image_id = _field(image, "id") image_base64 = _field(image, "image_base64") if isinstance(image_id, str) and isinstance(image_base64, str) and image_base64: - jobs.append((position, image_id, image_base64)) + found.append((position, page, image, image_id, image_base64)) - # The request's image_limit makes Mistral stop extracting at the cap, so - # reaching it means more images may exist that we never received; the slice - # is a safety net should more arrive anyway. - if 0 < MISTRAL_OCR_MAX_EMBEDDED_IMAGES <= len(jobs): + if len(found) > MISTRAL_OCR_MAX_EMBEDDED_IMAGES: logger.warning( - "%s reached MISTRAL_OCR_MAX_EMBEDDED_IMAGES=%d; some images may not have been extracted", + "%s has more than MISTRAL_OCR_MAX_EMBEDDED_IMAGES=%d embedded images; the rest keep their placeholders", source, MISTRAL_OCR_MAX_EMBEDDED_IMAGES, ) - jobs = jobs[:MISTRAL_OCR_MAX_EMBEDDED_IMAGES] - if not jobs: + found = found[:MISTRAL_OCR_MAX_EMBEDDED_IMAGES] + return [ + (position, image_id, image_base64) + for position, page, image, image_id, image_base64 in found + if not _below_min_page_fraction(image, page) + ] + + +def _inline_embedded_images(client, pages: list, markdowns: List[str], source: str) -> List[str]: + """Return per-page markdown with embedded images' OCR text in place of their placeholders. + + Image ids repeat across pages (``img-0.jpeg`` on each), so images are matched + to their own page. Identical images (a logo on every page) are OCR'd once and + their text reused. Images are OCR'd concurrently but applied in page order, so + the result doesn't depend on which call finishes first. + """ + placements = [] # (page position, image id, content digest) + distinct = {} # content digest -> (image id, image base64) of its first occurrence + for position, image_id, image_base64 in _collect_embedded_images(pages, source): + digest = hashlib.sha256(image_base64.encode()).hexdigest() + distinct.setdefault(digest, (image_id, image_base64)) + placements.append((position, image_id, digest)) + if not placements: return markdowns with ThreadPoolExecutor(max_workers=MISTRAL_OCR_IMAGE_CONCURRENCY) as executor: - texts = list( - executor.map( - lambda job: _ocr_embedded_image(client, job[1], job[2], source), jobs - ) + results = executor.map( + lambda image: _ocr_embedded_image(client, image[0], image[1], source), + distinct.values(), ) + texts = dict(zip(distinct, results)) result = list(markdowns) - for (position, image_id, _), text in zip(jobs, texts): - if text: - result[position] = _replace_placeholder(result[position], image_id, text) + for position, image_id, digest in placements: + if texts[digest]: + result[position] = _replace_placeholder(result[position], image_id, texts[digest]) return result @@ -480,9 +563,9 @@ def _ocr_document( request_kwargs = { "include_image_base64": True, "image_min_size": MISTRAL_OCR_IMAGE_MIN_SIZE, + # One more than the cap, so exceeding it is detectable. + "image_limit": MISTRAL_OCR_MAX_EMBEDDED_IMAGES + 1, } - if MISTRAL_OCR_MAX_EMBEDDED_IMAGES > 0: - request_kwargs["image_limit"] = MISTRAL_OCR_MAX_EMBEDDED_IMAGES try: pages = _ocr_pages(client, document_payload, **request_kwargs) diff --git a/tests/utils/test_document_loader.py b/tests/utils/test_document_loader.py index dca582db..b2d2979a 100644 --- a/tests/utils/test_document_loader.py +++ b/tests/utils/test_document_loader.py @@ -128,6 +128,8 @@ def test_safe_pdf_loader_ocr_maps_pages(): loader = SafePyPDFLoader("dummy.pdf") with patch.dict( "sys.modules", {"mistralai": fake_module, "mistralai.client": fake_module} + ), patch.object( + document_loader, "MISTRAL_OCR_EMBEDDED_IMAGES", False ), patch.object( document_loader, "MISTRAL_API_KEY", "test-key" ), patch.object(SafePyPDFLoader, "_encode_pdf_b64", return_value="b64data"): @@ -161,6 +163,8 @@ def test_safe_pdf_loader_empty_pages_returns_placeholder(): loader = SafePyPDFLoader("dummy.pdf") with patch.dict( "sys.modules", {"mistralai": fake_module, "mistralai.client": fake_module} + ), patch.object( + document_loader, "MISTRAL_OCR_EMBEDDED_IMAGES", False ), patch.object( document_loader, "MISTRAL_API_KEY", "test-key" ), patch.object(SafePyPDFLoader, "_encode_pdf_b64", return_value="b64data"): @@ -181,6 +185,8 @@ def test_safe_pdf_loader_ocr_error_propagates(): loader = SafePyPDFLoader("dummy.pdf") with patch.dict( "sys.modules", {"mistralai": fake_module, "mistralai.client": fake_module} + ), patch.object( + document_loader, "MISTRAL_OCR_EMBEDDED_IMAGES", False ), patch.object( document_loader, "MISTRAL_API_KEY", "test-key" ), patch.object(SafePyPDFLoader, "_encode_pdf_b64", return_value="b64data"): @@ -231,6 +237,8 @@ def test_mistral_client_uses_configured_server_url(): loader = SafePyPDFLoader("dummy.pdf") with patch.dict( "sys.modules", {"mistralai": fake_module, "mistralai.client": fake_module} + ), patch.object( + document_loader, "MISTRAL_OCR_EMBEDDED_IMAGES", False ), patch.object( document_loader, "MISTRAL_API_KEY", "test-key" ), patch.object( @@ -247,12 +255,16 @@ def test_mistral_client_uses_configured_server_url(): # --- Embedded-image OCR inside PDFs ----------------------------------------- -def _ocr_page(markdown, images=()): - return SimpleNamespace(markdown=markdown, images=list(images)) +def _ocr_page(markdown, images=(), dimensions=None): + return SimpleNamespace(markdown=markdown, images=list(images), dimensions=dimensions) -def _ocr_image(image_id, image_base64): - return SimpleNamespace(id=image_id, image_base64=image_base64) +def _ocr_image(image_id, image_base64, bbox=None): + """``bbox`` is ``(top_left_x, top_left_y, bottom_right_x, bottom_right_y)``.""" + corners = dict( + zip(("top_left_x", "top_left_y", "bottom_right_x", "bottom_right_y"), bbox or ()) + ) + return SimpleNamespace(id=image_id, image_base64=image_base64, **corners) def _load_pdf_with_images(pages, image_results=None, **settings): @@ -269,6 +281,7 @@ def _load_pdf_with_images(pages, image_results=None, **settings): settings = { "MISTRAL_OCR_EMBEDDED_IMAGES": True, "MISTRAL_OCR_IMAGE_MIN_SIZE": 100, + "MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION": 0.1, "MISTRAL_OCR_MAX_EMBEDDED_IMAGES": 50, "MISTRAL_OCR_IMAGE_CONCURRENCY": 4, **settings, @@ -342,19 +355,7 @@ def test_pdf_embedded_images_request_kwargs_when_enabled(): assert kwargs["document"]["type"] == "document_url" assert kwargs["include_image_base64"] is True assert kwargs["image_min_size"] == 123 - assert kwargs["image_limit"] == 7 - - -def test_pdf_embedded_images_no_image_limit_when_uncapped(): - _, client = _load_pdf_with_images( - [_ocr_page("text")], - MISTRAL_OCR_EMBEDDED_IMAGES=True, - MISTRAL_OCR_MAX_EMBEDDED_IMAGES=0, - ) - - kwargs = client.ocr.process.call_args.kwargs - assert kwargs["include_image_base64"] is True - assert "image_limit" not in kwargs + assert kwargs["image_limit"] == 8 # cap + 1, so exceeding the cap is detectable def test_pdf_embedded_images_disabled_keeps_original_request(): @@ -445,30 +446,36 @@ def test_pdf_main_ocr_failure_still_raises_and_logs_error(caplog): def test_pdf_embedded_image_base64_without_prefix_gets_data_uri(): pages = [ _ocr_page( - "![a.png](a.png) ![b.jpg](b.jpg) ![c.avif](c.avif) ![d.jpeg](d.jpeg)", + "![a.png](a.png) ![b.jpg](b.jpg) ![c.bin](c.bin) ![d.jpeg](d.jpeg)" + " ![e.gif](e.gif) ![f.pdf](f.pdf)", [ _ocr_image("a.png", "AAA"), _ocr_image("b.jpg", "BBB"), - _ocr_image("c.avif", "CCC"), + _ocr_image("c.bin", "CCC"), _ocr_image("d.jpeg", "data:image/webp;base64,DDD"), + _ocr_image("e.gif", "EEE"), + _ocr_image("f.pdf", "FFF"), ], ) ] _, client = _load_pdf_with_images( - pages, {"AAA": "a", "BBB": "b", "CCC": "c", "DDD": "d"} + pages, {"AAA": "a", "BBB": "b", "CCC": "c", "DDD": "d", "EEE": "e", "FFF": "f"} ) sent = sorted(c["document"]["image_url"] for c in _image_calls(client)) assert sent == [ + "data:image/gif;base64,EEE", "data:image/jpeg;base64,BBB", - "data:image/jpeg;base64,CCC", + "data:image/jpeg;base64,CCC", # unknown extension + "data:image/jpeg;base64,FFF", # a non-image mime is not trusted "data:image/png;base64,AAA", "data:image/webp;base64,DDD", ] def test_pdf_embedded_image_cap_limits_image_ocr_calls(caplog): + """Mistral hands back cap + 1 images; only ``cap`` are OCR'd, with a warning.""" pages = [ _ocr_page( "![img-0.jpeg](img-0.jpeg) ![img-1.jpeg](img-1.jpeg)", @@ -483,13 +490,14 @@ def test_pdf_embedded_image_cap_limits_image_ocr_calls(caplog): MISTRAL_OCR_MAX_EMBEDDED_IMAGES=1, ) + assert client.ocr.process.call_args_list[0].kwargs["image_limit"] == 2 assert len(_image_calls(client)) == 1 assert "first" in docs[0].page_content assert "![img-1.jpeg](img-1.jpeg)" in docs[0].page_content assert "MISTRAL_OCR_MAX_EMBEDDED_IMAGES" in caplog.text -def test_pdf_embedded_images_reaching_cap_warns_that_more_may_exist(caplog): +def test_pdf_embedded_images_exactly_at_cap_do_not_warn(caplog): pages = [_ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "AAA")])] with caplog.at_level("WARNING"): @@ -498,7 +506,25 @@ def test_pdf_embedded_images_reaching_cap_warns_that_more_may_exist(caplog): ) assert len(_image_calls(client)) == 1 - assert "may not have been extracted" in caplog.text + assert docs[0].page_content == "text" + assert "MISTRAL_OCR_MAX_EMBEDDED_IMAGES" not in caplog.text + + +def test_pdf_embedded_images_cap_counts_images_before_dedupe(): + """Duplicates count toward the cap: the cap bounds what Mistral sends back.""" + pages = [ + _ocr_page( + "![img-0.jpeg](img-0.jpeg) ![img-1.jpeg](img-1.jpeg)", + [_ocr_image("img-0.jpeg", "AAA"), _ocr_image("img-1.jpeg", "AAA")], + ) + ] + + docs, client = _load_pdf_with_images( + pages, {"AAA": "same"}, MISTRAL_OCR_MAX_EMBEDDED_IMAGES=1 + ) + + assert len(_image_calls(client)) == 1 + assert docs[0].page_content == "same ![img-1.jpeg](img-1.jpeg)" def test_pdf_embedded_image_only_normalizes_newlines_around_placeholder(): @@ -583,6 +609,261 @@ def test_pdf_embedded_image_without_placeholder_is_appended(): assert docs[0].page_content == "Page body\n\nhidden text" +@pytest.mark.parametrize("alt", ["", "image", "img-0.jpeg"]) +def test_pdf_embedded_image_placeholder_matches_any_alt_text(alt): + pages = [ + _ocr_page(f"Before\n\n![{alt}](img-0.jpeg)\n\nAfter", [_ocr_image("img-0.jpeg", "AAA")]) + ] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "scanned"}) + + assert docs[0].page_content == "Before\n\nscanned\n\nAfter" + + +def test_pdf_embedded_image_placeholder_does_not_match_other_image_ids(): + pages = [ + _ocr_page( + "![x](img-10.jpeg) and ![x](other-img-1.jpeg)", + [_ocr_image("img-1.jpeg", "AAA")], + ) + ] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "scanned"}) + + assert docs[0].page_content.startswith("![x](img-10.jpeg) and ![x](other-img-1.jpeg)") + assert docs[0].page_content.endswith("\n\nscanned") + + +def test_pdf_embedded_image_inline_in_sentence_is_not_split_into_paragraphs(): + pages = [ + _ocr_page( + "See ![img-0.jpeg](img-0.jpeg) for details.", [_ocr_image("img-0.jpeg", "AAA")] + ) + ] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "line one\n\nline two\nline three"}) + + assert docs[0].page_content == "See line one line two line three for details." + + +def test_pdf_embedded_image_in_table_cell_stays_in_its_row_with_pipes_escaped(): + markdown = "| Name | Scan |\n| --- | --- |\n| Ann | ![img-0.jpeg](img-0.jpeg) |\n| Bob | n/a |" + pages = [_ocr_page(markdown, [_ocr_image("img-0.jpeg", "AAA")])] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "a | b\nsecond line"}) + + assert docs[0].page_content == ( + "| Name | Scan |\n| --- | --- |\n| Ann | a \\| b second line |\n| Bob | n/a |" + ) + + +def test_pdf_embedded_image_in_list_item_stays_in_its_item(): + pages = [ + _ocr_page( + "- one\n- ![img-0.jpeg](img-0.jpeg)\n- three", [_ocr_image("img-0.jpeg", "AAA")] + ) + ] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "a\nb"}) + + assert docs[0].page_content == "- one\n- a b\n- three" + + +def test_pdf_embedded_image_alone_on_indented_line_becomes_paragraph(): + pages = [_ocr_page("Intro\n ![img-0.jpeg](img-0.jpeg) \nOutro", [_ocr_image("img-0.jpeg", "AAA")])] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "a\n\nb"}) + + assert docs[0].page_content == "Intro\n\na\n\nb\n\nOutro" + + +# --- Repeated images: content dedupe and page-fraction filter --- + + +def test_pdf_embedded_image_repeated_on_pages_is_ocrd_once_and_text_reused(): + """A letterhead on every page costs one OCR call; each page still gets its text.""" + pages = [ + _ocr_page("![img-0.jpeg](img-0.jpeg)\n\nPage one", [_ocr_image("img-0.jpeg", "LOGO")]), + _ocr_page("Page two\n\n![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "LOGO")]), + ] + + docs, client = _load_pdf_with_images(pages, {"LOGO": "ACME LAW LLP"}) + + assert len(_image_calls(client)) == 1 + assert docs[0].page_content == "ACME LAW LLP\n\nPage one" + assert docs[1].page_content == "Page two\n\nACME LAW LLP" + + +def test_pdf_embedded_image_dedupe_is_by_content_not_id(): + pages = [ + _ocr_page( + "![img-0.jpeg](img-0.jpeg) ![img-1.jpeg](img-1.jpeg)", + [_ocr_image("img-0.jpeg", "AAA"), _ocr_image("img-1.jpeg", "BBB")], + ) + ] + + docs, client = _load_pdf_with_images(pages, {"AAA": "first", "BBB": "second"}) + + assert len(_image_calls(client)) == 2 + assert docs[0].page_content == "first second" + + +PAGE = SimpleNamespace(width=1000, height=1000) # 1,000,000 px^2 + + +def _sized_page(bbox, dimensions=PAGE): + return [ + _ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "AAA", bbox)], dimensions) + ] + + +def test_pdf_embedded_image_below_page_fraction_is_skipped_without_a_call(): + # 200x200 = 4% of the page, below the 10% default. + docs, client = _load_pdf_with_images( + _sized_page((0, 0, 200, 200)), {"AAA": "never read"} + ) + + assert _image_calls(client) == [] + assert docs[0].page_content == "![img-0.jpeg](img-0.jpeg)" + + +def test_pdf_embedded_image_at_or_above_page_fraction_is_ocrd(): + # 400x400 = 16% of the page. + docs, client = _load_pdf_with_images(_sized_page((100, 100, 500, 500)), {"AAA": "big"}) + + assert len(_image_calls(client)) == 1 + assert docs[0].page_content == "big" + + +def test_pdf_embedded_image_page_fraction_zero_disables_the_filter(): + docs, client = _load_pdf_with_images( + _sized_page((0, 0, 10, 10)), + {"AAA": "tiny"}, + MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION=0, + ) + + assert len(_image_calls(client)) == 1 + assert docs[0].page_content == "tiny" + + +@pytest.mark.parametrize( + "bbox, dimensions", + [ + (None, PAGE), # no coordinates at all + ((0, 0, None, 10), PAGE), # one coordinate missing + ((0, 0, 10, 10), None), # no page dimensions + ((0, 0, 10, 10), SimpleNamespace(width=None, height=1000)), + ((0, 0, 10, 10), SimpleNamespace(width=0, height=1000)), + ((0, 0, 10, 10), SimpleNamespace(width=1000, height=-5)), + ((10, 10, 0, 0), PAGE), # inverted box + ], +) +def test_pdf_embedded_image_kept_when_page_fraction_cannot_be_computed(bbox, dimensions): + docs, client = _load_pdf_with_images(_sized_page(bbox, dimensions), {"AAA": "kept"}) + + assert len(_image_calls(client)) == 1 + assert docs[0].page_content == "kept" + + +def test_pdf_embedded_image_skipped_image_does_not_hide_a_repeat_that_is_large(): + """Fractions are judged per occurrence: the large copy is still OCR'd and reused.""" + pages = [ + _ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "AAA", (0, 0, 100, 100))], PAGE), + _ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "AAA", (0, 0, 900, 900))], PAGE), + ] + + docs, client = _load_pdf_with_images(pages, {"AAA": "text"}) + + assert len(_image_calls(client)) == 1 + assert docs[0].page_content == "![img-0.jpeg](img-0.jpeg)" + assert docs[1].page_content == "text" + + +# --- Retries --- + + +def test_safe_pdf_loader_passes_the_retry_config_for_main_and_image_requests(): + pages = [_ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "AAA")])] + + _, client = _load_pdf_with_images(pages, {"AAA": "text"}) + + calls = client.ocr.process.call_args_list + assert len(calls) == 2 + retries = [c.kwargs.get("retries") for c in calls] + assert all(r is not None for r in retries) + + +def test_mistral_ocr_retry_config_retries_transient_failures_with_backoff(): + from app.utils import document_loader + + config = document_loader._ocr_retries() + + assert config.strategy == "backoff" + assert config.retry_connection_errors is True + assert config.backoff.initial_interval == 500 + assert config.backoff.max_interval == 10_000 + assert config.backoff.exponent == 1.5 + assert config.backoff.max_elapsed_time == 60_000 + + +def test_mistral_ocr_call_survives_transient_http_errors_through_the_real_sdk(): + """429 and 503 are retried by the SDK itself; the third attempt succeeds.""" + import httpx + from mistralai.client import Mistral + + from app.utils import document_loader + + statuses = iter([429, 503, 200]) + requests = [] + + def handler(request): + requests.append(request) + status = next(statuses) + if status != 200: + return httpx.Response(status, json={"message": "transient"}) + page = {"index": 0, "markdown": "recovered", "images": [], "dimensions": None} + return httpx.Response( + 200, + json={"pages": [page], "model": "m", "usage_info": {"pages_processed": 1}}, + ) + + client = Mistral( + api_key="test-key", client=httpx.Client(transport=httpx.MockTransport(handler)) + ) + with patch("mistralai.client.utils.retries.time.sleep") as sleep: + pages = document_loader._ocr_pages( + client, {"type": "image_url", "image_url": "data:image/png;base64,AA"} + ) + + assert [p.markdown for p in pages] == ["recovered"] + assert len(requests) == 3 + assert sleep.call_count == 2 + + +def test_mistral_ocr_call_does_not_retry_client_errors_through_the_real_sdk(): + import httpx + from mistralai.client import Mistral + + from app.utils import document_loader + + requests = [] + + def handler(request): + requests.append(request) + return httpx.Response(400, json={"message": "bad request"}) + + client = Mistral( + api_key="test-key", client=httpx.Client(transport=httpx.MockTransport(handler)) + ) + with patch("mistralai.client.utils.retries.time.sleep") as sleep, pytest.raises(Exception): + document_loader._ocr_pages( + client, {"type": "image_url", "image_url": "data:image/png;base64,AA"} + ) + + assert len(requests) == 1 + sleep.assert_not_called() + + def test_pdf_embedded_images_support_dict_shaped_response(): response = { "pages": [ From 1f3cf12ddc02d14d29f02f3aff07369c2b0d3f5c Mon Sep 17 00:00:00 2001 From: Daniel Vejsada Date: Fri, 23 Oct 2026 21:16:45 +0200 Subject: [PATCH 4/4] fix(ocr): second review round for embedded-image OCR - Split the image caps: MISTRAL_OCR_MAX_EXTRACTED_IMAGES (default 200) bounds how many images Mistral extracts, while MISTRAL_OCR_MAX_EMBEDDED_IMAGES (50) now caps distinct images actually OCR'd, applied after the size filter and dedupe. Per-page logos no longer use up the cap and push real scans out. - Log each size-filtered image at DEBUG with its page fraction; keep images whose box extends beyond the page (units mismatch) and warn. - Send include_blocks=False; the SDK defaults it on and we don't use it. - A placeholder behind a block marker (#, -, >, 1.) counts as alone on its line, so a page's text no longer becomes one huge heading or item. - Bound image OCR calls in flight across all documents with a shared semaphore sized by MISTRAL_OCR_IMAGE_CONCURRENCY. - Insert the text at the first placeholder only; drop repeats. - Retry tests use a fast real RetryConfig instead of patching SDK internals. Co-Authored-By: Claude Opus 5.5 --- README.md | 7 +- app/config.py | 14 +- app/utils/document_loader.py | 159 +++++++++++---- tests/utils/test_document_loader.py | 300 ++++++++++++++++++++++++++-- 4 files changed, 413 insertions(+), 67 deletions(-) diff --git a/README.md b/README.md index 7ca6ec74..f4990af6 100644 --- a/README.md +++ b/README.md @@ -218,9 +218,10 @@ The following environment variables are required to run the application: - `IMAGE_OCR_MAX_PAGES`: (Optional) Max number of frames/pages of a single image upload (e.g. animated GIF or large multi-page TIFF) sent to OCR; frames beyond the cap are skipped with a warning. Default value is `100`. - `MISTRAL_OCR_EMBEDDED_IMAGES`: (Optional) Boolean. Mistral OCR returns only a placeholder (`![img-0.jpeg](img-0.jpeg)`) for a raster image embedded in a PDF page, such as a photo or scan of a page, and does not read the text inside it. When "True" (default), those images are extracted, OCR'd in a second Mistral call each, and the recognized text replaces the placeholder. **Each embedded image is billed as an extra OCR page**; set "False" to OCR only the PDF itself. If an image's OCR call fails, its placeholder is kept and the document still loads. - `MISTRAL_OCR_IMAGE_MIN_SIZE`: (Optional) Minimum image size in pixels for an embedded image to be extracted, so logos and icons are skipped. Default value is `100`. -- `MISTRAL_OCR_MAX_EMBEDDED_IMAGES`: (Optional) Max number of embedded images extracted and OCR'd per document (minimum `1`); a warning is logged when the document has more, and the extra images are left as placeholders. Default value is `50`. -- `MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION`: (Optional) Embedded images covering less than this fraction of their page (`0` to `1`) are not OCR'd and keep their placeholder, so letterheads and logos repeated on every page cost nothing. Set to `0` to OCR every extracted image. Identical images are always OCR'd only once per document. Default value is `0.1`. -- `MISTRAL_OCR_IMAGE_CONCURRENCY`: (Optional) Number of embedded images OCR'd in parallel (minimum `1`). Default value is `4`. +- `MISTRAL_OCR_MAX_EXTRACTED_IMAGES`: (Optional) Max number of embedded images Mistral extracts per document (minimum `1`), before any filtering; this bounds the size of the OCR response. A warning is logged when the document has more, and the later images are not extracted. Default value is `200`. +- `MISTRAL_OCR_MAX_EMBEDDED_IMAGES`: (Optional) Max number of distinct embedded images OCR'd per document (minimum `1`), i.e. the number of extra billed OCR calls. It is applied after the page-fraction filter and the removal of identical images, in page order; a warning is logged when the document has more, and the extra images are left as placeholders. Default value is `50`. +- `MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION`: (Optional) Embedded images covering less than this fraction of their page (`0` to `1`) are not OCR'd and keep their placeholder, so letterheads and logos repeated on every page cost nothing. Set to `0` to OCR every extracted image. Skipped images are logged at debug level. Identical images are always OCR'd only once per document. Default value is `0.1`. +- `MISTRAL_OCR_IMAGE_CONCURRENCY`: (Optional) Max number of embedded-image OCR calls in flight at once across all documents and requests (minimum `1`); it is also the parallelism within a single document. Default value is `4`. - `DEBUG_RAG_API`: (Optional) Set to "True" to show more verbose logging output in the server console, and to enable postgresql database routes - `DEBUG_PGVECTOR_QUERIES`: (Optional) Set to "True" to enable detailed PostgreSQL query logging for pgvector operations. Useful for debugging performance issues with vector database queries. - `CONSOLE_JSON`: (Optional) Set to "True" to log as json for Cloud Logging aggregations diff --git a/app/config.py b/app/config.py index 8edad8c3..d58f17fb 100644 --- a/app/config.py +++ b/app/config.py @@ -235,18 +235,26 @@ async def dispatch(self, request, call_next): MISTRAL_OCR_IMAGE_MIN_SIZE = max( 0, int(get_env_variable("MISTRAL_OCR_IMAGE_MIN_SIZE", "100")) ) -# Max embedded images extracted (and OCR'd) per document (minimum 1), which also -# bounds the size of the OCR response. +# Max embedded images Mistral extracts per document (minimum 1), before any +# filtering. This bounds the size of the OCR response; images beyond it are not +# extracted at all. +MISTRAL_OCR_MAX_EXTRACTED_IMAGES = max( + 1, int(get_env_variable("MISTRAL_OCR_MAX_EXTRACTED_IMAGES", "200")) +) +# Max distinct embedded images OCR'd per document (minimum 1), counted after the +# page-fraction filter and the dedupe. Each one is a billed OCR call. MISTRAL_OCR_MAX_EMBEDDED_IMAGES = max( 1, int(get_env_variable("MISTRAL_OCR_MAX_EMBEDDED_IMAGES", "50")) ) # Embedded images covering less than this fraction of their page (0-1) are not # OCR'd, so letterheads and logos repeated on every page cost nothing. 0 disables. +# Applied before MISTRAL_OCR_MAX_EMBEDDED_IMAGES, so skipped images don't count. MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION = min( 1.0, max(0.0, float(get_env_variable("MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION", "0.1"))), ) -# Number of embedded images OCR'd in parallel (minimum 1). +# Max embedded-image OCR calls in flight at once across all documents and +# requests (minimum 1); also the parallelism within a single document. MISTRAL_OCR_IMAGE_CONCURRENCY = max( 1, int(get_env_variable("MISTRAL_OCR_IMAGE_CONCURRENCY", "4")) ) diff --git a/app/utils/document_loader.py b/app/utils/document_loader.py index ee697354..de80a48f 100644 --- a/app/utils/document_loader.py +++ b/app/utils/document_loader.py @@ -6,6 +6,7 @@ import mimetypes import re import tempfile +import threading from concurrent.futures import ThreadPoolExecutor from typing import Iterator, List, Optional @@ -23,6 +24,7 @@ MISTRAL_OCR_EMBEDDED_IMAGES, MISTRAL_OCR_IMAGE_MIN_SIZE, MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION, + MISTRAL_OCR_MAX_EXTRACTED_IMAGES, MISTRAL_OCR_MAX_EMBEDDED_IMAGES, MISTRAL_OCR_IMAGE_CONCURRENCY, DOCX_TEXT_USE_PANDOC, @@ -369,6 +371,7 @@ def _ocr_pages(client, document_payload: dict, include_image_base64: bool = Fals model=MISTRAL_OCR_MODEL, document=document_payload, include_image_base64=include_image_base64, + include_blocks=False, # paragraph bounding boxes are unused; keep responses small retries=_ocr_retries(), **request_kwargs, ) @@ -396,6 +399,10 @@ def _image_data_uri(image_id: str, image_base64: str) -> str: return f"data:{mime};base64,{image_base64}" +# Bounds embedded-image OCR calls in flight across all documents and requests. +_IMAGE_OCR_SLOTS = threading.BoundedSemaphore(MISTRAL_OCR_IMAGE_CONCURRENCY) + + def _ocr_embedded_image(client, image_id: str, image_base64: str, source: str) -> str: """OCR one embedded image and return its text, or "" if there is none or the call fails. @@ -403,7 +410,8 @@ def _ocr_embedded_image(client, image_id: str, image_base64: str, source: str) - """ data_uri = _image_data_uri(image_id, image_base64) try: - pages = _ocr_pages(client, {"type": "image_url", "image_url": data_uri}) + with _IMAGE_OCR_SLOTS: + pages = _ocr_pages(client, {"type": "image_url", "image_url": data_uri}) except Exception as e: logger.warning( "OCR of embedded image %s in %s failed; keeping placeholder: %s", @@ -419,13 +427,33 @@ def _ocr_embedded_image(client, image_id: str, image_base64: str, source: str) - return text.strip() +# What may precede a placeholder that still stands alone on its line: indentation +# and one markdown block marker (heading, bullet, quote or ordered-list number). +_BLOCK_MARKER_PREFIX_RE = re.compile(r"[ \t]*(?:(?:#{1,6}|[-*+>]|\d+[.)])[ \t]+)?") + + +def _replace_line(markdown: str, line_start: int, line_end: int, text: str) -> str: + """Replace a whole line with ``text`` as its own paragraph (or drop it if ``text`` is empty). + + Adjacent newlines are folded into one blank line of separation; the rest of + the page is untouched. + """ + while line_start > 0 and markdown[line_start - 1] == "\n": + line_start -= 1 + while line_end < len(markdown) and markdown[line_end] == "\n": + line_end += 1 + parts = (markdown[:line_start], text, markdown[line_end:]) + return "\n\n".join(part for part in parts if part) + + def _replace_placeholder(markdown: str, image_id: str, text: str) -> str: """Swap the image's placeholder (``![alt](id)``) for ``text``; append it if there is none. - A placeholder alone on its line becomes its own paragraph: adjacent newlines - are folded into one blank line of separation and the rest of the page is - untouched. Elsewhere (a table cell, list item or sentence) the text goes - inline with its newlines collapsed, and ``|`` escaped inside a table row. + A placeholder alone on its line (optionally behind a markdown block marker + such as ``#``, ``-``, ``>`` or ``1.``) becomes its own paragraph, dropping + the marker. Elsewhere (a table cell or a sentence) the text goes inline with + its newlines collapsed, and ``|`` escaped inside a table row. The text goes + at the first occurrence only; repeated placeholders of the same image are removed. """ pattern = re.compile(rf"!\[[^\]]*\]\({re.escape(image_id)}\)") matches = list(pattern.finditer(markdown)) @@ -433,23 +461,22 @@ def _replace_placeholder(markdown: str, image_id: str, text: str) -> str: return f"{markdown.rstrip()}\n\n{text}" if markdown.strip() else text inline_text = re.sub(r"\s*\n\s*", " ", text) - for match in reversed(matches): # later first, so earlier offsets stay valid - start, end = match.span() + for index in reversed(range(len(matches))): # later first, so earlier offsets stay valid + start, end = matches[index].span() line_start = markdown.rfind("\n", 0, start) + 1 line_end = markdown.find("\n", end) line_end = len(markdown) if line_end == -1 else line_end - if markdown[line_start:start].strip() or markdown[end:line_end].strip(): + keep_text = index == 0 + if _BLOCK_MARKER_PREFIX_RE.fullmatch(markdown[line_start:start]) and not markdown[end:line_end].strip(): + markdown = _replace_line(markdown, line_start, line_end, text if keep_text else "") + continue + replacement = "" + if keep_text: is_table_row = markdown[line_start:line_end].lstrip().startswith("|") replacement = inline_text.replace("|", "\\|") if is_table_row else inline_text - markdown = markdown[:start] + replacement + markdown[end:] - continue - while line_start > 0 and markdown[line_start - 1] == "\n": - line_start -= 1 - while line_end < len(markdown) and markdown[line_end] == "\n": - line_end += 1 - lead = "\n\n" if line_start > 0 else "" - trail = "\n\n" if line_end < len(markdown) else "" - markdown = markdown[:line_start] + lead + text + trail + markdown[line_end:] + elif markdown[start - 1 : start] in (" ", "\t") and markdown[end : end + 1] in (" ", "\t"): + end += 1 # don't leave a double space where the placeholder was + markdown = markdown[:start] + replacement + markdown[end:] return markdown @@ -457,15 +484,15 @@ def _is_number(value) -> bool: return isinstance(value, (int, float)) and not isinstance(value, bool) -def _below_min_page_fraction(image, page) -> bool: - """Whether ``image`` covers less of ``page`` than MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION. +def _page_fraction(image, page): + """``(fraction, fits)``: the share of ``page`` that ``image`` covers, and whether its box lies on the page. The bounding box and the page dimensions are both taken to be in the pixel - space of the rendered page (UNVERIFIED against the live API). Images are kept - whenever the fraction can't be computed (missing coordinates or dimensions). + space of the rendered page (UNVERIFIED against the live API). ``fraction`` is + None when it can't be computed (missing coordinates or dimensions). ``fits`` + is False when the box extends beyond the page (with 1% tolerance), which + suggests the two are in different units. """ - if MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION <= 0: - return False dimensions = _field(page, "dimensions") page_width, page_height = _field(dimensions, "width"), _field(dimensions, "height") corners = [ @@ -473,20 +500,56 @@ def _below_min_page_fraction(image, page) -> bool: for name in ("top_left_x", "top_left_y", "bottom_right_x", "bottom_right_y") ] if not all(_is_number(v) for v in [page_width, page_height, *corners]): - return False + return None, True left, top, right, bottom = corners if page_width <= 0 or page_height <= 0 or right <= left or bottom <= top: - return False - page_area = page_width * page_height - return (right - left) * (bottom - top) < MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION * page_area + return None, True + if right > page_width * 1.01 or bottom > page_height * 1.01: + return None, False + return (right - left) * (bottom - top) / (page_width * page_height), True + + +def _drop_small_images(found: list, source: str) -> list: + """Drop images covering less than MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION of their page. + + ``found`` holds ``(page position, page, image, image id, image base64)``. + Images whose box doesn't fit their page are kept: the filter can't be trusted + for them. + """ + if MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION <= 0: + return found + kept = [] + units_warned = False + for entry in found: + position, page, image, image_id, _ = entry + fraction, fits = _page_fraction(image, page) + if not fits and not units_warned: + units_warned = True + logger.warning( + "Embedded image %s in %s has a bounding box outside its page; " + "page coordinates may be in other units, so such images are not filtered by size", + image_id, + source, + ) + if fraction is not None and fraction < MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION: + logger.debug( + "Skipping embedded image %s on page %d of %s: covers %.1f%% of the page", + image_id, + position + 1, + source, + fraction * 100, + ) + continue + kept.append(entry) + return kept def _collect_embedded_images(pages: list, source: str) -> list: """Embedded images worth OCR'ing as ``(page position, image id, image base64)``. - Mistral is asked for one image more than MISTRAL_OCR_MAX_EMBEDDED_IMAGES, so - receiving more than the cap means the document has images beyond it; those - are dropped with a warning. Images too small relative to their page are skipped. + Mistral is asked for one image more than MISTRAL_OCR_MAX_EXTRACTED_IMAGES, so + receiving more than that means the document has images beyond it; those are + dropped with a warning. Images too small relative to their page are skipped. """ found = [] # (page position, page, image, image id, image base64) for position, page in enumerate(pages): @@ -496,17 +559,16 @@ def _collect_embedded_images(pages: list, source: str) -> list: if isinstance(image_id, str) and isinstance(image_base64, str) and image_base64: found.append((position, page, image, image_id, image_base64)) - if len(found) > MISTRAL_OCR_MAX_EMBEDDED_IMAGES: + if len(found) > MISTRAL_OCR_MAX_EXTRACTED_IMAGES: logger.warning( - "%s has more than MISTRAL_OCR_MAX_EMBEDDED_IMAGES=%d embedded images; the rest keep their placeholders", + "%s has more than MISTRAL_OCR_MAX_EXTRACTED_IMAGES=%d embedded images; later images were not extracted", source, - MISTRAL_OCR_MAX_EMBEDDED_IMAGES, + MISTRAL_OCR_MAX_EXTRACTED_IMAGES, ) - found = found[:MISTRAL_OCR_MAX_EMBEDDED_IMAGES] + found = found[:MISTRAL_OCR_MAX_EXTRACTED_IMAGES] return [ (position, image_id, image_base64) - for position, page, image, image_id, image_base64 in found - if not _below_min_page_fraction(image, page) + for position, _, _, image_id, image_base64 in _drop_small_images(found, source) ] @@ -515,16 +577,25 @@ def _inline_embedded_images(client, pages: list, markdowns: List[str], source: s Image ids repeat across pages (``img-0.jpeg`` on each), so images are matched to their own page. Identical images (a logo on every page) are OCR'd once and - their text reused. Images are OCR'd concurrently but applied in page order, so - the result doesn't depend on which call finishes first. + their text reused; at most MISTRAL_OCR_MAX_EMBEDDED_IMAGES distinct images are + OCR'd, in page order. Images are OCR'd concurrently but applied in page order, + so the result doesn't depend on which call finishes first. """ - placements = [] # (page position, image id, content digest) + placements = {} # (page position, image id) -> content digest distinct = {} # content digest -> (image id, image base64) of its first occurrence for position, image_id, image_base64 in _collect_embedded_images(pages, source): digest = hashlib.sha256(image_base64.encode()).hexdigest() distinct.setdefault(digest, (image_id, image_base64)) - placements.append((position, image_id, digest)) - if not placements: + placements.setdefault((position, image_id), digest) + + if len(distinct) > MISTRAL_OCR_MAX_EMBEDDED_IMAGES: + logger.warning( + "%s has more than MISTRAL_OCR_MAX_EMBEDDED_IMAGES=%d distinct embedded images; the rest keep their placeholders", + source, + MISTRAL_OCR_MAX_EMBEDDED_IMAGES, + ) + distinct = dict(list(distinct.items())[:MISTRAL_OCR_MAX_EMBEDDED_IMAGES]) + if not distinct: return markdowns with ThreadPoolExecutor(max_workers=MISTRAL_OCR_IMAGE_CONCURRENCY) as executor: @@ -535,8 +606,8 @@ def _inline_embedded_images(client, pages: list, markdowns: List[str], source: s texts = dict(zip(distinct, results)) result = list(markdowns) - for position, image_id, digest in placements: - if texts[digest]: + for (position, image_id), digest in placements.items(): + if texts.get(digest): # absent when the image was over the OCR cap result[position] = _replace_placeholder(result[position], image_id, texts[digest]) return result @@ -564,7 +635,7 @@ def _ocr_document( "include_image_base64": True, "image_min_size": MISTRAL_OCR_IMAGE_MIN_SIZE, # One more than the cap, so exceeding it is detectable. - "image_limit": MISTRAL_OCR_MAX_EMBEDDED_IMAGES + 1, + "image_limit": MISTRAL_OCR_MAX_EXTRACTED_IMAGES + 1, } try: diff --git a/tests/utils/test_document_loader.py b/tests/utils/test_document_loader.py index b2d2979a..2c845123 100644 --- a/tests/utils/test_document_loader.py +++ b/tests/utils/test_document_loader.py @@ -282,6 +282,7 @@ def _load_pdf_with_images(pages, image_results=None, **settings): "MISTRAL_OCR_EMBEDDED_IMAGES": True, "MISTRAL_OCR_IMAGE_MIN_SIZE": 100, "MISTRAL_OCR_IMAGE_MIN_PAGE_FRACTION": 0.1, + "MISTRAL_OCR_MAX_EXTRACTED_IMAGES": 200, "MISTRAL_OCR_MAX_EMBEDDED_IMAGES": 50, "MISTRAL_OCR_IMAGE_CONCURRENCY": 4, **settings, @@ -347,13 +348,14 @@ def test_pdf_embedded_images_request_kwargs_when_enabled(): [_ocr_page("text")], MISTRAL_OCR_EMBEDDED_IMAGES=True, MISTRAL_OCR_IMAGE_MIN_SIZE=123, - MISTRAL_OCR_MAX_EMBEDDED_IMAGES=7, + MISTRAL_OCR_MAX_EXTRACTED_IMAGES=7, ) client.ocr.process.assert_called_once() kwargs = client.ocr.process.call_args.kwargs assert kwargs["document"]["type"] == "document_url" assert kwargs["include_image_base64"] is True + assert kwargs["include_blocks"] is False assert kwargs["image_min_size"] == 123 assert kwargs["image_limit"] == 8 # cap + 1, so exceeding the cap is detectable @@ -369,6 +371,7 @@ def test_pdf_embedded_images_disabled_keeps_original_request(): client.ocr.process.assert_called_once() kwargs = client.ocr.process.call_args.kwargs assert kwargs["include_image_base64"] is False + assert kwargs["include_blocks"] is False assert "image_min_size" not in kwargs assert "image_limit" not in kwargs assert docs[0].page_content == "![img-0.jpeg](img-0.jpeg)" @@ -475,7 +478,7 @@ def test_pdf_embedded_image_base64_without_prefix_gets_data_uri(): def test_pdf_embedded_image_cap_limits_image_ocr_calls(caplog): - """Mistral hands back cap + 1 images; only ``cap`` are OCR'd, with a warning.""" + """More distinct images than the cap: the first ``cap`` are OCR'd, with a warning.""" pages = [ _ocr_page( "![img-0.jpeg](img-0.jpeg) ![img-1.jpeg](img-1.jpeg)", @@ -490,7 +493,6 @@ def test_pdf_embedded_image_cap_limits_image_ocr_calls(caplog): MISTRAL_OCR_MAX_EMBEDDED_IMAGES=1, ) - assert client.ocr.process.call_args_list[0].kwargs["image_limit"] == 2 assert len(_image_calls(client)) == 1 assert "first" in docs[0].page_content assert "![img-1.jpeg](img-1.jpeg)" in docs[0].page_content @@ -510,8 +512,8 @@ def test_pdf_embedded_images_exactly_at_cap_do_not_warn(caplog): assert "MISTRAL_OCR_MAX_EMBEDDED_IMAGES" not in caplog.text -def test_pdf_embedded_images_cap_counts_images_before_dedupe(): - """Duplicates count toward the cap: the cap bounds what Mistral sends back.""" +def test_pdf_embedded_images_cap_counts_distinct_images_not_duplicates(caplog): + """Identical images are one OCR call, so they don't use up the cap.""" pages = [ _ocr_page( "![img-0.jpeg](img-0.jpeg) ![img-1.jpeg](img-1.jpeg)", @@ -519,12 +521,85 @@ def test_pdf_embedded_images_cap_counts_images_before_dedupe(): ) ] + with caplog.at_level("WARNING"): + docs, client = _load_pdf_with_images( + pages, {"AAA": "same"}, MISTRAL_OCR_MAX_EMBEDDED_IMAGES=1 + ) + + assert len(_image_calls(client)) == 1 + assert docs[0].page_content == "same same" + assert "MISTRAL_OCR_MAX_EMBEDDED_IMAGES" not in caplog.text + + +def test_pdf_embedded_images_truncated_by_cap_keep_placeholders_for_all_their_copies(): + pages = [ + _ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "AAA")]), + _ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "BBB")]), + _ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "BBB")]), + ] + + docs, client = _load_pdf_with_images( + pages, {"AAA": "first", "BBB": "second"}, MISTRAL_OCR_MAX_EMBEDDED_IMAGES=1 + ) + + assert len(_image_calls(client)) == 1 + assert [d.page_content for d in docs] == [ + "first", + "![img-0.jpeg](img-0.jpeg)", + "![img-0.jpeg](img-0.jpeg)", + ] + + +def test_pdf_embedded_images_skipped_logos_do_not_use_up_the_extraction_cap(): + """A small logo on each of 60 pages is filtered; the big photo on page 55 is still OCR'd.""" + logo_bbox, photo_bbox = (0, 0, 100, 100), (0, 0, 800, 800) + pages = [ + _ocr_page( + "![img-0.jpeg](img-0.jpeg)", + [_ocr_image("img-0.jpeg", "LOGO", logo_bbox)] + + ([_ocr_image("img-1.jpeg", "PHOTO", photo_bbox)] if number == 55 else []), + PAGE, + ) + for number in range(1, 61) + ] + docs, client = _load_pdf_with_images( - pages, {"AAA": "same"}, MISTRAL_OCR_MAX_EMBEDDED_IMAGES=1 + pages, {"PHOTO": "photo text"}, MISTRAL_OCR_MAX_EMBEDDED_IMAGES=2 ) + assert [c["document"]["image_url"] for c in _image_calls(client)] == [ + "data:image/jpeg;base64,PHOTO" + ] + assert client.ocr.process.call_args_list[0].kwargs["image_limit"] == 201 + assert "photo text" in docs[54].page_content + + +def test_pdf_embedded_images_over_the_extraction_cap_warn_and_are_dropped(caplog): + pages = [ + _ocr_page( + "![img-0.jpeg](img-0.jpeg) ![img-1.jpeg](img-1.jpeg)", + [_ocr_image("img-0.jpeg", "AAA"), _ocr_image("img-1.jpeg", "BBB")], + ) + ] + + with caplog.at_level("WARNING"): + docs, client = _load_pdf_with_images( + pages, {"AAA": "first", "BBB": "second"}, MISTRAL_OCR_MAX_EXTRACTED_IMAGES=1 + ) + assert len(_image_calls(client)) == 1 - assert docs[0].page_content == "same ![img-1.jpeg](img-1.jpeg)" + assert docs[0].page_content == "first ![img-1.jpeg](img-1.jpeg)" + assert "MISTRAL_OCR_MAX_EXTRACTED_IMAGES=1" in caplog.text + assert "not extracted" in caplog.text + + +def test_pdf_embedded_images_at_the_extraction_cap_do_not_warn(caplog): + pages = [_ocr_page("![img-0.jpeg](img-0.jpeg)", [_ocr_image("img-0.jpeg", "AAA")])] + + with caplog.at_level("WARNING"): + _load_pdf_with_images(pages, {"AAA": "text"}, MISTRAL_OCR_MAX_EXTRACTED_IMAGES=1) + + assert "MISTRAL_OCR_MAX_EXTRACTED_IMAGES" not in caplog.text def test_pdf_embedded_image_only_normalizes_newlines_around_placeholder(): @@ -593,6 +668,8 @@ def process(**kwargs): document_loader, "MISTRAL_OCR_EMBEDDED_IMAGES", True ), patch.object( document_loader, "MISTRAL_OCR_IMAGE_CONCURRENCY", 2 + ), patch.object( + document_loader, "_IMAGE_OCR_SLOTS", threading.BoundedSemaphore(2) ), patch.object( SafePyPDFLoader, "_encode_pdf_b64", return_value="b64data" ): @@ -601,6 +678,71 @@ def process(**kwargs): assert docs[0].page_content == "first text\n\nsecond text" +def test_image_ocr_calls_are_bounded_across_documents_loaded_concurrently(): + """The slot limit is global: two documents never run more image OCR calls at once than it allows.""" + import threading + + from app.utils import document_loader + from app.utils.document_loader import SafePyPDFLoader + + lock = threading.Lock() + overlap = threading.Event() + in_flight = max_in_flight = 0 + first_call = True + + def process(**kwargs): + nonlocal in_flight, max_in_flight, first_call + document = kwargs["document"] + if document["type"] == "document_url": + tag = document["document_url"].rsplit("-", 1)[1] + images = [_ocr_image(f"img-{n}.jpeg", f"{tag}{n}") for n in (0, 1)] + markdown = "\n\n".join(f"![img-{n}.jpeg](img-{n}.jpeg)" for n in (0, 1)) + return SimpleNamespace(pages=[_ocr_page(markdown, images)]) + with lock: + in_flight += 1 + max_in_flight = max(max_in_flight, in_flight) + if in_flight > 1: + overlap.set() + waits, first_call = first_call, False + if waits: + # Hold the slot until a second call overlaps; with a limit of 1 none can, + # so this times out and the call carries on. + overlap.wait(timeout=1) + with lock: + in_flight -= 1 + return SimpleNamespace(pages=[SimpleNamespace(markdown="text")]) + + fake_module, client = _make_mistral_module(pages=[]) + client.ocr.process.side_effect = process + results = {} + + def load(name): + results[name] = SafePyPDFLoader(name).load() + + with patch.dict( + "sys.modules", {"mistralai": fake_module, "mistralai.client": fake_module} + ), patch.object(document_loader, "MISTRAL_API_KEY", "test-key"), patch.object( + document_loader, "MISTRAL_OCR_EMBEDDED_IMAGES", True + ), patch.object( + document_loader, "MISTRAL_OCR_IMAGE_CONCURRENCY", 4 + ), patch.object( + document_loader, "_IMAGE_OCR_SLOTS", threading.BoundedSemaphore(1) + ), patch.object( + SafePyPDFLoader, "_encode_pdf_b64", lambda self: f"pdf-{self.filepath}" + ): + threads = [threading.Thread(target=load, args=(name,)) for name in ("A", "B")] + for thread in threads: + thread.start() + for thread in threads: + thread.join(timeout=30) + + assert max_in_flight == 1 + assert {name: docs[0].page_content for name, docs in results.items()} == { + "A": "text\n\ntext", + "B": "text\n\ntext", + } + + def test_pdf_embedded_image_without_placeholder_is_appended(): pages = [_ocr_page("Page body", [_ocr_image("img-0.jpeg", "AAA")])] @@ -657,24 +799,90 @@ def test_pdf_embedded_image_in_table_cell_stays_in_its_row_with_pipes_escaped(): ) -def test_pdf_embedded_image_in_list_item_stays_in_its_item(): +def test_pdf_embedded_image_alone_on_indented_line_becomes_paragraph(): + pages = [_ocr_page("Intro\n ![img-0.jpeg](img-0.jpeg) \nOutro", [_ocr_image("img-0.jpeg", "AAA")])] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "a\n\nb"}) + + assert docs[0].page_content == "Intro\n\na\n\nb\n\nOutro" + + +@pytest.mark.parametrize( + "marker", ["#", "######", "-", "*", "+", ">", "1.", "12)", " -", "\t>"] +) +def test_pdf_embedded_image_behind_a_block_marker_becomes_a_paragraph_without_the_marker(marker): + pages = [ + _ocr_page( + f"Intro\n{marker} ![img-0.jpeg](img-0.jpeg)\nOutro", [_ocr_image("img-0.jpeg", "AAA")] + ) + ] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "a\nb"}) + + assert docs[0].page_content == "Intro\n\na\nb\n\nOutro" + + +def test_pdf_embedded_image_in_the_middle_of_a_list_leaves_the_other_items(): pages = [ _ocr_page( "- one\n- ![img-0.jpeg](img-0.jpeg)\n- three", [_ocr_image("img-0.jpeg", "AAA")] ) ] + docs, _ = _load_pdf_with_images(pages, {"AAA": "scan"}) + + assert docs[0].page_content == "- one\n\nscan\n\n- three" + + +def test_pdf_embedded_image_after_a_marker_but_followed_by_text_stays_inline(): + pages = [ + _ocr_page( + "- ![img-0.jpeg](img-0.jpeg) caption", [_ocr_image("img-0.jpeg", "AAA")] + ) + ] + docs, _ = _load_pdf_with_images(pages, {"AAA": "a\nb"}) - assert docs[0].page_content == "- one\n- a b\n- three" + assert docs[0].page_content == "- a b caption" -def test_pdf_embedded_image_alone_on_indented_line_becomes_paragraph(): - pages = [_ocr_page("Intro\n ![img-0.jpeg](img-0.jpeg) \nOutro", [_ocr_image("img-0.jpeg", "AAA")])] +def test_pdf_embedded_image_repeated_on_a_page_is_inserted_once_at_the_first_occurrence(): + pages = [ + _ocr_page( + "Top\n\n![img-0.jpeg](img-0.jpeg)\n\nMiddle\n\n![img-0.jpeg](img-0.jpeg)\n\nEnd", + [_ocr_image("img-0.jpeg", "AAA")], + ) + ] - docs, _ = _load_pdf_with_images(pages, {"AAA": "a\n\nb"}) + docs, _ = _load_pdf_with_images(pages, {"AAA": "scan"}) - assert docs[0].page_content == "Intro\n\na\n\nb\n\nOutro" + assert docs[0].page_content == "Top\n\nscan\n\nMiddle\n\nEnd" + + +def test_pdf_embedded_image_repeated_at_the_end_of_a_page_drops_its_line_cleanly(): + pages = [ + _ocr_page( + "![img-0.jpeg](img-0.jpeg)\n\nBody\n\n![img-0.jpeg](img-0.jpeg)\n", + [_ocr_image("img-0.jpeg", "AAA")], + ) + ] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "scan"}) + + assert docs[0].page_content == "scan\n\nBody" + + +def test_pdf_embedded_image_repeated_inline_leaves_no_double_space(): + pages = [ + _ocr_page( + "See ![img-0.jpeg](img-0.jpeg) and ![img-0.jpeg](img-0.jpeg) now", + [_ocr_image("img-0.jpeg", "AAA")], + ) + ] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "scan"}) + + assert docs[0].page_content == "See scan and now" # --- Repeated images: content dedupe and page-fraction filter --- @@ -735,6 +943,52 @@ def test_pdf_embedded_image_at_or_above_page_fraction_is_ocrd(): assert docs[0].page_content == "big" +def test_pdf_embedded_image_skipped_by_page_fraction_is_logged_at_debug(caplog): + with caplog.at_level("DEBUG"): + _load_pdf_with_images(_sized_page((0, 0, 200, 200)), {"AAA": "never read"}) + + skipped = [r for r in caplog.records if "Skipping embedded image" in r.getMessage()] + assert [r.levelname for r in skipped] == ["DEBUG"] + message = skipped[0].getMessage() + assert "img-0.jpeg" in message and "scan.pdf" in message and "4.0%" in message + + +def test_pdf_embedded_images_outside_the_page_are_kept_with_one_units_warning(caplog): + """A box larger than its page suggests mismatched units, so the size filter is not trusted.""" + pages = [ + _ocr_page( + "![img-0.jpeg](img-0.jpeg) ![img-1.jpeg](img-1.jpeg)", + [ + _ocr_image("img-0.jpeg", "AAA", (0, 0, 20, 20)), + _ocr_image("img-1.jpeg", "BBB", (0, 0, 2000, 2000)), + ], + PAGE, + ), + _ocr_page( + "![img-0.jpeg](img-0.jpeg)", + [_ocr_image("img-0.jpeg", "CCC", (0, 0, 1500, 30))], + PAGE, + ), + ] + + with caplog.at_level("DEBUG"): + docs, client = _load_pdf_with_images(pages, {"BBB": "kept b", "CCC": "kept c"}) + + units = [r for r in caplog.records if "bounding box outside its page" in r.getMessage()] + assert [r.levelname for r in units] == ["WARNING"] # once per document + assert "scan.pdf" in units[0].getMessage() + assert len(_image_calls(client)) == 2 + assert docs[0].page_content == "![img-0.jpeg](img-0.jpeg) kept b" + assert docs[1].page_content == "kept c" + + +def test_pdf_embedded_image_within_one_percent_of_the_page_edge_is_not_a_units_mismatch(caplog): + with caplog.at_level("WARNING"): + _load_pdf_with_images(_sized_page((0, 0, 1005, 1000)), {"AAA": "text"}) + + assert "bounding box outside its page" not in caplog.text + + def test_pdf_embedded_image_page_fraction_zero_disables_the_filter(): docs, client = _load_pdf_with_images( _sized_page((0, 0, 10, 10)), @@ -806,6 +1060,19 @@ def test_mistral_ocr_retry_config_retries_transient_failures_with_backoff(): assert config.backoff.max_elapsed_time == 60_000 +def _fast_retries(): + """The loader's retry policy with intervals of a few milliseconds, so tests don't wait.""" + from mistralai.client import utils + + return utils.RetryConfig( + "backoff", + utils.BackoffStrategy( + initial_interval=1, max_interval=5, exponent=1.1, max_elapsed_time=2000 + ), + retry_connection_errors=True, + ) + + def test_mistral_ocr_call_survives_transient_http_errors_through_the_real_sdk(): """429 and 503 are retried by the SDK itself; the third attempt succeeds.""" import httpx @@ -830,14 +1097,13 @@ def handler(request): client = Mistral( api_key="test-key", client=httpx.Client(transport=httpx.MockTransport(handler)) ) - with patch("mistralai.client.utils.retries.time.sleep") as sleep: + with patch.object(document_loader, "_ocr_retries", _fast_retries): pages = document_loader._ocr_pages( client, {"type": "image_url", "image_url": "data:image/png;base64,AA"} ) assert [p.markdown for p in pages] == ["recovered"] assert len(requests) == 3 - assert sleep.call_count == 2 def test_mistral_ocr_call_does_not_retry_client_errors_through_the_real_sdk(): @@ -855,13 +1121,12 @@ def handler(request): client = Mistral( api_key="test-key", client=httpx.Client(transport=httpx.MockTransport(handler)) ) - with patch("mistralai.client.utils.retries.time.sleep") as sleep, pytest.raises(Exception): + with patch.object(document_loader, "_ocr_retries", _fast_retries), pytest.raises(Exception): document_loader._ocr_pages( client, {"type": "image_url", "image_url": "data:image/png;base64,AA"} ) assert len(requests) == 1 - sleep.assert_not_called() def test_pdf_embedded_images_support_dict_shaped_response(): @@ -909,6 +1174,7 @@ def test_image_ocr_loader_does_not_request_embedded_images(tmp_path): kwargs = client.ocr.process.call_args.kwargs assert kwargs["include_image_base64"] is False + assert kwargs["include_blocks"] is False assert "image_min_size" not in kwargs