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/README.md b/README.md index 0d81a103..f4990af6 100644 --- a/README.md +++ b/README.md @@ -214,8 +214,14 @@ 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_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 2a34257c..d58f17fb 100644 --- a/app/config.py +++ b/app/config.py @@ -224,6 +224,40 @@ 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 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"))), +) +# 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")) +) # 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..de80a48f 100644 --- a/app/utils/document_loader.py +++ b/app/utils/document_loader.py @@ -2,9 +2,13 @@ import os import codecs +import hashlib +import mimetypes import re import tempfile +import threading +from concurrent.futures import ThreadPoolExecutor from typing import Iterator, List, Optional import chardet @@ -17,6 +21,12 @@ MISTRAL_API_KEY, MISTRAL_OCR_MODEL, MISTRAL_SERVER_URL, + 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, DOCX_TEXT_TRACK_CHANGES, DOCX_TEXT_INCLUDE_HEADERS_FOOTERS, @@ -328,45 +338,326 @@ 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_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). + + 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, + include_blocks=False, # paragraph bounding boxes are unused; keep responses small + retries=_ocr_retries(), + **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"!\[[^\]]*\]\([^)]*\)") + + +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 + 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}" + + +# 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. + + 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: + 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", + image_id, + source, + e, + ) + return "" + + text = "\n\n".join( + _MARKDOWN_IMAGE_RE.sub("", md).strip() for md in _page_markdowns(pages) + ) + 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 (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)) + 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 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 + 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 + 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 + + +def _is_number(value) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) + + +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). ``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. + """ + 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 None, True + left, top, right, bottom = corners + if page_width <= 0 or page_height <= 0 or right <= left or bottom <= top: + 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_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): + 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: + found.append((position, page, image, image_id, image_base64)) + + if len(found) > MISTRAL_OCR_MAX_EXTRACTED_IMAGES: + logger.warning( + "%s has more than MISTRAL_OCR_MAX_EXTRACTED_IMAGES=%d embedded images; later images were not extracted", + source, + MISTRAL_OCR_MAX_EXTRACTED_IMAGES, + ) + found = found[:MISTRAL_OCR_MAX_EXTRACTED_IMAGES] + return [ + (position, image_id, image_base64) + for position, _, _, image_id, image_base64 in _drop_small_images(found, source) + ] + + +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; 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 + 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.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: + 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), 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 + + +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, + # One more than the cap, so exceeding it is detectable. + "image_limit": MISTRAL_OCR_MAX_EXTRACTED_IMAGES + 1, + } + 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 +669,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 +690,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/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 diff --git a/tests/utils/test_document_loader.py b/tests/utils/test_document_loader.py index b1b81a21..2c845123 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 @@ -126,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"): @@ -159,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"): @@ -179,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"): @@ -229,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( @@ -242,6 +252,932 @@ def test_mistral_client_uses_configured_server_url(): client.ocr.process.assert_called_once() +# --- Embedded-image OCR inside PDFs ----------------------------------------- + + +def _ocr_page(markdown, images=(), dimensions=None): + return SimpleNamespace(markdown=markdown, images=list(images), dimensions=dimensions) + + +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): + """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_IMAGE_MIN_PAGE_FRACTION": 0.1, + "MISTRAL_OCR_MAX_EXTRACTED_IMAGES": 200, + "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_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 + + +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 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)" + + +@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.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.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", "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", # 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): + """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)", + [_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_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"): + docs, client = _load_pdf_with_images( + pages, {"AAA": "text"}, MISTRAL_OCR_MAX_EMBEDDED_IMAGES=1 + ) + + assert len(_image_calls(client)) == 1 + assert docs[0].page_content == "text" + assert "MISTRAL_OCR_MAX_EMBEDDED_IMAGES" not in caplog.text + + +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)", + [_ocr_image("img-0.jpeg", "AAA"), _ocr_image("img-1.jpeg", "AAA")], + ) + ] + + 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, {"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 == "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(): + 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( + document_loader, "_IMAGE_OCR_SLOTS", threading.BoundedSemaphore(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_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")])] + + docs, _ = _load_pdf_with_images(pages, {"AAA": "hidden text"}) + + 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_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 == "- a b caption" + + +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": "scan"}) + + 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 --- + + +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_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)), + {"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 _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 + 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.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 + + +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.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 + + +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 kwargs["include_blocks"] 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"