From 6c23ac9ada0df72897b08476d13e082cf90d2870 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Sat, 25 Jul 2026 20:28:06 +0530 Subject: [PATCH 1/3] UN-2646 [FEAT] Add LLMWhisperer image output mode to the v2 adapter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds an "image" output mode to the LLMWhisperer V2 X2Text adapter: it converts a PDF to per-page images via the LLMWhisperer pdf-to-images API and returns them as PageImageReference objects (persisted to FileStorage), never smuggled into extracted_text, so text-mode consumers are unaffected. - dto.py: PageImageReference + additive TextExtractionMetadata.page_images. - constants.py: OutputModes.IMAGE, ImageOutputConfig (PDF-only), OUTPUT_MODE key. - helper.py: get_page_images() — submit / poll / retrieve+unzip the pdf-to-images job and persist the page images. - llm_whisperer_v2.py: _process_image_mode() + output-mode branch in process(), with PDF-only validation. - json_schema.json: "image" enum + "Image (PDF only)" label + description. Recovered from the earlier implementation (MFBT phase llmwhisperer-image-output-mode-adapter) and rebased onto main, kept independent of the document_insights/signature feature (PR #1967). 45 tests pass. MUNS-193/194/195 complete; MUNS-196 UI validation (UNS-757/758/759) remains. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/unstract/sdk1/adapters/x2text/dto.py | 62 +++ .../x2text/llm_whisperer_v2/src/constants.py | 82 ++++ .../x2text/llm_whisperer_v2/src/helper.py | 397 +++++++++++++++++- .../llm_whisperer_v2/src/llm_whisperer_v2.py | 69 +++ .../src/static/json_schema.json | 10 +- unstract/sdk1/tests/llmw_image_fixtures.py | 149 +++++++ .../tests/test_llm_whisperer_v2_constants.py | 46 ++ unstract/sdk1/tests/test_llmw_image_helper.py | 193 +++++++++ .../sdk1/tests/test_llmw_v2_process_image.py | 130 ++++++ unstract/sdk1/tests/test_x2text_dto.py | 127 ++++++ 10 files changed, 1256 insertions(+), 9 deletions(-) create mode 100644 unstract/sdk1/tests/llmw_image_fixtures.py create mode 100644 unstract/sdk1/tests/test_llm_whisperer_v2_constants.py create mode 100644 unstract/sdk1/tests/test_llmw_image_helper.py create mode 100644 unstract/sdk1/tests/test_llmw_v2_process_image.py create mode 100644 unstract/sdk1/tests/test_x2text_dto.py diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/dto.py b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/dto.py index 95c60bbe8c..23c22a5833 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/dto.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/dto.py @@ -1,14 +1,76 @@ +from __future__ import annotations + from dataclasses import dataclass from typing import Any +from unstract.sdk1.file_storage import FileStorageProvider + @dataclass class TextExtractionMetadata: whisper_hash: str line_metadata: dict[Any, Any] | None = None + # Optional, additive field populated only in image output mode. Defaults to + # None so existing text-mode consumers are entirely unaffected (the field is + # never encoded into ``extracted_text``). See PageImageReference below. + page_images: list[PageImageReference] | None = None @dataclass class TextExtractionResult: extracted_text: str extraction_metadata: TextExtractionMetadata | None = None + + +@dataclass +class PageImageReference: + """Per-page image reference for image-mode extraction results. + + Produced by the LLMWhisperer image output mode: each entry points to a + single page image that has been persisted to Unstract's FileStorage. This + is a dedicated value object so image references are never smuggled inside + the string ``extracted_text`` field used by text-mode consumers. + + Attributes: + page_number: 1-based index of the page this image represents. + path: FileStorage path / reference string to the stored page image. + filename: Stored image filename (e.g. ``page_001.png``). Optional. + size_bytes: Size of the stored image file in bytes. Optional. + provider: FileStorageProvider backend (LOCAL/S3/...) holding the + image. Optional. + """ + + page_number: int + path: str + filename: str | None = None + size_bytes: int | None = None + provider: FileStorageProvider | None = None + + def to_dict(self) -> dict[str, Any]: + """Serialize to a plain, JSON-friendly dictionary. + + The ``provider`` enum is stored as its string value so the result is + directly serializable; ``from_dict`` reverses this. + """ + return { + "page_number": self.page_number, + "path": self.path, + "filename": self.filename, + "size_bytes": self.size_bytes, + "provider": self.provider.value if self.provider is not None else None, + } + + @classmethod + def from_dict(cls, data: dict[str, Any]) -> PageImageReference: + """Reconstruct a PageImageReference from ``to_dict`` output. + + Round-trips with ``to_dict``: ``from_dict(ref.to_dict()) == ref``. + """ + provider = data.get("provider") + return cls( + page_number=data["page_number"], + path=data["path"], + filename=data.get("filename"), + size_bytes=data.get("size_bytes"), + provider=FileStorageProvider(provider) if provider is not None else None, + ) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/constants.py b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/constants.py index 090a3bf6f4..85b4cf8900 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/constants.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/constants.py @@ -12,6 +12,7 @@ class Modes(Enum): class OutputModes(Enum): LAYOUT_PRESERVING = "layout_preserving" TEXT = "text" + IMAGE = "image" class HTTPMethod(Enum): @@ -31,6 +32,12 @@ class WhispererEndpoint: STATUS = "whisper-status" RETRIEVE = "whisper-retrieve" HIGHLIGHTS = "highlights" + # Image output mode (pdf-to-images) endpoints. These are NOT exposed by the + # llmwhisperer-client package, so the adapter calls them via raw requests + # (decision 2A). See ImageOutputConfig for the assumed service contract. + PDF_TO_IMAGES = "pdf-to-images" + PDF_TO_IMAGES_STATUS = "pdf-to-images-status" + PDF_TO_IMAGES_RETRIEVE = "pdf-to-images-retrieve" class WhispererEnv: @@ -47,6 +54,16 @@ class WhispererEnv: MAX_RETRIES = "ADAPTER_LLMW_MAX_RETRIES" RETRY_MIN_WAIT = "ADAPTER_LLMW_RETRY_MIN_WAIT" RETRY_MAX_WAIT = "ADAPTER_LLMW_RETRY_MAX_WAIT" + # Max retry attempts for per-page FileStorage writes when persisting page + # images (image output mode). Applies to Unstract-side storage writes only, + # not to calls made to the LLMWhisperer service. + PAGE_STORE_MAX_RETRIES = "ADAPTER_LLMW_PAGE_STORE_MAX_RETRIES" + # Image output mode HTTP tuning. Submit/status calls use a short timeout; + # the ZIP download uses a distinct, longer timeout (large multi-page PDFs). + IMAGE_REQUEST_TIMEOUT = "ADAPTER_LLMW_IMAGE_REQUEST_TIMEOUT" + IMAGE_DOWNLOAD_TIMEOUT = "ADAPTER_LLMW_IMAGE_DOWNLOAD_TIMEOUT" + IMAGE_POLL_INTERVAL = "ADAPTER_LLMW_IMAGE_POLL_INTERVAL" + IMAGE_POLL_MAX_ATTEMPTS = "ADAPTER_LLMW_IMAGE_POLL_MAX_ATTEMPTS" LOG_LEVEL = "LOG_LEVEL" @@ -114,3 +131,68 @@ class WhispererDefaults: MAX_RETRIES = int(os.getenv(WhispererEnv.MAX_RETRIES, 3)) RETRY_MIN_WAIT = float(os.getenv(WhispererEnv.RETRY_MIN_WAIT, 1.0)) RETRY_MAX_WAIT = float(os.getenv(WhispererEnv.RETRY_MAX_WAIT, 60.0)) + PAGE_STORE_MAX_RETRIES = int(os.getenv(WhispererEnv.PAGE_STORE_MAX_RETRIES, 3)) + IMAGE_REQUEST_TIMEOUT = int(os.getenv(WhispererEnv.IMAGE_REQUEST_TIMEOUT, 30)) + IMAGE_DOWNLOAD_TIMEOUT = int(os.getenv(WhispererEnv.IMAGE_DOWNLOAD_TIMEOUT, 300)) + IMAGE_POLL_INTERVAL = float(os.getenv(WhispererEnv.IMAGE_POLL_INTERVAL, 3.0)) + IMAGE_POLL_MAX_ATTEMPTS = int(os.getenv(WhispererEnv.IMAGE_POLL_MAX_ATTEMPTS, 100)) + + +class ImageOutputConfig: + """Config and service contract for LLMWhisperer image output mode. + + CONTRACT SOURCE: verified against LLMWhisperer Service **PR #536** (branch + ``image-output``; PR #647 is a sub-fix). The endpoints are NOT exposed by + the installed ``llmwhisperer-client``, so the adapter calls them via raw + ``requests`` (decision 2A). Everything the adapter relies on is centralised + here. + + Flow (raw ``requests``, base = ``{url}/api/v2``): + + - Submit: ``POST {base}/pdf-to-images?format=png`` with the PDF bytes + -> JSON ``{"message": "...", "status": "processing", + "whisper_hash": "|"}`` (HTTP 202) + - Status: ``GET {base}/pdf-to-images-status?whisper_hash=`` + -> JSON ``{"status": "accepted|processing|processed|...", + "message": "..."}``. NOTE: no page count is exposed + (page count is billing-internal only). + - Retrieve: ``GET {base}/pdf-to-images-retrieve?whisper_hash=`` + -> ``application/zip`` stream of ``page_001.png``, ... + ONE-TIME by default: the service flips status to ``RETRIEVED`` + before streaming and rejects a second retrieve unless the + deployment sets ``RESULT_PERSISTENCE=true``. Hence the adapter + downloads exactly once and never retries the retrieve. + """ + + # --- Response field names --- + STATUS = "status" + # Not currently returned by pdf-to-images-status (billing-internal). Kept as + # a forward-compatible hook for verify_page_count(). + PROCESSED_PAGE_COUNT = "processed_page_count" + MESSAGE = "message" + + # Terminal service states. Ready-to-retrieve == PROCESSED (WhisperStatus). + STATUS_SUCCESS = frozenset({"processed"}) + STATUS_FAILURE = frozenset({"error", "failed", "unknown"}) + + # --- Submit query params --- + IMAGE_FORMAT_PARAM = "format" + DEFAULT_IMAGE_FORMAT = "png" + FILE_NAME_PARAM = "file_name" + + # --- Per-page image naming / storage layout --- + PAGE_IMAGE_PREFIX = "page_" + PAGE_IMAGE_EXTENSION = ".png" + PAGE_NUMBER_PADDING = 3 + PAGES_SUBFOLDER = "pages" + + # --- UI / validation (single source of truth) --- + # Display label for the image output mode option (UNS-754). + IMAGE_MODE_LABEL = "Image (PDF only)" + PDF_EXTENSION = ".pdf" + # Shared by runtime (process) and UI validation so the message is identical + # regardless of where the PDF-only check fires (UNS-757). + PDF_ONLY_ERROR = ( + "Image output mode supports PDF input only. " + "Please provide a PDF file or select a text output mode." + ) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/helper.py b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/helper.py index ade89f7cba..e9b8b6f04e 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/helper.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/helper.py @@ -1,5 +1,8 @@ import json import logging +import re +import time +import zipfile from io import BytesIO from pathlib import Path from typing import Any @@ -11,14 +14,18 @@ LLMWhispererClientException, LLMWhispererClientV2, ) + from unstract.sdk1.adapters.exceptions import ExtractorError from unstract.sdk1.adapters.utils import AdapterUtils from unstract.sdk1.adapters.x2text.constants import X2TextConstants +from unstract.sdk1.adapters.x2text.dto import PageImageReference from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.constants import ( + ImageOutputConfig, Modes, OutputModes, WhispererConfig, WhispererDefaults, + WhispererEndpoint, WhispererHeader, WhisperStatus, ) @@ -26,7 +33,9 @@ WhispererRequestParams, ) from unstract.sdk1.constants import MimeType +from unstract.sdk1.exceptions import FileOperationError from unstract.sdk1.file_storage import FileStorage, FileStorageProvider +from unstract.sdk1.utils.retry_utils import retry_with_exponential_backoff logger = logging.getLogger(__name__) @@ -45,17 +54,41 @@ def get_request_headers(config: dict[str, Any]) -> dict[str, Any]: } @staticmethod - def test_connection_request( - config: dict[str, Any], request_endpoint: str + def _send_raw_request( + config: dict[str, Any], + method: str, + endpoint: str, + *, + params: dict[str, Any] | None = None, + data: BytesIO | None = None, + headers: dict[str, Any] | None = None, + timeout: float | None = None, + stream: bool = False, ) -> Response: - llm_whisperer_svc_url = f"{config.get(WhispererConfig.URL)}/api/v2" - headers = LLMWhispererHelper.get_request_headers(config=config) + """Single outbound raw-``requests`` code path for the adapter (UNS-743). + Resolves the service base URL and auth headers from ``config`` so that + no caller constructs URLs or headers itself, issues the request with an + explicit timeout, and maps transport / HTTP failures to ``ExtractorError`` + with the same semantics used across the adapter. Both ``test_connection`` + and the ``pdf-to-images`` image-mode calls go through here. + """ + llm_whisperer_svc_url = f"{config.get(WhispererConfig.URL)}/api/v2" + url = f"{llm_whisperer_svc_url}/{endpoint}" + if headers is None: + headers = LLMWhispererHelper.get_request_headers(config=config) try: - response: Response - url = f"{llm_whisperer_svc_url}/{request_endpoint}" - response = requests.get(url=url, headers=headers) + response = requests.request( + method=method, + url=url, + headers=headers, + params=params, + data=data, + timeout=timeout, + stream=stream, + ) response.raise_for_status() + return response except ConnectionError as e: logger.error(f"Adapter error: {e}") raise ExtractorError( @@ -77,6 +110,16 @@ def test_connection_request( msg, status_code=e.response.status_code, actual_err=e ) from e + @staticmethod + def test_connection_request( + config: dict[str, Any], request_endpoint: str + ) -> Response: + return LLMWhispererHelper._send_raw_request( + config=config, + method="GET", + endpoint=request_endpoint, + ) + @staticmethod def make_request( config: dict[str, Any], @@ -390,3 +433,343 @@ def write_output_to_file( ) except Exception as e: logger.warn(f"Error while writing metadata to {metadata_file_path}: {e}") + + # ------------------------------------------------------------------ # + # Image output mode (pdf-to-images). # + # # + # These call the LLMWhisperer `pdf-to-images` endpoints via raw # + # `requests` (decision 2A). The exact endpoint/response contract is # + # centralised in ImageOutputConfig — see its docstring; it is an # + # ASSUMED contract (Service PR #647 is not available in this repo) # + # and is the single place to reconcile once the real API is known. # + # ------------------------------------------------------------------ # + + # Matches service page files like `page_001.png` / `page-1.png`. + _PAGE_IMAGE_RE = re.compile(r"page[_-]?0*(\d+)\.png$", re.IGNORECASE) + + @staticmethod + def _safe_json(response: Response) -> dict[str, Any]: + """Parse a JSON object body, tolerating non-JSON / non-object bodies.""" + try: + parsed = response.json() + except ValueError: + return {} + return parsed if isinstance(parsed, dict) else {} + + @staticmethod + def submit_pdf_to_images( + config: dict[str, Any], + file_data: BytesIO, + tag: str | list[str] | None = None, + file_name: str | None = None, + ) -> str: + """Submit a ``pdf-to-images`` job; returns the job id (whisper_hash). + + The image ``format``, ``tag`` (usage-report label) and ``file_name`` are + sent as query params — consistent with the ``/whisper`` endpoint so the + service attributes usage correctly (verified against Service PR #536). + ``tag`` falls back to the adapter config, then the default. + """ + resolved_tag = WhispererRequestParams(tag=tag).tag or config.get( + WhispererConfig.TAG, WhispererDefaults.TAG + ) + params: dict[str, Any] = { + ImageOutputConfig.IMAGE_FORMAT_PARAM: ImageOutputConfig.DEFAULT_IMAGE_FORMAT, + WhispererConfig.TAG: resolved_tag, + } + if file_name: + params[ImageOutputConfig.FILE_NAME_PARAM] = file_name + response = LLMWhispererHelper._send_raw_request( + config=config, + method="POST", + endpoint=WhispererEndpoint.PDF_TO_IMAGES, + params=params, + data=file_data, + timeout=WhispererDefaults.IMAGE_REQUEST_TIMEOUT, + ) + body = LLMWhispererHelper._safe_json(response) + whisper_hash = body.get(X2TextConstants.WHISPER_HASH_V2, "") + if not whisper_hash: + raise ExtractorError( + "LLMWhisperer pdf-to-images submit did not return a job id " + f"(whisper_hash). Response: {body}", + status_code=502, + ) + logger.info("Image mode: submitted pdf-to-images job %s", whisper_hash) + return whisper_hash + + @staticmethod + def poll_pdf_to_images_status( + config: dict[str, Any], whisper_hash: str + ) -> dict[str, Any]: + """Poll the status endpoint until a terminal state is reached. + + Returns the terminal status payload on success (``status`` reaches + ``PROCESSED``); raises ``ExtractorError`` on a failed/unknown state or + once the poll budget is exhausted. Mirrors the submit-then-poll pattern + already used for text extraction. + """ + headers = LLMWhispererHelper.get_request_headers(config) + params = {WhisperStatus.WHISPER_HASH: whisper_hash} + for attempt in range(WhispererDefaults.IMAGE_POLL_MAX_ATTEMPTS): + response = LLMWhispererHelper._send_raw_request( + config=config, + method="GET", + endpoint=WhispererEndpoint.PDF_TO_IMAGES_STATUS, + params=params, + headers=headers, + timeout=WhispererDefaults.IMAGE_REQUEST_TIMEOUT, + ) + body = LLMWhispererHelper._safe_json(response) + status = str(body.get(ImageOutputConfig.STATUS, "")).lower() + logger.info( + "Image mode: job %s status=%s (attempt %d/%d)", + whisper_hash, + status, + attempt + 1, + WhispererDefaults.IMAGE_POLL_MAX_ATTEMPTS, + ) + if status in ImageOutputConfig.STATUS_SUCCESS: + return body + if status in ImageOutputConfig.STATUS_FAILURE: + msg = body.get(ImageOutputConfig.MESSAGE, "unknown error") + raise ExtractorError( + f"LLMWhisperer pdf-to-images job {whisper_hash} failed: {msg}", + status_code=500, + ) + # Intermediate states (processing / queued / empty) -> keep polling. + time.sleep(WhispererDefaults.IMAGE_POLL_INTERVAL) + raise ExtractorError( + f"LLMWhisperer pdf-to-images job {whisper_hash} did not reach a " + f"terminal state within {WhispererDefaults.IMAGE_POLL_MAX_ATTEMPTS} " + "poll attempts", + status_code=504, + ) + + @staticmethod + def download_pdf_to_images_zip(config: dict[str, Any], whisper_hash: str) -> BytesIO: + """Stream the page-image ZIP into an in-memory buffer via chunked reads. + + Uses a distinct, longer download timeout (large multi-page PDFs) and + avoids a single ``response.content`` load. + """ + response = LLMWhispererHelper._send_raw_request( + config=config, + method="GET", + endpoint=WhispererEndpoint.PDF_TO_IMAGES_RETRIEVE, + params={WhisperStatus.WHISPER_HASH: whisper_hash}, + timeout=WhispererDefaults.IMAGE_DOWNLOAD_TIMEOUT, + stream=True, + ) + buffer = BytesIO() + for chunk in response.iter_content(chunk_size=1024 * 1024): + if chunk: + buffer.write(chunk) + buffer.seek(0) + return buffer + + @staticmethod + def extract_page_images_from_zip( + zip_buffer: BytesIO, + ) -> list[tuple[int, bytes]]: + """Extract page images from the ZIP, ordered ascending by page number. + + Returns ``[(page_number, image_bytes), ...]``. Raises ``ExtractorError`` + on a corrupt/invalid archive. + """ + pages: list[tuple[int, bytes]] = [] + try: + with zipfile.ZipFile(zip_buffer) as archive: + for name in archive.namelist(): + match = LLMWhispererHelper._PAGE_IMAGE_RE.search(name) + if not match: + continue + page_number = int(match.group(1)) + pages.append((page_number, archive.read(name))) + except zipfile.BadZipFile as e: + raise ExtractorError( + f"Corrupt or invalid ZIP received from pdf-to-images: {e}", + status_code=502, + actual_err=e, + ) from e + pages.sort(key=lambda item: item[0]) + return pages + + @staticmethod + def verify_page_count( + pages: list[tuple[int, bytes]], processed_page_count: int | None + ) -> None: + """Enforce a matching page count against the service's authority. + + The service-reported ``processed_page_count`` is authoritative + (UNS-746); any mismatch with the extracted count raises. + """ + if processed_page_count is None: + # Expected today: the pdf-to-images-status response does not expose a + # page count (verified vs PR #536). Kept as a forward-compatible hook. + logger.debug( + "Image mode: no processed_page_count in status response; " + "skipping page-count verification" + ) + return + actual = len(pages) + if actual != processed_page_count: + raise ExtractorError( + "Page count mismatch in image output mode: service reported " + f"processed_page_count={processed_page_count} but the extracted " + f"ZIP contained {actual} page image(s)", + status_code=502, + ) + + @staticmethod + def build_page_store_dir( + output_file_path: str | None, input_file_path: str, run_key: str + ) -> str: + """Collision-safe per-document folder for page images (UNS-747). + + ``run_key`` (the unique per-run whisper_hash) isolates every extraction, + so concurrent documents never share a prefix. Layout: + ``{base_dir}/{run_key}/pages``. + """ + reference = output_file_path or input_file_path + base_dir = str(Path(reference).parent) if reference else "." + return str(Path(base_dir) / run_key / ImageOutputConfig.PAGES_SUBFOLDER) + + @staticmethod + def _page_image_filename(page_number: int) -> str: + padded = str(page_number).zfill(ImageOutputConfig.PAGE_NUMBER_PADDING) + return ( + f"{ImageOutputConfig.PAGE_IMAGE_PREFIX}{padded}" + f"{ImageOutputConfig.PAGE_IMAGE_EXTENSION}" + ) + + @staticmethod + def _write_single_page(fs: FileStorage, path: str, data: bytes) -> None: + fs.write(path=path, mode="wb", data=data, encoding="utf-8") + + @staticmethod + def persist_page_images( + fs: FileStorage, + page_store_dir: str, + pages: list[tuple[int, bytes]], + ) -> list[PageImageReference]: + """Write every page image to FileStorage with per-page retry. + + All-or-nothing (fail-closed): the full ``PageImageReference`` list is + only returned once EVERY page is written. If any page exhausts its + retries, a hard ``ExtractorError`` propagates and no partial set is + returned (UNS-738 / UNS-739 / UNS-745). Works transparently for LOCAL + and S3 via the passed ``fs``. + """ + fs.mkdir(create_parents=True, path=page_store_dir) + + write_with_retry = retry_with_exponential_backoff( + max_retries=WhispererDefaults.PAGE_STORE_MAX_RETRIES, + base_delay=WhispererDefaults.RETRY_MIN_WAIT, + multiplier=2.0, + jitter=True, + exceptions=(FileOperationError, OSError), + logger_instance=logger, + prefix="LLMW_PAGE_STORE", + )(LLMWhispererHelper._write_single_page) + + references: list[PageImageReference] = [] + for page_number, data in pages: + filename = LLMWhispererHelper._page_image_filename(page_number) + path = str(Path(page_store_dir) / filename) + try: + write_with_retry(fs=fs, path=path, data=data) + except Exception as e: + raise ExtractorError( + "Failed to persist page image after retries: " + f"page={page_number}, provider={fs.provider.value}, " + f"path={path}", + status_code=500, + actual_err=e, + ) from e + references.append( + PageImageReference( + page_number=page_number, + path=path, + filename=filename, + size_bytes=len(data), + provider=fs.provider, + ) + ) + references.sort(key=lambda ref: ref.page_number) + logger.info( + "Image mode: persisted %d page image(s) under %s (provider=%s)", + len(references), + page_store_dir, + fs.provider.value, + ) + return references + + @staticmethod + def _download_and_extract( + config: dict[str, Any], whisper_hash: str + ) -> list[tuple[int, bytes]]: + zip_buffer = LLMWhispererHelper.download_pdf_to_images_zip(config, whisper_hash) + return LLMWhispererHelper.extract_page_images_from_zip(zip_buffer) + + @staticmethod + def get_page_images( + config: dict[str, Any], + input_file_path: str, + output_file_path: str | None, + fs: FileStorage | None = None, + tag: str | list[str] | None = None, + ) -> list[PageImageReference]: + """End-to-end image output flow (orchestrator). + + submit -> poll -> download+extract (ONCE) -> verify page count -> + persist per-page (retried). Returns the ordered ``PageImageReference`` + list, or raises (fail-closed — never partial). + + Retrieval is intentionally NOT retried. Verified against Service + PR #536: ``pdf-to-images-retrieve`` flips the job to ``RETRIEVED`` + *before* streaming and, with the service default + ``RESULT_PERSISTENCE=false``, a second retrieve returns + 400 "Result already retrieved". Re-downloading is therefore impossible + (and re-submitting would double-bill), so a mid-download failure is a + hard error — the job must be resubmitted by the caller. Per-page + FileStorage writes (Unstract-side) are still retried. + """ + if fs is None: + fs = FileStorage(provider=FileStorageProvider.LOCAL) + + input_data = BytesIO(fs.read(path=input_file_path, mode="rb")) + whisper_hash = LLMWhispererHelper.submit_pdf_to_images( + config, + input_data, + tag=tag, + file_name=Path(input_file_path).name, + ) + status_payload = LLMWhispererHelper.poll_pdf_to_images_status( + config, whisper_hash + ) + # NOTE (verified vs PR #536): the status response does NOT expose a page + # count today — it is billing-internal (pdfToImagesPageCount column). + # verify_page_count() therefore no-ops unless/until the service adds it. + processed_page_count = status_payload.get(ImageOutputConfig.PROCESSED_PAGE_COUNT) + + pages = LLMWhispererHelper._download_and_extract( + config=config, whisper_hash=whisper_hash + ) + + # Verify BEFORE persisting so nothing is written on a count mismatch. + LLMWhispererHelper.verify_page_count(pages, processed_page_count) + + page_store_dir = LLMWhispererHelper.build_page_store_dir( + output_file_path=output_file_path, + input_file_path=input_file_path, + run_key=whisper_hash, + ) + references = LLMWhispererHelper.persist_page_images(fs, page_store_dir, pages) + logger.info( + "Image mode: completed job=%s pages=%d processed_page_count=%s", + whisper_hash, + len(references), + processed_page_count, + ) + return references diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/llm_whisperer_v2.py b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/llm_whisperer_v2.py index 3a48a57647..07edcf1e84 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/llm_whisperer_v2.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/llm_whisperer_v2.py @@ -2,14 +2,19 @@ import logging import os +from pathlib import Path from typing import TYPE_CHECKING, Any +from unstract.sdk1.adapters.exceptions import ExtractorError from unstract.sdk1.adapters.x2text.constants import X2TextConstants from unstract.sdk1.adapters.x2text.dto import ( TextExtractionMetadata, TextExtractionResult, ) from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.constants import ( + ImageOutputConfig, + OutputModes, + WhispererConfig, WhispererEndpoint, ) from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.dto import ( @@ -61,6 +66,56 @@ def test_connection(self) -> bool: ) return True + @staticmethod + def _validate_pdf_only(input_file_path: str) -> None: + """Enforce the PDF-only constraint for image output mode (v1). + + The message is sourced from ``ImageOutputConfig`` so it stays identical + to the UI-layer validation surfaced in ``adapter_processor_v2``. + """ + if Path(input_file_path).suffix.lower() != ImageOutputConfig.PDF_EXTENSION: + raise ExtractorError( + ImageOutputConfig.PDF_ONLY_ERROR, + status_code=400, + ) + + def _process_image_mode( + self, + input_file_path: str, + output_file_path: str | None, + fs: FileStorage, + tag: str | list[str] | None = None, + ) -> TextExtractionResult: + """Image output mode branch of ``process()``. + + Validates PDF-only input, delegates the submit/download/persist flow to + the helper, and returns a ``TextExtractionResult`` whose ``page_images`` + metadata carries the per-page references. ``extracted_text`` is an empty + plain string (never JSON / never image data) so text-mode consumers + remain unaffected. ``tag`` is forwarded for service-side usage reporting. + """ + logger.info("Image mode: processing %s in image output mode", input_file_path) + self._validate_pdf_only(input_file_path) + page_images = LLMWhispererHelper.get_page_images( + config=self.config, + input_file_path=input_file_path, + output_file_path=output_file_path, + fs=fs, + tag=tag, + ) + logger.info( + "Image mode: returning %d page image reference(s) for %s", + len(page_images), + input_file_path, + ) + return TextExtractionResult( + extracted_text="", + extraction_metadata=TextExtractionMetadata( + whisper_hash="", + page_images=page_images, + ), + ) + def process( self, input_file_path: str, @@ -81,6 +136,20 @@ def process( """ if fs is None: fs = FileStorage(provider=FileStorageProvider.LOCAL) + + # Branch on the configured output mode. Image mode routes to a dedicated + # path (PDF-only); every other mode follows the unchanged text path. + output_mode = self.config.get( + WhispererConfig.OUTPUT_MODE, OutputModes.LAYOUT_PRESERVING.value + ) + if output_mode == OutputModes.IMAGE.value: + return self._process_image_mode( + input_file_path, + output_file_path, + fs, + tag=kwargs.get(X2TextConstants.TAGS), + ) + enable_highlight = kwargs.get(X2TextConstants.ENABLE_HIGHLIGHT, False) logger.info( "HIGHLIGHT_DEBUG LLMWhispererV2.process: enable_highlight=%s", diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/static/json_schema.json b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/static/json_schema.json index ef0a036d4d..05d0c66d32 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/static/json_schema.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/static/json_schema.json @@ -44,10 +44,16 @@ "title": "Output Mode", "enum": [ "layout_preserving", - "text" + "text", + "image" + ], + "enumNames": [ + "Layout Preserving", + "Text", + "Image (PDF only)" ], "default": "layout_preserving", - "description": "Output format, described in the [LLMWhisperer documentation](https://docs.unstract.com/llmwhisperer/llm_whisperer/apis/llm_whisperer_text_extraction_api/#output-modes)" + "description": "Output format, described in the [LLMWhisperer documentation](https://docs.unstract.com/llmwhisperer/llm_whisperer/apis/llm_whisperer_text_extraction_api/#output-modes). Note: **Image** mode returns per-page images instead of text and supports **PDF input only**." }, "line_splitter_tolerance": { "type": "number", diff --git a/unstract/sdk1/tests/llmw_image_fixtures.py b/unstract/sdk1/tests/llmw_image_fixtures.py new file mode 100644 index 0000000000..10133eb525 --- /dev/null +++ b/unstract/sdk1/tests/llmw_image_fixtures.py @@ -0,0 +1,149 @@ +"""Shared test fixtures and stubs for LLMWhisperer image output mode (UNS-762). + +Importable from multiple test modules:: + + from tests.llmw_image_fixtures import ( + make_page_zip, + CORRUPT_ZIP, + minimal_png, + InMemoryFileStorage, + FlakyFileStorage, + ) + +Provides: +- A happy-path ZIP builder producing ``page_00N.png`` entries with valid PNGs. +- Corrupt / non-ZIP byte fixtures for error-path testing. +- In-memory ``FileStorage`` doubles (S3-like) needing no network/credentials. +""" + +from __future__ import annotations + +import binascii +import io +import struct +import zipfile +import zlib + +from unstract.sdk1.exceptions import FileOperationError +from unstract.sdk1.file_storage import FileStorageProvider + + +def _png_chunk(tag: bytes, data: bytes) -> bytes: + crc = binascii.crc32(tag + data) & 0xFFFFFFFF + return struct.pack(">I", len(data)) + tag + data + struct.pack(">I", crc) + + +def minimal_png() -> bytes: + """Return the bytes of a valid 1x1 RGB PNG.""" + signature = b"\x89PNG\r\n\x1a\n" + ihdr = struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0) # 1x1, 8-bit RGB + raw = b"\x00\xff\x00\x00" # one scanline: filter byte 0 + red pixel + idat = zlib.compress(raw) + return ( + signature + + _png_chunk(b"IHDR", ihdr) + + _png_chunk(b"IDAT", idat) + + _png_chunk(b"IEND", b"") + ) + + +def make_page_zip( + num_pages: int, *, padding: int = 3, ext: str = ".png", shuffle: bool = False +) -> bytes: + """Build a ZIP of ``page_00N.png`` entries (valid PNGs). + + Args: + num_pages: Number of page images to include. + padding: Zero-padding width for the page number. + ext: File extension for each page entry. + shuffle: If True, write entries in reverse order (to prove the + extractor sorts, not relies on archive order). + """ + order = range(num_pages, 0, -1) if shuffle else range(1, num_pages + 1) + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + for i in order: + archive.writestr(f"page_{str(i).zfill(padding)}{ext}", minimal_png()) + return buffer.getvalue() + + +# A byte sequence that is not a valid ZIP archive. +CORRUPT_ZIP = b"this is definitely not a zip archive" + + +class InMemoryFileStorage: + """Minimal in-memory FileStorage double (S3-like), no network/credentials. + + Implements only the surface the image-mode helper uses: ``provider``, + ``mkdir``, ``write``, ``read``, ``exists``. + """ + + def __init__(self, provider: FileStorageProvider = FileStorageProvider.S3) -> None: + """Create an empty in-memory store for the given provider.""" + self.provider = provider + self._files: dict[str, bytes] = {} + self._dirs: set[str] = set() + self.write_calls = 0 + + def mkdir(self, path: str, create_parents: bool = True) -> None: + self._dirs.add(str(path)) + + def write( + self, + path: str, + mode: str = "wb", + encoding: str = "utf-8", + data: bytes | str = b"", + **_: object, + ) -> int: + self.write_calls += 1 + payload = data.encode(encoding) if isinstance(data, str) else bytes(data) + self._files[str(path)] = payload + return len(payload) + + def read( + self, path: str, mode: str = "rb", encoding: str = "utf-8", **_: object + ) -> bytes | str: + payload = self._files[str(path)] + return payload if "b" in mode else payload.decode(encoding) + + def exists(self, path: str) -> bool: + key = str(path) + return key in self._files or key in self._dirs + + @property + def stored_paths(self) -> list[str]: + return sorted(self._files) + + +class FlakyFileStorage(InMemoryFileStorage): + """In-memory double whose writes fail a configurable number of times. + + Used to exercise the per-page write retry loop and the fail-closed policy. + """ + + def __init__( + self, fail_times: int = 1, fail_always: bool = False, **kwargs: object + ) -> None: + """Configure how many writes per path fail before succeeding.""" + super().__init__(**kwargs) + self.fail_times = fail_times + self.fail_always = fail_always + self._attempts: dict[str, int] = {} + + def write( + self, + path: str, + mode: str = "wb", + encoding: str = "utf-8", + data: bytes | str = b"", + **kwargs: object, + ) -> int: + key = str(path) + self._attempts[key] = self._attempts.get(key, 0) + 1 + if self.fail_always or self._attempts[key] <= self.fail_times: + raise FileOperationError(f"simulated write failure for {key}") + return super().write(path, mode, encoding, data, **kwargs) + + def attempts_for(self, path: str) -> int: + return self._attempts.get(str(path), 0) diff --git a/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py b/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py new file mode 100644 index 0000000000..4eaafd4bd4 --- /dev/null +++ b/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py @@ -0,0 +1,46 @@ +"""Unit tests for LLMWhisperer v2 adapter constants (MUNS-193). + +Covers: +- UNS-732: OutputModes.IMAGE enum value. +- UNS-733: ADAPTER_LLMW_PAGE_STORE_MAX_RETRIES env-var-backed constant. +""" + +import importlib + +from _pytest.monkeypatch import MonkeyPatch + +from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src import constants as c + +_ENV_VAR = "ADAPTER_LLMW_PAGE_STORE_MAX_RETRIES" + + +class TestOutputModesImage: + def test_image_mode_value(self) -> None: + assert c.OutputModes.IMAGE.value == "image" + + def test_existing_modes_unchanged(self) -> None: + assert c.OutputModes.TEXT.value == "text" + assert c.OutputModes.LAYOUT_PRESERVING.value == "layout_preserving" + + +class TestPageStoreMaxRetries: + def test_env_var_name(self) -> None: + assert c.WhispererEnv.PAGE_STORE_MAX_RETRIES == _ENV_VAR + + def test_default_is_three(self, monkeypatch: MonkeyPatch) -> None: + monkeypatch.delenv(_ENV_VAR, raising=False) + reloaded = importlib.reload(c) + try: + assert reloaded.WhispererDefaults.PAGE_STORE_MAX_RETRIES == 3 + assert isinstance(reloaded.WhispererDefaults.PAGE_STORE_MAX_RETRIES, int) + finally: + importlib.reload(c) + + def test_reads_from_env(self, monkeypatch: MonkeyPatch) -> None: + monkeypatch.setenv(_ENV_VAR, "5") + reloaded = importlib.reload(c) + try: + assert reloaded.WhispererDefaults.PAGE_STORE_MAX_RETRIES == 5 + finally: + monkeypatch.delenv(_ENV_VAR, raising=False) + importlib.reload(c) diff --git a/unstract/sdk1/tests/test_llmw_image_helper.py b/unstract/sdk1/tests/test_llmw_image_helper.py new file mode 100644 index 0000000000..eb8fe9fec6 --- /dev/null +++ b/unstract/sdk1/tests/test_llmw_image_helper.py @@ -0,0 +1,193 @@ +"""Unit tests for the LLMWhisperer v2 image-output helper (MUNS-194 / 196). + +Covers ZIP extraction/ordering, corrupt-ZIP handling, page-count verification, +collision-safe folder keys, zero-padded naming, FileStorage persistence with +retry + fail-closed semantics, and write/read round-trip content fidelity. + +All tests are pure in-memory / temp-dir units: no network, no live service. +""" + +import io + +import pytest +from _pytest.monkeypatch import MonkeyPatch + +from tests.llmw_image_fixtures import ( + CORRUPT_ZIP, + FlakyFileStorage, + InMemoryFileStorage, + make_page_zip, + minimal_png, +) +from unstract.sdk1.adapters.exceptions import ExtractorError +from unstract.sdk1.adapters.x2text.dto import PageImageReference +from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src import constants as c +from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.helper import ( + LLMWhispererHelper, +) +from unstract.sdk1.file_storage import FileStorage, FileStorageProvider + +H = LLMWhispererHelper + + +class TestZipExtraction: + def test_extracts_all_pages_ordered(self) -> None: + pages = H.extract_page_images_from_zip(io.BytesIO(make_page_zip(3))) + assert [p for p, _ in pages] == [1, 2, 3] + assert all(data.startswith(b"\x89PNG") for _, data in pages) + + def test_orders_even_when_archive_unordered(self) -> None: + pages = H.extract_page_images_from_zip(io.BytesIO(make_page_zip(4, shuffle=True))) + assert [p for p, _ in pages] == [1, 2, 3, 4] + + def test_ignores_non_page_entries(self) -> None: + import zipfile + + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr("page_001.png", minimal_png()) + archive.writestr("readme.txt", b"not a page") + buffer.seek(0) + pages = H.extract_page_images_from_zip(buffer) + assert [p for p, _ in pages] == [1] + + def test_corrupt_zip_raises_extractor_error(self) -> None: + with pytest.raises(ExtractorError, match="Corrupt or invalid ZIP"): + H.extract_page_images_from_zip(io.BytesIO(CORRUPT_ZIP)) + + +class TestPageCountVerification: + def test_matching_count_passes(self) -> None: + pages = H.extract_page_images_from_zip(io.BytesIO(make_page_zip(2))) + H.verify_page_count(pages, processed_page_count=2) # no raise + + def test_fewer_pages_raises(self) -> None: + pages = H.extract_page_images_from_zip(io.BytesIO(make_page_zip(2))) + with pytest.raises(ExtractorError, match="Page count mismatch"): + H.verify_page_count(pages, processed_page_count=3) + + def test_more_pages_raises(self) -> None: + pages = H.extract_page_images_from_zip(io.BytesIO(make_page_zip(3))) + with pytest.raises(ExtractorError, match="Page count mismatch"): + H.verify_page_count(pages, processed_page_count=2) + + def test_none_count_skips_check(self) -> None: + pages = H.extract_page_images_from_zip(io.BytesIO(make_page_zip(2))) + H.verify_page_count(pages, processed_page_count=None) # no raise + + +class TestFolderKeyAndNaming: + def test_folder_key_isolates_runs(self) -> None: + dir_a = H.build_page_store_dir("/data/out.txt", "/data/in.pdf", "run-A") + dir_b = H.build_page_store_dir("/data/out.txt", "/data/in.pdf", "run-B") + assert dir_a != dir_b + assert "run-A" in dir_a and "run-B" in dir_b + assert dir_a.endswith("pages") + + def test_folder_key_deterministic_for_same_run(self) -> None: + assert H.build_page_store_dir( + "/data/out.txt", "/data/in.pdf", "run-A" + ) == H.build_page_store_dir("/data/out.txt", "/data/in.pdf", "run-A") + + def test_folder_falls_back_to_input_dir(self) -> None: + result = H.build_page_store_dir(None, "/docs/in.pdf", "job1") + assert result.startswith("/docs/") + assert "job1" in result + + @pytest.mark.parametrize( + ("page", "expected"), + [ + (1, "page_001.png"), + (9, "page_009.png"), + (42, "page_042.png"), + (100, "page_100.png"), + (1234, "page_1234.png"), + ], + ) + def test_zero_padding_consistency(self, page: int, expected: str) -> None: + assert H._page_image_filename(page) == expected + + +class TestPersistence: + def test_persists_all_pages_as_ordered_references(self) -> None: + fs = InMemoryFileStorage(provider=FileStorageProvider.S3) + pages = [(2, b"two"), (1, b"one"), (3, b"three")] + refs = H.persist_page_images(fs, "doc/pages", pages) + + assert [r.page_number for r in refs] == [1, 2, 3] + assert all(isinstance(r, PageImageReference) for r in refs) + assert refs[0].filename == "page_001.png" + assert refs[0].path == "doc/pages/page_001.png" + assert refs[0].size_bytes == len(b"one") + assert refs[0].provider is FileStorageProvider.S3 + assert len(fs.stored_paths) == 3 + + def test_retry_then_success(self, monkeypatch: MonkeyPatch) -> None: + monkeypatch.setattr(c.WhispererDefaults, "RETRY_MIN_WAIT", 0.0) + monkeypatch.setattr(c.WhispererDefaults, "PAGE_STORE_MAX_RETRIES", 3) + fs = FlakyFileStorage(fail_times=2) # succeeds on 3rd attempt + refs = H.persist_page_images(fs, "doc/pages", [(1, b"data")]) + assert len(refs) == 1 + assert fs.attempts_for("doc/pages/page_001.png") == 3 + + def test_fail_closed_when_retries_exhausted(self, monkeypatch: MonkeyPatch) -> None: + monkeypatch.setattr(c.WhispererDefaults, "RETRY_MIN_WAIT", 0.0) + monkeypatch.setattr(c.WhispererDefaults, "PAGE_STORE_MAX_RETRIES", 2) + fs = FlakyFileStorage(fail_always=True) + with pytest.raises(ExtractorError, match="Failed to persist page image"): + H.persist_page_images(fs, "doc/pages", [(1, b"a"), (2, b"b")]) + # Fail-closed: the second page is never attempted after the first fails. + assert fs.stored_paths == [] + + def test_local_write_read_round_trip(self, tmp_path) -> None: # noqa: ANN001 + fs = FileStorage(provider=FileStorageProvider.LOCAL) + page_dir = H.build_page_store_dir( + output_file_path=str(tmp_path / "out.txt"), + input_file_path=str(tmp_path / "in.pdf"), + run_key="job-xyz", + ) + original = [(1, minimal_png()), (2, b"second-page-bytes")] + refs = H.persist_page_images(fs, page_dir, original) + + for (page_number, data), ref in zip(original, refs, strict=True): + assert ref.page_number == page_number + round_tripped = fs.read(path=ref.path, mode="rb") + assert round_tripped == data + + +class TestSubmitParams: + """submit_pdf_to_images sends tag + file_name for service-side usage reports.""" + + _CONFIG = {"url": "u", "unstract_key": "k", "tag": "cfgtag"} + + def _patch(self, monkeypatch: MonkeyPatch) -> dict: + captured: dict = {} + monkeypatch.setattr( + H, "_send_raw_request", lambda **kw: captured.update(kw) or object() + ) + monkeypatch.setattr(H, "_safe_json", lambda _r: {"whisper_hash": "wh1"}) + return captured + + def test_explicit_tag_and_file_name_are_sent(self, monkeypatch: MonkeyPatch) -> None: + captured = self._patch(monkeypatch) + wh = H.submit_pdf_to_images( + self._CONFIG, io.BytesIO(b"pdf"), tag="mytag", file_name="doc.pdf" + ) + assert wh == "wh1" + params = captured["params"] + assert params["tag"] == "mytag" + assert params["file_name"] == "doc.pdf" + assert params["format"] == "png" + + def test_tag_falls_back_to_config_and_no_filename( + self, monkeypatch: MonkeyPatch + ) -> None: + captured = self._patch(monkeypatch) + H.submit_pdf_to_images(self._CONFIG, io.BytesIO(b"pdf")) + assert captured["params"]["tag"] == "cfgtag" + assert "file_name" not in captured["params"] + + def test_list_tag_is_normalized(self, monkeypatch: MonkeyPatch) -> None: + captured = self._patch(monkeypatch) + H.submit_pdf_to_images(self._CONFIG, io.BytesIO(b"pdf"), tag=["first", "second"]) + assert captured["params"]["tag"] == "first" diff --git a/unstract/sdk1/tests/test_llmw_v2_process_image.py b/unstract/sdk1/tests/test_llmw_v2_process_image.py new file mode 100644 index 0000000000..0a304f8d40 --- /dev/null +++ b/unstract/sdk1/tests/test_llmw_v2_process_image.py @@ -0,0 +1,130 @@ +"""Tests for LLMWhispererV2.process() output-mode branching (MUNS-195). + +Covers image/text branching (UNS-749), PDF-only validation (UNS-749/757), +image-mode result population (UNS-751), and the text-mode regression guarantee +that image logic is never triggered in text mode (UNS-753). + +Network and the image helper flow are stubbed — no live service is contacted. +""" + +import pytest +from _pytest.monkeypatch import MonkeyPatch + +from unstract.sdk1.adapters.exceptions import ExtractorError +from unstract.sdk1.adapters.x2text.dto import PageImageReference +from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.helper import ( + LLMWhispererHelper, +) +from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.llm_whisperer_v2 import ( + LLMWhispererV2, +) + +_BASE_CONFIG = {"url": "https://svc.example.com", "unstract_key": "key"} + + +def _adapter(**overrides: object) -> LLMWhispererV2: + return LLMWhispererV2({**_BASE_CONFIG, **overrides}) + + +class TestTextModeRegression: + def test_text_mode_follows_existing_path(self, monkeypatch: MonkeyPatch) -> None: + image_called = {"hit": False} + monkeypatch.setattr( + LLMWhispererHelper, + "send_whisper_request", + lambda **_: {"whisper_hash": "wh1", "line_metadata": [[1, 0, 10, 100]]}, + ) + monkeypatch.setattr( + LLMWhispererHelper, + "extract_text_from_response", + lambda *_a, **_k: "hello text", + ) + monkeypatch.setattr( + LLMWhispererHelper, + "get_page_images", + lambda **_: image_called.__setitem__("hit", True), + ) + + result = _adapter().process("in.pdf") + + assert result.extracted_text == "hello text" + assert result.extraction_metadata.whisper_hash == "wh1" + assert result.extraction_metadata.page_images is None + assert image_called["hit"] is False # image path never touched + + def test_text_mode_error_path_propagates(self, monkeypatch: MonkeyPatch) -> None: + def _boom(**_: object) -> None: + raise ExtractorError("service error", status_code=500) + + monkeypatch.setattr(LLMWhispererHelper, "send_whisper_request", _boom) + with pytest.raises(ExtractorError, match="service error"): + _adapter().process("in.pdf") + + +class TestImageModeBranch: + def test_populates_page_images_and_empty_text(self, monkeypatch: MonkeyPatch) -> None: + refs = [ + PageImageReference(page_number=1, path="d/pages/page_001.png"), + PageImageReference(page_number=2, path="d/pages/page_002.png"), + ] + captured: dict[str, object] = {} + + def _fake_get_page_images(**kwargs: object) -> list[PageImageReference]: + captured.update(kwargs) + return refs + + monkeypatch.setattr(LLMWhispererHelper, "get_page_images", _fake_get_page_images) + monkeypatch.setattr( + LLMWhispererHelper, + "send_whisper_request", + lambda **_: pytest.fail("text path must not run in image mode"), + ) + + result = _adapter(output_mode="image").process("in.pdf", "out.txt") + + assert result.extracted_text == "" # plain string, never JSON + assert result.extraction_metadata.page_images == refs + assert captured["input_file_path"] == "in.pdf" + assert captured["output_file_path"] == "out.txt" + + def test_empty_page_list_is_safe(self, monkeypatch: MonkeyPatch) -> None: + monkeypatch.setattr(LLMWhispererHelper, "get_page_images", lambda **_: []) + result = _adapter(output_mode="image").process("in.pdf") + assert result.extraction_metadata.page_images == [] + assert result.extracted_text == "" + + def test_tag_forwarded_to_helper(self, monkeypatch: MonkeyPatch) -> None: + captured: dict[str, object] = {} + + def _capture(**kwargs: object) -> list: + captured.update(kwargs) + return [] + + monkeypatch.setattr(LLMWhispererHelper, "get_page_images", _capture) + _adapter(output_mode="image").process("in.pdf", tags=["cust-42"]) + assert captured["tag"] == ["cust-42"] + + def test_pdf_extension_is_case_insensitive(self, monkeypatch: MonkeyPatch) -> None: + monkeypatch.setattr(LLMWhispererHelper, "get_page_images", lambda **_: []) + # Should not raise for an uppercase .PDF extension. + _adapter(output_mode="image").process("SCAN.PDF") + + +class TestPdfOnlyValidation: + def test_non_pdf_rejected_before_helper_runs(self, monkeypatch: MonkeyPatch) -> None: + image_called = {"hit": False} + monkeypatch.setattr( + LLMWhispererHelper, + "get_page_images", + lambda **_: image_called.__setitem__("hit", True) or [], + ) + with pytest.raises(ExtractorError, match="PDF input only"): + _adapter(output_mode="image").process("in.png") + assert image_called["hit"] is False + + def test_validate_pdf_only_accepts_pdf(self) -> None: + LLMWhispererV2._validate_pdf_only("/tmp/doc.pdf") # no raise + + def test_validate_pdf_only_rejects_other(self) -> None: + with pytest.raises(ExtractorError, match="PDF input only"): + LLMWhispererV2._validate_pdf_only("/tmp/doc.tiff") diff --git a/unstract/sdk1/tests/test_x2text_dto.py b/unstract/sdk1/tests/test_x2text_dto.py new file mode 100644 index 0000000000..ffea22a6d9 --- /dev/null +++ b/unstract/sdk1/tests/test_x2text_dto.py @@ -0,0 +1,127 @@ +"""Unit tests for x2text DTOs — image output mode extension (MUNS-193). + +Covers: +- UNS-730: PageImageReference dataclass shape. +- UNS-731: additive, non-breaking ``page_images`` field on + TextExtractionMetadata. +- UNS-734: non-breaking serialization + round-trip guarantees. +- UNS-735: PageImageReference.to_dict / from_dict helpers. + +All tests are pure in-memory unit tests: no live services, file storage, or +network calls. +""" + +from dataclasses import asdict + +from unstract.sdk1.adapters.x2text.dto import ( + PageImageReference, + TextExtractionMetadata, + TextExtractionResult, +) +from unstract.sdk1.file_storage import FileStorageProvider + + +def _serialize(obj: object) -> dict: + """Serialize a dataclass, omitting None-valued fields. + + Mirrors a None-omitting wire convention: optional fields left unset never + introduce new keys, which is precisely the non-breaking guarantee under + test for existing (text-mode) consumers. + """ + return {k: v for k, v in asdict(obj).items() if v is not None} + + +class TestNonBreakingSerialization: + """The additive ``page_images`` field must not change text-mode output.""" + + def test_text_mode_metadata_matches_baseline(self) -> None: + # Baseline = the exact key set produced before page_images existed. + baseline = {"whisper_hash": "abc123"} + meta = TextExtractionMetadata(whisper_hash="abc123") + + assert meta.page_images is None + assert _serialize(meta) == baseline + assert "page_images" not in _serialize(meta) + + def test_text_mode_metadata_full_fields_unchanged(self) -> None: + meta = TextExtractionMetadata( + whisper_hash="h", + line_metadata={"1": "x"}, + ) + assert _serialize(meta) == { + "whisper_hash": "h", + "line_metadata": {"1": "x"}, + } + + def test_result_default_serialization_unchanged(self) -> None: + result = TextExtractionResult(extracted_text="hello") + assert _serialize(result) == {"extracted_text": "hello"} + + +class TestImageModeRoundTrip: + """Metadata carrying page_images must round-trip losslessly.""" + + def test_metadata_with_page_images_round_trips(self) -> None: + original = TextExtractionMetadata( + whisper_hash="h", + page_images=[ + PageImageReference(page_number=1, path="doc/page_001.png"), + PageImageReference( + page_number=2, + path="doc/page_002.png", + filename="page_002.png", + size_bytes=2048, + provider=FileStorageProvider.S3, + ), + ], + ) + # Serialize to a wire form using the per-page to_dict helper... + wire = { + "whisper_hash": original.whisper_hash, + "page_images": [pi.to_dict() for pi in original.page_images], + } + # ...then deserialize back into an equivalent object. + restored = TextExtractionMetadata( + whisper_hash=wire["whisper_hash"], + page_images=[PageImageReference.from_dict(d) for d in wire["page_images"]], + ) + assert restored == original + + def test_metadata_page_images_none_round_trips(self) -> None: + original = TextExtractionMetadata(whisper_hash="h") + restored = TextExtractionMetadata(**asdict(original)) + assert restored == original + assert restored.page_images is None + + +class TestPageImageReferenceSerialization: + """to_dict/from_dict coverage for minimal and full field sets.""" + + def test_construct_with_required_fields_only(self) -> None: + ref = PageImageReference(page_number=5, path="p") + assert ref.page_number == 5 + assert ref.path == "p" + assert ref.filename is None + assert ref.size_bytes is None + assert ref.provider is None + + def test_minimal_round_trip(self) -> None: + ref = PageImageReference(page_number=1, path="doc/page_001.png") + restored = PageImageReference.from_dict(ref.to_dict()) + assert restored == ref + + def test_full_round_trip_serializes_provider_to_value(self) -> None: + ref = PageImageReference( + page_number=3, + path="doc/page_003.png", + filename="page_003.png", + size_bytes=4096, + provider=FileStorageProvider.LOCAL, + ) + as_dict = ref.to_dict() + # Enum is serialized to its string value for JSON-friendliness. + assert as_dict["provider"] == "local" + + restored = PageImageReference.from_dict(as_dict) + assert restored == ref + assert restored.provider is FileStorageProvider.LOCAL From bfc473dcf801b7c145d584979b9ac42b715de12b Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Sat, 25 Jul 2026 21:21:44 +0530 Subject: [PATCH 2/3] UN-2646 [FEAT] Conditional PDF-only guidance on image mode (UNS-759) Convert the adapter json_schema's single top-level if/then into an allOf so a second conditional can coexist with the existing low_cost one. Adds a PDF-only guidance note (type: null field, RJSF renders it as a labelled callout) shown only when output_mode == "image", guarded by required:[output_mode]. Purely UI/UX; no effect on submission, validation, or backend behaviour. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/static/json_schema.json | 77 ++++++++++++------- 1 file changed, 51 insertions(+), 26 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/static/json_schema.json b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/static/json_schema.json index 05d0c66d32..0f07626849 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/static/json_schema.json +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/static/json_schema.json @@ -117,35 +117,60 @@ "description": "Any metadata which should be sent to the webhook. This data is sent verbatim to the callback endpoint." } }, - "if": { - "anyOf": [ - { + "allOf": [ + { + "if": { + "anyOf": [ + { + "properties": { + "mode": { + "const": "low_cost" + } + } + } + ] + }, + "then": { "properties": { - "mode": { - "const": "low_cost" + "median_filter_size": { + "type": "integer", + "title": "Median Filter Size", + "default": 0, + "description": "The size of the median filter to use for pre-processing the image during OCR based extraction. Useful to eliminate scanning artifacts and low quality JPEG artifacts. Default is 0 if the value is not explicitly set. Available only in the Enterprise version." + }, + "gaussian_blur_radius": { + "type": "number", + "title": "Gaussian Blur Radius", + "default": 0.0, + "description": "The radius of the gaussian blur to use for pre-processing the image during OCR based extraction. Useful to eliminate noise from the image. Default is 0.0 if the value is not explicitly set. Available only in the Enterprise version." } - } + }, + "required": [ + "median_filter_size", + "gaussian_blur_radius" + ] } - ] - }, - "then": { - "properties": { - "median_filter_size": { - "type": "integer", - "title": "Median Filter Size", - "default": 0, - "description": "The size of the median filter to use for pre-processing the image during OCR based extraction. Useful to eliminate scanning artifacts and low quality JPEG artifacts. Default is 0 if the value is not explicitly set. Available only in the Enterprise version." + }, + { + "if": { + "properties": { + "output_mode": { + "const": "image" + } + }, + "required": [ + "output_mode" + ] }, - "gaussian_blur_radius": { - "type": "number", - "title": "Gaussian Blur Radius", - "default": 0.0, - "description": "The radius of the gaussian blur to use for pre-processing the image during OCR based extraction. Useful to eliminate noise from the image. Default is 0.0 if the value is not explicitly set. Available only in the Enterprise version." + "then": { + "properties": { + "image_pdf_only_notice": { + "type": "null", + "title": "Image output mode - PDF only", + "description": "Image output mode returns per-page images and supports **PDF input files only**. Non-PDF inputs are rejected before processing." + } + } } - }, - "required": [ - "median_filter_size", - "gaussian_blur_radius" - ] - } + } + ] } From 2a505bd739299135285c408012b58df6b17b9eb2 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Sat, 25 Jul 2026 15:52:30 +0000 Subject: [PATCH 3/3] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../adapters/x2text/llm_whisperer_v2/src/helper.py | 1 - .../sdk1/tests/test_llm_whisperer_v2_constants.py | 1 - unstract/sdk1/tests/test_llmw_image_helper.py | 14 +++++++------- unstract/sdk1/tests/test_llmw_v2_process_image.py | 1 - 4 files changed, 7 insertions(+), 10 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/helper.py b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/helper.py index e9b8b6f04e..c978d92abc 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/helper.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/helper.py @@ -14,7 +14,6 @@ LLMWhispererClientException, LLMWhispererClientV2, ) - from unstract.sdk1.adapters.exceptions import ExtractorError from unstract.sdk1.adapters.utils import AdapterUtils from unstract.sdk1.adapters.x2text.constants import X2TextConstants diff --git a/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py b/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py index 4eaafd4bd4..3c54e4154f 100644 --- a/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py +++ b/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py @@ -8,7 +8,6 @@ import importlib from _pytest.monkeypatch import MonkeyPatch - from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src import constants as c _ENV_VAR = "ADAPTER_LLMW_PAGE_STORE_MAX_RETRIES" diff --git a/unstract/sdk1/tests/test_llmw_image_helper.py b/unstract/sdk1/tests/test_llmw_image_helper.py index eb8fe9fec6..cf258cb5b8 100644 --- a/unstract/sdk1/tests/test_llmw_image_helper.py +++ b/unstract/sdk1/tests/test_llmw_image_helper.py @@ -11,6 +11,13 @@ import pytest from _pytest.monkeypatch import MonkeyPatch +from unstract.sdk1.adapters.exceptions import ExtractorError +from unstract.sdk1.adapters.x2text.dto import PageImageReference +from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src import constants as c +from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.helper import ( + LLMWhispererHelper, +) +from unstract.sdk1.file_storage import FileStorage, FileStorageProvider from tests.llmw_image_fixtures import ( CORRUPT_ZIP, @@ -19,13 +26,6 @@ make_page_zip, minimal_png, ) -from unstract.sdk1.adapters.exceptions import ExtractorError -from unstract.sdk1.adapters.x2text.dto import PageImageReference -from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src import constants as c -from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.helper import ( - LLMWhispererHelper, -) -from unstract.sdk1.file_storage import FileStorage, FileStorageProvider H = LLMWhispererHelper diff --git a/unstract/sdk1/tests/test_llmw_v2_process_image.py b/unstract/sdk1/tests/test_llmw_v2_process_image.py index 0a304f8d40..e4eaea1bfa 100644 --- a/unstract/sdk1/tests/test_llmw_v2_process_image.py +++ b/unstract/sdk1/tests/test_llmw_v2_process_image.py @@ -9,7 +9,6 @@ import pytest from _pytest.monkeypatch import MonkeyPatch - from unstract.sdk1.adapters.exceptions import ExtractorError from unstract.sdk1.adapters.x2text.dto import PageImageReference from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.helper import (