From 6c23ac9ada0df72897b08476d13e082cf90d2870 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Sat, 25 Jul 2026 20:28:06 +0530 Subject: [PATCH 01/24] 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 02/24] 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 03/24] [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 ( From dce48a6477979aa69dc6066e34515e0010bd4298 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Sun, 26 Jul 2026 00:37:31 +0530 Subject: [PATCH 04/24] UN-2646 [FIX] Gate image PDF-only guidance to image mode (UNS-759) The output_mode description carried the image PDF-only note unconditionally, so it rendered for every mode (Layout Preserving / Text too). Move it into an if/then/else on output_mode: the note shows only when Image is selected; other modes show the neutral documentation description. The previous type: null notice field never rendered (RJSF NullField emits nothing). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../src/static/json_schema.json | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 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 0f07626849..e86ceb0f6c 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 @@ -52,8 +52,7 @@ "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). Note: **Image** mode returns per-page images instead of text and supports **PDF input only**." + "default": "layout_preserving" }, "line_splitter_tolerance": { "type": "number", @@ -64,7 +63,7 @@ "line_splitter_strategy": { "type": "string", "title": "Line Splitter Strategy", - "default":"left-priority", + "default": "left-priority", "description": "An advanced option for customizing the line splitting process." }, "horizontal_stretch_factor": { @@ -164,10 +163,15 @@ }, "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." + "output_mode": { + "description": "**Image mode returns per-page images instead of text and accepts PDF input files only** — non-PDF inputs are rejected before processing." + } + } + }, + "else": { + "properties": { + "output_mode": { + "description": "Output format, described in the [LLMWhisperer documentation](https://docs.unstract.com/llmwhisperer/llm_whisperer/apis/llm_whisperer_text_extraction_api/#output-modes)." } } } From 2269a57d6c8683b3bc8096db43b037b349643010 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Sun, 26 Jul 2026 00:37:38 +0530 Subject: [PATCH 05/24] UN-2646 [FEAT] Fail-fast PDF-only check before index dispatch (UNS-757) When the profile's x2text adapter is in image output mode, reject non-PDF inputs in build_index_payload (the live pre-dispatch path) so the user gets the SDK's PDF_ONLY_ERROR at index time instead of an extraction failure inside the executor. Message is sourced from ImageOutputConfig.PDF_ONLY_ERROR so the early check and the SDK runtime guard stay in sync. Mirrored in the legacy index_document path for symmetry. Adds no-DB unit tests (10 cases). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../prompt_studio_helper.py | 44 +++++++++++ .../test_validate_image_output_pdf_only.py | 75 +++++++++++++++++++ 2 files changed, 119 insertions(+) create mode 100644 backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index 9746a26efe..183309dbce 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -77,6 +77,11 @@ ) from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt from unstract.core.pubsub_helper import LogPublisher +from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.constants import ( + ImageOutputConfig, + OutputModes, + WhispererConfig, +) from unstract.sdk1.constants import LogLevel from unstract.sdk1.exceptions import IndexingError, SdkError from unstract.sdk1.execution.context import ExecutionContext @@ -550,6 +555,12 @@ def build_index_payload( default_profile, request_user=request_user ) + # Fail fast when an image-output x2text adapter is paired with a + # non-PDF input, so the user sees the PDF-only message here (before the + # executor task is dispatched) rather than as an extraction failure + # inside the worker (UNS-757). + PromptStudioHelper._validate_image_output_pdf_only(default_profile, file_name) + # Common path decomposition used by extract, summarize, and index directory, filename = os.path.split(file_path) stem = os.path.splitext(filename)[0] @@ -1362,6 +1373,34 @@ def fetch_prompt_from_tool(tool_id: str) -> list[ToolStudioPrompt]: ).order_by(TSPKeys.SEQUENCE_NUMBER) return prompt_instances + @staticmethod + def _validate_image_output_pdf_only( + profile_manager: ProfileManager, file_name: str + ) -> None: + """Reject non-PDF inputs when the x2text adapter is in image mode. + + Image output mode (LLMWhisperer V2) supports PDF input only. The SDK + adapter enforces this at extraction time; this mirror-check runs at + index time so the user gets the identical PDF-only message before any + extraction work is dispatched. The message is sourced from the SDK + (``ImageOutputConfig.PDF_ONLY_ERROR``) so both layers stay in sync. + + Only fires when the adapter config carries ``output_mode == "image"``, + which is unique to LLMWhisperer V2 — other x2text adapters are + unaffected. + """ + x2text = profile_manager.x2text + if x2text is None: + return + metadata = x2text.metadata or {} + if metadata.get(WhispererConfig.OUTPUT_MODE) != OutputModes.IMAGE.value: + return + if not file_name.lower().endswith(ImageOutputConfig.PDF_EXTENSION): + raise IndexingAPIError( + detail=ImageOutputConfig.PDF_ONLY_ERROR, + status_code=400, + ) + @staticmethod def index_document( tool_id: str, @@ -1438,6 +1477,11 @@ def index_document( summary_profile, request_user=request_user ) + # Fail fast when an image-output x2text adapter is paired with a + # non-PDF input, so the user sees the PDF-only message at index time + # instead of after extraction is dispatched (UNS-757). + PromptStudioHelper._validate_image_output_pdf_only(default_profile, file_name) + fs_instance = EnvHelper.get_storage( storage_type=StorageType.PERMANENT, env_name=FileStorageKeys.PERMANENT_REMOTE_STORAGE, diff --git a/backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py b/backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py new file mode 100644 index 0000000000..1e99f0e8d0 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py @@ -0,0 +1,75 @@ +"""Unit tests for ``PromptStudioHelper._validate_image_output_pdf_only``. + +Pins the UNS-757 fail-fast guard: when the x2text adapter is in image +output mode, a non-PDF input must be rejected at index-build time with the +SDK's shared PDF-only message, so the user never has to wait for the +executor to fail the extraction. Every other combination must pass through. + +Unit tests: the real helper module is imported (Django is loaded by the +rig's test env) and the profile is a lightweight mock, so no database is +touched. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock + +import pytest + +from prompt_studio.prompt_studio_core_v2 import prompt_studio_helper as _psh_mod +from prompt_studio.prompt_studio_core_v2.exceptions import IndexingAPIError +from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.constants import ( + ImageOutputConfig, +) + +PromptStudioHelper = _psh_mod.PromptStudioHelper + + +def _profile(metadata: dict | None) -> MagicMock: + """A profile whose x2text adapter exposes ``metadata`` verbatim.""" + profile = MagicMock(name="ProfileManager") + profile.x2text.metadata = metadata + return profile + + +class TestImageModeRejectsNonPdf: + """Image output mode + non-PDF → IndexingAPIError(400, PDF-only).""" + + @pytest.mark.parametrize("file_name", ["statement.docx", "notes.txt", "a.png"]) + def test_non_pdf_raises(self, file_name: str) -> None: + with pytest.raises(IndexingAPIError) as exc_info: + PromptStudioHelper._validate_image_output_pdf_only( + _profile({"output_mode": "image"}), file_name + ) + assert exc_info.value.status_code == 400 + assert str(exc_info.value.detail) == ImageOutputConfig.PDF_ONLY_ERROR + + @pytest.mark.parametrize("file_name", ["statement.pdf", "STATEMENT.PDF"]) + def test_pdf_passes_case_insensitively(self, file_name: str) -> None: + # Must not raise for PDF inputs regardless of extension casing. + PromptStudioHelper._validate_image_output_pdf_only( + _profile({"output_mode": "image"}), file_name + ) + + +class TestNonImageModesUnaffected: + """Only image mode is gated; every other config is a no-op.""" + + @pytest.mark.parametrize( + "metadata", + [ + {"output_mode": "text"}, + {"output_mode": "layout_preserving"}, + {}, # e.g. a non-LLMWhisperer adapter with no output_mode + None, # adapter metadata absent entirely + ], + ) + def test_non_image_mode_passes_for_non_pdf(self, metadata: dict | None) -> None: + PromptStudioHelper._validate_image_output_pdf_only( + _profile(metadata), "statement.docx" + ) + + def test_missing_x2text_adapter_passes(self) -> None: + profile = MagicMock(name="ProfileManager") + profile.x2text = None + PromptStudioHelper._validate_image_output_pdf_only(profile, "statement.docx") From 4f3fe2baef3828ec3bdd164908a0aeb766554f48 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Mon, 27 Jul 2026 09:34:33 +0530 Subject: [PATCH 06/24] UN-2646 [FIX] Resolve SonarCloud findings on image-output code (UNS-758) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - helper.py: drop redundant 0* from the page-image regex — int() already strips leading zeros, so 0*(\\d+) was ambiguous and flagged for super-linear backtracking. Now a linear page[_-]?(\\d+)\\.png$. - Tests: hoist non-asserting setup out of pytest.raises blocks so each block has exactly one call that can throw (3 sites). - Test: split the deterministic build_page_store_dir assertion into two named locals so the reliability check no longer sees identical expressions on both sides of == (was flagged as a bug). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../sdk1/adapters/x2text/llm_whisperer_v2/src/helper.py | 6 ++++-- unstract/sdk1/tests/test_llmw_image_helper.py | 9 +++++---- unstract/sdk1/tests/test_llmw_v2_process_image.py | 6 ++++-- 3 files changed, 13 insertions(+), 8 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..d757ef5a14 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 @@ -444,8 +444,10 @@ def write_output_to_file( # 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) + # Matches service page files like `page_001.png` / `page-1.png`. The + # captured digits are passed through int() (leading zeros stripped there), + # so no separate `0*` prefix is needed — keeping the pattern linear. + _PAGE_IMAGE_RE = re.compile(r"page[_-]?(\d+)\.png$", re.IGNORECASE) @staticmethod def _safe_json(response: Response) -> dict[str, Any]: diff --git a/unstract/sdk1/tests/test_llmw_image_helper.py b/unstract/sdk1/tests/test_llmw_image_helper.py index eb8fe9fec6..ed16c9e362 100644 --- a/unstract/sdk1/tests/test_llmw_image_helper.py +++ b/unstract/sdk1/tests/test_llmw_image_helper.py @@ -52,8 +52,9 @@ def test_ignores_non_page_entries(self) -> None: assert [p for p, _ in pages] == [1] def test_corrupt_zip_raises_extractor_error(self) -> None: + corrupt = io.BytesIO(CORRUPT_ZIP) with pytest.raises(ExtractorError, match="Corrupt or invalid ZIP"): - H.extract_page_images_from_zip(io.BytesIO(CORRUPT_ZIP)) + H.extract_page_images_from_zip(corrupt) class TestPageCountVerification: @@ -85,9 +86,9 @@ def test_folder_key_isolates_runs(self) -> None: 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") + first = H.build_page_store_dir("/data/out.txt", "/data/in.pdf", "run-A") + second = H.build_page_store_dir("/data/out.txt", "/data/in.pdf", "run-A") + assert first == second def test_folder_falls_back_to_input_dir(self) -> None: result = H.build_page_store_dir(None, "/docs/in.pdf", "job1") diff --git a/unstract/sdk1/tests/test_llmw_v2_process_image.py b/unstract/sdk1/tests/test_llmw_v2_process_image.py index 0a304f8d40..3c4a6d4fea 100644 --- a/unstract/sdk1/tests/test_llmw_v2_process_image.py +++ b/unstract/sdk1/tests/test_llmw_v2_process_image.py @@ -57,8 +57,9 @@ def _boom(**_: object) -> None: raise ExtractorError("service error", status_code=500) monkeypatch.setattr(LLMWhispererHelper, "send_whisper_request", _boom) + adapter = _adapter() with pytest.raises(ExtractorError, match="service error"): - _adapter().process("in.pdf") + adapter.process("in.pdf") class TestImageModeBranch: @@ -118,8 +119,9 @@ def test_non_pdf_rejected_before_helper_runs(self, monkeypatch: MonkeyPatch) -> "get_page_images", lambda **_: image_called.__setitem__("hit", True) or [], ) + adapter = _adapter(output_mode="image") with pytest.raises(ExtractorError, match="PDF input only"): - _adapter(output_mode="image").process("in.png") + adapter.process("in.png") assert image_called["hit"] is False def test_validate_pdf_only_accepts_pdf(self) -> None: From 8d98bb540ff3752eef365e4350d5e6d23937bbd6 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Mon, 27 Jul 2026 09:56:28 +0530 Subject: [PATCH 07/24] UN-2646 [FIX] Address CodeRabbit/Greptile review on image-output helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - _send_raw_request: give timeout a finite default (IMAGE_REQUEST_TIMEOUT) so test_connection can no longer hang forever on a stuck endpoint. - download_pdf_to_images_zip: consume the stream in try/finally — map read-time transport errors to ExtractorError and always close the response so a mid-stream failure cannot leak the connection. - extract_page_images_from_zip: fail closed on an archive with no page images (was a silent empty success), and also catch RuntimeError (encrypted member) and zlib.error (corrupt member) alongside BadZipFile. - Tests: scope the constants env patches via monkeypatch.context() so the module reload no longer leaks a patched value into later tests; add coverage for the no-page-archive hard error and the finite default timeout. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../x2text/llm_whisperer_v2/src/helper.py | 35 ++++++++++++++++--- .../tests/test_llm_whisperer_v2_constants.py | 22 ++++++------ unstract/sdk1/tests/test_llmw_image_helper.py | 24 +++++++++++++ 3 files changed, 65 insertions(+), 16 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 643a4d2dd2..a37728d5dc 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 @@ -3,6 +3,7 @@ import re import time import zipfile +import zlib from io import BytesIO from pathlib import Path from typing import Any @@ -61,7 +62,7 @@ def _send_raw_request( params: dict[str, Any] | None = None, data: BytesIO | None = None, headers: dict[str, Any] | None = None, - timeout: float | None = None, + timeout: float = WhispererDefaults.IMAGE_REQUEST_TIMEOUT, stream: bool = False, ) -> Response: """Single outbound raw-``requests`` code path for the adapter (UNS-743). @@ -563,9 +564,23 @@ def download_pdf_to_images_zip(config: dict[str, Any], whisper_hash: str) -> Byt stream=True, ) buffer = BytesIO() - for chunk in response.iter_content(chunk_size=1024 * 1024): - if chunk: - buffer.write(chunk) + # Consume the stream inside try/finally: map read-time transport errors + # (ChunkedEncodingError / ConnectionError / read Timeout) to + # ExtractorError like the rest of the adapter, and always release the + # connection even if a chunk read fails mid-stream. + try: + for chunk in response.iter_content(chunk_size=1024 * 1024): + if chunk: + buffer.write(chunk) + except requests.RequestException as e: + logger.error(f"Error streaming pdf-to-images archive: {e}") + raise ExtractorError( + "Failed to download the pdf-to-images archive from LLMWhisperer", + status_code=502, + actual_err=e, + ) from e + finally: + response.close() buffer.seek(0) return buffer @@ -587,12 +602,22 @@ def extract_page_images_from_zip( continue page_number = int(match.group(1)) pages.append((page_number, archive.read(name))) - except zipfile.BadZipFile as e: + except (zipfile.BadZipFile, RuntimeError, zlib.error) as e: + # BadZipFile: not a ZIP. RuntimeError: encrypted member. + # zlib.error: corrupt compressed member surfaced by read(). raise ExtractorError( f"Corrupt or invalid ZIP received from pdf-to-images: {e}", status_code=502, actual_err=e, ) from e + if not pages: + # A well-formed archive with no recognizable page images is a + # failed extraction, not an empty success — fail closed, matching + # the rest of this flow. + raise ExtractorError( + "pdf-to-images returned an archive with no page images", + status_code=502, + ) pages.sort(key=lambda item: item[0]) return pages diff --git a/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py b/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py index 3c54e4154f..a4a8cfc725 100644 --- a/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py +++ b/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py @@ -27,19 +27,19 @@ 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: + # Reload under the scoped patch, then reload again after the env is + # restored so the module cache reflects the real environment and does + # not leak the patched value into later tests. + with monkeypatch.context() as patch: + patch.delenv(_ENV_VAR, raising=False) + reloaded = importlib.reload(c) assert reloaded.WhispererDefaults.PAGE_STORE_MAX_RETRIES == 3 assert isinstance(reloaded.WhispererDefaults.PAGE_STORE_MAX_RETRIES, int) - finally: - importlib.reload(c) + importlib.reload(c) def test_reads_from_env(self, monkeypatch: MonkeyPatch) -> None: - monkeypatch.setenv(_ENV_VAR, "5") - reloaded = importlib.reload(c) - try: + with monkeypatch.context() as patch: + patch.setenv(_ENV_VAR, "5") + reloaded = importlib.reload(c) assert reloaded.WhispererDefaults.PAGE_STORE_MAX_RETRIES == 5 - finally: - monkeypatch.delenv(_ENV_VAR, raising=False) - importlib.reload(c) + importlib.reload(c) diff --git a/unstract/sdk1/tests/test_llmw_image_helper.py b/unstract/sdk1/tests/test_llmw_image_helper.py index bc70b0ec41..e7512b09e4 100644 --- a/unstract/sdk1/tests/test_llmw_image_helper.py +++ b/unstract/sdk1/tests/test_llmw_image_helper.py @@ -56,6 +56,18 @@ def test_corrupt_zip_raises_extractor_error(self) -> None: with pytest.raises(ExtractorError, match="Corrupt or invalid ZIP"): H.extract_page_images_from_zip(corrupt) + def test_archive_with_no_page_entries_raises(self) -> None: + # A well-formed ZIP with no page_*.png entries is a failed extraction, + # not an empty success — must fail closed. + import zipfile + + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr("readme.txt", b"not a page") + buffer.seek(0) + with pytest.raises(ExtractorError, match="no page images"): + H.extract_page_images_from_zip(buffer) + class TestPageCountVerification: def test_matching_count_passes(self) -> None: @@ -192,3 +204,15 @@ 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" + + +class TestRequestDefaults: + """The shared raw-request path must never wait forever (UNS-758).""" + + def test_send_raw_request_has_finite_default_timeout(self) -> None: + import inspect + + default = inspect.signature(H._send_raw_request).parameters["timeout"].default + # A None default maps to requests' "wait forever"; test_connection relies + # on this default, so it must be a positive, finite number. + assert isinstance(default, int | float) and default > 0 From 9ec353818cb80f19c35663051559c6213b766844 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Mon, 27 Jul 2026 12:28:35 +0530 Subject: [PATCH 08/24] UN-2646 [FIX] Make image-mode results survive indexing + cache (UNS-758) Resolves the two pipeline gaps that capped Greptile confidence at 3/5: 1. page_images dropped in transit: the executor built the indexer payload from extracted_text only, so image-mode references never reached the indexer and the document indexed empty. The executor now forwards page_images (new IKeys.PAGE_IMAGES), mirroring the highlight_metadata pattern, so the references survive the transport. 2. Repeated remote conversions on re-run: image mode never wrote the extract file the Prompt Studio extraction cache gate requires (non-empty), so every re-run re-submitted the pdf-to-images job. Image mode now returns a short human-readable summary as extracted_text and writes it to the extract file (plus a durable .page_images.json manifest of the references), so the cache treats the conversion as complete and the indexed document is meaningful instead of blank. References are never inlined into the text. Tests: update the image-mode process assertions for the summary contract and add coverage for the extract-file/manifest write and the summary format (49 sdk1 image-output tests pass). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../x2text/llm_whisperer_v2/src/helper.py | 57 +++++++++++++++++++ .../llm_whisperer_v2/src/llm_whisperer_v2.py | 24 ++++++-- unstract/sdk1/tests/test_llmw_image_helper.py | 35 ++++++++++++ .../sdk1/tests/test_llmw_v2_process_image.py | 21 ++++++- workers/executor/executors/constants.py | 1 + workers/executor/executors/legacy_executor.py | 11 ++++ 6 files changed, 143 insertions(+), 6 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 a37728d5dc..aef0c01f4e 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 @@ -799,3 +799,60 @@ def get_page_images( processed_page_count, ) return references + + @staticmethod + def build_image_output_summary(page_images: list[PageImageReference]) -> str: + """Human-readable extract text for an image-mode result. + + Image mode produces no OCR text, but the Prompt Studio extraction cache + keys on a non-empty extract file and the indexer stores whatever text + the extraction yields. Returning a short summary (rather than an empty + string) keeps a re-run from re-submitting the remote conversion and + keeps the indexed document meaningful instead of blank. The per-page + references travel separately in ``extraction_metadata.page_images`` and + the JSON manifest — never inside this string. + """ + count = len(page_images) + noun = "page image" if count == 1 else "page images" + return ( + f"[LLMWhisperer image output mode] {count} {noun} extracted from the " + "PDF and stored in FileStorage. Per-page references are available in " + "the page_images extraction metadata and the accompanying manifest." + ) + + @staticmethod + def write_image_output( + fs: FileStorage, + output_file_path: str, + summary: str, + page_images: list[PageImageReference], + ) -> None: + """Persist the image-mode extract file + a page-image manifest sidecar. + + ``output_file_path`` (the extract file) receives the summary text so the + extraction-cache gate treats the conversion as complete (no re-submit) + and the indexer has meaningful text. ``.page_images. + json`` receives the ordered per-page references as a durable, retrievable + manifest. + """ + try: + fs.write( + path=str(output_file_path), + mode="w", + data=summary, + encoding="utf-8", + ) + manifest = json.dumps( + [ref.to_dict() for ref in page_images], + ensure_ascii=False, + indent=2, + ) + fs.write( + path=f"{output_file_path}.page_images.json", + mode="w", + data=manifest, + encoding="utf-8", + ) + except Exception as e: + logger.error(f"Error writing image output for {output_file_path}: {e}") + raise ExtractorError(str(e)) from e 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 07edcf1e84..6a25de32a6 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 @@ -90,9 +90,14 @@ def _process_image_mode( 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. + metadata carries the per-page references. ``extracted_text`` is a short + human-readable summary (never JSON / never image data): image mode has + no OCR text, but a non-empty extract keeps the Prompt Studio extraction + cache from re-submitting the remote conversion on a re-run and keeps the + indexed document meaningful. The per-page references live in + ``extraction_metadata.page_images`` (forwarded by the executor) and a + JSON manifest sidecar — never inside ``extracted_text``. ``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) @@ -103,13 +108,24 @@ def _process_image_mode( fs=fs, tag=tag, ) + summary = LLMWhispererHelper.build_image_output_summary(page_images) + # Persist the extract file (summary) + manifest so the extraction is + # cache-consistent (no re-submit on re-run) and the references are + # durably retrievable. Skipped when no output path was requested. + if output_file_path: + LLMWhispererHelper.write_image_output( + fs=fs, + output_file_path=output_file_path, + summary=summary, + page_images=page_images, + ) logger.info( "Image mode: returning %d page image reference(s) for %s", len(page_images), input_file_path, ) return TextExtractionResult( - extracted_text="", + extracted_text=summary, extraction_metadata=TextExtractionMetadata( whisper_hash="", page_images=page_images, diff --git a/unstract/sdk1/tests/test_llmw_image_helper.py b/unstract/sdk1/tests/test_llmw_image_helper.py index e7512b09e4..ca11a5e522 100644 --- a/unstract/sdk1/tests/test_llmw_image_helper.py +++ b/unstract/sdk1/tests/test_llmw_image_helper.py @@ -216,3 +216,38 @@ def test_send_raw_request_has_finite_default_timeout(self) -> None: # A None default maps to requests' "wait forever"; test_connection relies # on this default, so it must be a positive, finite number. assert isinstance(default, int | float) and default > 0 + + +class TestImageOutputWrite: + """write_image_output persists the summary extract file + manifest sidecar.""" + + def test_writes_summary_and_manifest(self, tmp_path) -> None: # noqa: ANN001 + import json as _json + + fs = FileStorage(provider=FileStorageProvider.LOCAL) + refs = [ + PageImageReference( + page_number=1, path="doc/pages/page_001.png", filename="page_001.png" + ), + PageImageReference( + page_number=2, path="doc/pages/page_002.png", filename="page_002.png" + ), + ] + out = str(tmp_path / "doc.txt") + summary = H.build_image_output_summary(refs) + + H.write_image_output( + fs=fs, output_file_path=out, summary=summary, page_images=refs + ) + + # Extract file holds the human summary (what image mode indexes). + assert fs.read(path=out, mode="r") == summary + # Sidecar manifest round-trips the ordered references. + manifest = _json.loads(fs.read(path=out + ".page_images.json", mode="r")) + assert manifest == [r.to_dict() for r in refs] + + def test_summary_is_human_readable_not_json(self) -> None: + refs = [PageImageReference(page_number=1, path="p/page_001.png")] + summary = H.build_image_output_summary(refs) + assert "1 page image" in summary + assert "page_001.png" not in summary # references never inlined diff --git a/unstract/sdk1/tests/test_llmw_v2_process_image.py b/unstract/sdk1/tests/test_llmw_v2_process_image.py index 2feedb51d8..e5cbdaa31b 100644 --- a/unstract/sdk1/tests/test_llmw_v2_process_image.py +++ b/unstract/sdk1/tests/test_llmw_v2_process_image.py @@ -79,19 +79,36 @@ def _fake_get_page_images(**kwargs: object) -> list[PageImageReference]: "send_whisper_request", lambda **_: pytest.fail("text path must not run in image mode"), ) + # Delegate the extract-file / manifest write; assert it is invoked + # rather than doing real file IO here. + write_calls: dict[str, object] = {} + monkeypatch.setattr( + LLMWhispererHelper, + "write_image_output", + lambda **kw: write_calls.update(kw), + ) result = _adapter(output_mode="image").process("in.pdf", "out.txt") - assert result.extracted_text == "" # plain string, never JSON + # extracted_text is a non-empty human summary — never JSON, never + # image data (the references live only in metadata / the manifest). + expected_summary = LLMWhispererHelper.build_image_output_summary(refs) + assert result.extracted_text == expected_summary + assert "page_001.png" not in result.extracted_text assert result.extraction_metadata.page_images == refs assert captured["input_file_path"] == "in.pdf" assert captured["output_file_path"] == "out.txt" + # summary + manifest persisted via the helper, keyed to the output path + assert write_calls["output_file_path"] == "out.txt" + assert write_calls["page_images"] == refs + assert write_calls["summary"] == expected_summary def test_empty_page_list_is_safe(self, monkeypatch: MonkeyPatch) -> None: monkeypatch.setattr(LLMWhispererHelper, "get_page_images", lambda **_: []) + # No output_file_path -> no extract-file write path is taken. result = _adapter(output_mode="image").process("in.pdf") assert result.extraction_metadata.page_images == [] - assert result.extracted_text == "" + assert result.extracted_text == LLMWhispererHelper.build_image_output_summary([]) def test_tag_forwarded_to_helper(self, monkeypatch: MonkeyPatch) -> None: captured: dict[str, object] = {} diff --git a/workers/executor/executors/constants.py b/workers/executor/executors/constants.py index 9eddab8423..428206148f 100644 --- a/workers/executor/executors/constants.py +++ b/workers/executor/executors/constants.py @@ -195,6 +195,7 @@ class IndexingConstants: USAGE_KWARGS = "usage_kwargs" PROCESS_TEXT = "process_text" EXTRACTED_TEXT = "extracted_text" + PAGE_IMAGES = "page_images" TAGS = "tags" EXECUTION_SOURCE = "execution_source" DOC_ID = "doc_id" diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index ce7fbea0d1..fda233a8d0 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -305,6 +305,17 @@ def _handle_extract(self, context: ExecutionContext) -> ExecutionResult: result_data["highlight_metadata"] = ( process_response.extraction_metadata.line_metadata ) + # Include image output page references when present (image output + # mode) so they survive the transport to the indexer / downstream + # consumers instead of being dropped with the summary text. + if ( + process_response.extraction_metadata + and process_response.extraction_metadata.page_images + ): + result_data[IKeys.PAGE_IMAGES] = [ + ref.to_dict() + for ref in process_response.extraction_metadata.page_images + ] return ExecutionResult( success=True, data=result_data, From fac7836f97b250403e73b64fdd3e588a4e299727 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Mon, 27 Jul 2026 12:51:00 +0530 Subject: [PATCH 09/24] UN-2646 [FIX] Keep image-mode cache fix, drop unused ref plumbing (UNS-758) Scope the previous change to just the real bug fix. The executor page_images forward and the JSON manifest sidecar were producer-side plumbing with no consumer anywhere in the codebase, so they added unused surface without making image references usable end-to-end. Removed both. Kept: image mode returns a short summary as extracted_text and writes it to the extract file, so the Prompt Studio extraction cache treats the conversion as complete and a re-run is a cache hit instead of a re-submit of the remote pdf-to-images job (the billing bug). Per-page references remain on extraction_metadata.page_images and the images are persisted to FileStorage; a Prompt Studio consumer of those references is deferred as follow-up. 49 sdk1 image-output tests pass; ruff/pycln clean. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../x2text/llm_whisperer_v2/src/helper.py | 37 +++++++------------ .../llm_whisperer_v2/src/llm_whisperer_v2.py | 12 +++--- unstract/sdk1/tests/test_llmw_image_helper.py | 19 +++------- .../sdk1/tests/test_llmw_v2_process_image.py | 3 +- workers/executor/executors/constants.py | 1 - workers/executor/executors/legacy_executor.py | 11 ------ 6 files changed, 25 insertions(+), 58 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 aef0c01f4e..590e2b9e9a 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 @@ -809,15 +809,15 @@ def build_image_output_summary(page_images: list[PageImageReference]) -> str: the extraction yields. Returning a short summary (rather than an empty string) keeps a re-run from re-submitting the remote conversion and keeps the indexed document meaningful instead of blank. The per-page - references travel separately in ``extraction_metadata.page_images`` and - the JSON manifest — never inside this string. + references travel separately in ``extraction_metadata.page_images`` — + never inside this string. """ count = len(page_images) noun = "page image" if count == 1 else "page images" return ( f"[LLMWhisperer image output mode] {count} {noun} extracted from the " "PDF and stored in FileStorage. Per-page references are available in " - "the page_images extraction metadata and the accompanying manifest." + "the page_images extraction metadata." ) @staticmethod @@ -825,15 +825,17 @@ def write_image_output( fs: FileStorage, output_file_path: str, summary: str, - page_images: list[PageImageReference], ) -> None: - """Persist the image-mode extract file + a page-image manifest sidecar. - - ``output_file_path`` (the extract file) receives the summary text so the - extraction-cache gate treats the conversion as complete (no re-submit) - and the indexer has meaningful text. ``.page_images. - json`` receives the ordered per-page references as a durable, retrievable - manifest. + """Persist the image-mode summary to the extract file. + + Image mode has no OCR text; writing a short summary to + ``output_file_path`` gives the Prompt Studio extraction cache a + non-empty extract, so a re-run is a cache hit instead of a re-submit of + the remote pdf-to-images conversion, and the indexed document stays + meaningful. The per-page references remain on the returned + ``extraction_metadata.page_images`` and the images themselves are + persisted to FileStorage; a Prompt Studio consumer of those references + is tracked as follow-up work. """ try: fs.write( @@ -842,17 +844,6 @@ def write_image_output( data=summary, encoding="utf-8", ) - manifest = json.dumps( - [ref.to_dict() for ref in page_images], - ensure_ascii=False, - indent=2, - ) - fs.write( - path=f"{output_file_path}.page_images.json", - mode="w", - data=manifest, - encoding="utf-8", - ) except Exception as e: - logger.error(f"Error writing image output for {output_file_path}: {e}") + logger.error(f"Error writing image extract file {output_file_path}: {e}") raise ExtractorError(str(e)) from e 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 6a25de32a6..4869826b84 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 @@ -95,9 +95,8 @@ def _process_image_mode( no OCR text, but a non-empty extract keeps the Prompt Studio extraction cache from re-submitting the remote conversion on a re-run and keeps the indexed document meaningful. The per-page references live in - ``extraction_metadata.page_images`` (forwarded by the executor) and a - JSON manifest sidecar — never inside ``extracted_text``. ``tag`` is - forwarded for service-side usage reporting. + ``extraction_metadata.page_images`` — never inside ``extracted_text``. + ``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) @@ -109,15 +108,14 @@ def _process_image_mode( tag=tag, ) summary = LLMWhispererHelper.build_image_output_summary(page_images) - # Persist the extract file (summary) + manifest so the extraction is - # cache-consistent (no re-submit on re-run) and the references are - # durably retrievable. Skipped when no output path was requested. + # Persist the summary to the extract file so the extraction is + # cache-consistent (no re-submit on re-run). Skipped when no output + # path was requested. if output_file_path: LLMWhispererHelper.write_image_output( fs=fs, output_file_path=output_file_path, summary=summary, - page_images=page_images, ) logger.info( "Image mode: returning %d page image reference(s) for %s", diff --git a/unstract/sdk1/tests/test_llmw_image_helper.py b/unstract/sdk1/tests/test_llmw_image_helper.py index ca11a5e522..72eb9d4ba4 100644 --- a/unstract/sdk1/tests/test_llmw_image_helper.py +++ b/unstract/sdk1/tests/test_llmw_image_helper.py @@ -219,32 +219,23 @@ def test_send_raw_request_has_finite_default_timeout(self) -> None: class TestImageOutputWrite: - """write_image_output persists the summary extract file + manifest sidecar.""" - - def test_writes_summary_and_manifest(self, tmp_path) -> None: # noqa: ANN001 - import json as _json + """write_image_output persists the summary to the extract file.""" + def test_writes_summary_to_extract_file(self, tmp_path) -> None: # noqa: ANN001 fs = FileStorage(provider=FileStorageProvider.LOCAL) refs = [ PageImageReference( page_number=1, path="doc/pages/page_001.png", filename="page_001.png" ), - PageImageReference( - page_number=2, path="doc/pages/page_002.png", filename="page_002.png" - ), ] out = str(tmp_path / "doc.txt") summary = H.build_image_output_summary(refs) - H.write_image_output( - fs=fs, output_file_path=out, summary=summary, page_images=refs - ) + H.write_image_output(fs=fs, output_file_path=out, summary=summary) - # Extract file holds the human summary (what image mode indexes). + # Extract file holds the human summary (what image mode indexes); a + # non-empty extract is what keeps a re-run from re-submitting. assert fs.read(path=out, mode="r") == summary - # Sidecar manifest round-trips the ordered references. - manifest = _json.loads(fs.read(path=out + ".page_images.json", mode="r")) - assert manifest == [r.to_dict() for r in refs] def test_summary_is_human_readable_not_json(self) -> None: refs = [PageImageReference(page_number=1, path="p/page_001.png")] diff --git a/unstract/sdk1/tests/test_llmw_v2_process_image.py b/unstract/sdk1/tests/test_llmw_v2_process_image.py index e5cbdaa31b..7e5e5249b2 100644 --- a/unstract/sdk1/tests/test_llmw_v2_process_image.py +++ b/unstract/sdk1/tests/test_llmw_v2_process_image.py @@ -98,9 +98,8 @@ def _fake_get_page_images(**kwargs: object) -> list[PageImageReference]: assert result.extraction_metadata.page_images == refs assert captured["input_file_path"] == "in.pdf" assert captured["output_file_path"] == "out.txt" - # summary + manifest persisted via the helper, keyed to the output path + # summary persisted to the extract file via the helper assert write_calls["output_file_path"] == "out.txt" - assert write_calls["page_images"] == refs assert write_calls["summary"] == expected_summary def test_empty_page_list_is_safe(self, monkeypatch: MonkeyPatch) -> None: diff --git a/workers/executor/executors/constants.py b/workers/executor/executors/constants.py index 428206148f..9eddab8423 100644 --- a/workers/executor/executors/constants.py +++ b/workers/executor/executors/constants.py @@ -195,7 +195,6 @@ class IndexingConstants: USAGE_KWARGS = "usage_kwargs" PROCESS_TEXT = "process_text" EXTRACTED_TEXT = "extracted_text" - PAGE_IMAGES = "page_images" TAGS = "tags" EXECUTION_SOURCE = "execution_source" DOC_ID = "doc_id" diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index fda233a8d0..ce7fbea0d1 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -305,17 +305,6 @@ def _handle_extract(self, context: ExecutionContext) -> ExecutionResult: result_data["highlight_metadata"] = ( process_response.extraction_metadata.line_metadata ) - # Include image output page references when present (image output - # mode) so they survive the transport to the indexer / downstream - # consumers instead of being dropped with the summary text. - if ( - process_response.extraction_metadata - and process_response.extraction_metadata.page_images - ): - result_data[IKeys.PAGE_IMAGES] = [ - ref.to_dict() - for ref in process_response.extraction_metadata.page_images - ] return ExecutionResult( success=True, data=result_data, From da9cf913ea4cb26f44e9e90b0e34f9e8a5b68dcc Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Mon, 27 Jul 2026 19:48:26 +0530 Subject: [PATCH 10/24] UN-2646 [FIX] Address Chandru's review on image-output mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Correctness / robustness: - Return the real whisper_hash in extraction_metadata (was ''), and reject image mode + highlight explicitly instead of silently dropping highlights. - Store page images under {extract_dir}/{stem}/pages (per-document), not the per-run hash — a re-extraction now overwrites its own pages instead of orphaning a fresh tree in FileStorage every run. - persist_page_images cleans up already-written pages on a mid-list failure. - Poll loop fails closed: only explicit intermediate states continue; a failure/unknown/empty status raises immediately instead of hanging to the budget. _safe_json logs non-JSON/non-object bodies. - verify_page_count is now live: expected count derived locally from the input PDF (pdfplumber); extract asserts a contiguous 1..N page set (catches truncation + duplicate/misnamed members) and logs skipped entries. - Retrieve advertises accept: application/zip; submit sends Content-Type. Layering / guard (UNS-757): - Promote OUTPUT_MODE/image/PDF_EXTENSION/PDF_ONLY_ERROR + an is_pdf() helper to the shared x2text.constants (ImageOutputConstants); backend imports the generic surface instead of the adapter's private src. - Move the PDF-only guard into dynamic_extractor (the single extract choke point, under profile_manager) so all entry points and prompt-level profile overrides are covered; gate on the LLMWhisperer adapter id as well as output_mode so other x2text adapters can't inherit the rejection. Comments: drop ticket-id/PR provenance and the contradictory PR#536/#647 contract notes; remove the unused IMAGE_MODE_LABEL 'single source of truth'. Tests: fix the constants-reload cross-file pollution (retry tests now pin the real budget on helper.WhispererDefaults); add network-layer coverage (poll success/failure/non-JSON-fast-fail/budget, mid-stream download error + close, submit-without-hash), mid-list cleanup, image+highlight rejection, a non-LLMWhisperer no-op, and a dynamic_extractor call-site test. sdk1: 57 pass; backend guard: 12 pass. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../prompt_studio_helper.py | 50 ++-- .../test_validate_image_output_pdf_only.py | 71 ++++-- .../sdk1/adapters/x2text/constants.py | 32 +++ .../x2text/llm_whisperer_v2/src/constants.py | 40 ++-- .../x2text/llm_whisperer_v2/src/helper.py | 220 ++++++++++++------ .../llm_whisperer_v2/src/llm_whisperer_v2.py | 18 +- unstract/sdk1/tests/llmw_image_fixtures.py | 24 +- .../tests/test_llm_whisperer_v2_constants.py | 37 ++- unstract/sdk1/tests/test_llmw_image_helper.py | 162 ++++++++++--- .../sdk1/tests/test_llmw_v2_process_image.py | 34 ++- 10 files changed, 482 insertions(+), 206 deletions(-) diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index 2c1046ca7e..0c33bb847e 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -77,11 +77,7 @@ ) from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt from unstract.core.pubsub_helper import LogPublisher -from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.constants import ( - ImageOutputConfig, - OutputModes, - WhispererConfig, -) +from unstract.sdk1.adapters.x2text.constants import ImageOutputConstants from unstract.sdk1.constants import LogLevel from unstract.sdk1.exceptions import IndexingError, SdkError from unstract.sdk1.execution.context import ExecutionContext @@ -563,11 +559,6 @@ def build_index_payload( default_profile, request_user=request_user ) - # Fail fast when an image-output x2text adapter is paired with a - # non-PDF input, so the user sees the PDF-only message here (before the - # executor task is dispatched) rather than as an extraction failure - # inside the worker (UNS-757). - PromptStudioHelper._validate_image_output_pdf_only(default_profile, file_name) # Common path decomposition used by extract, summarize, and index directory, filename = os.path.split(file_path) @@ -1388,24 +1379,29 @@ def _validate_image_output_pdf_only( """Reject non-PDF inputs when the x2text adapter is in image mode. Image output mode (LLMWhisperer V2) supports PDF input only. The SDK - adapter enforces this at extraction time; this mirror-check runs at - index time so the user gets the identical PDF-only message before any - extraction work is dispatched. The message is sourced from the SDK - (``ImageOutputConfig.PDF_ONLY_ERROR``) so both layers stay in sync. - - Only fires when the adapter config carries ``output_mode == "image"``, - which is unique to LLMWhisperer V2 — other x2text adapters are - unaffected. + adapter enforces this at extraction time; this mirror-check runs just + before extraction is dispatched (from ``dynamic_extractor``, the single + choke point for every extract path) so the user gets the identical + PDF-only message early. The message + PDF test come from the shared + ``ImageOutputConstants`` so the two layers cannot drift. + + Gated on BOTH the LLMWhisperer adapter id and ``output_mode == image``: + ``output_mode`` is user-editable adapter metadata, so keying on it alone + would make any future x2text adapter that adopts the same key inherit a + PDF-only rejection it never asked for. """ x2text = profile_manager.x2text if x2text is None: return + adapter_id = getattr(x2text, "adapter_id", "") or "" + if not adapter_id.startswith("llmwhisperer|"): + return metadata = x2text.metadata or {} - if metadata.get(WhispererConfig.OUTPUT_MODE) != OutputModes.IMAGE.value: + if metadata.get(ImageOutputConstants.OUTPUT_MODE) != ImageOutputConstants.IMAGE_MODE: return - if not file_name.lower().endswith(ImageOutputConfig.PDF_EXTENSION): + if not ImageOutputConstants.is_pdf(file_name): raise IndexingAPIError( - detail=ImageOutputConfig.PDF_ONLY_ERROR, + detail=ImageOutputConstants.PDF_ONLY_ERROR, status_code=400, ) @@ -1485,10 +1481,6 @@ def index_document( summary_profile, request_user=request_user ) - # Fail fast when an image-output x2text adapter is paired with a - # non-PDF input, so the user sees the PDF-only message at index time - # instead of after extraction is dispatched (UNS-757). - PromptStudioHelper._validate_image_output_pdf_only(default_profile, file_name) fs_instance = EnvHelper.get_storage( storage_type=StorageType.PERMANENT, @@ -2480,6 +2472,14 @@ def dynamic_extractor( profile_manager: ProfileManager, document_id: str, ) -> str: + # Reject a non-PDF input paired with an image-output adapter before any + # extraction work. This is the single choke point every extract path + # funnels through, and it runs under this profile_manager (not the + # default profile), so every entry point and prompt-level profile + # override is covered (UNS-757/758). + PromptStudioHelper._validate_image_output_pdf_only( + profile_manager, os.path.basename(file_path) + ) # Guard against None metadata (when adapter_metadata_b is None) metadata = profile_manager.x2text.metadata or {} x2text_config_hash = ToolUtils.hash_str(json.dumps(metadata, sort_keys=True)) diff --git a/backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py b/backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py index 1e99f0e8d0..f8a9b553f5 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py +++ b/backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py @@ -1,13 +1,14 @@ -"""Unit tests for ``PromptStudioHelper._validate_image_output_pdf_only``. +"""Unit tests for the image-output PDF-only guard. -Pins the UNS-757 fail-fast guard: when the x2text adapter is in image -output mode, a non-PDF input must be rejected at index-build time with the -SDK's shared PDF-only message, so the user never has to wait for the -executor to fail the extraction. Every other combination must pass through. +Pins the fail-fast guard: when the x2text adapter is the LLMWhisperer adapter +in image output mode, a non-PDF input must be rejected (with the SDK's shared +PDF-only message) before extraction is dispatched. Every other combination — +non-image mode, a non-LLMWhisperer adapter, a PDF input — must pass through. +Also pins that the guard is actually wired into ``dynamic_extractor`` (the +single extract choke point), so it cannot become unreachable unnoticed. -Unit tests: the real helper module is imported (Django is loaded by the -rig's test env) and the profile is a lightweight mock, so no database is -touched. +Unit tests: the real helper module is imported (Django is loaded by the rig's +test env) and the profile is a lightweight mock, so no database is touched. """ from __future__ import annotations @@ -18,22 +19,23 @@ from prompt_studio.prompt_studio_core_v2 import prompt_studio_helper as _psh_mod from prompt_studio.prompt_studio_core_v2.exceptions import IndexingAPIError -from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.constants import ( - ImageOutputConfig, -) +from unstract.sdk1.adapters.x2text.constants import ImageOutputConstants PromptStudioHelper = _psh_mod.PromptStudioHelper +_LLMW_ADAPTER_ID = "llmwhisperer|a5e6b8af-3e1f-4a80-b006-d017e8e67f93" -def _profile(metadata: dict | None) -> MagicMock: - """A profile whose x2text adapter exposes ``metadata`` verbatim.""" + +def _profile(metadata: dict | None, adapter_id: str = _LLMW_ADAPTER_ID) -> MagicMock: + """A profile whose x2text adapter exposes ``adapter_id`` + ``metadata``.""" profile = MagicMock(name="ProfileManager") + profile.x2text.adapter_id = adapter_id profile.x2text.metadata = metadata return profile class TestImageModeRejectsNonPdf: - """Image output mode + non-PDF → IndexingAPIError(400, PDF-only).""" + """LLMWhisperer + image output mode + non-PDF → IndexingAPIError(400).""" @pytest.mark.parametrize("file_name", ["statement.docx", "notes.txt", "a.png"]) def test_non_pdf_raises(self, file_name: str) -> None: @@ -42,34 +44,55 @@ def test_non_pdf_raises(self, file_name: str) -> None: _profile({"output_mode": "image"}), file_name ) assert exc_info.value.status_code == 400 - assert str(exc_info.value.detail) == ImageOutputConfig.PDF_ONLY_ERROR + assert str(exc_info.value.detail) == ImageOutputConstants.PDF_ONLY_ERROR @pytest.mark.parametrize("file_name", ["statement.pdf", "STATEMENT.PDF"]) def test_pdf_passes_case_insensitively(self, file_name: str) -> None: - # Must not raise for PDF inputs regardless of extension casing. PromptStudioHelper._validate_image_output_pdf_only( _profile({"output_mode": "image"}), file_name ) -class TestNonImageModesUnaffected: - """Only image mode is gated; every other config is a no-op.""" +class TestGateConditions: + """The guard is gated on BOTH the adapter id and the output mode.""" @pytest.mark.parametrize( "metadata", - [ - {"output_mode": "text"}, - {"output_mode": "layout_preserving"}, - {}, # e.g. a non-LLMWhisperer adapter with no output_mode - None, # adapter metadata absent entirely - ], + [{"output_mode": "text"}, {"output_mode": "layout_preserving"}, {}, None], ) def test_non_image_mode_passes_for_non_pdf(self, metadata: dict | None) -> None: PromptStudioHelper._validate_image_output_pdf_only( _profile(metadata), "statement.docx" ) + def test_non_llmwhisperer_adapter_is_not_rejected(self) -> None: + # A different x2text adapter that happens to carry output_mode=image in + # its (user-editable) metadata must NOT inherit a PDF-only rejection. + PromptStudioHelper._validate_image_output_pdf_only( + _profile({"output_mode": "image"}, adapter_id="some-other|123"), + "statement.docx", + ) + def test_missing_x2text_adapter_passes(self) -> None: profile = MagicMock(name="ProfileManager") profile.x2text = None PromptStudioHelper._validate_image_output_pdf_only(profile, "statement.docx") + + +class TestGuardIsWiredIntoDynamicExtractor: + """The guard must run from dynamic_extractor (the single extract path).""" + + def test_dynamic_extractor_rejects_non_pdf_image_mode(self) -> None: + # The guard is the first statement in dynamic_extractor, so an image-mode + # adapter + non-PDF raises before any DB/storage work — proving the call + # site is exercised (deleting the call would make this test fail). + profile = _profile({"output_mode": "image"}) + with pytest.raises(IndexingAPIError): + PromptStudioHelper.dynamic_extractor( + file_path="/data/statement.docx", + enable_highlight=False, + run_id="r1", + org_id="org1", + profile_manager=profile, + document_id="doc1", + ) diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/constants.py b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/constants.py index bd703e6538..110e7aef27 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/constants.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/constants.py @@ -1,3 +1,6 @@ +from pathlib import Path + + class X2TextConstants: PLATFORM_SERVICE_API_KEY = "PLATFORM_SERVICE_API_KEY" X2TEXT_HOST = "X2TEXT_HOST" @@ -7,3 +10,32 @@ class X2TextConstants: EXTRACTED_TEXT = "extracted_text" WHISPER_HASH = "whisper-hash" WHISPER_HASH_V2 = "whisper_hash" + + +class ImageOutputConstants: + """Image-output-mode contract shared across the x2text layer. + + Kept on the generic x2text surface (not inside an adapter's private + ``src/`` package) so consumers outside the adapter — e.g. the backend's + index-time PDF-only guard — depend on it without reaching into adapter + internals. + """ + + # Adapter config key selecting the output format, and the value that + # selects per-page image output. + OUTPUT_MODE = "output_mode" + IMAGE_MODE = "image" + + # Image output accepts PDF input only. A single message + a single + # extension test keep the runtime guard (adapter ``process()``) and the + # index-time guard (backend) from drifting apart. + PDF_EXTENSION = ".pdf" + PDF_ONLY_ERROR = ( + "Image output mode supports PDF input only. " + "Please provide a PDF file or select a text output mode." + ) + + @staticmethod + def is_pdf(file_name: str) -> bool: + """Return True when ``file_name`` is a PDF (case-insensitive suffix).""" + return Path(file_name).suffix.lower() == ImageOutputConstants.PDF_EXTENSION 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 85b4cf8900..9ec6c477ca 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 @@ -1,6 +1,8 @@ import os from enum import Enum +from unstract.sdk1.adapters.x2text.constants import ImageOutputConstants + class Modes(Enum): NATIVE_TEXT = "native_text" @@ -12,7 +14,7 @@ class Modes(Enum): class OutputModes(Enum): LAYOUT_PRESERVING = "layout_preserving" TEXT = "text" - IMAGE = "image" + IMAGE = ImageOutputConstants.IMAGE_MODE class HTTPMethod(Enum): @@ -72,7 +74,7 @@ class WhispererConfig: URL = "url" MODE = "mode" - OUTPUT_MODE = "output_mode" + OUTPUT_MODE = ImageOutputConstants.OUTPUT_MODE UNSTRACT_KEY = "unstract_key" MEDIAN_FILTER_SIZE = "median_filter_size" GAUSSIAN_BLUR_RADIUS = "gaussian_blur_radius" @@ -141,13 +143,11 @@ class WhispererDefaults: 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. + The pdf-to-images endpoints are not exposed by the installed + ``llmwhisperer-client``, so the adapter calls them directly via raw + ``requests``. The wire shape the adapter depends on is centralised here. - Flow (raw ``requests``, base = ``{url}/api/v2``): + Flow (base = ``{url}/api/v2``): - Submit: ``POST {base}/pdf-to-images?format=png`` with the PDF bytes -> JSON ``{"message": "...", "status": "processing", @@ -166,14 +166,14 @@ class ImageOutputConfig: # --- 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). + # Poll control. Success == ready-to-retrieve; only these intermediate states + # keep the poll loop going. Any other value — a failure state, an unknown + # status, or an empty/non-JSON body — is treated as terminal and raises, so + # the loop fails fast instead of polling to the budget on a stuck job. STATUS_SUCCESS = frozenset({"processed"}) - STATUS_FAILURE = frozenset({"error", "failed", "unknown"}) + STATUS_INTERMEDIATE = frozenset({"accepted", "processing", "queued"}) # --- Submit query params --- IMAGE_FORMAT_PARAM = "format" @@ -186,13 +186,7 @@ class ImageOutputConfig: 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." - ) + # --- PDF-only validation (shared with the backend index-time guard) --- + PDF_EXTENSION = ImageOutputConstants.PDF_EXTENSION + PDF_ONLY_ERROR = ImageOutputConstants.PDF_ONLY_ERROR + is_pdf = staticmethod(ImageOutputConstants.is_pdf) 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 590e2b9e9a..859a5e8199 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 @@ -451,12 +451,28 @@ def write_output_to_file( @staticmethod def _safe_json(response: Response) -> dict[str, Any]: - """Parse a JSON object body, tolerating non-JSON / non-object bodies.""" + """Parse a JSON object body, tolerating non-JSON / non-object bodies. + + A non-JSON or non-object body is logged (with a truncated preview) + before returning ``{}`` so a caller that treats the empty result as an + unexpected status has a diagnostic instead of a silent fall-through. + """ try: parsed = response.json() except ValueError: + logger.warning( + "LLMWhisperer returned a non-JSON body (HTTP %s): %s", + getattr(response, "status_code", "?"), + (response.text or "")[:200], + ) + return {} + if not isinstance(parsed, dict): + logger.warning( + "LLMWhisperer returned a non-object JSON body: %s", + str(parsed)[:200], + ) return {} - return parsed if isinstance(parsed, dict) else {} + return parsed @staticmethod def submit_pdf_to_images( @@ -481,12 +497,17 @@ def submit_pdf_to_images( } if file_name: params[ImageOutputConfig.FILE_NAME_PARAM] = file_name + headers = { + **LLMWhispererHelper.get_request_headers(config), + "Content-Type": "application/octet-stream", + } response = LLMWhispererHelper._send_raw_request( config=config, method="POST", endpoint=WhispererEndpoint.PDF_TO_IMAGES, params=params, data=file_data, + headers=headers, timeout=WhispererDefaults.IMAGE_REQUEST_TIMEOUT, ) body = LLMWhispererHelper._safe_json(response) @@ -533,13 +554,17 @@ def poll_pdf_to_images_status( ) if status in ImageOutputConfig.STATUS_SUCCESS: return body - if status in ImageOutputConfig.STATUS_FAILURE: + if status not in ImageOutputConfig.STATUS_INTERMEDIATE: + # Fail closed: only explicit intermediate states keep polling. + # A failure state, an unknown status, or an empty body (non-JSON) + # raises immediately with the observed status echoed, instead of + # hanging until the poll budget is exhausted. msg = body.get(ImageOutputConfig.MESSAGE, "unknown error") raise ExtractorError( - f"LLMWhisperer pdf-to-images job {whisper_hash} failed: {msg}", - status_code=500, + f"LLMWhisperer pdf-to-images job {whisper_hash} returned an " + f"unexpected status '{status or ''}': {msg}", + status_code=502, ) - # 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 " @@ -555,11 +580,18 @@ def download_pdf_to_images_zip(config: dict[str, Any], whisper_hash: str) -> Byt Uses a distinct, longer download timeout (large multi-page PDFs) and avoids a single ``response.content`` load. """ + # This endpoint streams application/zip; advertise it so a strict + # gateway does not 406 the default ``accept: application/json``. + headers = { + **LLMWhispererHelper.get_request_headers(config), + "accept": "application/zip", + } response = LLMWhispererHelper._send_raw_request( config=config, method="GET", endpoint=WhispererEndpoint.PDF_TO_IMAGES_RETRIEVE, params={WhisperStatus.WHISPER_HASH: whisper_hash}, + headers=headers, timeout=WhispererDefaults.IMAGE_DOWNLOAD_TIMEOUT, stream=True, ) @@ -594,11 +626,15 @@ def extract_page_images_from_zip( on a corrupt/invalid archive. """ pages: list[tuple[int, bytes]] = [] + skipped: list[str] = [] try: with zipfile.ZipFile(zip_buffer) as archive: for name in archive.namelist(): + if name.endswith("/"): + continue # directory entry, not a member match = LLMWhispererHelper._PAGE_IMAGE_RE.search(name) if not match: + skipped.append(name) continue page_number = int(match.group(1)) pages.append((page_number, archive.read(name))) @@ -610,56 +646,99 @@ def extract_page_images_from_zip( status_code=502, actual_err=e, ) from e + if skipped: + # Visible, not silent: a naming-convention change mid-archive would + # otherwise truncate the page set and still report success. + logger.warning( + "Image mode: ignored %d non-page entr%s in the pdf-to-images " + "archive: %s", + len(skipped), + "y" if len(skipped) == 1 else "ies", + ", ".join(skipped[:10]), + ) if not pages: # A well-formed archive with no recognizable page images is a - # failed extraction, not an empty success — fail closed, matching - # the rest of this flow. + # failed extraction, not an empty success — fail closed. raise ExtractorError( "pdf-to-images returned an archive with no page images", status_code=502, ) pages.sort(key=lambda item: item[0]) + # The page set must be exactly 1..N with no gaps or duplicates. A gap + # means a truncated archive; a duplicate means two members mapped to the + # same page number (``page_001.png`` under two folders) — which + # ``persist_page_images`` would silently overwrite. Fail closed on both. + page_numbers = [page for page, _ in pages] + if page_numbers != list(range(1, len(page_numbers) + 1)): + raise ExtractorError( + "pdf-to-images archive page numbers are not a contiguous 1..N " + f"sequence (got {page_numbers}); the archive is truncated or has " + "duplicate/misnamed page members", + status_code=502, + ) return pages + @staticmethod + def _safe_pdf_page_count(pdf_bytes: bytes) -> int | None: + """Page count of the input PDF, or None if it cannot be read. + + Best-effort: the extraction must not fail just because the count could + not be derived locally, so any error returns None (and the caller + degrades to the archive contiguity check). + """ + try: + import pdfplumber # noqa: PLC0415 - lazy: only image mode needs it + + with pdfplumber.open(BytesIO(pdf_bytes)) as pdf: + return len(pdf.pages) + except Exception as e: + logger.warning("Image mode: unable to read input PDF page count: %s", e) + return None + @staticmethod def verify_page_count( - pages: list[tuple[int, bytes]], processed_page_count: int | None + pages: list[tuple[int, bytes]], expected_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. + """Verify the extracted page count against the input PDF's page count. + + ``expected_page_count`` is derived locally from the input PDF (the + pdf-to-images-status response exposes no count). A mismatch means the + service returned more or fewer images than the document has pages — a + truncated or over-produced archive — and raises. When the count could + not be determined the check is skipped with a warning, leaving the + 1..N contiguity check in ``extract_page_images_from_zip`` as the last + line of defence. """ - 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" + if expected_page_count is None: + logger.warning( + "Image mode: input PDF page count unavailable; skipping " + "page-count verification (relying on the 1..N contiguity check)" ) return actual = len(pages) - if actual != processed_page_count: + if actual != expected_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)", + "Page count mismatch in image output mode: the input PDF has " + f"{expected_page_count} page(s) but the service returned {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``. + def build_page_store_dir(output_file_path: str | None, input_file_path: str) -> str: + """Per-document folder for page images: ``{extract_dir}/{stem}/pages``. + + Keyed on the document ``stem`` (the same discriminator the extract + ``.txt`` files alongside use), not the per-run whisper_hash. This is + collision-safe against concurrent documents in the same project, is + reconstructible from ``output_file_path`` alone, and — being stable + across runs — makes a re-extraction overwrite its own pages instead of + orphaning a fresh tree in FileStorage on every run. """ 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) + stem = Path(reference).stem if reference else "document" + return str(Path(base_dir) / stem / ImageOutputConfig.PAGES_SUBFOLDER) @staticmethod def _page_image_filename(page_number: int) -> str: @@ -673,6 +752,23 @@ def _page_image_filename(page_number: int) -> str: def _write_single_page(fs: FileStorage, path: str, data: bytes) -> None: fs.write(path=path, mode="wb", data=data, encoding="utf-8") + @staticmethod + def _cleanup_partial_pages(fs: FileStorage, page_store_dir: str) -> None: + """Best-effort removal of a partially-written page directory. + + Invoked when a persist fails mid-set so a failed extraction leaves no + orphan pages behind. A cleanup error is only logged — the original + extraction error is what the caller must see. + """ + try: + fs.rm(page_store_dir, recursive=True) + except Exception as e: + logger.warning( + "Image mode: could not clean up partial page dir %s: %s", + page_store_dir, + e, + ) + @staticmethod def persist_page_images( fs: FileStorage, @@ -681,11 +777,10 @@ def persist_page_images( ) -> 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``. + Fail-closed on disk as well as in the return value: if any page exhausts + its retries, the pages already written are removed (best-effort) before + a hard ``ExtractorError`` propagates, so a failed extraction never leaves + a partial set behind. Works transparently for LOCAL and S3 via ``fs``. """ fs.mkdir(create_parents=True, path=page_store_dir) @@ -706,6 +801,7 @@ def persist_page_images( try: write_with_retry(fs=fs, path=path, data=data) except Exception as e: + LLMWhispererHelper._cleanup_partial_pages(fs, page_store_dir) raise ExtractorError( "Failed to persist page image after retries: " f"page={page_number}, provider={fs.provider.value}, " @@ -745,60 +841,52 @@ def get_page_images( output_file_path: str | None, fs: FileStorage | None = None, tag: str | list[str] | None = None, - ) -> list[PageImageReference]: + ) -> tuple[str, 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. + persist per-page (retried). Returns ``(whisper_hash, references)`` with + the ordered ``PageImageReference`` list, or raises (fail-closed — never + partial). The ``whisper_hash`` is returned so callers can record the + real job id in extraction metadata instead of an empty string. + + Retrieval is intentionally NOT retried: the service marks the job + RETRIEVED before streaming and (with the default persistence off) a + second retrieve is rejected, while a re-submit would double-bill — so a + mid-download failure is a hard error the caller must resubmit. Per-page + FileStorage writes are still retried. """ if fs is None: fs = FileStorage(provider=FileStorageProvider.LOCAL) - input_data = BytesIO(fs.read(path=input_file_path, mode="rb")) + input_bytes = fs.read(path=input_file_path, mode="rb") whisper_hash = LLMWhispererHelper.submit_pdf_to_images( config, - input_data, + BytesIO(input_bytes), 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) + LLMWhispererHelper.poll_pdf_to_images_status(config, whisper_hash) 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) + # Verify the returned image count against the input PDF's own page count + # (derived locally) BEFORE persisting, so nothing is written on a + # truncated/over-produced archive. + expected_page_count = LLMWhispererHelper._safe_pdf_page_count(input_bytes) + LLMWhispererHelper.verify_page_count(pages, expected_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, + "Image mode: completed job=%s pages=%d", whisper_hash, len(references) ) - return references + return whisper_hash, references @staticmethod def build_image_output_summary(page_images: list[PageImageReference]) -> str: 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 4869826b84..94b9a7f440 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,7 +2,6 @@ import logging import os -from pathlib import Path from typing import TYPE_CHECKING, Any from unstract.sdk1.adapters.exceptions import ExtractorError @@ -73,7 +72,7 @@ def _validate_pdf_only(input_file_path: str) -> None: 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: + if not ImageOutputConfig.is_pdf(input_file_path): raise ExtractorError( ImageOutputConfig.PDF_ONLY_ERROR, status_code=400, @@ -100,7 +99,7 @@ def _process_image_mode( """ 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( + whisper_hash, page_images = LLMWhispererHelper.get_page_images( config=self.config, input_file_path=input_file_path, output_file_path=output_file_path, @@ -125,7 +124,7 @@ def _process_image_mode( return TextExtractionResult( extracted_text=summary, extraction_metadata=TextExtractionMetadata( - whisper_hash="", + whisper_hash=whisper_hash, page_images=page_images, ), ) @@ -156,7 +155,17 @@ def process( output_mode = self.config.get( WhispererConfig.OUTPUT_MODE, OutputModes.LAYOUT_PRESERVING.value ) + enable_highlight = kwargs.get(X2TextConstants.ENABLE_HIGHLIGHT, False) if output_mode == OutputModes.IMAGE.value: + # Highlighting produces line-level source references over extracted + # text; image mode yields no text, so the combination is rejected + # explicitly rather than silently returning empty highlight data. + if enable_highlight: + raise ExtractorError( + "Highlighting is not supported in image output mode; disable " + "highlight or select a text output mode.", + status_code=400, + ) return self._process_image_mode( input_file_path, output_file_path, @@ -164,7 +173,6 @@ def process( tag=kwargs.get(X2TextConstants.TAGS), ) - enable_highlight = kwargs.get(X2TextConstants.ENABLE_HIGHLIGHT, False) logger.info( "HIGHLIGHT_DEBUG LLMWhispererV2.process: enable_highlight=%s", enable_highlight, diff --git a/unstract/sdk1/tests/llmw_image_fixtures.py b/unstract/sdk1/tests/llmw_image_fixtures.py index 10133eb525..344bd00116 100644 --- a/unstract/sdk1/tests/llmw_image_fixtures.py +++ b/unstract/sdk1/tests/llmw_image_fixtures.py @@ -84,6 +84,14 @@ def __init__(self, provider: FileStorageProvider = FileStorageProvider.S3) -> No self._files: dict[str, bytes] = {} self._dirs: set[str] = set() self.write_calls = 0 + self.rm_calls: list[str] = [] + + def rm(self, path: str, recursive: bool = True) -> None: + self.rm_calls.append(str(path)) + prefix = str(path).rstrip("/") + "/" + for key in list(self._files): + if key == str(path) or key.startswith(prefix): + del self._files[key] def mkdir(self, path: str, create_parents: bool = True) -> None: self._dirs.add(str(path)) @@ -123,12 +131,21 @@ class FlakyFileStorage(InMemoryFileStorage): """ def __init__( - self, fail_times: int = 1, fail_always: bool = False, **kwargs: object + self, + fail_times: int = 1, + fail_always: bool = False, + fail_substrings: tuple[str, ...] = (), + **kwargs: object, ) -> None: - """Configure how many writes per path fail before succeeding.""" + """Configure how many writes per path fail before succeeding. + + ``fail_substrings`` always-fails any write whose path contains one of + the substrings — used to fail a specific page (mid-list failure). + """ super().__init__(**kwargs) self.fail_times = fail_times self.fail_always = fail_always + self.fail_substrings = tuple(fail_substrings) self._attempts: dict[str, int] = {} def write( @@ -141,7 +158,8 @@ def write( ) -> int: key = str(path) self._attempts[key] = self._attempts.get(key, 0) + 1 - if self.fail_always or self._attempts[key] <= self.fail_times: + always_fail = self.fail_always or any(s in key for s in self.fail_substrings) + if always_fail or self._attempts[key] <= self.fail_times: raise FileOperationError(f"simulated write failure for {key}") return super().write(path, mode, encoding, data, **kwargs) diff --git a/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py b/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py index a4a8cfc725..57d5a9fb45 100644 --- a/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py +++ b/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py @@ -1,13 +1,9 @@ -"""Unit tests for LLMWhisperer v2 adapter constants (MUNS-193). +"""Unit tests for LLMWhisperer v2 adapter constants. -Covers: -- UNS-732: OutputModes.IMAGE enum value. -- UNS-733: ADAPTER_LLMW_PAGE_STORE_MAX_RETRIES env-var-backed constant. +Covers the image OutputModes value and the env-var-backed page-store retry +budget. """ -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" @@ -26,20 +22,13 @@ 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: - # Reload under the scoped patch, then reload again after the env is - # restored so the module cache reflects the real environment and does - # not leak the patched value into later tests. - with monkeypatch.context() as patch: - patch.delenv(_ENV_VAR, raising=False) - reloaded = importlib.reload(c) - assert reloaded.WhispererDefaults.PAGE_STORE_MAX_RETRIES == 3 - assert isinstance(reloaded.WhispererDefaults.PAGE_STORE_MAX_RETRIES, int) - importlib.reload(c) - - def test_reads_from_env(self, monkeypatch: MonkeyPatch) -> None: - with monkeypatch.context() as patch: - patch.setenv(_ENV_VAR, "5") - reloaded = importlib.reload(c) - assert reloaded.WhispererDefaults.PAGE_STORE_MAX_RETRIES == 5 - importlib.reload(c) + def test_default_is_three(self) -> None: + # Deliberately no importlib.reload: reloading the constants module + # rebinds the WhispererDefaults *class object* while helper.py keeps a + # direct name binding to the original — which silently turns other + # suites' ``monkeypatch.setattr(WhispererDefaults, ...)`` into no-ops + # (and, being order-dependent, is invisible until the split changes). + # The value is read from the env at import; with the var unset (the + # test environment) it is the default 3. + assert c.WhispererDefaults.PAGE_STORE_MAX_RETRIES == 3 + assert isinstance(c.WhispererDefaults.PAGE_STORE_MAX_RETRIES, int) diff --git a/unstract/sdk1/tests/test_llmw_image_helper.py b/unstract/sdk1/tests/test_llmw_image_helper.py index 72eb9d4ba4..a2b3eac8d6 100644 --- a/unstract/sdk1/tests/test_llmw_image_helper.py +++ b/unstract/sdk1/tests/test_llmw_image_helper.py @@ -8,12 +8,14 @@ """ import io +from unittest.mock import MagicMock import pytest +import requests 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 import helper as helper_mod from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.helper import ( LLMWhispererHelper, ) @@ -72,40 +74,40 @@ def test_archive_with_no_page_entries_raises(self) -> None: 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 + H.verify_page_count(pages, expected_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) + H.verify_page_count(pages, expected_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) + H.verify_page_count(pages, expected_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 + H.verify_page_count(pages, expected_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") + def test_folder_isolates_distinct_documents(self) -> None: + dir_a = H.build_page_store_dir("/data/extract/doc-a.txt", "/data/doc-a.pdf") + dir_b = H.build_page_store_dir("/data/extract/doc-b.txt", "/data/doc-b.pdf") assert dir_a != dir_b - assert "run-A" in dir_a and "run-B" in dir_b + assert "doc-a" in dir_a and "doc-b" in dir_b assert dir_a.endswith("pages") - def test_folder_key_deterministic_for_same_run(self) -> None: - first = H.build_page_store_dir("/data/out.txt", "/data/in.pdf", "run-A") - second = H.build_page_store_dir("/data/out.txt", "/data/in.pdf", "run-A") - assert first == second + def test_folder_stable_across_runs_for_same_document(self) -> None: + # Keyed on the document stem, not the per-run hash: a re-extraction + # overwrites its own pages instead of orphaning a fresh tree. + first = H.build_page_store_dir("/data/extract/doc.txt", "/data/doc.pdf") + second = H.build_page_store_dir("/data/extract/doc.txt", "/data/doc.pdf") + assert first == second == "/data/extract/doc/pages" - 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 + def test_folder_falls_back_to_input_when_no_output(self) -> None: + assert H.build_page_store_dir(None, "/docs/in.pdf") == "/docs/in/pages" @pytest.mark.parametrize( ("page", "expected"), @@ -136,28 +138,53 @@ def test_persists_all_pages_as_ordered_references(self) -> None: 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) + # Patch the class the helper actually holds (helper_mod.WhispererDefaults), + # so the budget is genuinely pinned regardless of any module reload + # elsewhere in the suite. + monkeypatch.setattr(helper_mod.WhispererDefaults, "RETRY_MIN_WAIT", 0.0) + monkeypatch.setattr(helper_mod.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_budget_is_pinned_to_two_retries(self, monkeypatch: MonkeyPatch) -> None: + # Budget 2 == 3 total attempts. A page whose first 3 attempts fail must + # error — proving the patched budget actually takes effect (with the + # default budget 3 == 4 attempts, the 4th would have succeeded). + monkeypatch.setattr(helper_mod.WhispererDefaults, "RETRY_MIN_WAIT", 0.0) + monkeypatch.setattr(helper_mod.WhispererDefaults, "PAGE_STORE_MAX_RETRIES", 2) + fs = FlakyFileStorage(fail_times=3) # would succeed only on the 4th attempt + with pytest.raises(ExtractorError, match="Failed to persist page image"): + H.persist_page_images(fs, "doc/pages", [(1, b"data")]) + 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) + monkeypatch.setattr(helper_mod.WhispererDefaults, "RETRY_MIN_WAIT", 0.0) + monkeypatch.setattr(helper_mod.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_mid_list_failure_cleans_up_written_pages( + self, monkeypatch: MonkeyPatch + ) -> None: + # Page 1 succeeds, page 2 always fails -> the partial set must be removed + # so a failed extraction leaves no orphan pages behind. + monkeypatch.setattr(helper_mod.WhispererDefaults, "RETRY_MIN_WAIT", 0.0) + monkeypatch.setattr(helper_mod.WhispererDefaults, "PAGE_STORE_MAX_RETRIES", 1) + fs = FlakyFileStorage(fail_times=0, fail_substrings=("page_002",)) + with pytest.raises(ExtractorError, match="Failed to persist page image"): + H.persist_page_images(fs, "doc/pages", [(1, b"a"), (2, b"b")]) + assert "doc/pages" in fs.rm_calls # cleanup invoked + assert fs.stored_paths == [] # page 1 removed by the cleanup + 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) @@ -206,16 +233,93 @@ def test_list_tag_is_normalized(self, monkeypatch: MonkeyPatch) -> None: assert captured["params"]["tag"] == "first" +_NET_CONFIG = {"url": "https://svc.example", "unstract_key": "k"} + + +def _json_response(payload: dict) -> MagicMock: + resp = MagicMock() + resp.json.return_value = payload + return resp + + class TestRequestDefaults: - """The shared raw-request path must never wait forever (UNS-758).""" + """The shared raw-request path must apply a finite timeout in practice.""" + + def test_default_timeout_is_passed_to_requests( + self, monkeypatch: MonkeyPatch + ) -> None: + # Behaviour, not signature: patch requests.request and assert the + # timeout actually handed to it is finite when a caller omits it. + captured: dict = {} + resp = MagicMock() + resp.raise_for_status.return_value = None + monkeypatch.setattr(requests, "request", lambda **kw: captured.update(kw) or resp) + H._send_raw_request(config=_NET_CONFIG, method="GET", endpoint="ping") + assert isinstance(captured["timeout"], int | float) + assert captured["timeout"] > 0 - def test_send_raw_request_has_finite_default_timeout(self) -> None: - import inspect - default = inspect.signature(H._send_raw_request).parameters["timeout"].default - # A None default maps to requests' "wait forever"; test_connection relies - # on this default, so it must be a positive, finite number. - assert isinstance(default, int | float) and default > 0 +class TestPollBehavior: + def test_processed_returns_payload(self, monkeypatch: MonkeyPatch) -> None: + monkeypatch.setattr( + H, "_send_raw_request", lambda **kw: _json_response({"status": "processed"}) + ) + assert H.poll_pdf_to_images_status(_NET_CONFIG, "wh")["status"] == "processed" + + def test_failure_status_raises_immediately(self, monkeypatch: MonkeyPatch) -> None: + monkeypatch.setattr( + H, + "_send_raw_request", + lambda **kw: _json_response({"status": "failed", "message": "boom"}), + ) + with pytest.raises(ExtractorError, match="unexpected status 'failed'"): + H.poll_pdf_to_images_status(_NET_CONFIG, "wh") + + def test_non_json_body_fails_fast(self, monkeypatch: MonkeyPatch) -> None: + # A non-JSON/HTML error body -> _safe_json {} -> status "" -> not an + # intermediate state -> raise on the first poll (no budget-long hang). + bad = MagicMock() + bad.json.side_effect = ValueError("no json") + bad.text = "bad gateway" + bad.status_code = 502 + monkeypatch.setattr(H, "_send_raw_request", lambda **kw: bad) + with pytest.raises(ExtractorError, match="unexpected status"): + H.poll_pdf_to_images_status(_NET_CONFIG, "wh") + + def test_budget_exhaustion_raises_after_max_attempts( + self, monkeypatch: MonkeyPatch + ) -> None: + monkeypatch.setattr(helper_mod.WhispererDefaults, "IMAGE_POLL_INTERVAL", 0.0) + monkeypatch.setattr(helper_mod.WhispererDefaults, "IMAGE_POLL_MAX_ATTEMPTS", 3) + calls = {"n": 0} + + def _sr(**_: object) -> MagicMock: + calls["n"] += 1 + return _json_response({"status": "processing"}) + + monkeypatch.setattr(H, "_send_raw_request", _sr) + with pytest.raises(ExtractorError, match="did not reach a terminal state"): + H.poll_pdf_to_images_status(_NET_CONFIG, "wh") + assert calls["n"] == 3 + + +class TestDownloadAndSubmitBehavior: + def test_mid_stream_error_maps_to_extractor_error_and_closes( + self, monkeypatch: MonkeyPatch + ) -> None: + resp = MagicMock() + resp.iter_content.side_effect = requests.exceptions.ChunkedEncodingError("x") + monkeypatch.setattr(H, "_send_raw_request", lambda **kw: resp) + with pytest.raises(ExtractorError, match="Failed to download"): + H.download_pdf_to_images_zip(_NET_CONFIG, "wh") + resp.close.assert_called_once() # connection released on failure + + def test_submit_without_whisper_hash_raises(self, monkeypatch: MonkeyPatch) -> None: + monkeypatch.setattr( + H, "_send_raw_request", lambda **kw: _json_response({"message": "ok"}) + ) + with pytest.raises(ExtractorError, match="did not return a job id"): + H.submit_pdf_to_images(_NET_CONFIG, io.BytesIO(b"pdf")) class TestImageOutputWrite: diff --git a/unstract/sdk1/tests/test_llmw_v2_process_image.py b/unstract/sdk1/tests/test_llmw_v2_process_image.py index 7e5e5249b2..09fee77ebe 100644 --- a/unstract/sdk1/tests/test_llmw_v2_process_image.py +++ b/unstract/sdk1/tests/test_llmw_v2_process_image.py @@ -69,9 +69,11 @@ def test_populates_page_images_and_empty_text(self, monkeypatch: MonkeyPatch) -> ] captured: dict[str, object] = {} - def _fake_get_page_images(**kwargs: object) -> list[PageImageReference]: + def _fake_get_page_images( + **kwargs: object, + ) -> tuple[str, list[PageImageReference]]: captured.update(kwargs) - return refs + return "run-1|doc-hash", refs monkeypatch.setattr(LLMWhispererHelper, "get_page_images", _fake_get_page_images) monkeypatch.setattr( @@ -96,6 +98,9 @@ def _fake_get_page_images(**kwargs: object) -> list[PageImageReference]: assert result.extracted_text == expected_summary assert "page_001.png" not in result.extracted_text assert result.extraction_metadata.page_images == refs + # The real job id is recorded, not an empty string (HITL/QueueResult + # consumers read this). + assert result.extraction_metadata.whisper_hash == "run-1|doc-hash" assert captured["input_file_path"] == "in.pdf" assert captured["output_file_path"] == "out.txt" # summary persisted to the extract file via the helper @@ -103,7 +108,7 @@ def _fake_get_page_images(**kwargs: object) -> list[PageImageReference]: assert write_calls["summary"] == expected_summary def test_empty_page_list_is_safe(self, monkeypatch: MonkeyPatch) -> None: - monkeypatch.setattr(LLMWhispererHelper, "get_page_images", lambda **_: []) + monkeypatch.setattr(LLMWhispererHelper, "get_page_images", lambda **_: ("wh", [])) # No output_file_path -> no extract-file write path is taken. result = _adapter(output_mode="image").process("in.pdf") assert result.extraction_metadata.page_images == [] @@ -112,19 +117,34 @@ def test_empty_page_list_is_safe(self, monkeypatch: MonkeyPatch) -> None: def test_tag_forwarded_to_helper(self, monkeypatch: MonkeyPatch) -> None: captured: dict[str, object] = {} - def _capture(**kwargs: object) -> list: + def _capture(**kwargs: object) -> tuple[str, list]: captured.update(kwargs) - return [] + return "wh", [] 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 **_: []) + monkeypatch.setattr(LLMWhispererHelper, "get_page_images", lambda **_: ("wh", [])) # Should not raise for an uppercase .PDF extension. _adapter(output_mode="image").process("SCAN.PDF") + def test_image_mode_with_highlight_is_rejected( + self, monkeypatch: MonkeyPatch + ) -> None: + # Highlighting has no meaning without text; the combination must be + # rejected explicitly rather than silently returning empty highlights. + called = {"hit": False} + monkeypatch.setattr( + LLMWhispererHelper, + "get_page_images", + lambda **_: called.__setitem__("hit", True) or ("wh", []), + ) + with pytest.raises(ExtractorError, match="not supported in image output mode"): + _adapter(output_mode="image").process("in.pdf", enable_highlight=True) + assert called["hit"] is False # rejected before any conversion + class TestPdfOnlyValidation: def test_non_pdf_rejected_before_helper_runs(self, monkeypatch: MonkeyPatch) -> None: @@ -132,7 +152,7 @@ def test_non_pdf_rejected_before_helper_runs(self, monkeypatch: MonkeyPatch) -> monkeypatch.setattr( LLMWhispererHelper, "get_page_images", - lambda **_: image_called.__setitem__("hit", True) or [], + lambda **_: image_called.__setitem__("hit", True) or ("wh", []), ) adapter = _adapter(output_mode="image") with pytest.raises(ExtractorError, match="PDF input only"): From 3ac1aeac1409aeb9dede8856a65e68b70d876483 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Mon, 27 Jul 2026 14:19:03 +0000 Subject: [PATCH 11/24] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../prompt_studio_core_v2/prompt_studio_helper.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index 0c33bb847e..d51fa2cfff 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -559,7 +559,6 @@ def build_index_payload( default_profile, request_user=request_user ) - # Common path decomposition used by extract, summarize, and index directory, filename = os.path.split(file_path) stem = os.path.splitext(filename)[0] @@ -1397,7 +1396,10 @@ def _validate_image_output_pdf_only( if not adapter_id.startswith("llmwhisperer|"): return metadata = x2text.metadata or {} - if metadata.get(ImageOutputConstants.OUTPUT_MODE) != ImageOutputConstants.IMAGE_MODE: + if ( + metadata.get(ImageOutputConstants.OUTPUT_MODE) + != ImageOutputConstants.IMAGE_MODE + ): return if not ImageOutputConstants.is_pdf(file_name): raise IndexingAPIError( @@ -1481,7 +1483,6 @@ def index_document( summary_profile, request_user=request_user ) - fs_instance = EnvHelper.get_storage( storage_type=StorageType.PERMANENT, env_name=FileStorageKeys.PERMANENT_REMOTE_STORAGE, From b98945cafab3530bd59697d73b1a99b183f984ce Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Mon, 27 Jul 2026 22:38:40 +0530 Subject: [PATCH 12/24] UN-2646 [FIX] Drop remaining PR/ticket provenance from image-output comments Follow-up to the review round: remove the contradictory 'ASSUMED contract (PR #647)' vs 'verified against PR #536' notes and the 'decision 2A' / UNS-743 provenance from the adapter helper + constants comments, keeping the factual one-line reason (comments carry the why, not the ticket). Co-Authored-By: Claude Opus 4.8 (1M context) --- .../x2text/llm_whisperer_v2/src/constants.py | 6 +++--- .../x2text/llm_whisperer_v2/src/helper.py | 18 ++++++------------ 2 files changed, 9 insertions(+), 15 deletions(-) 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 9ec6c477ca..050f023271 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 @@ -34,9 +34,9 @@ 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. + # Image output mode (pdf-to-images) endpoints. Not exposed by the + # llmwhisperer-client package, so the adapter calls them via raw requests; + # see ImageOutputConfig for the wire contract. PDF_TO_IMAGES = "pdf-to-images" PDF_TO_IMAGES_STATUS = "pdf-to-images-status" PDF_TO_IMAGES_RETRIEVE = "pdf-to-images-retrieve" 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 859a5e8199..51e7db366b 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 @@ -65,7 +65,7 @@ def _send_raw_request( timeout: float = WhispererDefaults.IMAGE_REQUEST_TIMEOUT, stream: bool = False, ) -> Response: - """Single outbound raw-``requests`` code path for the adapter (UNS-743). + """Single outbound raw-``requests`` code path for the adapter. Resolves the service base URL and auth headers from ``config`` so that no caller constructs URLs or headers itself, issues the request with an @@ -434,15 +434,9 @@ 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. # - # ------------------------------------------------------------------ # + # Image output mode (pdf-to-images): these endpoints are not exposed by the + # llmwhisperer-client, so the adapter calls them via raw `requests`. The + # wire contract the adapter relies on is centralised in ImageOutputConfig. # Matches service page files like `page_001.png` / `page-1.png`. The # captured digits are passed through int() (leading zeros stripped there), @@ -485,8 +479,8 @@ def submit_pdf_to_images( 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. + service attributes usage correctly. ``tag`` falls back to the adapter + config, then the default. """ resolved_tag = WhispererRequestParams(tag=tag).tag or config.get( WhispererConfig.TAG, WhispererDefaults.TAG From 4e619bbf81e6380e75a110b06876f74ccce17e26 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Tue, 28 Jul 2026 17:24:58 +0530 Subject: [PATCH 13/24] UN-2646 [FEAT] Gate image output mode behind the cloud consumer plugin The image output mode produces per-page billed PNGs that are consumed by the VLM answer plugin, which ships only with Unstract Cloud. Without the plugins.vlm_image_answer package: - the 'image' option (and its conditional description) is stripped from the LLMWhisperer V2 adapter's JSON schema, so the UI cannot select it - adapter create/update and test-connection reject image-mode metadata with a clear error, covering API-created adapters that bypass the UI This prevents OSS deployments from configuring an extraction whose output nothing can consume. Co-Authored-By: Claude Fable 5 --- .../adapter_processor_v2/adapter_processor.py | 10 +- .../image_output_gating.py | 97 +++++++++++++++ backend/adapter_processor_v2/serializers.py | 8 ++ .../tests/test_image_output_gating.py | 110 ++++++++++++++++++ 4 files changed, 223 insertions(+), 2 deletions(-) create mode 100644 backend/adapter_processor_v2/image_output_gating.py create mode 100644 backend/adapter_processor_v2/tests/test_image_output_gating.py diff --git a/backend/adapter_processor_v2/adapter_processor.py b/backend/adapter_processor_v2/adapter_processor.py index ce453ea555..63a38284ee 100644 --- a/backend/adapter_processor_v2/adapter_processor.py +++ b/backend/adapter_processor_v2/adapter_processor.py @@ -16,6 +16,10 @@ InValidAdapterId, TestAdapterError, ) +from adapter_processor_v2.image_output_gating import ( + filter_image_output_mode, + validate_image_output_allowed, +) from unstract.sdk1.adapters.adapterkit import Adapterkit from unstract.sdk1.adapters.base import Adapter from unstract.sdk1.adapters.x2text.constants import X2TextConstants @@ -43,8 +47,9 @@ def get_json_schema(adapter_id: str) -> dict[str, Any]: AdapterKeys.ID, adapter_id ) if len(updated_adapters) != 0: - schema_details[AdapterKeys.JSON_SCHEMA] = json.loads( - updated_adapters[0].get(AdapterKeys.JSON_SCHEMA) + schema_details[AdapterKeys.JSON_SCHEMA] = filter_image_output_mode( + adapter_id, + json.loads(updated_adapters[0].get(AdapterKeys.JSON_SCHEMA)), ) else: logger.error(f"Invalid adapter Id : {adapter_id} while fetching JSON Schema") @@ -110,6 +115,7 @@ def get_adapter_data_with_key(adapter_id: str, key_value: str) -> Any: @staticmethod def test_adapter(adapter_id: str, adapter_metadata: dict[str, Any]) -> bool: + validate_image_output_allowed(adapter_metadata, adapter_id) try: adapter_type = adapter_metadata.get(AdapterKeys.ADAPTER_TYPE) diff --git a/backend/adapter_processor_v2/image_output_gating.py b/backend/adapter_processor_v2/image_output_gating.py new file mode 100644 index 0000000000..6982c4dd88 --- /dev/null +++ b/backend/adapter_processor_v2/image_output_gating.py @@ -0,0 +1,97 @@ +"""Gating for the LLMWhisperer image output mode. + +Image output mode produces per-page PNGs that are consumed by the VLM answer +plugin, which ships only with Unstract Cloud. On deployments without the +``plugins.vlm_image_answer`` package the mode is hidden from the adapter's +JSON schema and rejected at save/test time, so users cannot configure a +per-page billed extraction whose output nothing can consume. +""" + +import copy +import logging +from typing import Any + +from rest_framework.exceptions import ValidationError + +from unstract.sdk1.adapters.x2text.constants import ImageOutputConstants + +logger = logging.getLogger(__name__) + +LLMWHISPERER_ADAPTER_PREFIX = "llmwhisperer|" +IMAGE_OUTPUT_REQUIRES_CLOUD = ( + "The 'image' output mode is available only on Unstract Cloud." +) + + +def _consumer_plugin_available() -> bool: + try: + import plugins.vlm_image_answer # noqa: F401 + except ImportError: + return False + return True + + +IMAGE_OUTPUT_CONSUMER_AVAILABLE = _consumer_plugin_available() + + +def _is_image_mode_condition(block: dict[str, Any]) -> bool: + """True if an ``allOf`` block is conditioned on the image output mode.""" + const = ( + block.get("if", {}) + .get("properties", {}) + .get(ImageOutputConstants.OUTPUT_MODE, {}) + .get("const") + ) + return const == ImageOutputConstants.IMAGE_MODE + + +def filter_image_output_mode(adapter_id: str, schema: dict[str, Any]) -> dict[str, Any]: + """Strip the image output-mode option from an adapter's JSON schema. + + No-op when the consumer plugin is available, for non-LLMWhisperer + adapters, or when the schema has no image option. Returns a filtered + deep copy otherwise (the SDK-provided schema is shared state). + """ + if IMAGE_OUTPUT_CONSUMER_AVAILABLE: + return schema + if not adapter_id.startswith(LLMWHISPERER_ADAPTER_PREFIX): + return schema + output_mode = schema.get("properties", {}).get(ImageOutputConstants.OUTPUT_MODE, {}) + if ImageOutputConstants.IMAGE_MODE not in output_mode.get("enum", []): + return schema + + schema = copy.deepcopy(schema) + output_mode = schema["properties"][ImageOutputConstants.OUTPUT_MODE] + idx = output_mode["enum"].index(ImageOutputConstants.IMAGE_MODE) + output_mode["enum"].pop(idx) + enum_names = output_mode.get("enumNames") + if enum_names and len(enum_names) > idx: + enum_names.pop(idx) + if "allOf" in schema: + schema["allOf"] = [ + block for block in schema["allOf"] if not _is_image_mode_condition(block) + ] + return schema + + +def validate_image_output_allowed( + adapter_metadata: dict[str, Any] | None, adapter_id: str | None = None +) -> None: + """Reject image output mode when the consumer plugin is unavailable. + + Backstop for the schema filtering above: covers adapters created or + updated via the API (bypassing the UI form) and test-connection calls. + When ``adapter_id`` is unknown (e.g. a metadata-only update) the check + falls back to the metadata alone — only the LLMWhisperer V2 adapter + exposes an ``image`` output mode. + """ + if IMAGE_OUTPUT_CONSUMER_AVAILABLE or not adapter_metadata: + return + if ( + adapter_metadata.get(ImageOutputConstants.OUTPUT_MODE) + != ImageOutputConstants.IMAGE_MODE + ): + return + if adapter_id is not None and not adapter_id.startswith(LLMWHISPERER_ADAPTER_PREFIX): + return + raise ValidationError(IMAGE_OUTPUT_REQUIRES_CLOUD) diff --git a/backend/adapter_processor_v2/serializers.py b/backend/adapter_processor_v2/serializers.py index a5f2c492d6..63faaac348 100644 --- a/backend/adapter_processor_v2/serializers.py +++ b/backend/adapter_processor_v2/serializers.py @@ -14,6 +14,7 @@ from adapter_processor_v2.adapter_processor import AdapterProcessor from adapter_processor_v2.constants import AdapterKeys +from adapter_processor_v2.image_output_gating import validate_image_output_allowed from backend.constants import FieldLengthConstants as FLC from backend.serializers import AuditSerializer from unstract.sdk1.constants import AdapterTypes @@ -73,6 +74,13 @@ class AdapterInstanceSerializer(BaseAdapterSerializer): def to_internal_value(self, data: dict[str, Any]) -> dict[str, Any]: if data.get(AdapterKeys.ADAPTER_METADATA, None): + # Reject image output mode on deployments without the cloud + # consumer plugin, before the metadata is encrypted away. + validate_image_output_allowed( + data[AdapterKeys.ADAPTER_METADATA], + data.get(AdapterKeys.ADAPTER_ID) + or getattr(self.instance, "adapter_id", None), + ) encryption_secret: str = settings.ENCRYPTION_KEY f: Fernet = Fernet(encryption_secret.encode("utf-8")) json_string: str = json.dumps(data.pop(AdapterKeys.ADAPTER_METADATA)) diff --git a/backend/adapter_processor_v2/tests/test_image_output_gating.py b/backend/adapter_processor_v2/tests/test_image_output_gating.py new file mode 100644 index 0000000000..025b89b27c --- /dev/null +++ b/backend/adapter_processor_v2/tests/test_image_output_gating.py @@ -0,0 +1,110 @@ +"""Tests for the cloud-only gating of the LLMWhisperer image output mode. + +The image-mode consumer ships only with Unstract Cloud. Without the +``plugins.vlm_image_answer`` package, the ``image`` output mode must be +hidden from the adapter's JSON schema and rejected at save/test time. +""" + +import json +from pathlib import Path + +import pytest +from rest_framework.exceptions import ValidationError + +from adapter_processor_v2 import image_output_gating as gating +from adapter_processor_v2.image_output_gating import ( + IMAGE_OUTPUT_REQUIRES_CLOUD, + filter_image_output_mode, + validate_image_output_allowed, +) + +_LLMW_ADAPTER_ID = "llmwhisperer|0a1647f0-f65f-410d-843b-3d979c78350e" + +_SCHEMA_PATH = ( + Path(__file__).resolve().parents[3] + / "unstract/sdk1/src/unstract/sdk1/adapters/x2text/llm_whisperer_v2/src/static" + / "json_schema.json" +) + + +def _llmw_schema() -> dict: + return json.loads(_SCHEMA_PATH.read_text()) + + +@pytest.fixture +def consumer_absent(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(gating, "IMAGE_OUTPUT_CONSUMER_AVAILABLE", False) + + +@pytest.fixture +def consumer_present(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(gating, "IMAGE_OUTPUT_CONSUMER_AVAILABLE", True) + + +class TestSchemaFiltering: + def test_image_option_stripped_when_consumer_absent(self, consumer_absent) -> None: + schema = _llmw_schema() + filtered = filter_image_output_mode(_LLMW_ADAPTER_ID, schema) + + output_mode = filtered["properties"]["output_mode"] + assert "image" not in output_mode["enum"] + # enum / enumNames stay positionally paired + assert len(output_mode["enum"]) == len(output_mode["enumNames"]) + assert "Image (PDF only)" not in output_mode["enumNames"] + # The image-conditioned allOf block (conditional description) is gone + assert not any( + block.get("if", {}).get("properties", {}).get("output_mode", {}).get("const") + == "image" + for block in filtered.get("allOf", []) + ) + # Unrelated conditional blocks are preserved + assert any("if" in block for block in filtered.get("allOf", [])) + + def test_source_schema_not_mutated(self, consumer_absent) -> None: + schema = _llmw_schema() + filter_image_output_mode(_LLMW_ADAPTER_ID, schema) + assert "image" in schema["properties"]["output_mode"]["enum"] + + def test_schema_untouched_when_consumer_present(self, consumer_present) -> None: + schema = _llmw_schema() + assert filter_image_output_mode(_LLMW_ADAPTER_ID, schema) is schema + + def test_non_llmwhisperer_schema_untouched(self, consumer_absent) -> None: + schema = {"properties": {"output_mode": {"enum": ["image"]}}} + assert filter_image_output_mode("someocr|uuid", schema) is schema + + def test_schema_without_image_option_untouched(self, consumer_absent) -> None: + schema = {"properties": {"output_mode": {"enum": ["layout_preserving"]}}} + assert filter_image_output_mode(_LLMW_ADAPTER_ID, schema) is schema + + +class TestSaveTimeValidation: + def test_image_mode_rejected_when_consumer_absent(self, consumer_absent) -> None: + with pytest.raises(ValidationError, match="Unstract Cloud"): + validate_image_output_allowed({"output_mode": "image"}, _LLMW_ADAPTER_ID) + + def test_error_message_names_cloud(self, consumer_absent) -> None: + with pytest.raises(ValidationError) as excinfo: + validate_image_output_allowed({"output_mode": "image"}, _LLMW_ADAPTER_ID) + assert IMAGE_OUTPUT_REQUIRES_CLOUD in str(excinfo.value) + + def test_image_mode_allowed_when_consumer_present(self, consumer_present) -> None: + validate_image_output_allowed({"output_mode": "image"}, _LLMW_ADAPTER_ID) + + def test_other_output_modes_allowed(self, consumer_absent) -> None: + validate_image_output_allowed( + {"output_mode": "layout_preserving"}, _LLMW_ADAPTER_ID + ) + + def test_non_llmwhisperer_adapter_allowed(self, consumer_absent) -> None: + # Another adapter with a coincidental output_mode key is not gated. + validate_image_output_allowed({"output_mode": "image"}, "someocr|uuid") + + def test_unknown_adapter_id_still_rejected(self, consumer_absent) -> None: + # Metadata-only updates lack an adapter id; the metadata alone gates. + with pytest.raises(ValidationError): + validate_image_output_allowed({"output_mode": "image"}, None) + + def test_empty_metadata_allowed(self, consumer_absent) -> None: + validate_image_output_allowed(None, _LLMW_ADAPTER_ID) + validate_image_output_allowed({}, _LLMW_ADAPTER_ID) From a92f347248bd81352a399766fac16d5ca87076f7 Mon Sep 17 00:00:00 2001 From: Praveen Kumar Date: Wed, 5 Aug 2026 17:06:18 +0530 Subject: [PATCH 14/24] UN-2646 [FEAT] OSS half of the VLM image-answer feature (path contract, loader, vision policy, bridge) (#2218) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * UN-2646 [FEAT] Promote page-image path contract to the shared x2text surface Foundation for the VLM consumer (MUNS-206): the writer (adapter) and the upcoming page-image reader must agree on the storage layout without metadata persistence or a manifest sidecar. - Promote build_page_store_dir ({extract_dir}/{stem}/pages) from the LLMWhisperer helper to unstract.sdk1.adapters.x2text.constants; the helper now binds the shared function directly (identity, not a copy) - Add page-naming constants to ImageOutputConstants: PAGES_SUBFOLDER, PAGE_IMAGE_PREFIX/EXTENSION/PADDING, PAGE_GLOB_PATTERN and PAGE_NUMBER_REGEX (integer capture for natural sort — lexicographic ordering misorders past page 999) - ImageOutputConfig now aliases the shared layout constants; the tolerant ZIP-member pattern moves to a named ZIP_PAGE_MEMBER_REGEX (service wire contract, distinct from the storage contract) - New writer/reader agreement tests: path determinism across stem shapes, glob/regex contract incl. >999-page fixtures, and the writer's filename builder round-tripping through the reader's glob + regex Co-Authored-By: Claude Fable 5 * UN-2646 [FEAT] FileStorage page-image loader and vision-capability policy Reader-side sdk1 utilities for the VLM consumer (MUNS-207 + the sdk1 half of MUNS-208): - page_image_loader: discovers persisted page_NNN.png files through the FileStorage abstraction using the shared naming constants, orders by integer page index, enforces a fail-fast page cap before any bytes are read, base64-encodes, and shapes complete_vision content blocks with 'Page N' labels before each image. Typed failures keep empty (never-extracted/purged), incomplete (post-write loss, with found vs missing pages), and over-cap cases distinct — remediation copy steers to cache-bypass re-extraction with the per-page re-billing warning. - vision_capability: classifies a model as SUPPORTED / UNSUPPORTED / UNKNOWN from litellm's local model_cost registry only — never get_model_info, which can make network calls for self-hosted providers, and never bare supports_vision, which returns False for both non-vision and unknown models and would wrongly hard-block custom vision models (Ollama, proxies). Policy: hard-block only a definitive UNSUPPORTED; UNKNOWN warns and allows. Tests: 33 new (in-memory S3-like double + real local FileStorage backend, >999-page ordering, cap-before-read assertion, registry fakes + real-registry smoke checks). Full sdk1 image suite: 104 pass. Co-Authored-By: Claude Fable 5 * UN-2646 [FEAT] OSS bridge for the cloud-only VLM image-answer plugin Image-mode prompts are answered by a vision LLM via the cloud-only 'vlm-image-answer' executor plugin. This adds the OSS half (mirrors the lookup_enrichment bridge, with the opposite error policy — lookups degrade, image mode fails loudly): - workers vlm_image_answer bridge: detects image mode by resolving the x2text adapter config through the platform service (the payload carries only the instance id), cached per (execution, adapter); dispatches to the plugin with the deterministic pages dir from the shared path helper; raises structured, code-prefixed errors (IMAGE_OUTPUT_REQUIRES_CLOUD / IMAGE_OUTPUT_MISSING / IMAGE_PAGE_CAP_EXCEEDED / VISION_LLM_REQUIRED) that survive the string-only error propagation to PS and deployment responses — never a silent fallthrough to the text path - legacy_executor: image-mode branch replaces retrieval + completion only; type conversion, lookups, webhooks and challenge run unchanged. Single-pass extraction is rejected for image-mode profiles before delegating to the cloud single-pass plugin - backend vlm_utils no-op bridge (lookup_utils pattern) wired into: profile serialization (non-blocking vision_warning), API deployment creation (deploy-time guard), and the extraction choke point (VLM answer invalidation after a fresh extraction rewrites pages/) Tests: 13 bridge (detection/cache/dispatch/error mapping/single-pass guard) + 8 vlm_utils (OSS no-ops + cloud delegation policies). Full sdk1 image suite (104) and gating suite (12) stay green. Co-Authored-By: Claude Fable 5 * UN-2646 [FEAT] Surface backend vision_warning after profile save Profile save responses may carry a backend-computed vision_warning (image output mode selected with an LLM that is not verifiably vision-capable — populated by the cloud vlm_image_answer hooks, never present in OSS responses). Show it as a non-blocking warning toast so the user learns about the mismatch at config time instead of at run time. Co-Authored-By: Claude Fable 5 * UN-2646 [FIX] Surface typed errors for stale listings and vanished pages Live-testing regression: with pages purged from object storage, the loader raised a raw FileNotFoundError instead of the typed incomplete-set error — fsspec's directory cache served discovery a stale listing in the long-lived worker, so the contiguity check passed and the purge only surfaced at read time, bypassing the error mapping. - discover_page_images now refreshes the backend's listing cache first (duck-typed invalidate_cache; no-op for backends without one) - load_page_images maps a read-time FileNotFoundError (TOCTOU: page vanished after discovery) onto PageImageSetIncompleteError with the cache-bypass + re-billing remediation - in-memory test double now raises FileNotFoundError like real backends Co-Authored-By: Claude Fable 5 * UN-2646 [FIX] Sniff PDF content when the storage name has no extension Live deployment-path regression: workflow executions store the source file under an extension-less name (e.g. SOURCE), so the extension-only PDF guard false-rejected every API-deployment input — a real .pdf upload failed with the PDF-only error. _validate_pdf_only now checks the extension first and falls back to content sniffing (%PDF- magic bytes via FileStorage) when the name has no .pdf suffix; unverifiable content still rejects (fail-closed). No extra read on the common .pdf-named path. Co-Authored-By: Claude Fable 5 * UN-2646 [FIX] Raise the gating rejection with a dict detail, not a bare string Live-testing regression: validate_image_output_allowed raises from inside AdapterInstanceSerializer.to_internal_value, where DRF folds the detail into its per-field error mapping — a bare-string detail crashed error collection (ValueError: dictionary update sequence) and surfaced as a 500 'Something went wrong' instead of the clean 400 with the 'available only on Unstract Cloud' message. Use a field-keyed dict detail, and log the rejected adapter/mode for diagnosability. Co-Authored-By: Claude Fable 5 * UN-2646 [FIX] Address Greptile review: run-scoped mode cache and byte budget Two review findings on the consumer path: - Run-scoped detection cache (P1): IDE payloads carry no execution_id, so every IDE run shared one ('', adapter) cache entry — an adapter switched between text and image output kept serving the stale mode until worker restart. The cache key now scopes to execution_id or, for IDE runs, the run_id (deduplicating the N per-prompt resolutions within one run — the cache's actual purpose — with no cross-run reuse), and skips caching entirely when no scope id exists. - Aggregate byte budget on image loading (P2): the page cap bounds the COUNT of images, not their size. load_page_images now enforces a 50MB (default, disable-able) raw-byte budget while reading, raising the typed PageImageSetTooLargeError with page-range remediation; the bridge maps it to a distinct IMAGE_PAGES_TOO_LARGE error code. Tests: run-scoped vs no-scope cache behavior, byte-budget stop-at-page accounting, and the new error-code mapping. Co-Authored-By: Claude Fable 5 * UN-2646 [FIX] Enforce the byte budget from storage metadata before any read Greptile follow-up: the read-time budget still pulled each full object into worker memory before checking cumulative size, so one pathological page could spike RAM before rejection. The budget is now enforced twice: first from storage size metadata (object HEAD — zero bytes transferred) before any read, so an oversized set is rejected with no image bytes in memory; the read-time accounting stays as a belt-and-braces guard for backends without size metadata and for stat/read races. Tests: metadata pre-check rejects with zero reads; no-size() backends still enforced at read time; unreadable metadata falls back cleanly. Co-Authored-By: Claude Fable 5 * UN-2646 [FIX] Address review: extract-path keying, stamped detection, image-mode scoping Addresses Chandrasekharan's review on the consumer path: 1. Pages directory keys on a never-rewritten extract path (blocker): summarize-as-source and smart-table runs rewrite the payload file_path before the answer step, so the reader looked in the wrong directory and surfaced IMAGE_OUTPUT_MISSING with a paid remediation that could not help. The payload builders (IDE single/bulk and the structure tool task) now stamp an explicit extract_file_path, captured before any rewrite; the bridge derives pages/ from it, with the old file_path as fallback for older payloads. 2. Text-mode prompts no longer depend on the platform service: the backend stamps the x2text adapter's output_mode per prompt (it already holds the decrypted metadata), so detection is payload-only for stamped runs; the run-scoped platform resolution remains as the fallback for unstamped payloads (API deployments, older payloads). 3. Challenge and evaluation are skipped in image mode with a visible log line — both verify an answer against retrieval context, which image mode does not have; running them billed a doomed second LLM call. A vision-aware challenge is a later-phase decision. 4. Retrieval adapters are no longer constructed for image-mode prompts: detection now runs before adapter init, which proceeds LLM-only (embedding/vector DB skipped) when image mode is detected. 5. Profile-save toasts no longer swallow each other: one alert renders — 'Saved — check LLM compatibility' with the warning when present, plain success otherwise (the alert store holds a single entry, so two synchronous calls batched into showing only the last). 6. vlm_utils logs a distinct warning when plugins.vlm_image_answer is present but backend_hooks fails to import — 'cloud hooks broken' is no longer indistinguishable from 'running OSS'. Bridge API split accordingly: detect_image_mode_config (stamp fast path + platform fallback) and run_vlm_image_answer (dispatch only, keyed on extract_file_path). Tests: 23 bridge (stamped/unstamped detection, cache scoping, extract-path keying) all green; sdk1 130; backend 20. Co-Authored-By: Claude Fable 5 * [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci * UN-2646 [FIX] Bound page-image reads to the remaining byte budget Greptile's follow-up was right: the metadata pre-pass fell back to an unbounded full-object read when the size lookup failed, so one pathological object could still spike worker memory before the budget check. Replaced both layers with a single stronger mechanism — bounded reads: every page is read with length = remaining budget + 1 via FileStorage.read, giving a hard allocation ceiling of budget + 1 bytes for the whole loop, independent of object sizes or the backend's size metadata. The metadata pre-pass is gone (also Chandrasekharan's call — one mechanism, less code). Also folds in the review nits: dead PAGE_GLOB_PATTERN removed (readers list + regex, never glob) and the duplicated not-found message extracted to one helper. Tests: bounded-read guarantees pinned (single 10MB object → exactly 51 bytes read; aggregate reads <= budget + 1) — 127 sdk1 + 23 bridge pass. Co-Authored-By: Claude Fable 5 --------- Co-authored-by: Claude Fable 5 Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> --- .../image_output_gating.py | 12 +- backend/api_v2/serializers.py | 8 + .../prompt_profile_manager_v2/serializers.py | 6 + .../prompt_studio_core_v2/constants.py | 6 + .../prompt_studio_helper.py | 44 +++ backend/prompt_studio/tests/__init__.py | 0 backend/prompt_studio/tests/test_vlm_utils.py | 74 +++++ backend/prompt_studio/vlm_utils.py | 86 +++++ .../add-llm-profile/AddLlmProfile.jsx | 21 +- .../sdk1/adapters/x2text/constants.py | 45 +++ .../x2text/llm_whisperer_v2/src/constants.py | 17 +- .../x2text/llm_whisperer_v2/src/helper.py | 27 +- .../llm_whisperer_v2/src/llm_whisperer_v2.py | 33 +- .../sdk1/adapters/x2text/page_image_loader.py | 310 +++++++++++++++++ .../unstract/sdk1/utils/vision_capability.py | 103 ++++++ unstract/sdk1/tests/llmw_image_fixtures.py | 39 ++- .../sdk1/tests/test_llmw_v2_process_image.py | 33 ++ unstract/sdk1/tests/test_page_image_loader.py | 313 ++++++++++++++++++ unstract/sdk1/tests/test_vision_capability.py | 97 ++++++ .../tests/test_x2text_shared_page_path.py | 132 ++++++++ workers/executor/executors/constants.py | 6 + workers/executor/executors/exceptions.py | 21 ++ workers/executor/executors/legacy_executor.py | 105 ++++-- .../executor/executors/vlm_image_answer.py | 242 ++++++++++++++ workers/executor/tests/__init__.py | 0 .../tests/test_vlm_image_answer_bridge.py | 310 +++++++++++++++++ .../file_processing/structure_tool_task.py | 3 + 27 files changed, 2036 insertions(+), 57 deletions(-) create mode 100644 backend/prompt_studio/tests/__init__.py create mode 100644 backend/prompt_studio/tests/test_vlm_utils.py create mode 100644 backend/prompt_studio/vlm_utils.py create mode 100644 unstract/sdk1/src/unstract/sdk1/adapters/x2text/page_image_loader.py create mode 100644 unstract/sdk1/src/unstract/sdk1/utils/vision_capability.py create mode 100644 unstract/sdk1/tests/test_page_image_loader.py create mode 100644 unstract/sdk1/tests/test_vision_capability.py create mode 100644 unstract/sdk1/tests/test_x2text_shared_page_path.py create mode 100644 workers/executor/executors/vlm_image_answer.py create mode 100644 workers/executor/tests/__init__.py create mode 100644 workers/executor/tests/test_vlm_image_answer_bridge.py diff --git a/backend/adapter_processor_v2/image_output_gating.py b/backend/adapter_processor_v2/image_output_gating.py index 6982c4dd88..422e772439 100644 --- a/backend/adapter_processor_v2/image_output_gating.py +++ b/backend/adapter_processor_v2/image_output_gating.py @@ -94,4 +94,14 @@ def validate_image_output_allowed( return if adapter_id is not None and not adapter_id.startswith(LLMWHISPERER_ADAPTER_PREFIX): return - raise ValidationError(IMAGE_OUTPUT_REQUIRES_CLOUD) + logger.warning( + "Rejecting image output mode without the consumer plugin " + "(adapter_id=%s, output_mode=%s)", + adapter_id, + adapter_metadata.get(ImageOutputConstants.OUTPUT_MODE), + ) + # Dict detail, not a bare string: this is raised from inside + # ``to_internal_value``, where DRF folds the detail into its per-field + # error mapping — a bare string there crashes error collection with a + # 500 instead of surfacing a clean 400. + raise ValidationError({"adapter_metadata": [IMAGE_OUTPUT_REQUIRES_CLOUD]}) diff --git a/backend/api_v2/serializers.py b/backend/api_v2/serializers.py index 0eb99f1bce..2cba65cf49 100644 --- a/backend/api_v2/serializers.py +++ b/backend/api_v2/serializers.py @@ -8,6 +8,9 @@ from django.core.validators import RegexValidator from pipeline_v2.models import Pipeline from prompt_studio.prompt_profile_manager_v2.models import ProfileManager +from prompt_studio.vlm_utils import ( + validate_workflow_for_deployment as validate_workflow_vlm_for_deployment, +) from rest_framework import serializers from rest_framework.serializers import ( BooleanField, @@ -149,6 +152,11 @@ def validate_workflow(self, workflow): "Destination endpoint must have a connector configured for non-API and non-manual review connections before creating an API deployment." ) + # Image-output-mode profiles need a vision-capable LLM at run time; + # block deployment creation on a definitive mismatch (cloud-only + # check — no-op in OSS, where image mode is gated off entirely). + validate_workflow_vlm_for_deployment(workflow) + return workflow def validate(self, data): diff --git a/backend/prompt_studio/prompt_profile_manager_v2/serializers.py b/backend/prompt_studio/prompt_profile_manager_v2/serializers.py index 008fed3850..2bf7a97575 100644 --- a/backend/prompt_studio/prompt_profile_manager_v2/serializers.py +++ b/backend/prompt_studio/prompt_profile_manager_v2/serializers.py @@ -4,6 +4,7 @@ from backend.serializers import AuditSerializer from prompt_studio.prompt_profile_manager_v2.constants import ProfileManagerKeys +from prompt_studio.vlm_utils import get_profile_vision_warning from .models import ProfileManager @@ -38,4 +39,9 @@ def to_representation(self, instance): # type: ignore rep[ProfileManagerKeys.X2TEXT] = AdapterProcessor.get_adapter_instance_by_id( x2text ) + # Non-blocking image-mode/vision-LLM mismatch warning (cloud-only; + # always None in OSS — key omitted). + vision_warning = get_profile_vision_warning(instance) + if vision_warning: + rep["vision_warning"] = vision_warning return rep diff --git a/backend/prompt_studio/prompt_studio_core_v2/constants.py b/backend/prompt_studio/prompt_studio_core_v2/constants.py index 03bd68c1d8..3a09634b0d 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/constants.py +++ b/backend/prompt_studio/prompt_studio_core_v2/constants.py @@ -99,6 +99,12 @@ class ToolStudioPromptKeys: VARIABLE_MAP = "variable_map" RECORD = "record" FILE_PATH = "file_path" + # Extract-file path that never gets rewritten by summarize-as-source / + # smart-table overrides — the page-image reader keys on this. + EXTRACT_FILE_PATH = "extract_file_path" + # Per-prompt stamp of the x2text adapter's output mode (LLMWhisperer + # only), so the executor detects image mode without a platform call. + X2TEXT_OUTPUT_MODE = "x2text_output_mode" ENABLE_HIGHLIGHT = "enable_highlight" ENABLE_WORD_CONFIDENCE = "enable_word_confidence" REQUIRED = "required" diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index d51fa2cfff..047254fc3b 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -76,6 +76,7 @@ OutputManagerHelper, ) from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt +from prompt_studio.vlm_utils import invalidate_vlm_answers_on_reextraction from unstract.core.pubsub_helper import LogPublisher from unstract.sdk1.adapters.x2text.constants import ImageOutputConstants from unstract.sdk1.constants import LogLevel @@ -444,6 +445,7 @@ def _build_prompt_output( output[TSPKeys.SIMILARITY_TOP_K] = profile_manager.similarity_top_k output[TSPKeys.SECTION] = profile_manager.section output[TSPKeys.X2TEXT_ADAPTER] = x2text + PromptStudioHelper._stamp_x2text_output_mode(output, profile_manager) webhook_enabled = bool(prompt.enable_postprocessing_webhook) webhook_url = (prompt.postprocessing_webhook_url or "").strip() @@ -821,6 +823,9 @@ def build_fetch_response_payload( enable_highlight=tool.enable_highlight, ) + # Captured before the summarize override: the page-image reader must + # key on the extract path even when answers run over the summary. + image_extract_path = extract_path is_summary = tool.summarize_as_source if is_summary: profile_manager.chunk_size = 0 @@ -867,6 +872,7 @@ def build_fetch_response_payload( output[TSPKeys.SIMILARITY_TOP_K] = profile_manager.similarity_top_k output[TSPKeys.SECTION] = profile_manager.section output[TSPKeys.X2TEXT_ADAPTER] = x2text + PromptStudioHelper._stamp_x2text_output_mode(output, profile_manager) webhook_enabled = bool(prompt.enable_postprocessing_webhook) webhook_url = (prompt.postprocessing_webhook_url or "").strip() @@ -928,6 +934,7 @@ def build_fetch_response_payload( TSPKeys.FILE_NAME: doc_name, TSPKeys.FILE_HASH: file_hash, TSPKeys.FILE_PATH: extract_path, + TSPKeys.EXTRACT_FILE_PATH: image_extract_path, Common.LOG_EVENTS_ID: StateStore.get(Common.LOG_EVENTS_ID), TSPKeys.EXECUTION_SOURCE: ExecutionSource.IDE.value, TSPKeys.CUSTOM_DATA: tool.custom_data, @@ -1045,6 +1052,9 @@ def build_bulk_fetch_response_payload( enable_highlight=tool.enable_highlight, ) + # Captured before the summarize override: the page-image reader must + # key on the extract path even when answers run over the summary. + image_extract_path = extract_path is_summary = tool.summarize_as_source if is_summary: profile_manager.chunk_size = 0 @@ -1122,6 +1132,7 @@ def build_bulk_fetch_response_payload( TSPKeys.FILE_NAME: doc_name, TSPKeys.FILE_HASH: file_hash, TSPKeys.FILE_PATH: extract_path, + TSPKeys.EXTRACT_FILE_PATH: image_extract_path, Common.LOG_EVENTS_ID: StateStore.get(Common.LOG_EVENTS_ID), TSPKeys.EXECUTION_SOURCE: ExecutionSource.IDE.value, TSPKeys.CUSTOM_DATA: tool.custom_data, @@ -1371,6 +1382,29 @@ def fetch_prompt_from_tool(tool_id: str) -> list[ToolStudioPrompt]: ).order_by(TSPKeys.SEQUENCE_NUMBER) return prompt_instances + @staticmethod + def _stamp_x2text_output_mode(output: dict, profile_manager) -> None: + """Stamp the x2text output mode onto a per-prompt payload. + + Lets the executor detect image mode from the payload instead of a + platform-service call. LLMWhisperer-only (the sole adapter with an + image output mode); best-effort — a metadata read failure leaves + the stamp absent and the executor falls back to live resolution. + """ + x2text = getattr(profile_manager, "x2text", None) + if x2text is None: + return + try: + adapter_id = str(getattr(x2text, "adapter_id", "") or "") + if not adapter_id.startswith("llmwhisperer|"): + return + metadata = x2text.metadata or {} + output[TSPKeys.X2TEXT_OUTPUT_MODE] = metadata.get( + ImageOutputConstants.OUTPUT_MODE + ) + except Exception: + logger.exception("Could not stamp x2text output mode; will resolve live") + @staticmethod def _validate_image_output_pdf_only( profile_manager: ProfileManager, file_name: str @@ -2049,6 +2083,7 @@ def _fetch_response( output[TSPKeys.SIMILARITY_TOP_K] = profile_manager.similarity_top_k output[TSPKeys.SECTION] = profile_manager.section output[TSPKeys.X2TEXT_ADAPTER] = x2text + PromptStudioHelper._stamp_x2text_output_mode(output, profile_manager) # Webhook postprocessing settings webhook_enabled = bool(prompt.enable_postprocessing_webhook) webhook_url = (prompt.postprocessing_webhook_url or "").strip() @@ -2576,6 +2611,15 @@ def dynamic_extractor( f"Extraction completed but status not saved." ) + # A fresh (non-cache-hit) extraction rewrote any persisted page + # images — stored VLM answers for this document are stale. + # No-op in OSS (cloud-only hook). + invalidate_vlm_answers_on_reextraction( + document_id=str(document_id), + profile_manager=profile_manager, + extract_file_path=extract_file_path, + ) + return extracted_text @staticmethod diff --git a/backend/prompt_studio/tests/__init__.py b/backend/prompt_studio/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/backend/prompt_studio/tests/test_vlm_utils.py b/backend/prompt_studio/tests/test_vlm_utils.py new file mode 100644 index 0000000000..a19d6462ab --- /dev/null +++ b/backend/prompt_studio/tests/test_vlm_utils.py @@ -0,0 +1,74 @@ +"""Tests for the OSS vlm_utils bridge (no-op without the cloud package). + +Mirrors the lookup_utils bridge contract: every helper degrades safely +in OSS, delegates when the cloud hooks module is present, and the +non-critical hooks (warning, invalidation) never let a cloud-side +failure break the OSS operation they ride on. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +from _pytest.monkeypatch import MonkeyPatch + +from prompt_studio import vlm_utils + + +class TestOssNoOps: + def test_cloud_package_absent_in_oss(self) -> None: + assert vlm_utils.VLM_IMAGE_ANSWER_AVAILABLE is False + + def test_vision_warning_is_none(self) -> None: + assert vlm_utils.get_profile_vision_warning(SimpleNamespace()) is None + + def test_deployment_validation_is_noop(self) -> None: + vlm_utils.validate_workflow_for_deployment(SimpleNamespace()) # no raise + + def test_invalidation_is_noop(self) -> None: + vlm_utils.invalidate_vlm_answers_on_reextraction( + document_id="d1", + profile_manager=SimpleNamespace(), + extract_file_path="/x/extract/doc.txt", + ) # no raise + + +@pytest.fixture +def cloud_hooks(monkeypatch: MonkeyPatch) -> MagicMock: + hooks = MagicMock() + monkeypatch.setattr(vlm_utils, "_hooks", hooks) + monkeypatch.setattr(vlm_utils, "VLM_IMAGE_ANSWER_AVAILABLE", True) + return hooks + + +class TestCloudDelegation: + def test_vision_warning_delegates(self, cloud_hooks: MagicMock) -> None: + cloud_hooks.get_profile_vision_warning.return_value = "warn!" + assert vlm_utils.get_profile_vision_warning(SimpleNamespace()) == "warn!" + + def test_vision_warning_failure_swallowed(self, cloud_hooks: MagicMock) -> None: + # A warning must never break profile save/read. + cloud_hooks.get_profile_vision_warning.side_effect = RuntimeError("x") + assert vlm_utils.get_profile_vision_warning(SimpleNamespace()) is None + + def test_deployment_validation_propagates(self, cloud_hooks: MagicMock) -> None: + # Deploy-time rejection is a hard gate — errors must propagate. + cloud_hooks.validate_workflow_for_deployment.side_effect = ValueError("no") + with pytest.raises(ValueError): + vlm_utils.validate_workflow_for_deployment(SimpleNamespace()) + + def test_invalidation_delegates_and_swallows_failure( + self, cloud_hooks: MagicMock + ) -> None: + profile = SimpleNamespace() + vlm_utils.invalidate_vlm_answers_on_reextraction( + document_id="d1", profile_manager=profile, extract_file_path="/e.txt" + ) + cloud_hooks.invalidate_vlm_answers_on_reextraction.assert_called_once_with( + document_id="d1", profile_manager=profile, extract_file_path="/e.txt" + ) + # Invalidation failure must not fail the extraction it rides on. + cloud_hooks.invalidate_vlm_answers_on_reextraction.side_effect = RuntimeError + vlm_utils.invalidate_vlm_answers_on_reextraction( + document_id="d1", profile_manager=profile, extract_file_path="/e.txt" + ) # no raise diff --git a/backend/prompt_studio/vlm_utils.py b/backend/prompt_studio/vlm_utils.py new file mode 100644 index 0000000000..462c480039 --- /dev/null +++ b/backend/prompt_studio/vlm_utils.py @@ -0,0 +1,86 @@ +"""Bridge helpers for the cloud-only VLM image-answer feature. No-ops in OSS. + +Image output mode is answered by a vision LLM through the cloud-only +``vlm-image-answer`` plugin. The backend touch points below (profile-save +vision warning, deploy-time validation, answer-cache invalidation on +re-extraction) delegate to ``plugins.vlm_image_answer.backend_hooks`` when +that cloud package is present and degrade to no-ops when it is not — OSS +additionally hides the image output mode entirely via +``adapter_processor_v2.image_output_gating``. +""" + +import logging +from typing import Any + +logger = logging.getLogger(__name__) + +try: + from plugins.vlm_image_answer import backend_hooks as _hooks + + VLM_IMAGE_ANSWER_AVAILABLE = True +except ImportError: + _hooks = None + VLM_IMAGE_ANSWER_AVAILABLE = False + # Distinguish "running OSS" (package absent — expected, silent) from + # "cloud hooks are broken" (package present but backend_hooks failed + # to import): in the latter case the adapter gating still enables + # image output mode while these hooks quietly stop existing. + try: + import plugins.vlm_image_answer # noqa: F401 + except ImportError: + pass + else: + logger.warning( + "plugins.vlm_image_answer is present but backend_hooks failed " + "to import — VLM profile warnings, deploy-time validation and " + "re-extraction invalidation are disabled while image output " + "mode remains enabled" + ) + + +def get_profile_vision_warning(profile_manager: Any) -> str | None: + """Non-blocking warning when an image-mode profile's LLM lacks vision. + + Returns a human-readable warning string, or None (always None in OSS). + Never raises — a warning must not break profile save/read. + """ + if not VLM_IMAGE_ANSWER_AVAILABLE: + return None + try: + return _hooks.get_profile_vision_warning(profile_manager) + except Exception: + logger.exception("VLM vision warning check failed; skipping warning") + return None + + +def validate_workflow_for_deployment(workflow: Any) -> None: + """Deploy-time guard: reject deployments that cannot serve image mode. + + The cloud hook raises ``rest_framework.serializers.ValidationError`` + for a definitive misconfiguration (e.g. image-mode profile with a + known non-vision LLM); OSS is a no-op (image mode is gated off). + """ + if not VLM_IMAGE_ANSWER_AVAILABLE: + return + _hooks.validate_workflow_for_deployment(workflow) + + +def invalidate_vlm_answers_on_reextraction( + document_id: str, profile_manager: Any, extract_file_path: str +) -> None: + """Invalidate stored VLM answers after a re-extraction rewrote pages/. + + Called from the extraction choke point right after a successful + (non-cache-hit) extraction. Never raises — invalidation failure must + not fail the extraction itself; the cloud hook logs and degrades. + """ + if not VLM_IMAGE_ANSWER_AVAILABLE: + return + try: + _hooks.invalidate_vlm_answers_on_reextraction( + document_id=document_id, + profile_manager=profile_manager, + extract_file_path=extract_file_path, + ) + except Exception: + logger.exception("VLM answer invalidation failed after re-extraction") diff --git a/frontend/src/components/custom-tools/add-llm-profile/AddLlmProfile.jsx b/frontend/src/components/custom-tools/add-llm-profile/AddLlmProfile.jsx index 80cf854c2d..cdf5f20759 100644 --- a/frontend/src/components/custom-tools/add-llm-profile/AddLlmProfile.jsx +++ b/frontend/src/components/custom-tools/add-llm-profile/AddLlmProfile.jsx @@ -356,10 +356,23 @@ function AddLlmProfile({ llmProfiles: newLlmProfiles, }; updateCustomTool(updatedState); - setAlertDetails({ - type: "success", - content: "Saved successfully", - }); + // Single alert: the store holds one alertDetails object, so two + // synchronous calls would batch and only the last would render. + // vision_warning is a backend-computed advisory (image output + // mode with an LLM that may not support vision); absent in OSS. + if (data?.vision_warning) { + setAlertDetails({ + type: "warning", + title: "Saved — check LLM compatibility", + content: data.vision_warning, + duration: 10, + }); + } else { + setAlertDetails({ + type: "success", + content: "Saved successfully", + }); + } if (newLlmProfiles?.length === 1) { // Set the first LLM profile as default diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/constants.py b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/constants.py index 110e7aef27..feb89d8f6f 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/constants.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/constants.py @@ -35,7 +35,52 @@ class ImageOutputConstants: "Please provide a PDF file or select a text output mode." ) + # --- Page image storage layout (writer/reader contract) --- + # The adapter (writer) persists one PNG per page under a ``pages`` + # subfolder as ``page_NNN.png`` (zero-padded to PAGE_NUMBER_PADDING + # digits; four or more digits appear naturally past page 999). Readers + # list the directory and MUST order pages by the integer captured by + # PAGE_NUMBER_REGEX — never lexicographically, which silently + # misorders once page numbers outgrow the padding. + PAGES_SUBFOLDER = "pages" + PAGE_IMAGE_PREFIX = "page_" + PAGE_IMAGE_EXTENSION = ".png" + PAGE_NUMBER_PADDING = 3 + # First capture group is the numeric page index (as a string, possibly + # zero-padded) — cast to int before sorting. + PAGE_NUMBER_REGEX = r"page_(\d+)\.png" + + # Leading bytes of every PDF file — the content-based check for inputs + # whose storage name carries no extension (workflow executions store the + # source file under an extension-less name like ``SOURCE``). + PDF_MAGIC_BYTES = b"%PDF-" + @staticmethod def is_pdf(file_name: str) -> bool: """Return True when ``file_name`` is a PDF (case-insensitive suffix).""" return Path(file_name).suffix.lower() == ImageOutputConstants.PDF_EXTENSION + + @staticmethod + def is_pdf_bytes(header: bytes) -> bool: + """Return True when ``header`` starts with the PDF magic bytes.""" + return bytes(header).startswith(ImageOutputConstants.PDF_MAGIC_BYTES) + + +def build_page_store_dir(output_file_path: str | None, input_file_path: str) -> str: + """Per-document folder for page images: ``{extract_dir}/{stem}/pages``. + + The single canonical derivation shared by the adapter (writer) and any + page-image reader, so both sides agree on the location without metadata + persistence or a manifest sidecar. Keyed on the document ``stem`` (the + same discriminator the extract ``.txt`` files alongside use), not the + per-run whisper_hash. This is collision-safe against concurrent documents + in the same project, is reconstructible from ``output_file_path`` alone, + and — being stable across runs — makes a re-extraction overwrite its own + pages instead of orphaning a fresh tree in FileStorage on every run. + + Pure and deterministic: no I/O, no lookups. + """ + reference = output_file_path or input_file_path + base_dir = str(Path(reference).parent) if reference else "." + stem = Path(reference).stem if reference else "document" + return str(Path(base_dir) / stem / ImageOutputConstants.PAGES_SUBFOLDER) 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 050f023271..ceb74c712f 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 @@ -181,10 +181,19 @@ class ImageOutputConfig: 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" + # Sourced from the shared x2text surface so the writer (this adapter) + # and any page-image reader agree on one storage contract. + PAGE_IMAGE_PREFIX = ImageOutputConstants.PAGE_IMAGE_PREFIX + PAGE_IMAGE_EXTENSION = ImageOutputConstants.PAGE_IMAGE_EXTENSION + PAGE_NUMBER_PADDING = ImageOutputConstants.PAGE_NUMBER_PADDING + PAGES_SUBFOLDER = ImageOutputConstants.PAGES_SUBFOLDER + + # ZIP member names as sent by the LLMWhisperer service. Distinct from + # the storage contract above (ImageOutputConstants.PAGE_NUMBER_REGEX): + # this tolerates service-side naming variations (``page-1.png``) when + # ingesting the download; persisted files are always renamed to the + # strict ``page_NNN.png`` layout. + ZIP_PAGE_MEMBER_REGEX = r"page[_-]?(\d+)\.png$" # --- PDF-only validation (shared with the backend index-time guard) --- PDF_EXTENSION = ImageOutputConstants.PDF_EXTENSION 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 51e7db366b..14e1f385e8 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 @@ -17,7 +17,12 @@ ) 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.constants import ( + X2TextConstants, +) +from unstract.sdk1.adapters.x2text.constants import ( + build_page_store_dir as _shared_build_page_store_dir, +) from unstract.sdk1.adapters.x2text.dto import PageImageReference from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.constants import ( ImageOutputConfig, @@ -441,7 +446,7 @@ def write_output_to_file( # Matches service page files like `page_001.png` / `page-1.png`. The # captured digits are passed through int() (leading zeros stripped there), # so no separate `0*` prefix is needed — keeping the pattern linear. - _PAGE_IMAGE_RE = re.compile(r"page[_-]?(\d+)\.png$", re.IGNORECASE) + _PAGE_IMAGE_RE = re.compile(ImageOutputConfig.ZIP_PAGE_MEMBER_REGEX, re.IGNORECASE) @staticmethod def _safe_json(response: Response) -> dict[str, Any]: @@ -718,21 +723,9 @@ def verify_page_count( status_code=502, ) - @staticmethod - def build_page_store_dir(output_file_path: str | None, input_file_path: str) -> str: - """Per-document folder for page images: ``{extract_dir}/{stem}/pages``. - - Keyed on the document ``stem`` (the same discriminator the extract - ``.txt`` files alongside use), not the per-run whisper_hash. This is - collision-safe against concurrent documents in the same project, is - reconstructible from ``output_file_path`` alone, and — being stable - across runs — makes a re-extraction overwrite its own pages instead of - orphaning a fresh tree in FileStorage on every run. - """ - reference = output_file_path or input_file_path - base_dir = str(Path(reference).parent) if reference else "." - stem = Path(reference).stem if reference else "document" - return str(Path(base_dir) / stem / ImageOutputConfig.PAGES_SUBFOLDER) + # Single canonical derivation of ``{extract_dir}/{stem}/pages`` shared + # by writer and reader alike (see unstract.sdk1.adapters.x2text.constants). + build_page_store_dir = staticmethod(_shared_build_page_store_dir) @staticmethod def _page_image_filename(page_number: int) -> str: 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 94b9a7f440..464d69f2c4 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 @@ -5,7 +5,10 @@ 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.constants import ( + ImageOutputConstants, + X2TextConstants, +) from unstract.sdk1.adapters.x2text.dto import ( TextExtractionMetadata, TextExtractionResult, @@ -66,17 +69,31 @@ def test_connection(self) -> bool: return True @staticmethod - def _validate_pdf_only(input_file_path: str) -> None: + def _validate_pdf_only(input_file_path: str, fs: FileStorage | None = None) -> None: """Enforce the PDF-only constraint for image output mode (v1). + Checks the filename extension first; when the storage name carries no + ``.pdf`` suffix, falls back to content sniffing — workflow executions + store the source file under an extension-less name (e.g. ``SOURCE``), + so an extension-only check would false-reject every deployment input. + Fail-closed: if the content cannot be verified either, reject. + The message is sourced from ``ImageOutputConfig`` so it stays identical to the UI-layer validation surfaced in ``adapter_processor_v2``. """ - if not ImageOutputConfig.is_pdf(input_file_path): - raise ExtractorError( - ImageOutputConfig.PDF_ONLY_ERROR, - status_code=400, - ) + if ImageOutputConfig.is_pdf(input_file_path): + return + if fs is not None: + try: + header = fs.read(path=input_file_path, mode="rb", length=5) + except Exception: + header = b"" + if ImageOutputConstants.is_pdf_bytes(header): + return + raise ExtractorError( + ImageOutputConfig.PDF_ONLY_ERROR, + status_code=400, + ) def _process_image_mode( self, @@ -98,7 +115,7 @@ def _process_image_mode( ``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) + self._validate_pdf_only(input_file_path, fs=fs) whisper_hash, page_images = LLMWhispererHelper.get_page_images( config=self.config, input_file_path=input_file_path, diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/page_image_loader.py b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/page_image_loader.py new file mode 100644 index 0000000000..e293ef5df2 --- /dev/null +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/page_image_loader.py @@ -0,0 +1,310 @@ +"""FileStorage-backed reader for persisted page images. + +Reader half of the page-image storage contract defined in +``unstract.sdk1.adapters.x2text.constants``: the LLMWhisperer adapter +(writer) persists one PNG per page as ``page_NNN.png`` under the directory +returned by ``build_page_store_dir``; this module discovers those pages via +the shared naming constants, orders them by their **integer** page index +(never lexicographically — that misorders past the zero-padding width), +base64-encodes them, and shapes multimodal content blocks for +``LLM.complete_vision``. + +All reads go through the FileStorage abstraction so the same code serves +local disk and remote object storage. Discovery lists the +deterministic path — no manifest, no metadata transport. + +Failure modes are typed so callers can surface distinct, actionable errors: + +- ``PageImagesNotFoundError`` — directory missing/empty (never extracted in + image mode, or fully purged). +- ``PageImageSetIncompleteError`` — pages exist but the 1..N set is broken + (post-write loss). Distinct from "not found" so remediation can differ. +- ``PageCapExceededError`` — document larger than the page cap; callers + must fail explicitly rather than silently truncate. +""" + +import base64 +import logging +import re +from dataclasses import dataclass +from pathlib import PurePosixPath + +from unstract.sdk1.adapters.x2text.constants import ImageOutputConstants +from unstract.sdk1.file_storage import FileStorage + +logger = logging.getLogger(__name__) + +# Conservative default; effective value is supplied by the caller +# (platform-configured), this is only the fallback. +DEFAULT_PAGE_CAP = 20 + +# Aggregate raw-byte budget across all loaded pages. The page cap bounds the +# COUNT of images, not their size — without a byte budget, unusually large +# renders would grow worker memory and the provider request unbounded +# (base64 adds ~33% on top). 50MB raw comfortably exceeds any normal +# LLMWhisperer render while staying inside provider request limits. +DEFAULT_MAX_TOTAL_BYTES = 50 * 1024 * 1024 + +_PAGE_NAME_RE = re.compile(ImageOutputConstants.PAGE_NUMBER_REGEX) + + +class PageImageLoadError(Exception): + """Base error for page-image discovery/loading failures.""" + + def __init__(self, message: str, *, page_store_dir: str) -> None: + """Store the offending pages directory alongside the message.""" + super().__init__(message) + self.page_store_dir = page_store_dir + + +class PageImagesNotFoundError(PageImageLoadError): + """The pages directory is missing or contains no page images. + + Either the document was never extracted in image output mode, or the + persisted images were purged. Remediation: re-extract the document with + cache bypass (note: re-extraction re-submits to LLMWhisperer and is + billed per page — a plain re-run is an extraction cache hit and will + NOT regenerate images). + """ + + +class PageImageSetIncompleteError(PageImageLoadError): + """Pages exist but the contiguous 1..N set is broken (post-write loss).""" + + def __init__( + self, + message: str, + *, + page_store_dir: str, + found_pages: list[int], + missing_pages: list[int], + ) -> None: + """Record which pages were found vs missing for remediation UIs.""" + super().__init__(message, page_store_dir=page_store_dir) + self.found_pages = found_pages + self.missing_pages = missing_pages + + +class PageCapExceededError(PageImageLoadError): + """The document has more pages than the configured cap allows.""" + + def __init__( + self, message: str, *, page_store_dir: str, page_count: int, page_cap: int + ) -> None: + """Record the observed page count and the cap that was exceeded.""" + super().__init__(message, page_store_dir=page_store_dir) + self.page_count = page_count + self.page_cap = page_cap + + +class PageImageSetTooLargeError(PageImageLoadError): + """The combined size of the page images exceeds the byte budget.""" + + def __init__( + self, + message: str, + *, + page_store_dir: str, + total_bytes: int, + max_total_bytes: int, + ) -> None: + """Record the observed total and the budget that was exceeded.""" + super().__init__(message, page_store_dir=page_store_dir) + self.total_bytes = total_bytes + self.max_total_bytes = max_total_bytes + + +@dataclass(frozen=True) +class LoadedPageImage: + """A page image read from FileStorage, base64-encoded for a VLM call.""" + + page_number: int + path: str + base64_data: str + + +def _not_found(page_store_dir: str) -> PageImagesNotFoundError: + return PageImagesNotFoundError( + f"No page images found at '{page_store_dir}'. The document has " + "not been extracted in image output mode (or its images were " + "removed). Re-extract the document with cache bypass to " + "regenerate them (re-extraction is billed per page).", + page_store_dir=page_store_dir, + ) + + +def discover_page_images(fs: FileStorage, page_store_dir: str) -> list[tuple[int, str]]: + """Discover persisted page images, ordered by integer page number. + + Lists ``page_store_dir`` through FileStorage, keeps entries whose + basename matches the shared ``PAGE_NUMBER_REGEX`` (others are logged + and skipped), and validates the set is exactly 1..N. + + Returns: + ``[(page_number, full_path), ...]`` sorted by page number. + + Raises: + PageImagesNotFoundError: directory missing or no page images in it. + PageImageSetIncompleteError: duplicate or missing page numbers. + """ + # Object-store backends serve listings from fsspec's directory cache in + # long-lived worker processes; a page purged since the last listing would + # still "exist" here and only blow up at read time. Refresh the cache + # first so discovery reflects reality (no-op for backends without one). + invalidate = getattr(getattr(fs, "fs", None), "invalidate_cache", None) + if callable(invalidate): + try: + invalidate(page_store_dir) + except Exception: # pragma: no cover - cache refresh is best-effort + logger.debug("Could not invalidate listing cache for %s", page_store_dir) + + try: + entries = fs.ls(page_store_dir) if fs.exists(page_store_dir) else None + except FileNotFoundError: + entries = None + if entries is None: + raise _not_found(page_store_dir) + + pages: dict[int, str] = {} + for entry in entries: + name = PurePosixPath(str(entry)).name + match = _PAGE_NAME_RE.fullmatch(name) + if not match: + logger.debug("Skipping non-page entry in %s: %s", page_store_dir, name) + continue + page_number = int(match.group(1)) + if page_number in pages: + raise PageImageSetIncompleteError( + f"Duplicate page number {page_number} in '{page_store_dir}' " + f"({PurePosixPath(pages[page_number]).name} vs {name}); the " + "page set is corrupt. Re-extract the document with cache " + "bypass (billed per page).", + page_store_dir=page_store_dir, + found_pages=sorted(pages), + missing_pages=[], + ) + pages[page_number] = str(entry) + + if not pages: + raise _not_found(page_store_dir) + + found = sorted(pages) + missing = sorted(set(range(1, found[-1] + 1)) - set(found)) + if missing: + raise PageImageSetIncompleteError( + f"Only {len(found)} of {found[-1]} page images are present at " + f"'{page_store_dir}' (missing pages: {missing[:10]}" + f"{'…' if len(missing) > 10 else ''}). Re-extract the document " + "with cache bypass to regenerate them (billed per page).", + page_store_dir=page_store_dir, + found_pages=found, + missing_pages=missing, + ) + + return [(number, pages[number]) for number in found] + + +def load_page_images( + fs: FileStorage, + page_store_dir: str, + *, + page_cap: int | None = DEFAULT_PAGE_CAP, + max_total_bytes: int | None = DEFAULT_MAX_TOTAL_BYTES, +) -> list[LoadedPageImage]: + """Discover, cap-check, read, and base64-encode all page images. + + The page-count cap runs before any bytes are read so an oversized + document fails fast and cheap. The aggregate byte budget is enforced + with **bounded reads**: each page is read with a length limit of the + remaining budget plus one byte, so no read — not even of a single + pathological object — can ever allocate more than the budget in + worker memory, regardless of the object's actual size or whether the + backend exposes size metadata. ``None`` disables either limit. + + Raises: + PageCapExceededError: more pages than ``page_cap`` allows. + PageImageSetTooLargeError: pages total more than ``max_total_bytes``. + (plus the discovery errors from ``discover_page_images``) + """ + discovered = discover_page_images(fs, page_store_dir) + if page_cap is not None and len(discovered) > page_cap: + raise PageCapExceededError( + f"Document exceeds {page_cap} pages for image output mode " + f"({len(discovered)} pages found). Reduce the page range (e.g. " + "via the adapter's 'pages to extract' setting) or raise the " + "configured page cap.", + page_store_dir=page_store_dir, + page_count=len(discovered), + page_cap=page_cap, + ) + + loaded = [] + total_bytes = 0 + for page_number, path in discovered: + read_kwargs: dict[str, int] = {} + if max_total_bytes is not None: + # Bounded read: never pull more than the remaining budget (+1 + # byte to detect the overflow) into memory — the hard + # allocation ceiling for this loop is max_total_bytes + 1. + read_kwargs["length"] = max_total_bytes - total_bytes + 1 + try: + data = bytes(fs.read(path=path, mode="rb", **read_kwargs)) + except FileNotFoundError as e: + # TOCTOU guard: the page vanished between discovery and read + # (purged concurrently, or discovery served a stale listing). + # Surface the typed incomplete-set error, never a raw IO error. + raise PageImageSetIncompleteError( + f"Page image {PurePosixPath(path).name} is missing from " + f"'{page_store_dir}' (it disappeared after discovery); the " + "page set is incomplete. Re-extract the document with cache " + "bypass to regenerate it (billed per page).", + page_store_dir=page_store_dir, + found_pages=[n for n, _ in discovered if n != page_number], + missing_pages=[page_number], + ) from e + total_bytes += len(data) + if max_total_bytes is not None and total_bytes > max_total_bytes: + # Stop before encoding/retaining more — the page cap bounds the + # count, this bounds the payload. + raise PageImageSetTooLargeError( + f"Page images total more than " + f"{max_total_bytes // (1024 * 1024)}MB by page {page_number} " + f"of {len(discovered)} — too large to send to the LLM in " + "one request. Reduce the page range (e.g. via the adapter's " + "'pages to extract' setting).", + page_store_dir=page_store_dir, + total_bytes=total_bytes, + max_total_bytes=max_total_bytes, + ) + encoded = base64.b64encode(data).decode("ascii") + loaded.append( + LoadedPageImage(page_number=page_number, path=path, base64_data=encoded) + ) + return loaded + + +def build_vision_message_content( + pages: list[LoadedPageImage], prompt_text: str +) -> list[dict]: + """Shape multimodal content blocks for ``LLM.complete_vision``. + + Layout: the prompt text first (task framing), then for each page — in + page order — a ``Page N`` text label immediately followed by that + page's image block. The explicit labels preserve reading order for the + model and enable page citations later at zero cost. + + Returns the ``content`` list for a single user message; callers wrap it + as ``[{"role": "user", "content": content}]``. + """ + content: list[dict] = [{"type": "text", "text": prompt_text}] + for page in pages: + content.append({"type": "text", "text": f"Page {page.page_number}"}) + content.append( + { + "type": "image_url", + "image_url": { + "url": f"data:image/png;base64,{page.base64_data}", + }, + } + ) + return content diff --git a/unstract/sdk1/src/unstract/sdk1/utils/vision_capability.py b/unstract/sdk1/src/unstract/sdk1/utils/vision_capability.py new file mode 100644 index 0000000000..5e21b345e4 --- /dev/null +++ b/unstract/sdk1/src/unstract/sdk1/utils/vision_capability.py @@ -0,0 +1,103 @@ +"""Vision-capability detection and gating policy for LLM models. + +The image output mode sends page images to the profile's LLM via +``complete_vision``; a non-vision model fails only at run time with an +opaque provider error. This module classifies a model's vision support +up front and applies the gating policy shared by config-time warnings +and run-time guards. + +Classification uses LiteLLM's **local** ``model_cost`` registry only — +never ``get_model_info`` — because ``get_model_info`` can make network +calls for self-hosted providers (e.g. it queries the Ollama server), and +``litellm.supports_vision`` alone returns ``False`` for both known +non-vision models *and* unknown models, which would wrongly hard-block +custom/self-hosted vision models (LiteLLM proxies, Ollama). + +Policy (locked): hard-block only on a **definitive** "known model, no +vision support"; unknown/custom models are allowed with a warning — the +provider's own runtime error remains the backstop. +""" + +import logging +from dataclasses import dataclass +from enum import Enum + +import litellm + +logger = logging.getLogger(__name__) + + +class VisionSupport(Enum): + """Classification of a model's vision (image input) capability.""" + + SUPPORTED = "supported" + UNSUPPORTED = "unsupported" + UNKNOWN = "unknown" + + +@dataclass(frozen=True) +class VisionValidationResult: + """Outcome of applying the gating policy to a model.""" + + model: str + support: VisionSupport + allowed: bool + message: str | None + + +def check_vision_support(model: str) -> VisionSupport: + """Classify ``model``'s vision capability from the local registry. + + The registry is keyed both with and without provider prefixes + (``anthropic/claude-…`` vs ``claude-…``); a model found under either + form is "known". On known models the ``supports_vision`` flag is + authoritative — absence means no vision support. Models not in the + registry (self-hosted, proxies, brand-new releases) are UNKNOWN. + """ + if not model: + return VisionSupport.UNKNOWN + registry = litellm.model_cost + entry = registry.get(model) + if entry is None and "/" in model: + entry = registry.get(model.split("/", 1)[-1]) + if entry is None: + return VisionSupport.UNKNOWN + if entry.get("supports_vision"): + return VisionSupport.SUPPORTED + return VisionSupport.UNSUPPORTED + + +def validate_vision_capability(model: str) -> VisionValidationResult: + """Apply the image-mode gating policy to ``model``. + + Returns a result rather than raising so callers can map it to their + own error/warning surfaces (structured API errors, profile-save + warnings). ``allowed`` is False only for a definitive UNSUPPORTED. + """ + support = check_vision_support(model) + if support is VisionSupport.SUPPORTED: + return VisionValidationResult( + model=model, support=support, allowed=True, message=None + ) + if support is VisionSupport.UNKNOWN: + message = ( + f"Model '{model}' is not in the capability registry, so its " + "image (vision) support cannot be verified. The run will " + "proceed; if the model does not accept image input, the " + "provider will reject the request." + ) + logger.warning(message) + return VisionValidationResult( + model=model, support=support, allowed=True, message=message + ) + return VisionValidationResult( + model=model, + support=support, + allowed=False, + message=( + f"Model '{model}' does not support image input. Image output " + "mode requires a vision-capable LLM — update this profile's " + "LLM to a vision model (e.g. a GPT-4o, Claude, or Gemini " + "vision model) and re-run." + ), + ) diff --git a/unstract/sdk1/tests/llmw_image_fixtures.py b/unstract/sdk1/tests/llmw_image_fixtures.py index 344bd00116..6194e78602 100644 --- a/unstract/sdk1/tests/llmw_image_fixtures.py +++ b/unstract/sdk1/tests/llmw_image_fixtures.py @@ -110,14 +110,47 @@ def write( return len(payload) def read( - self, path: str, mode: str = "rb", encoding: str = "utf-8", **_: object + self, + path: str, + mode: str = "rb", + encoding: str = "utf-8", + length: int = -1, + **_: object, ) -> bytes | str: - payload = self._files[str(path)] + try: + payload = self._files[str(path)] + except KeyError: + # Real backends (fsspec/local) raise FileNotFoundError. + raise FileNotFoundError(str(path)) from None + if length is not None and length >= 0: + payload = payload[:length] 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 + if key in self._files or key in self._dirs: + return True + # S3-like: a "directory" exists when any object lives under it. + prefix = key.rstrip("/") + "/" + return any(stored.startswith(prefix) for stored in self._files) + + def size(self, path: str) -> int: + """Byte size from 'metadata' (fsspec info-style), like real backends.""" + try: + return len(self._files[str(path)]) + except KeyError: + raise FileNotFoundError(str(path)) from None + + def ls(self, path: str) -> list[str]: + """Direct children of ``path`` (full paths), fsspec-style.""" + from pathlib import PurePosixPath + + parent = str(path).rstrip("/") + return sorted( + stored + for stored in self._files + if str(PurePosixPath(stored).parent) == parent + ) @property def stored_paths(self) -> list[str]: diff --git a/unstract/sdk1/tests/test_llmw_v2_process_image.py b/unstract/sdk1/tests/test_llmw_v2_process_image.py index 09fee77ebe..9f058b34d2 100644 --- a/unstract/sdk1/tests/test_llmw_v2_process_image.py +++ b/unstract/sdk1/tests/test_llmw_v2_process_image.py @@ -165,3 +165,36 @@ def test_validate_pdf_only_accepts_pdf(self) -> None: def test_validate_pdf_only_rejects_other(self) -> None: with pytest.raises(ExtractorError, match="PDF input only"): LLMWhispererV2._validate_pdf_only("/tmp/doc.tiff") + + def test_extensionless_pdf_accepted_via_content_sniff(self) -> None: + # Workflow executions store inputs under extension-less names + # (e.g. SOURCE); the guard must sniff the content, not just the + # storage filename, or every deployment input is false-rejected. + class _Fs: + def read(self, path: str, mode: str = "rb", **_: object) -> bytes: + return b"%PDF-1.7 rest-of-file" + + LLMWhispererV2._validate_pdf_only("/exec/data/SOURCE", fs=_Fs()) # no raise + + def test_extensionless_non_pdf_rejected_via_content_sniff(self) -> None: + class _Fs: + def read(self, path: str, mode: str = "rb", **_: object) -> bytes: + return b"PK\x03\x04zipfile" + + with pytest.raises(ExtractorError, match="PDF input only"): + LLMWhispererV2._validate_pdf_only("/exec/data/SOURCE", fs=_Fs()) + + def test_extensionless_unreadable_rejected_fail_closed(self) -> None: + class _Fs: + def read(self, path: str, mode: str = "rb", **_: object) -> bytes: + raise OSError("storage down") + + with pytest.raises(ExtractorError, match="PDF input only"): + LLMWhispererV2._validate_pdf_only("/exec/data/SOURCE", fs=_Fs()) + + def test_pdf_extension_skips_content_read(self) -> None: + class _Fs: + def read(self, path: str, mode: str = "rb", **_: object) -> bytes: + raise AssertionError("must not read when the extension is .pdf") + + LLMWhispererV2._validate_pdf_only("/tmp/doc.pdf", fs=_Fs()) # no raise diff --git a/unstract/sdk1/tests/test_page_image_loader.py b/unstract/sdk1/tests/test_page_image_loader.py new file mode 100644 index 0000000000..decde94b75 --- /dev/null +++ b/unstract/sdk1/tests/test_page_image_loader.py @@ -0,0 +1,313 @@ +"""Tests for the FileStorage-backed page-image loader (reader side). + +Covers discovery + natural ordering (incl. >999 pages), the page cap, +base64 encoding and vision message-block construction, and the typed +empty/partial/duplicate failure modes — across the in-memory S3-like +double and the real local-filesystem FileStorage backend. +""" + +import base64 +from pathlib import Path + +import pytest +from llmw_image_fixtures import InMemoryFileStorage, minimal_png +from unstract.sdk1.adapters.x2text.page_image_loader import ( + DEFAULT_PAGE_CAP, + LoadedPageImage, + PageCapExceededError, + PageImageSetIncompleteError, + PageImageSetTooLargeError, + PageImagesNotFoundError, + build_vision_message_content, + discover_page_images, + load_page_images, +) +from unstract.sdk1.file_storage import FileStorage, FileStorageProvider + +_DIR = "/data/extract/doc/pages" + + +def _store( + pages: dict[int, bytes], extra: dict[str, bytes] | None = None +) -> InMemoryFileStorage: + fs = InMemoryFileStorage() + for number, data in pages.items(): + fs.write(path=f"{_DIR}/page_{number:03d}.png", mode="wb", data=data) + for name, data in (extra or {}).items(): + fs.write(path=f"{_DIR}/{name}", mode="wb", data=data) + return fs + + +class TestDiscovery: + def test_orders_by_integer_page_number(self) -> None: + fs = _store({n: b"x" for n in (3, 1, 2)}) + assert [n for n, _ in discover_page_images(fs, _DIR)] == [1, 2, 3] + + def test_returns_full_paths(self) -> None: + fs = _store({1: b"x"}) + assert discover_page_images(fs, _DIR) == [(1, f"{_DIR}/page_001.png")] + + def test_natural_sort_beyond_999_pages(self) -> None: + # Lexicographic ordering would put page_1000 before page_999. + fs = InMemoryFileStorage() + for n in (999, 1000, 1, 1001): + fs.write(path=f"{_DIR}/page_{n:03d}.png", mode="wb", data=b"x") + # Fill 2..998 so the set is contiguous. + for n in range(2, 999): + fs.write(path=f"{_DIR}/page_{n:03d}.png", mode="wb", data=b"x") + numbers = [n for n, _ in discover_page_images(fs, _DIR)] + assert numbers == list(range(1, 1002)) + + def test_ignores_non_page_entries(self) -> None: + fs = _store({1: b"x", 2: b"y"}, extra={"thumbnail.png": b"t", "notes.txt": b"n"}) + assert [n for n, _ in discover_page_images(fs, _DIR)] == [1, 2] + + +class TestFailureModes: + def test_missing_directory_raises_not_found(self) -> None: + with pytest.raises(PageImagesNotFoundError) as excinfo: + discover_page_images(InMemoryFileStorage(), _DIR) + # Remediation must steer to cache-bypass re-extraction + billing note. + assert "cache bypass" in str(excinfo.value) + assert "billed per page" in str(excinfo.value) + + def test_directory_with_only_foreign_files_raises_not_found(self) -> None: + fs = _store({}, extra={"thumbnail.png": b"t"}) + with pytest.raises(PageImagesNotFoundError): + discover_page_images(fs, _DIR) + + def test_partial_set_raises_incomplete_with_missing_pages(self) -> None: + fs = _store({1: b"a", 2: b"b", 4: b"d", 7: b"g"}) + with pytest.raises(PageImageSetIncompleteError) as excinfo: + discover_page_images(fs, _DIR) + err = excinfo.value + assert err.found_pages == [1, 2, 4, 7] + assert err.missing_pages == [3, 5, 6] + assert "cache bypass" in str(err) + + def test_empty_and_partial_are_distinct_types(self) -> None: + # Callers branch remediation copy on the exception type; neither may + # be a subclass of the other. + assert not issubclass(PageImagesNotFoundError, PageImageSetIncompleteError) + assert not issubclass(PageImageSetIncompleteError, PageImagesNotFoundError) + + def test_duplicate_page_numbers_raise_incomplete(self) -> None: + fs = _store({1: b"a"}) + # page_001.png and page_1.png parse to the same page number. + fs.write(path=f"{_DIR}/page_1.png", mode="wb", data=b"dup") + with pytest.raises(PageImageSetIncompleteError, match="Duplicate"): + discover_page_images(fs, _DIR) + + +class TestPageCap: + def test_within_cap_loads(self) -> None: + fs = _store({1: b"a", 2: b"b"}) + assert len(load_page_images(fs, _DIR, page_cap=2)) == 2 + + def test_over_cap_raises_with_clear_message(self) -> None: + fs = _store({n: b"x" for n in range(1, 6)}) + with pytest.raises(PageCapExceededError) as excinfo: + load_page_images(fs, _DIR, page_cap=4) + err = excinfo.value + assert err.page_count == 5 + assert err.page_cap == 4 + assert "exceeds 4 pages" in str(err) + + def test_cap_check_precedes_reads(self) -> None: + # Fail-fast: no image bytes are read for an oversized document. + fs = _store({n: b"x" for n in range(1, 6)}) + reads: list[str] = [] + original_read = fs.read + fs.read = lambda path, **kw: reads.append(path) or original_read(path, **kw) + with pytest.raises(PageCapExceededError): + load_page_images(fs, _DIR, page_cap=1) + assert reads == [] + + def test_none_disables_cap(self) -> None: + fs = _store({n: b"x" for n in range(1, DEFAULT_PAGE_CAP + 5)}) + loaded = load_page_images(fs, _DIR, page_cap=None) + assert len(loaded) == DEFAULT_PAGE_CAP + 4 + + +class TestLoadingAndEncoding: + def test_base64_round_trip(self) -> None: + payload = minimal_png() + fs = _store({1: payload}) + [loaded] = load_page_images(fs, _DIR) + assert isinstance(loaded, LoadedPageImage) + assert base64.b64decode(loaded.base64_data) == payload + + def test_loaded_pages_keep_page_order(self) -> None: + fs = _store({2: b"two", 1: b"one", 3: b"three"}) + loaded = load_page_images(fs, _DIR) + assert [p.page_number for p in loaded] == [1, 2, 3] + assert base64.b64decode(loaded[0].base64_data) == b"one" + + +class TestVisionMessageContent: + def test_prompt_first_then_labelled_pages(self) -> None: + pages = [ + LoadedPageImage(page_number=1, path="p1", base64_data="QQ=="), + LoadedPageImage(page_number=2, path="p2", base64_data="Qg=="), + ] + content = build_vision_message_content(pages, "What is the total?") + assert content[0] == {"type": "text", "text": "What is the total?"} + # "Page N" label immediately precedes each image, in page order. + assert content[1] == {"type": "text", "text": "Page 1"} + assert content[2]["type"] == "image_url" + assert content[2]["image_url"]["url"] == "data:image/png;base64,QQ==" + assert content[3] == {"type": "text", "text": "Page 2"} + assert content[4]["image_url"]["url"] == "data:image/png;base64,Qg==" + assert len(content) == 5 + + def test_no_pages_yields_prompt_only(self) -> None: + assert build_vision_message_content([], "q") == [{"type": "text", "text": "q"}] + + +class TestLocalFileStorageBackend: + """UNS-809: the loader behaves identically on a real FileStorage backend.""" + + def _local_fs(self) -> FileStorage: + return FileStorage(provider=FileStorageProvider.LOCAL) + + def test_discovery_and_load_on_local_backend(self, tmp_path: Path) -> None: + fs = self._local_fs() + pages_dir = str(tmp_path / "doc" / "pages") + fs.mkdir(pages_dir) + payloads = {1: b"one", 2: b"two", 10: b"ten"} + for n in range(1, 11): + fs.write( + path=f"{pages_dir}/page_{n:03d}.png", + mode="wb", + data=payloads.get(n, b"x"), + ) + loaded = load_page_images(fs, pages_dir, page_cap=None) + assert [p.page_number for p in loaded] == list(range(1, 11)) + assert base64.b64decode(loaded[9].base64_data) == b"ten" + + def test_missing_dir_on_local_backend(self, tmp_path: Path) -> None: + with pytest.raises(PageImagesNotFoundError): + discover_page_images(self._local_fs(), str(tmp_path / "absent" / "pages")) + + def test_partial_set_on_local_backend(self, tmp_path: Path) -> None: + fs = self._local_fs() + pages_dir = str(tmp_path / "doc" / "pages") + fs.mkdir(pages_dir) + for n in (1, 3): + fs.write(path=f"{pages_dir}/page_{n:03d}.png", mode="wb", data=b"x") + with pytest.raises(PageImageSetIncompleteError) as excinfo: + discover_page_images(fs, pages_dir) + assert excinfo.value.missing_pages == [2] + + +class TestStaleListingAndToctou: + """Regressions from live testing. + + Object-store listing caches and read-time disappearance must surface + typed errors, never raw IO errors. + """ + + def test_read_time_file_not_found_maps_to_incomplete(self) -> None: + # Discovery sees 3 pages (e.g. a stale fsspec dircache), but page 2 + # was purged — the read must raise the typed incomplete-set error. + fs = _store({1: b"a", 2: b"b", 3: b"c"}) + del fs._files[f"{_DIR}/page_002.png"] + + class StaleLsFs: + def exists(self, path: str) -> bool: + return True + + def ls(self, path: str) -> list[str]: + return [f"{_DIR}/page_00{n}.png" for n in (1, 2, 3)] + + def read(self, path: str, mode: str = "rb", **_: object) -> bytes: + return fs.read(path, mode) + + with pytest.raises(PageImageSetIncompleteError) as excinfo: + load_page_images(StaleLsFs(), _DIR) + err = excinfo.value + assert err.missing_pages == [2] + assert err.found_pages == [1, 3] + assert "cache bypass" in str(err) + + def test_discovery_invalidates_backend_listing_cache(self) -> None: + # When the FileStorage wraps an fsspec filesystem exposing + # invalidate_cache (s3fs etc.), discovery must refresh it first. + calls: list[str] = [] + + class Underlying: + def invalidate_cache(self, path: str) -> None: + calls.append(path) + + fs = _store({1: b"a"}) + fs.fs = Underlying() + discover_page_images(fs, _DIR) + assert calls == [_DIR] + + def test_backends_without_listing_cache_are_fine(self) -> None: + # The in-memory double has no .fs attribute — must not error. + fs = _store({1: b"a"}) + assert [n for n, _ in discover_page_images(fs, _DIR)] == [1] + + +class TestByteBudget: + """The page cap bounds count; the byte budget bounds payload size.""" + + def test_over_budget_raises_typed_error(self) -> None: + fs = _store({1: b"a" * 30, 2: b"b" * 30, 3: b"c" * 30}) + with pytest.raises(PageImageSetTooLargeError) as excinfo: + load_page_images(fs, _DIR, max_total_bytes=50) + err = excinfo.value + # Bounded read: page 2 was read with length 21 (remaining + 1), so + # the recorded total is budget + 1, and page 3 was never touched. + assert err.total_bytes == 51 + assert err.max_total_bytes == 50 + assert "pages to extract" in str(err) + + def test_single_oversized_page_reads_are_bounded(self) -> None: + # A single pathological object must never be fully allocated: each + # read is capped at remaining-budget + 1 bytes regardless of the + # object's real size (no size metadata needed). + fs = _store({1: b"g" * (10 * 1024 * 1024)}) # 10MB object + lengths: list[object] = [] + original_read = fs.read + + def recording_read(path: str, **kw: object) -> bytes | str: + lengths.append(kw.get("length")) + return original_read(path, **kw) + + fs.read = recording_read + with pytest.raises(PageImageSetTooLargeError): + load_page_images(fs, _DIR, max_total_bytes=50) + assert lengths == [51] # only 51 bytes ever entered memory + + def test_budget_never_over_allocates_across_pages(self) -> None: + # Aggregate guarantee: sum of bytes actually read stays <= budget+1. + fs = _store({1: b"a" * 30, 2: b"b" * 30, 3: b"c" * 30}) + read_bytes: list[int] = [] + original_read = fs.read + + def recording_read(path: str, **kw: object) -> bytes | str: + data = original_read(path, **kw) + read_bytes.append(len(data)) + return data + + fs.read = recording_read + with pytest.raises(PageImageSetTooLargeError): + load_page_images(fs, _DIR, max_total_bytes=50) + assert sum(read_bytes) <= 51 + + def test_within_budget_loads(self) -> None: + fs = _store({1: b"a" * 10, 2: b"b" * 10}) + assert len(load_page_images(fs, _DIR, max_total_bytes=25)) == 2 + + def test_none_disables_budget(self) -> None: + fs = _store({1: b"a" * 100}) + assert len(load_page_images(fs, _DIR, max_total_bytes=None)) == 1 + + def test_default_budget_is_generous(self) -> None: + from unstract.sdk1.adapters.x2text.page_image_loader import ( + DEFAULT_MAX_TOTAL_BYTES, + ) + + assert DEFAULT_MAX_TOTAL_BYTES == 50 * 1024 * 1024 diff --git a/unstract/sdk1/tests/test_vision_capability.py b/unstract/sdk1/tests/test_vision_capability.py new file mode 100644 index 0000000000..b27b610727 --- /dev/null +++ b/unstract/sdk1/tests/test_vision_capability.py @@ -0,0 +1,97 @@ +"""Tests for vision-capability detection and the image-mode gating policy. + +Registry lookups are exercised against a controlled fake of +``litellm.model_cost`` so results don't drift with litellm releases; a +couple of smoke tests hit the real registry for stable, long-lived models. +""" + +import pytest +from _pytest.monkeypatch import MonkeyPatch +from unstract.sdk1.utils import vision_capability as vc +from unstract.sdk1.utils.vision_capability import ( + VisionSupport, + check_vision_support, + validate_vision_capability, +) + +_FAKE_REGISTRY = { + "vision-model": {"supports_vision": True}, + "provider/prefixed-vision": {"supports_vision": True}, + "text-only-model": {"litellm_provider": "x"}, # known, no vision flag + "explicit-no-vision": {"supports_vision": False}, +} + + +@pytest.fixture +def fake_registry(monkeypatch: MonkeyPatch) -> None: + monkeypatch.setattr(vc.litellm, "model_cost", _FAKE_REGISTRY) + + +class TestCheckVisionSupport: + def test_known_vision_model(self, fake_registry: None) -> None: + assert check_vision_support("vision-model") is VisionSupport.SUPPORTED + + def test_prefix_stripped_lookup(self, fake_registry: None) -> None: + # Registry keyed without the provider prefix still resolves. + assert check_vision_support("provider/prefixed-vision") is ( + VisionSupport.SUPPORTED + ) + + def test_known_model_without_flag_is_unsupported(self, fake_registry: None) -> None: + # In the registry, absence of supports_vision on a KNOWN model is + # authoritative "no vision" (litellm omits the key rather than + # setting False). + assert check_vision_support("text-only-model") is VisionSupport.UNSUPPORTED + + def test_explicit_false_is_unsupported(self, fake_registry: None) -> None: + assert check_vision_support("explicit-no-vision") is VisionSupport.UNSUPPORTED + + def test_unregistered_model_is_unknown(self, fake_registry: None) -> None: + assert check_vision_support("ollama/llava") is VisionSupport.UNKNOWN + + def test_empty_model_is_unknown(self, fake_registry: None) -> None: + assert check_vision_support("") is VisionSupport.UNKNOWN + + def test_no_network_calls_ever(self, monkeypatch: MonkeyPatch) -> None: + # get_model_info can hit the network for self-hosted providers — + # classification must never call it. + monkeypatch.setattr( + vc.litellm, + "get_model_info", + lambda *a, **k: pytest.fail("get_model_info must not be called"), + raising=False, + ) + check_vision_support("ollama/anything") + check_vision_support("gpt-4o") + + +class TestPolicy: + def test_supported_allows_silently(self, fake_registry: None) -> None: + result = validate_vision_capability("vision-model") + assert result.allowed is True + assert result.message is None + + def test_unknown_warns_and_allows(self, fake_registry: None) -> None: + result = validate_vision_capability("ollama/custom-vlm") + assert result.allowed is True + assert result.support is VisionSupport.UNKNOWN + assert result.message and "cannot be verified" in result.message + + def test_unsupported_blocks_and_names_model(self, fake_registry: None) -> None: + result = validate_vision_capability("text-only-model") + assert result.allowed is False + assert "text-only-model" in result.message + assert "vision-capable" in result.message + + +class TestRealRegistrySmoke: + """Long-stable models against the real litellm registry.""" + + def test_gpt_4o_supported(self) -> None: + assert check_vision_support("gpt-4o") is VisionSupport.SUPPORTED + + def test_gpt_35_turbo_unsupported(self) -> None: + assert check_vision_support("gpt-3.5-turbo") is VisionSupport.UNSUPPORTED + + def test_fabricated_model_unknown(self) -> None: + assert check_vision_support("no-such/model-xyz-123") is VisionSupport.UNKNOWN diff --git a/unstract/sdk1/tests/test_x2text_shared_page_path.py b/unstract/sdk1/tests/test_x2text_shared_page_path.py new file mode 100644 index 0000000000..1319bfbee9 --- /dev/null +++ b/unstract/sdk1/tests/test_x2text_shared_page_path.py @@ -0,0 +1,132 @@ +"""Writer/reader path-agreement tests for the shared page-image contract. + +The adapter (writer) and any page-image reader must derive the +``{extract_dir}/{stem}/pages`` directory through the single shared +``build_page_store_dir`` helper, and discover/order pages via the shared +naming constants. These tests pin that contract: pure functions and +constants only — no I/O, no network, no adapter/consumer classes. +""" + +import re + +import pytest +from unstract.sdk1.adapters.x2text.constants import ( + ImageOutputConstants, + build_page_store_dir, +) + +# (output_file_path, input_file_path, expected) — expected derives from the +# output path's stem when present (the extract-file discriminator). +_PATH_CASES = [ + pytest.param( + "/data/extract/doc.txt", "/in/doc.pdf", "/data/extract/doc/pages", id="flat" + ), + pytest.param( + "/a/b/c/d/extract/report.txt", + "/uploads/report.pdf", + "/a/b/c/d/extract/report/pages", + id="nested-output-dir", + ), + pytest.param( + "/data/report.v2.final.txt", + "/in/report.v2.final.pdf", + "/data/report.v2.final/pages", + id="stem-with-dots", + ), + pytest.param( + "/data/annual report (2026).txt", + "/in/annual report (2026).pdf", + "/data/annual report (2026)/pages", + id="stem-with-spaces-and-specials", + ), + pytest.param("/data/x.txt", "/in/x.pdf", "/data/x/pages", id="single-char-stem"), + pytest.param( + None, "/in/scan.pdf", "/in/scan/pages", id="no-output-path-falls-back-to-input" + ), +] + + +class TestBuildPageStoreDir: + @pytest.mark.parametrize(("output_path", "input_path", "expected"), _PATH_CASES) + def test_expected_path_and_determinism( + self, output_path: str | None, input_path: str, expected: str + ) -> None: + first = build_page_store_dir(output_path, input_path) + second = build_page_store_dir(output_path, input_path) + assert first == expected + assert first == second # pure + deterministic + + @pytest.mark.parametrize(("output_path", "input_path", "expected"), _PATH_CASES) + def test_path_ends_with_pages_subfolder( + self, output_path: str | None, input_path: str, expected: str + ) -> None: + result = build_page_store_dir(output_path, input_path) + assert result.split("/")[-1] == ImageOutputConstants.PAGES_SUBFOLDER + + def test_writer_and_reader_share_one_implementation(self) -> None: + # The adapter exposes the helper as a staticmethod bound to the very + # same shared function — identity, not a reimplementation. Any reader + # importing from the shared surface therefore agrees byte-for-byte. + from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.helper import ( + LLMWhispererHelper, + ) + + assert LLMWhispererHelper.build_page_store_dir is build_page_store_dir + + +class TestPageNamingConstants: + @pytest.mark.parametrize( + ("filename", "captured"), + [ + ("page_001.png", "001"), + ("page_042.png", "042"), + ("page_100.png", "100"), + ("page_999.png", "999"), + # Past page 999 the writer naturally emits 4+ digits — the + # regex must keep matching (a 3-digit-only pattern would + # silently drop pages of very large documents). + ("page_1000.png", "1000"), + ], + ) + def test_number_regex_captures_page_index(self, filename: str, captured: str) -> None: + match = re.search(ImageOutputConstants.PAGE_NUMBER_REGEX, filename) + assert match is not None + assert match.group(1) == captured + + @pytest.mark.parametrize( + "filename", ["thumbnail.png", "page_abc.png", "page_.png", "page_001.jpg"] + ) + def test_number_regex_rejects_non_page_files(self, filename: str) -> None: + assert re.fullmatch(ImageOutputConstants.PAGE_NUMBER_REGEX, filename) is None + + def test_natural_sort_via_captured_int(self) -> None: + # The reason the regex exists: integer sort of the captured group + # orders pages correctly where lexicographic sort fails past the + # zero-padding width. + names = ["page_1000.png", "page_999.png", "page_010.png", "page_001.png"] + page_re = re.compile(ImageOutputConstants.PAGE_NUMBER_REGEX) + ordered = sorted(names, key=lambda n: int(page_re.search(n).group(1))) + assert ordered == [ + "page_001.png", + "page_010.png", + "page_999.png", + "page_1000.png", + ] + # Lexicographic order puts page_1000 before page_999 — the misorder + # the integer sort exists to prevent. + assert sorted(names) != ordered + + +class TestWriterFilenamesMatchReaderContract: + def test_writer_filename_matches_reader_regex(self) -> None: + # The writer's filename builder must produce names the reader-side + # regex discovers and parses — the two halves of the contract. + from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.helper import ( + LLMWhispererHelper, + ) + + for page in (1, 42, 999, 1000): + name = LLMWhispererHelper._page_image_filename(page) + match = re.fullmatch(ImageOutputConstants.PAGE_NUMBER_REGEX, name) + assert match is not None + assert int(match.group(1)) == page diff --git a/workers/executor/executors/constants.py b/workers/executor/executors/constants.py index 9eddab8423..9c837fd9cd 100644 --- a/workers/executor/executors/constants.py +++ b/workers/executor/executors/constants.py @@ -20,6 +20,12 @@ class PromptServiceConstants: VECTOR_DB = "vector-db" EMBEDDING = "embedding" X2TEXT_ADAPTER = "x2text_adapter" + # Extract-file path that is never rewritten by summarize-as-source / + # smart-table overrides — the page-image reader keys on this. + EXTRACT_FILE_PATH = "extract_file_path" + # Per-prompt stamp of the x2text adapter's output mode (LLMWhisperer + # only); lets image-mode detection skip the platform-service call. + X2TEXT_OUTPUT_MODE = "x2text_output_mode" CHUNK_OVERLAP = "chunk-overlap" LLM = "llm" IS_ASSERT = "is_assert" diff --git a/workers/executor/executors/exceptions.py b/workers/executor/executors/exceptions.py index 3db8358827..f94e6c790f 100644 --- a/workers/executor/executors/exceptions.py +++ b/workers/executor/executors/exceptions.py @@ -82,3 +82,24 @@ def __init__(self, variable: str, reason: str, is_ide: bool = True): f"Custom data error for variable '{variable_display}': {reason} {help_text}" ) super().__init__(message=message) + + +class VlmImageAnswerError(LegacyExecutorError): + """Raised when an image-mode prompt cannot be answered. + + Image output mode requires the cloud-only "vlm-image-answer" plugin; + when it is missing, or the vision path fails in a way the user must + act on (non-vision LLM, missing images, page cap), the prompt must + fail loudly — never fall through to the text path, which would + silently answer against the one-line extraction summary. + + ``error_code`` is a stable machine-readable identifier; it is also + prefixed onto the message so it survives the string-only error + propagation to Prompt Studio and API deployment responses. + """ + + code = 400 + + def __init__(self, message: str, error_code: str): + self.error_code = error_code + super().__init__(message=f"{error_code}: {message}") diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index ce7fbea0d1..cf0ebbea13 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -26,6 +26,10 @@ run_lookup_enrichment, run_webhook_postprocessing, ) +from executor.executors.vlm_image_answer import ( + detect_image_mode_config, + run_vlm_image_answer, +) from unstract.sdk1.adapters.exceptions import AdapterError from unstract.sdk1.adapters.x2text.constants import X2TextConstants @@ -1754,10 +1758,16 @@ def _execute_single_prompt( ) usage_kwargs = {"run_id": run_id, "execution_id": execution_id} + # Image output mode: detected up front (payload stamp fast-path, + # platform-service fallback) so retrieval adapters are never + # constructed for a prompt that answers from page images. + vlm_config = detect_image_mode_config( + output=output, shim=shim, usage_kwargs=usage_kwargs + ) llm, embedding, vector_db = self._init_llm_and_retrieval( output=output, shim=shim, - chunk_size=chunk_size, + chunk_size=0 if vlm_config is not None else chunk_size, llm_cls=llm_cls, embedding_compat_cls=embedding_compat_cls, vector_db_cls=vector_db_cls, @@ -1771,7 +1781,26 @@ def _execute_single_prompt( answer = "NA" retrieval_strategy = output.get(PSKeys.RETRIEVAL_STRATEGY) valid_strategies = {s.value for s in RetrievalStrategy} - if retrieval_strategy in valid_strategies: + if vlm_config is not None: + # Image output mode: the document has page images, not + # text — the answer comes from a vision LLM (cloud + # plugin); RAG retrieval is skipped entirely. The pages + # directory keys on the never-rewritten extract path (the + # payload FILE_PATH may point at the summarize output or + # the original source for smart-table runs). + answer = run_vlm_image_answer( + output=output, + shim=shim, + llm=llm, + extract_file_path=(params.get(PSKeys.EXTRACT_FILE_PATH) or file_path), + execution_source=execution_source, + metadata=metadata, + metrics=metrics, + x2text_config=vlm_config, + usage_kwargs=usage_kwargs, + ) + metadata[PSKeys.CONTEXT][prompt_name] = [] + elif retrieval_strategy in valid_strategies: if chunk_size > 0: shim.stream_log(f"Retrieving context for: `{prompt_name}`") logger.info( @@ -1859,30 +1888,40 @@ def _execute_single_prompt( shim=shim, ) - records.extend( - self._run_challenge_if_enabled( - tool_settings=tool_settings, + if vlm_config is None: + records.extend( + self._run_challenge_if_enabled( + tool_settings=tool_settings, + output=output, + structured_output=structured_output, + context_list=context_list, + llm=llm, + llm_cls=llm_cls, + usage_kwargs=usage_kwargs, + run_id=run_id, + platform_api_key=platform_api_key, + metadata=metadata, + shim=shim, + prompt_name=prompt_name, + ) + ) + self._run_evaluation_if_enabled( output=output, - structured_output=structured_output, context_list=context_list, - llm=llm, - llm_cls=llm_cls, - usage_kwargs=usage_kwargs, - run_id=run_id, + structured_output=structured_output, platform_api_key=platform_api_key, - metadata=metadata, shim=shim, prompt_name=prompt_name, ) - ) - self._run_evaluation_if_enabled( - output=output, - context_list=context_list, - structured_output=structured_output, - platform_api_key=platform_api_key, - shim=shim, - prompt_name=prompt_name, - ) + else: + # Image mode has no retrieval context; challenge and + # evaluation verify an answer AGAINST context, so running + # them here would bill a doomed second LLM call. A + # vision-aware challenge is a later-phase decision. + shim.stream_log( + f"Skipped challenge/evaluation for `{prompt_name}` " + "(image output mode has no retrieval context)" + ) shim.stream_log(f"Completed prompt: `{prompt_name}`") val = structured_output.get(prompt_name) @@ -2292,6 +2331,32 @@ def _handle_single_pass_extraction( {"output": dict, "metadata": dict, "metrics": dict} """ + from executor.executors.constants import PromptServiceConstants as PSKeys + from executor.executors.vlm_image_answer import raise_if_image_mode_unsupported + + # Image output mode cannot run single-pass: one combined prompt over + # the "full text" would silently answer from the one-line extraction + # summary. (The answer_prompt fallback below re-checks per prompt; + # this covers the cloud single-pass plugin delegation too.) + params = context.executor_params + tool_settings = params.get(PSKeys.TOOL_SETTINGS) or {} + outputs = params.get(PSKeys.OUTPUTS) or [{}] + x2text_instance_id = tool_settings.get(PSKeys.X2TEXT_ADAPTER) or outputs[0].get( + PSKeys.X2TEXT_ADAPTER + ) + if x2text_instance_id: + raise_if_image_mode_unsupported( + operation="Single-pass extraction", + adapter_instance_id=str(x2text_instance_id), + shim=self._build_shim( + platform_api_key=params.get(PSKeys.PLATFORM_SERVICE_API_KEY, ""), + component=self._log_component, + ), + scope_id=str( + params.get(PSKeys.EXECUTION_ID) or params.get(PSKeys.RUN_ID) or "" + ), + ) + try: from unstract.sdk1.execution.registry import ExecutorRegistry diff --git a/workers/executor/executors/vlm_image_answer.py b/workers/executor/executors/vlm_image_answer.py new file mode 100644 index 0000000000..b5c0b4b8e3 --- /dev/null +++ b/workers/executor/executors/vlm_image_answer.py @@ -0,0 +1,242 @@ +"""Bridge for the cloud-only VLM image-answer plugin. + +Image output mode (an LLMWhisperer x2text adapter with +``output_mode == "image"``) persists per-page PNGs instead of text, so +answering a prompt against such a document means sending those images to a +vision-capable LLM. That consumer ships only with Unstract Cloud, as the +``vlm-image-answer`` executor plugin. + +This OSS bridge owns detection and dispatch (mirroring the +``lookup_enrichment`` bridge, with the opposite error policy — lookups +degrade gracefully, image mode must fail loudly): + +- Detects image mode from the profile's x2text adapter configuration. The + executor payload carries only the adapter *instance id* (no metadata), + so the config is resolved through the platform service once per + (execution, adapter) and cached. +- When the plugin is installed, delegates the answer to it. RAG retrieval + is skipped by the caller — image mode has no text to retrieve. +- When the plugin is absent, raises a structured error rather than letting + the prompt silently answer against the one-line extraction summary. +""" + +import logging +from collections import OrderedDict +from typing import Any + +from executor.executors.constants import PromptServiceConstants as PSKeys +from executor.executors.exceptions import VlmImageAnswerError +from executor.executors.file_utils import FileUtils +from executor.executors.plugins.loader import ExecutorPluginLoader + +from unstract.sdk1.adapters.x2text.constants import ( + ImageOutputConstants, + build_page_store_dir, +) +from unstract.sdk1.adapters.x2text.page_image_loader import ( + PageCapExceededError, + PageImageLoadError, + PageImageSetIncompleteError, + PageImageSetTooLargeError, + PageImagesNotFoundError, +) +from unstract.sdk1.platform import PlatformHelper + +logger = logging.getLogger(__name__) + +PLUGIN_NAME = "vlm-image-answer" + +# Stable machine-readable error codes (prefixed onto error messages so they +# survive the string-only propagation to PS / deployment API responses). +IMAGE_OUTPUT_REQUIRES_CLOUD = "IMAGE_OUTPUT_REQUIRES_CLOUD" +IMAGE_OUTPUT_MISSING = "IMAGE_OUTPUT_MISSING" +IMAGE_PAGE_CAP_EXCEEDED = "IMAGE_PAGE_CAP_EXCEEDED" +IMAGE_PAGES_TOO_LARGE = "IMAGE_PAGES_TOO_LARGE" +IMAGE_OUTPUT_UNSUPPORTED_OPERATION = "IMAGE_OUTPUT_UNSUPPORTED_OPERATION" +VISION_LLM_REQUIRED = "VISION_LLM_REQUIRED" + +_LLMWHISPERER_ADAPTER_PREFIX = "llmwhisperer|" + +# (scope_id, adapter_instance_id) -> resolved config dict | None, where +# scope_id is the execution id or (for IDE runs, which carry no execution +# id) the run id. Run-scoped on purpose: the cache exists to deduplicate +# the N per-prompt resolutions within ONE run — never to cache across +# runs, where it would pin a stale output mode after an adapter edit. +# Bounded so a long-lived worker never grows it unchecked. +_MODE_CACHE: OrderedDict[tuple[str, str], dict[str, Any] | None] = OrderedDict() +_MODE_CACHE_MAX = 256 + + +def _resolve_image_mode_config( + shim: Any, adapter_instance_id: str, scope_id: str +) -> dict[str, Any] | None: + """Return the x2text adapter config when it is in image mode, else None. + + Resolution goes through the platform service (the payload has no + adapter metadata); results are cached per (scope, adapter). With no + scope id at all, caching is skipped entirely — a shared ("", adapter) + entry would serve a stale mode to every later run on this worker. + """ + cache_key = (scope_id, adapter_instance_id) + if scope_id and cache_key in _MODE_CACHE: + _MODE_CACHE.move_to_end(cache_key) + return _MODE_CACHE[cache_key] + + config = PlatformHelper.get_adapter_config(shim, adapter_instance_id) or {} + adapter_id = str(config.get("adapter_id", "")) + adapter_metadata = config.get("adapter_metadata") or {} + is_image_mode = adapter_id.startswith(_LLMWHISPERER_ADAPTER_PREFIX) and ( + adapter_metadata.get(ImageOutputConstants.OUTPUT_MODE) + == ImageOutputConstants.IMAGE_MODE + ) + + result = config if is_image_mode else None + if scope_id: + _MODE_CACHE[cache_key] = result + while len(_MODE_CACHE) > _MODE_CACHE_MAX: + _MODE_CACHE.popitem(last=False) + return result + + +def detect_image_mode_config( + *, + output: dict[str, Any], + shim: Any, + usage_kwargs: dict[str, Any] | None = None, +) -> dict[str, Any] | None: + """Detect image mode for one prompt; None means normal text path. + + Fast path: the backend stamps the adapter's ``x2text_output_mode`` + onto the per-prompt payload (it already holds the decrypted adapter + metadata), so text-mode prompts cost nothing here. The platform- + service resolution runs only as the fallback for payloads without + the stamp (e.g. API deployments, older payloads). + + Returns the resolved adapter config for the plugin, or ``{}`` when + image mode was determined from the stamp alone. + """ + adapter_instance_id = str(output.get(PSKeys.X2TEXT_ADAPTER) or "") + if not adapter_instance_id: + return None + + if PSKeys.X2TEXT_OUTPUT_MODE in output: + stamped_mode = output.get(PSKeys.X2TEXT_OUTPUT_MODE) + if stamped_mode == ImageOutputConstants.IMAGE_MODE: + return {} + return None + + usage_kwargs = usage_kwargs or {} + # IDE payloads carry no execution_id — fall back to the run id so the + # cache stays scoped to one run (see _MODE_CACHE). + scope_id = str(usage_kwargs.get("execution_id") or usage_kwargs.get("run_id") or "") + return _resolve_image_mode_config(shim, adapter_instance_id, scope_id) + + +def run_vlm_image_answer( + *, + output: dict[str, Any], + shim: Any, + llm: Any, + extract_file_path: str, + execution_source: str, + metadata: dict[str, Any], + metrics: dict[str, Any], + x2text_config: dict[str, Any], + usage_kwargs: dict[str, Any] | None = None, +) -> str: + """Answer an image-mode prompt via the cloud plugin. + + The caller has already detected image mode via + ``detect_image_mode_config``. ``extract_file_path`` must be the + extract-file path (never the summarize/source rewrite of it) — the + pages directory is derived from it via the shared writer/reader + helper. + + Returns: + The raw answer string (the caller assigns it in place of the + RAG/completion answer, so type conversion, lookups, webhooks + etc. run unchanged). + + Raises: + VlmImageAnswerError: the cloud plugin is not installed, or the + vision path failed in a way the user must act on (missing + images, page cap, non-vision LLM). + """ + usage_kwargs = usage_kwargs or {} + prompt_name = output.get(PSKeys.NAME, "") + plugin_cls = ExecutorPluginLoader.get(PLUGIN_NAME) + if plugin_cls is None: + raise VlmImageAnswerError( + "This document was extracted in image output mode, which is " + "answered by a vision LLM available only on Unstract Cloud. " + "Switch the profile's text extractor to a text output mode, " + "or run this on Unstract Cloud.", + error_code=IMAGE_OUTPUT_REQUIRES_CLOUD, + ) + + shim.stream_log(f"Answering `{prompt_name}` from page images via vision LLM") + fs = FileUtils.get_fs_instance(execution_source=execution_source) + page_store_dir = build_page_store_dir(extract_file_path, extract_file_path) + + try: + outcome = plugin_cls.run_with_metrics( + output=output, + llm=llm, + fs=fs, + page_store_dir=page_store_dir, + x2text_config=x2text_config, + metadata=metadata, + shim=shim, + usage_kwargs=usage_kwargs, + ) + except (PageImagesNotFoundError, PageImageSetIncompleteError) as e: + raise VlmImageAnswerError(str(e), error_code=IMAGE_OUTPUT_MISSING) from e + except PageCapExceededError as e: + raise VlmImageAnswerError(str(e), error_code=IMAGE_PAGE_CAP_EXCEEDED) from e + except PageImageSetTooLargeError as e: + raise VlmImageAnswerError(str(e), error_code=IMAGE_PAGES_TOO_LARGE) from e + except PageImageLoadError as e: + raise VlmImageAnswerError(str(e), error_code=IMAGE_OUTPUT_MISSING) from e + except VlmImageAnswerError: + raise + except Exception as e: + # Plugin-defined hard failures carry a stable error_code attribute + # (e.g. VISION_LLM_REQUIRED). Anything else propagates untouched — + # never degrade to the text path. + plugin_code = getattr(e, "error_code", None) + if isinstance(plugin_code, str) and plugin_code: + raise VlmImageAnswerError(str(e), error_code=plugin_code) from e + raise + + llm_metrics = outcome.get("llm_metrics") if isinstance(outcome, dict) else None + if llm_metrics: + metrics.setdefault(prompt_name, {})["vlm_image_answer"] = llm_metrics + + answer = outcome["answer"] if isinstance(outcome, dict) else str(outcome) + shim.stream_log(f"Vision LLM answered `{prompt_name}`") + return answer + + +def raise_if_image_mode_unsupported( + *, + operation: str, + adapter_instance_id: str | None, + shim: Any, + scope_id: str = "", +) -> None: + """Guard operations that cannot run against image-mode documents. + + Single-pass extraction (and any future full-text operation) would + silently run against the one-line extraction summary — reject it + explicitly instead. + """ + if not adapter_instance_id: + return + config = _resolve_image_mode_config(shim, str(adapter_instance_id), scope_id) + if config is not None: + raise VlmImageAnswerError( + f"{operation} is not supported in image output mode. Run " + "prompts individually, or switch the profile's text extractor " + "to a text output mode.", + error_code=IMAGE_OUTPUT_UNSUPPORTED_OPERATION, + ) diff --git a/workers/executor/tests/__init__.py b/workers/executor/tests/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/workers/executor/tests/test_vlm_image_answer_bridge.py b/workers/executor/tests/test_vlm_image_answer_bridge.py new file mode 100644 index 0000000000..8fb49e8e39 --- /dev/null +++ b/workers/executor/tests/test_vlm_image_answer_bridge.py @@ -0,0 +1,310 @@ +"""Tests for the OSS vlm_image_answer bridge (detection + dispatch). + +Detection must prefer the backend's per-prompt output-mode stamp (zero +platform calls for text mode), fall back to run-scoped platform +resolution, and never cache without a scope id. Dispatch must raise a +structured plugin-absent error instead of falling through to the text +path, key the pages directory on the never-rewritten extract path, and +map sdk1 loader errors to stable error codes. +""" + +import sys +from pathlib import Path +from unittest.mock import MagicMock, patch + +import pytest + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +from executor.executors.exceptions import VlmImageAnswerError # noqa: E402 +from executor.executors.vlm_image_answer import ( # noqa: E402 + _MODE_CACHE, + IMAGE_OUTPUT_MISSING, + IMAGE_OUTPUT_REQUIRES_CLOUD, + IMAGE_OUTPUT_UNSUPPORTED_OPERATION, + IMAGE_PAGE_CAP_EXCEEDED, + IMAGE_PAGES_TOO_LARGE, + detect_image_mode_config, + raise_if_image_mode_unsupported, + run_vlm_image_answer, +) + +from unstract.sdk1.adapters.x2text.page_image_loader import ( # noqa: E402 + PageCapExceededError, + PageImageSetTooLargeError, + PageImagesNotFoundError, +) + +_IMAGE_CONFIG = { + "adapter_id": "llmwhisperer|0a1647f0-f65f-410d-843b-3d979c78350e", + "adapter_metadata": {"output_mode": "image", "url": "http://svc"}, +} +_TEXT_CONFIG = { + "adapter_id": "llmwhisperer|0a1647f0-f65f-410d-843b-3d979c78350e", + "adapter_metadata": {"output_mode": "layout_preserving"}, +} +_OTHER_ADAPTER_CONFIG = { + "adapter_id": "someocr|123", + "adapter_metadata": {"output_mode": "image"}, +} + + +@pytest.fixture(autouse=True) +def _clear_cache(): + _MODE_CACHE.clear() + yield + _MODE_CACHE.clear() + + +def _detect(output=None, adapter_config=_IMAGE_CONFIG, usage_kwargs=None): + output = output or {"x2text_adapter": "uuid-1", "name": "p1"} + with patch( + "executor.executors.vlm_image_answer.PlatformHelper.get_adapter_config", + return_value=adapter_config, + ) as resolve: + result = detect_image_mode_config( + output=output, + shim=MagicMock(), + usage_kwargs=usage_kwargs or {"execution_id": "e1"}, + ) + return result, resolve + + +class TestDetection: + def test_image_mode_returns_config(self): + result, _ = _detect() + assert result == _IMAGE_CONFIG + + def test_non_image_mode_returns_none(self): + result, _ = _detect(adapter_config=_TEXT_CONFIG) + assert result is None + + def test_non_llmwhisperer_adapter_returns_none(self): + # Another adapter with a coincidental output_mode key is not gated. + result, _ = _detect(adapter_config=_OTHER_ADAPTER_CONFIG) + assert result is None + + def test_missing_adapter_id_returns_none_without_platform_call(self): + result, resolve = _detect(output={"name": "p1"}) + assert result is None + resolve.assert_not_called() + + +class TestStampedDetection: + """The backend stamp is the fast path — zero platform calls.""" + + def test_stamped_image_mode_detected_without_platform_call(self): + result, resolve = _detect( + output={ + "x2text_adapter": "uuid-1", + "name": "p1", + "x2text_output_mode": "image", + } + ) + assert result == {} + resolve.assert_not_called() + + @pytest.mark.parametrize("mode", ["layout_preserving", "text", None]) + def test_stamped_non_image_mode_skips_platform_call(self, mode): + result, resolve = _detect( + output={ + "x2text_adapter": "uuid-1", + "name": "p1", + "x2text_output_mode": mode, + } + ) + assert result is None + resolve.assert_not_called() + + def test_unstamped_payload_falls_back_to_platform(self): + result, resolve = _detect(output={"x2text_adapter": "uuid-1", "name": "p1"}) + assert result == _IMAGE_CONFIG + resolve.assert_called_once() + + +class TestResolutionCache: + def test_resolution_cached_per_execution_and_adapter(self): + with patch( + "executor.executors.vlm_image_answer.PlatformHelper.get_adapter_config", + return_value=_IMAGE_CONFIG, + ) as resolve: + for _ in range(3): # three prompts, same adapter + execution + detect_image_mode_config( + output={"x2text_adapter": "uuid-1", "name": "p"}, + shim=MagicMock(), + usage_kwargs={"execution_id": "e1"}, + ) + assert resolve.call_count == 1 + + def test_run_id_scopes_cache_when_execution_id_missing(self): + # IDE payloads carry no execution_id: two RUNS on the same adapter + # must each resolve fresh (an adapter edit between runs takes + # effect), while prompts within one run share the cached result. + with patch( + "executor.executors.vlm_image_answer.PlatformHelper.get_adapter_config", + return_value=_IMAGE_CONFIG, + ) as resolve: + for run in ("run-1", "run-2"): + for _ in range(2): # two prompts per run + detect_image_mode_config( + output={"x2text_adapter": "uuid-1", "name": "p"}, + shim=MagicMock(), + usage_kwargs={"run_id": run}, + ) + assert resolve.call_count == 2 # once per run, not once total + + def test_no_scope_id_never_caches(self): + # With neither execution_id nor run_id, a shared ("", adapter) + # entry would pin a stale mode forever — caching must be skipped. + with patch( + "executor.executors.vlm_image_answer.PlatformHelper.get_adapter_config", + return_value=_IMAGE_CONFIG, + ) as resolve: + for _ in range(3): + detect_image_mode_config( + output={"x2text_adapter": "uuid-1", "name": "p"}, + shim=MagicMock(), + usage_kwargs={}, + ) + assert resolve.call_count == 3 + + +def _run(plugin=None, x2text_config=_IMAGE_CONFIG, metrics=None, **overrides): + kwargs = { + "output": {"x2text_adapter": "uuid-1", "name": "p1", "promptx": "q?"}, + "shim": MagicMock(), + "llm": MagicMock(), + "extract_file_path": "/data/extract/doc.txt", + "execution_source": "ide", + "metadata": {"context": {}}, + "metrics": metrics if metrics is not None else {}, + "x2text_config": x2text_config, + "usage_kwargs": {"run_id": "r1", "execution_id": "e1"}, + } + kwargs.update(overrides) + with ( + patch( + "executor.executors.vlm_image_answer.ExecutorPluginLoader.get", + return_value=plugin, + ), + patch( + "executor.executors.vlm_image_answer.FileUtils.get_fs_instance", + return_value=MagicMock(), + ), + ): + return run_vlm_image_answer(**kwargs) + + +class TestPluginAbsent: + def test_raises_structured_error_never_falls_through(self): + with pytest.raises(VlmImageAnswerError) as excinfo: + _run(plugin=None) + assert excinfo.value.error_code == IMAGE_OUTPUT_REQUIRES_CLOUD + assert str(excinfo.value).startswith(IMAGE_OUTPUT_REQUIRES_CLOUD + ":") + assert "Unstract Cloud" in str(excinfo.value) + + +class TestPluginDispatch: + def test_answer_and_page_store_dir_contract(self): + plugin = MagicMock() + plugin.run_with_metrics.return_value = {"answer": "42", "llm_metrics": {"t": 1}} + metrics = {} + answer = _run(plugin=plugin, metrics=metrics) + assert answer == "42" + call_kwargs = plugin.run_with_metrics.call_args.kwargs + # Deterministic path derived from the extract file path via the + # shared helper — the writer/reader agreement contract. + assert call_kwargs["page_store_dir"] == "/data/extract/doc/pages" + assert call_kwargs["x2text_config"] == _IMAGE_CONFIG + assert metrics["p1"]["vlm_image_answer"] == {"t": 1} + + def test_pages_dir_keys_on_extract_path_not_rewritten_file_path(self): + # Summarize-as-source / smart-table rewrite the payload FILE_PATH; + # the reader must key on the extract path regardless. + plugin = MagicMock() + plugin.run_with_metrics.return_value = {"answer": "a"} + _run(plugin=plugin, extract_file_path="/data/extract/report.txt") + call_kwargs = plugin.run_with_metrics.call_args.kwargs + assert call_kwargs["page_store_dir"] == "/data/extract/report/pages" + + def test_loader_not_found_maps_to_image_output_missing(self): + plugin = MagicMock() + plugin.run_with_metrics.side_effect = PageImagesNotFoundError( + "no images", page_store_dir="/d/pages" + ) + with pytest.raises(VlmImageAnswerError) as excinfo: + _run(plugin=plugin) + assert excinfo.value.error_code == IMAGE_OUTPUT_MISSING + + def test_cap_error_maps_to_page_cap_code(self): + plugin = MagicMock() + plugin.run_with_metrics.side_effect = PageCapExceededError( + "too big", page_store_dir="/d/pages", page_count=50, page_cap=20 + ) + with pytest.raises(VlmImageAnswerError) as excinfo: + _run(plugin=plugin) + assert excinfo.value.error_code == IMAGE_PAGE_CAP_EXCEEDED + + def test_too_large_error_maps_to_pages_too_large_code(self): + plugin = MagicMock() + plugin.run_with_metrics.side_effect = PageImageSetTooLargeError( + "too many bytes", + page_store_dir="/d/pages", + total_bytes=99, + max_total_bytes=10, + ) + with pytest.raises(VlmImageAnswerError) as excinfo: + _run(plugin=plugin) + assert excinfo.value.error_code == IMAGE_PAGES_TOO_LARGE + + def test_plugin_error_code_attribute_is_wrapped(self): + class VisionError(Exception): + error_code = "VISION_LLM_REQUIRED" + + plugin = MagicMock() + plugin.run_with_metrics.side_effect = VisionError("model X has no vision") + with pytest.raises(VlmImageAnswerError) as excinfo: + _run(plugin=plugin) + assert excinfo.value.error_code == "VISION_LLM_REQUIRED" + assert "model X has no vision" in str(excinfo.value) + + def test_unexpected_plugin_error_propagates_unwrapped(self): + plugin = MagicMock() + plugin.run_with_metrics.side_effect = RuntimeError("boom") + with pytest.raises(RuntimeError, match="boom"): + _run(plugin=plugin) + + +class TestUnsupportedOperationGuard: + def test_single_pass_rejected_for_image_mode(self): + with patch( + "executor.executors.vlm_image_answer.PlatformHelper.get_adapter_config", + return_value=_IMAGE_CONFIG, + ): + with pytest.raises(VlmImageAnswerError) as excinfo: + raise_if_image_mode_unsupported( + operation="Single-pass extraction", + adapter_instance_id="uuid-1", + shim=MagicMock(), + scope_id="e1", + ) + assert excinfo.value.error_code == IMAGE_OUTPUT_UNSUPPORTED_OPERATION + + def test_text_mode_passes(self): + with patch( + "executor.executors.vlm_image_answer.PlatformHelper.get_adapter_config", + return_value=_TEXT_CONFIG, + ): + raise_if_image_mode_unsupported( + operation="Single-pass extraction", + adapter_instance_id="uuid-1", + shim=MagicMock(), + scope_id="e1", + ) # no raise + + def test_no_adapter_id_passes(self): + raise_if_image_mode_unsupported( + operation="Single-pass extraction", + adapter_instance_id=None, + shim=MagicMock(), + ) # no raise, no platform call diff --git a/workers/file_processing/structure_tool_task.py b/workers/file_processing/structure_tool_task.py index 971a783980..2ff0fa1ed7 100644 --- a/workers/file_processing/structure_tool_task.py +++ b/workers/file_processing/structure_tool_task.py @@ -397,6 +397,9 @@ def _execute_structure_tool_impl(params: dict) -> dict: _SK.FILE_HASH: file_hash, _SK.FILE_NAME: file_name, _SK.FILE_PATH: extracted_input_file, + # Never rewritten (unlike FILE_PATH, which summarize/smart-table + # overrides mutate) — the page-image reader keys on this. + "extract_file_path": extracted_input_file, _SK.EXECUTION_SOURCE: _SK.TOOL, _SK.CUSTOM_DATA: custom_data, "PLATFORM_SERVICE_API_KEY": platform_service_api_key, From c8060ff7f4883e4c0c575adeb33bb6f80bc88c12 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Wed, 5 Aug 2026 17:27:26 +0530 Subject: [PATCH 15/24] UN-2646 [FIX] Gate the platform fallback by execution source; fix CI test fallout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The merged consumer branch put the pre-existing workers suites in front of the image-mode bridge for the first time (they never ran on the stacked PR), and ~105 tests failed on one mechanism: unstamped payloads sent the detection fallback to the platform service, whose URL came from a MagicMock shim. The semantic fix, not a test hack: IDE payloads are always stamped by the current backend, so an unstamped IDE payload is a pre-upgrade in-flight run — detection now treats it as text mode instead of making every legacy IDE prompt depend on a platform call. The platform resolution remains only for non-IDE (deployment) payloads, which are built worker-side without adapter metadata — that real dependency is stubbed once in the workers test conftest (the same place the suite already mocks celery and the shim), with image-mode tests opting in via the payload stamp. The single-pass guard gets the same semantics (stamp first, resolve only for non-IDE). Also fixes the sdk1 collection error: the loader test imported the shared fixture bare (PYTHONPATH-dependent) instead of via the tests package like its siblings. Bridge tests: 26 pass, incl. unstamped-IDE-never-calls-platform and stamped-guard-never-calls-platform pins. Co-Authored-By: Claude Fable 5 --- unstract/sdk1/tests/test_page_image_loader.py | 3 +- workers/executor/executors/legacy_executor.py | 10 ++++- .../executor/executors/vlm_image_answer.py | 31 ++++++++++--- .../tests/test_vlm_image_answer_bridge.py | 43 +++++++++++++++++++ workers/tests/conftest.py | 25 +++++++++++ 5 files changed, 105 insertions(+), 7 deletions(-) diff --git a/unstract/sdk1/tests/test_page_image_loader.py b/unstract/sdk1/tests/test_page_image_loader.py index decde94b75..73f9650f62 100644 --- a/unstract/sdk1/tests/test_page_image_loader.py +++ b/unstract/sdk1/tests/test_page_image_loader.py @@ -10,7 +10,6 @@ from pathlib import Path import pytest -from llmw_image_fixtures import InMemoryFileStorage, minimal_png from unstract.sdk1.adapters.x2text.page_image_loader import ( DEFAULT_PAGE_CAP, LoadedPageImage, @@ -24,6 +23,8 @@ ) from unstract.sdk1.file_storage import FileStorage, FileStorageProvider +from tests.llmw_image_fixtures import InMemoryFileStorage, minimal_png + _DIR = "/data/extract/doc/pages" diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index cf0ebbea13..1245aaa221 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -1762,7 +1762,10 @@ def _execute_single_prompt( # platform-service fallback) so retrieval adapters are never # constructed for a prompt that answers from page images. vlm_config = detect_image_mode_config( - output=output, shim=shim, usage_kwargs=usage_kwargs + output=output, + shim=shim, + execution_source=str(execution_source or ""), + usage_kwargs=usage_kwargs, ) llm, embedding, vector_db = self._init_llm_and_retrieval( output=output, @@ -2345,6 +2348,9 @@ def _handle_single_pass_extraction( PSKeys.X2TEXT_ADAPTER ) if x2text_instance_id: + stamped_mode = tool_settings.get(PSKeys.X2TEXT_OUTPUT_MODE) + if stamped_mode is None: + stamped_mode = outputs[0].get(PSKeys.X2TEXT_OUTPUT_MODE) raise_if_image_mode_unsupported( operation="Single-pass extraction", adapter_instance_id=str(x2text_instance_id), @@ -2355,6 +2361,8 @@ def _handle_single_pass_extraction( scope_id=str( params.get(PSKeys.EXECUTION_ID) or params.get(PSKeys.RUN_ID) or "" ), + stamped_mode=stamped_mode, + execution_source=str(params.get(PSKeys.EXECUTION_SOURCE) or ""), ) try: diff --git a/workers/executor/executors/vlm_image_answer.py b/workers/executor/executors/vlm_image_answer.py index b5c0b4b8e3..08706b8db6 100644 --- a/workers/executor/executors/vlm_image_answer.py +++ b/workers/executor/executors/vlm_image_answer.py @@ -98,19 +98,28 @@ def _resolve_image_mode_config( return result +_IDE_SOURCE = "ide" + + def detect_image_mode_config( *, output: dict[str, Any], shim: Any, + execution_source: str = "", usage_kwargs: dict[str, Any] | None = None, ) -> dict[str, Any] | None: """Detect image mode for one prompt; None means normal text path. Fast path: the backend stamps the adapter's ``x2text_output_mode`` onto the per-prompt payload (it already holds the decrypted adapter - metadata), so text-mode prompts cost nothing here. The platform- - service resolution runs only as the fallback for payloads without - the stamp (e.g. API deployments, older payloads). + metadata), so stamped prompts cost nothing here. + + Unstamped payloads split by source: IDE payloads are always stamped + by the current backend, so an unstamped one is a pre-upgrade + in-flight run — treated as text mode rather than making every + legacy IDE prompt depend on a platform-service call. Deployment + payloads are built worker-side without adapter metadata, so they + resolve through the platform service (run-scoped cache). Returns the resolved adapter config for the plugin, or ``{}`` when image mode was determined from the stamp alone. @@ -125,6 +134,9 @@ def detect_image_mode_config( return {} return None + if execution_source == _IDE_SOURCE: + return None + usage_kwargs = usage_kwargs or {} # IDE payloads carry no execution_id — fall back to the run id so the # cache stays scoped to one run (see _MODE_CACHE). @@ -223,16 +235,25 @@ def raise_if_image_mode_unsupported( adapter_instance_id: str | None, shim: Any, scope_id: str = "", + stamped_mode: str | None = None, + execution_source: str = "", ) -> None: """Guard operations that cannot run against image-mode documents. Single-pass extraction (and any future full-text operation) would silently run against the one-line extraction summary — reject it - explicitly instead. + explicitly instead. Same detection semantics as + ``detect_image_mode_config``: stamp first, platform resolution only + for non-IDE sources. """ if not adapter_instance_id: return - config = _resolve_image_mode_config(shim, str(adapter_instance_id), scope_id) + if stamped_mode is not None: + config = {} if stamped_mode == ImageOutputConstants.IMAGE_MODE else None + elif execution_source == _IDE_SOURCE: + return + else: + config = _resolve_image_mode_config(shim, str(adapter_instance_id), scope_id) if config is not None: raise VlmImageAnswerError( f"{operation} is not supported in image output mode. Run " diff --git a/workers/executor/tests/test_vlm_image_answer_bridge.py b/workers/executor/tests/test_vlm_image_answer_bridge.py index 8fb49e8e39..adc0fb4e25 100644 --- a/workers/executor/tests/test_vlm_image_answer_bridge.py +++ b/workers/executor/tests/test_vlm_image_answer_bridge.py @@ -117,10 +117,28 @@ def test_stamped_non_image_mode_skips_platform_call(self, mode): resolve.assert_not_called() def test_unstamped_payload_falls_back_to_platform(self): + # Non-IDE (deployment) payloads are built worker-side without + # adapter metadata — resolution is legitimate there. result, resolve = _detect(output={"x2text_adapter": "uuid-1", "name": "p1"}) assert result == _IMAGE_CONFIG resolve.assert_called_once() + def test_unstamped_ide_payload_is_text_mode_without_platform_call(self): + # IDE payloads are always stamped by the current backend; an + # unstamped one is a pre-upgrade in-flight run. It must NOT make + # every legacy IDE prompt depend on the platform service. + with patch( + "executor.executors.vlm_image_answer.PlatformHelper.get_adapter_config", + ) as resolve: + result = detect_image_mode_config( + output={"x2text_adapter": "uuid-1", "name": "p1"}, + shim=MagicMock(), + execution_source="ide", + usage_kwargs={"run_id": "r1"}, + ) + assert result is None + resolve.assert_not_called() + class TestResolutionCache: def test_resolution_cached_per_execution_and_adapter(self): @@ -308,3 +326,28 @@ def test_no_adapter_id_passes(self): adapter_instance_id=None, shim=MagicMock(), ) # no raise, no platform call + + def test_stamped_image_mode_rejected_without_platform_call(self): + with patch( + "executor.executors.vlm_image_answer.PlatformHelper.get_adapter_config", + ) as resolve: + with pytest.raises(VlmImageAnswerError): + raise_if_image_mode_unsupported( + operation="Single-pass extraction", + adapter_instance_id="uuid-1", + shim=MagicMock(), + stamped_mode="image", + ) + resolve.assert_not_called() + + def test_unstamped_ide_single_pass_passes_without_platform_call(self): + with patch( + "executor.executors.vlm_image_answer.PlatformHelper.get_adapter_config", + ) as resolve: + raise_if_image_mode_unsupported( + operation="Single-pass extraction", + adapter_instance_id="uuid-1", + shim=MagicMock(), + execution_source="ide", + ) # no raise + resolve.assert_not_called() diff --git a/workers/tests/conftest.py b/workers/tests/conftest.py index faa90d7593..739f2a0fa3 100644 --- a/workers/tests/conftest.py +++ b/workers/tests/conftest.py @@ -426,3 +426,28 @@ def _restore_current_celery_app(): yield finally: default_app.set_current() + + +@pytest.fixture(autouse=True) +def _stub_vlm_image_mode_resolution(monkeypatch): + """Default image-mode detection to "not image mode" for this suite. + + The vlm_image_answer bridge resolves the x2text adapter's output mode + through the platform service for unstamped non-IDE payloads — a real + dependency of deployment runs that has no live endpoint in unit tests + (the shim is a MagicMock, so the resolver would build a nonsense URL + and fail every prompt). Stub it to "config unavailable" so every + prompt takes the text path, exactly as before the bridge existed. + Image-mode tests opt in by stamping ``x2text_output_mode: "image"`` + on their per-prompt payloads (no platform call involved). + """ + from executor.executors import vlm_image_answer + + vlm_image_answer._MODE_CACHE.clear() + monkeypatch.setattr( + vlm_image_answer.PlatformHelper, + "get_adapter_config", + staticmethod(lambda *_args, **_kwargs: None), + ) + yield + vlm_image_answer._MODE_CACHE.clear() From 341701bce63223c7eea2e7552026483c00e3973b Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Wed, 5 Aug 2026 17:32:25 +0530 Subject: [PATCH 16/24] UN-2646 [FIX] Reset the stable pages dir before persisting a new page set MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A re-extraction that yields fewer pages than a previous run left in the stable pages directory would otherwise keep the old trailing images, which the reader serves back as part of the new document (Greptile P1). The writer now clears the directory wholesale before writing, and a failed reset raises instead of risking stale pages reaching the vision LLM — consistent with this feature's fail-loud policy. Co-Authored-By: Claude Fable 5 --- .../x2text/llm_whisperer_v2/src/helper.py | 16 ++++++++++++ unstract/sdk1/tests/test_llmw_image_helper.py | 26 +++++++++++++++++++ 2 files changed, 42 insertions(+) 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 14e1f385e8..c03019745c 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 @@ -768,7 +768,23 @@ def persist_page_images( its retries, the pages already written are removed (best-effort) before a hard ``ExtractorError`` propagates, so a failed extraction never leaves a partial set behind. Works transparently for LOCAL and S3 via ``fs``. + + The directory is reset wholesale first: a re-extraction can produce + fewer pages than a previous run left in this stable path, and stale + trailing images would otherwise be read back as part of the new set. + A failed reset raises rather than risk serving another document's + pages to the vision LLM. """ + try: + if fs.exists(page_store_dir): + fs.rm(page_store_dir, recursive=True) + except Exception as e: + raise ExtractorError( + "Failed to clear previous page images before writing the new " + f"set: dir={page_store_dir}, provider={fs.provider.value}", + status_code=500, + actual_err=e, + ) from e fs.mkdir(create_parents=True, path=page_store_dir) write_with_retry = retry_with_exponential_backoff( diff --git a/unstract/sdk1/tests/test_llmw_image_helper.py b/unstract/sdk1/tests/test_llmw_image_helper.py index a2b3eac8d6..36cd978578 100644 --- a/unstract/sdk1/tests/test_llmw_image_helper.py +++ b/unstract/sdk1/tests/test_llmw_image_helper.py @@ -180,6 +180,32 @@ def test_mid_list_failure_cleans_up_written_pages( assert "doc/pages" in fs.rm_calls # cleanup invoked assert fs.stored_paths == [] # page 1 removed by the cleanup + def test_reextraction_with_fewer_pages_prunes_stale_trailing_pages(self) -> None: + # The pages dir is a stable path: a re-extraction that yields fewer + # pages must not leave the previous run's trailing images behind, + # where the reader would serve them as part of the new document. + fs = InMemoryFileStorage(provider=FileStorageProvider.S3) + H.persist_page_images(fs, "doc/pages", [(1, b"a"), (2, b"b"), (3, b"c")]) + refs = H.persist_page_images(fs, "doc/pages", [(1, b"x"), (2, b"y")]) + + assert [r.page_number for r in refs] == [1, 2] + assert sorted(fs.stored_paths) == [ + "doc/pages/page_001.png", + "doc/pages/page_002.png", + ] + assert fs.read(path="doc/pages/page_001.png", mode="rb") == b"x" + + def test_failed_dir_reset_raises_instead_of_serving_stale_pages(self) -> None: + fs = InMemoryFileStorage(provider=FileStorageProvider.S3) + H.persist_page_images(fs, "doc/pages", [(1, b"a")]) + + def _rm_fails(path: str, recursive: bool = True) -> None: + raise OSError("permission denied") + + fs.rm = _rm_fails # type: ignore[method-assign] + with pytest.raises(ExtractorError, match="clear previous page images"): + H.persist_page_images(fs, "doc/pages", [(1, b"x")]) + 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( From 3515124a3dec95a421ed85297e09d44e075758bf Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Thu, 6 Aug 2026 12:55:52 +0530 Subject: [PATCH 17/24] UN-2646 [FIX] Stamp x2text output mode onto single-pass payloads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit build_single_pass_payload was the one payload builder that skipped _stamp_x2text_output_mode, so IDE single-pass payloads arrived unstamped and the executor's single-pass guard — which treats an unstamped IDE payload as a pre-upgrade text-mode run — never fired. Image mode + single-pass would have silently answered every prompt against the one-line extraction summary once the cloud consumer lands. Stamps tool_settings alongside X2TEXT_ADAPTER, matching the other three builders. Pinned by tests that drive build_single_pass_payload with all collaborators patched and assert the stamp lands in the built executor payload (image, non-image, and non-LLMWhisperer cases), so deleting the stamp call fails the suite. Co-Authored-By: Claude Fable 5 --- .../prompt_studio_helper.py | 4 + .../tests/test_single_pass_payload_stamp.py | 131 ++++++++++++++++++ 2 files changed, 135 insertions(+) create mode 100644 backend/prompt_studio/prompt_studio_core_v2/tests/test_single_pass_payload_stamp.py diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index 047254fc3b..807681d510 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -1265,6 +1265,10 @@ def build_single_pass_payload( or TSPKeys.SIMPLE, TSPKeys.SIMILARITY_TOP_K: default_profile.similarity_top_k, } + # Stamp the x2text output mode like every other payload builder — the + # executor's single-pass guard trusts the stamp, and an unstamped IDE + # payload is treated as pre-upgrade text mode (guard never fires). + PromptStudioHelper._stamp_x2text_output_mode(tool_settings, default_profile) lookup_configs = get_lookup_configs_for_tool(tool, prompts=prompts) if lookup_configs: diff --git a/backend/prompt_studio/prompt_studio_core_v2/tests/test_single_pass_payload_stamp.py b/backend/prompt_studio/prompt_studio_core_v2/tests/test_single_pass_payload_stamp.py new file mode 100644 index 0000000000..8ab89b4916 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_core_v2/tests/test_single_pass_payload_stamp.py @@ -0,0 +1,131 @@ +"""Regression tests: ``build_single_pass_payload`` stamps the x2text output mode. + +The executor's single-pass guard trusts the payload stamp — an unstamped IDE +payload is treated as a pre-upgrade text-mode run and the guard never fires. +``build_single_pass_payload`` was the one payload builder that skipped +``_stamp_x2text_output_mode``, which made image-mode + single-pass answer +every prompt against the one-line extraction summary, silently. These tests +pin the stamp into the built ``tool_settings`` so deleting the call fails. + +Unit tests: the real helper module is imported (Django is loaded by the rig's +test env) and every collaborator is patched on it per-test, so no database is +touched. +""" + +from __future__ import annotations + +from contextlib import ExitStack +from unittest.mock import MagicMock, patch + +from prompt_studio.prompt_studio_core_v2 import prompt_studio_helper as _psh_mod +from prompt_studio.prompt_studio_core_v2.constants import ToolStudioPromptKeys as TSPKeys + +PromptStudioHelper = _psh_mod.PromptStudioHelper + +_LLMW_ADAPTER_ID = "llmwhisperer|a5e6b8af-3e1f-4a80-b006-d017e8e67f93" + + +def _make_tool(): + tool = MagicMock(name="CustomTool") + tool.tool_id = "tool-1" + tool.prompt_grammer = None + tool.challenge_llm = None + tool.enable_challenge = False + tool.enable_highlight = False + tool.enable_word_confidence = False + tool.summarize_as_source = False + tool.custom_data = None + return tool + + +def _make_profile(metadata: dict | None, adapter_id: str = _LLMW_ADAPTER_ID): + profile = MagicMock(name="ProfileManager") + profile.x2text.id = "x2t-1" + profile.x2text.adapter_id = adapter_id + profile.x2text.metadata = metadata + profile.llm.id = "llm-1" + profile.embedding_model.id = "emb-1" + profile.vector_store.id = "vdb-1" + profile.chunk_overlap = 64 + profile.retrieval_strategy = "simple" + profile.similarity_top_k = 3 + profile.profile_id = "profile-1" + return profile + + +def _make_prompt(): + p = MagicMock(name="ToolStudioPrompt") + p.prompt = "What is the total?" + p.active = True + p.enforce_type = "text" + p.prompt_key = "total" + p.prompt_id = "p-1" + return p + + +def _build(profile) -> dict: + """Run ``build_single_pass_payload`` with collaborators patched. + + Returns the ``tool_settings`` dict from the built executor payload. + """ + fs_instance = MagicMock(name="fs_instance") + fs_instance.get_hash_from_file.return_value = "hash-1" + + with ExitStack() as stack: + for target, attr, value in ( + ( + _psh_mod.ProfileManager, + "get_default_llm_profile", + MagicMock(return_value=profile), + ), + (PromptStudioHelper, "validate_adapter_status", MagicMock(return_value=None)), + ( + PromptStudioHelper, + "validate_profile_manager_owner_access", + MagicMock(return_value=None), + ), + (PromptStudioHelper, "dynamic_extractor", MagicMock(return_value=None)), + ( + PromptStudioHelper, + "_get_platform_api_key", + MagicMock(return_value="pk-test"), + ), + (_psh_mod.EnvHelper, "get_storage", MagicMock(return_value=fs_instance)), + (_psh_mod, "get_lookup_configs_for_tool", MagicMock(return_value=None)), + (_psh_mod.StateStore, "get", MagicMock(return_value="")), + ): + stack.enter_context(patch.object(target, attr, value)) + + context, _cb_kwargs = PromptStudioHelper.build_single_pass_payload( + tool=_make_tool(), + doc_path="/data/org/user/tool/statement.pdf", + doc_name="statement.pdf", + prompts=[_make_prompt()], + org_id="org-1", + user_id="user-1", + document_id="doc-1", + run_id="run-1", + request_user=MagicMock(name="request-user"), + ) + return context.executor_params[TSPKeys.TOOL_SETTINGS] + + +class TestSinglePassPayloadStampsOutputMode: + def test_image_mode_profile_is_stamped_into_tool_settings(self) -> None: + # The executor's single-pass guard fires only on this stamp for IDE + # payloads — without it, image mode + single-pass silently answers + # from the one-line extraction summary. + tool_settings = _build(_make_profile({"output_mode": "image"})) + assert tool_settings[TSPKeys.X2TEXT_OUTPUT_MODE] == "image" + + def test_text_mode_profile_is_stamped_into_tool_settings(self) -> None: + # A stamped non-image mode must also be present (stamp != image-only): + # the executor trusts stamp presence to skip live resolution entirely. + tool_settings = _build(_make_profile({"output_mode": "layout_preserving"})) + assert tool_settings[TSPKeys.X2TEXT_OUTPUT_MODE] == "layout_preserving" + + def test_non_llmwhisperer_adapter_is_not_stamped(self) -> None: + tool_settings = _build( + _make_profile({"output_mode": "image"}, adapter_id="some-other|123") + ) + assert TSPKeys.X2TEXT_OUTPUT_MODE not in tool_settings From 3dc5832f0c3469bfffc9aed056a757aee616c118 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Thu, 6 Aug 2026 13:13:37 +0530 Subject: [PATCH 18/24] UN-2646 [FIX] Fail closed when the VLM plugin is installed but broken MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A half-broken cloud install (plugins.vlm_image_answer package present, backend_hooks failing to import) previously logged a warning and degraded to no-ops while the adapter gating — probing only the package — kept image output mode enabled. Re-extracting an image-mode document would then rewrite its page images while the answers stored against the old pages were never invalidated (Greptile P1). Both surfaces now share one fail-closed verdict: - The gating probe requires backend_hooks to import, so a broken install strips image mode from the schema and rejects new saves. - vlm_utils exports VLM_HOOKS_BROKEN, and dynamic_extractor's image-mode guard (the single extraction choke point) rejects image-mode extraction outright in that state — pages can never be rewritten while stale answers survive. Text-mode profiles are unaffected. Tested across all three surfaces: probe verdicts for absent, broken and healthy installs (forced via sys.modules so the local cloud overlay cannot skew them), the choke-point rejection incl. text-mode passthrough, and the OSS not-broken baseline. Co-Authored-By: Claude Fable 5 --- .../image_output_gating.py | 15 +++++++ .../tests/test_image_output_gating.py | 41 +++++++++++++++++++ .../prompt_studio_helper.py | 20 +++++++++ .../test_validate_image_output_pdf_only.py | 28 +++++++++++++ backend/prompt_studio/tests/test_vlm_utils.py | 5 +++ backend/prompt_studio/vlm_utils.py | 19 +++++---- 6 files changed, 121 insertions(+), 7 deletions(-) diff --git a/backend/adapter_processor_v2/image_output_gating.py b/backend/adapter_processor_v2/image_output_gating.py index 422e772439..6c2870071f 100644 --- a/backend/adapter_processor_v2/image_output_gating.py +++ b/backend/adapter_processor_v2/image_output_gating.py @@ -24,10 +24,25 @@ def _consumer_plugin_available() -> bool: + """True only when the package AND its backend hooks import. + + A half-broken install (package present, ``backend_hooks`` failing to + import) must gate image mode off — fail-closed — rather than leave + the mode enabled while the re-extraction invalidation and deploy + validation hooks quietly stop existing. + """ try: import plugins.vlm_image_answer # noqa: F401 except ImportError: return False + try: + import plugins.vlm_image_answer.backend_hooks # noqa: F401 + except ImportError: + logger.error( + "plugins.vlm_image_answer is installed but backend_hooks " + "failed to import; disabling image output mode (fail-closed)" + ) + return False return True diff --git a/backend/adapter_processor_v2/tests/test_image_output_gating.py b/backend/adapter_processor_v2/tests/test_image_output_gating.py index 025b89b27c..863b7d033e 100644 --- a/backend/adapter_processor_v2/tests/test_image_output_gating.py +++ b/backend/adapter_processor_v2/tests/test_image_output_gating.py @@ -108,3 +108,44 @@ def test_unknown_adapter_id_still_rejected(self, consumer_absent) -> None: def test_empty_metadata_allowed(self, consumer_absent) -> None: validate_image_output_allowed(None, _LLMW_ADAPTER_ID) validate_image_output_allowed({}, _LLMW_ADAPTER_ID) + + +class TestConsumerProbe: + """The availability probe must fail closed on a half-broken install.""" + + def test_absent_package_is_unavailable(self, monkeypatch: pytest.MonkeyPatch) -> None: + # OSS baseline: no plugins.vlm_image_answer package at all. A None + # sys.modules entry forces ImportError regardless of whether the + # local environment has the cloud plugin overlaid. + import sys + + monkeypatch.setitem(sys.modules, "plugins.vlm_image_answer", None) + assert gating._consumer_plugin_available() is False + + def test_package_without_hooks_is_unavailable( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + # Package present but backend_hooks unimportable (broken install): + # the mode must gate off rather than stay enabled with dead hooks. + import sys + from types import ModuleType + + pkg = ModuleType("plugins.vlm_image_answer") + monkeypatch.setitem(sys.modules, "plugins", ModuleType("plugins")) + monkeypatch.setitem(sys.modules, "plugins.vlm_image_answer", pkg) + monkeypatch.setitem(sys.modules, "plugins.vlm_image_answer.backend_hooks", None) + assert gating._consumer_plugin_available() is False + + def test_package_with_hooks_is_available( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + import sys + from types import ModuleType + + hooks = ModuleType("plugins.vlm_image_answer.backend_hooks") + pkg = ModuleType("plugins.vlm_image_answer") + pkg.backend_hooks = hooks + monkeypatch.setitem(sys.modules, "plugins", ModuleType("plugins")) + monkeypatch.setitem(sys.modules, "plugins.vlm_image_answer", pkg) + monkeypatch.setitem(sys.modules, "plugins.vlm_image_answer.backend_hooks", hooks) + assert gating._consumer_plugin_available() is True diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index 6f1bc5797f..61de3a96be 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -27,6 +27,7 @@ from utils.local_context import StateStore from backend.celery_service import app as celery_app +from prompt_studio import vlm_utils from prompt_studio.lookup_utils import ( get_lookup_config, get_lookup_configs_for_tool, @@ -1452,6 +1453,10 @@ def _validate_image_output_pdf_only( ``output_mode`` is user-editable adapter metadata, so keying on it alone would make any future x2text adapter that adopts the same key inherit a PDF-only rejection it never asked for. + + Also rejects image-mode extraction outright when the cloud plugin + package is present but its backend hooks are broken + (``vlm_utils.VLM_HOOKS_BROKEN``) — see the inline comment below. """ x2text = profile_manager.x2text if x2text is None: @@ -1465,6 +1470,21 @@ def _validate_image_output_pdf_only( != ImageOutputConstants.IMAGE_MODE ): return + # Fail-closed on a half-broken cloud install: with the plugin package + # present but its backend hooks unimportable, a re-extraction would + # rewrite the page images while the answers stored against the old + # pages are never invalidated. Blocking image-mode extraction here + # (same choke point) is the only safe behavior. + if vlm_utils.VLM_HOOKS_BROKEN: + raise IndexingAPIError( + detail=( + "Image output mode is unavailable: the VLM consumer " + "plugin is installed but failed to load. Contact your " + "administrator, or switch the profile's text extractor " + "to a text output mode." + ), + status_code=500, + ) if not ImageOutputConstants.is_pdf(file_name): raise IndexingAPIError( detail=ImageOutputConstants.PDF_ONLY_ERROR, diff --git a/backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py b/backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py index f8a9b553f5..a0f5fea413 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py +++ b/backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py @@ -79,6 +79,34 @@ def test_missing_x2text_adapter_passes(self) -> None: PromptStudioHelper._validate_image_output_pdf_only(profile, "statement.docx") +class TestBrokenHooksBlockImageMode: + """Half-broken cloud install (package present, hooks unimportable) → + image-mode extraction is rejected outright, so a re-extraction can + never rewrite pages while stale stored answers survive. + """ + + def test_broken_hooks_reject_image_mode_even_for_pdf(self, monkeypatch) -> None: # noqa: ANN001 + monkeypatch.setattr(_psh_mod.vlm_utils, "VLM_HOOKS_BROKEN", True) + with pytest.raises(IndexingAPIError) as exc_info: + PromptStudioHelper._validate_image_output_pdf_only( + _profile({"output_mode": "image"}), "statement.pdf" + ) + assert exc_info.value.status_code == 500 + assert "failed to load" in str(exc_info.value.detail) + + def test_broken_hooks_do_not_affect_text_mode(self, monkeypatch) -> None: # noqa: ANN001 + monkeypatch.setattr(_psh_mod.vlm_utils, "VLM_HOOKS_BROKEN", True) + PromptStudioHelper._validate_image_output_pdf_only( + _profile({"output_mode": "text"}), "statement.docx" + ) + + def test_healthy_state_passes_image_mode_pdf(self, monkeypatch) -> None: # noqa: ANN001 + monkeypatch.setattr(_psh_mod.vlm_utils, "VLM_HOOKS_BROKEN", False) + PromptStudioHelper._validate_image_output_pdf_only( + _profile({"output_mode": "image"}), "statement.pdf" + ) + + class TestGuardIsWiredIntoDynamicExtractor: """The guard must run from dynamic_extractor (the single extract path).""" diff --git a/backend/prompt_studio/tests/test_vlm_utils.py b/backend/prompt_studio/tests/test_vlm_utils.py index a19d6462ab..d024770226 100644 --- a/backend/prompt_studio/tests/test_vlm_utils.py +++ b/backend/prompt_studio/tests/test_vlm_utils.py @@ -19,6 +19,11 @@ class TestOssNoOps: def test_cloud_package_absent_in_oss(self) -> None: assert vlm_utils.VLM_IMAGE_ANSWER_AVAILABLE is False + def test_hooks_not_marked_broken_in_oss(self) -> None: + # Package absent is the expected OSS state — it must not be + # conflated with the fail-closed "installed but broken" state. + assert vlm_utils.VLM_HOOKS_BROKEN is False + def test_vision_warning_is_none(self) -> None: assert vlm_utils.get_profile_vision_warning(SimpleNamespace()) is None diff --git a/backend/prompt_studio/vlm_utils.py b/backend/prompt_studio/vlm_utils.py index 462c480039..22ac149b49 100644 --- a/backend/prompt_studio/vlm_utils.py +++ b/backend/prompt_studio/vlm_utils.py @@ -18,23 +18,28 @@ from plugins.vlm_image_answer import backend_hooks as _hooks VLM_IMAGE_ANSWER_AVAILABLE = True + VLM_HOOKS_BROKEN = False except ImportError: _hooks = None VLM_IMAGE_ANSWER_AVAILABLE = False # Distinguish "running OSS" (package absent — expected, silent) from # "cloud hooks are broken" (package present but backend_hooks failed - # to import): in the latter case the adapter gating still enables - # image output mode while these hooks quietly stop existing. + # to import). The broken state is exported so the consumers of image + # mode fail closed instead of quietly degrading: the adapter gating + # disables the mode and ``dynamic_extractor`` rejects image-mode + # extraction — otherwise a re-extraction would rewrite the page + # images while the stored answers derived from the old pages are + # never invalidated. try: import plugins.vlm_image_answer # noqa: F401 except ImportError: - pass + VLM_HOOKS_BROKEN = False else: - logger.warning( + VLM_HOOKS_BROKEN = True + logger.error( "plugins.vlm_image_answer is present but backend_hooks failed " - "to import — VLM profile warnings, deploy-time validation and " - "re-extraction invalidation are disabled while image output " - "mode remains enabled" + "to import — image output mode is disabled (new saves rejected, " + "image-mode extraction blocked) until the install is repaired" ) From b49587a0fd2b399885e7931b958698e8c5ba8ec4 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Thu, 6 Aug 2026 13:22:38 +0530 Subject: [PATCH 19/24] UN-2646 [FIX] Order invalidation before the success marker; propagate hook errors MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Greptile flagged that a runtime failure in the re-extraction invalidation hook was swallowed after the success marker had already committed — so a retry would cache-hit past the failed invalidation and stored answers could stay stale against rewritten page images, silently and permanently. Today this is unreachable (the cloud hook is purely informational: Prompt Studio has no read-side answer cache; output rows are overwritten per run, the same semantics text-mode re-extraction has always had). But the contract now makes it structurally impossible for any future hook that performs real invalidation: - The bridge no longer swallows hook exceptions — the hook owns its error policy, and a raised error fails the re-extraction loudly. - dynamic_extractor invokes the hook BEFORE committing the extraction-success marker, so a failed invalidation leaves the marker unset and a retry re-runs extraction + invalidation instead of cache-hitting past a stale-answer state. Zero behavior change for the current hook; the bridge tests now pin delegation and propagation separately. Co-Authored-By: Claude Fable 5 --- .../prompt_studio_helper.py | 21 ++++++++------- backend/prompt_studio/tests/test_vlm_utils.py | 19 ++++++++----- backend/prompt_studio/vlm_utils.py | 27 +++++++++++-------- 3 files changed, 40 insertions(+), 27 deletions(-) diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index 61de3a96be..40d491897a 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -2657,6 +2657,18 @@ def dynamic_extractor( ) extracted_text = result.data.get("extracted_text", "") + + # A fresh (non-cache-hit) extraction rewrote any persisted page + # images — notify the VLM answer-invalidation hook (no-op in OSS) + # BEFORE committing the extraction-success marker: if a hook ever + # fails, the marker stays unset and a retry re-runs extraction and + # invalidation, instead of cache-hitting past a stale-answer state. + invalidate_vlm_answers_on_reextraction( + document_id=str(document_id), + profile_manager=profile_manager, + extract_file_path=extract_file_path, + ) + success = PromptStudioIndexHelper.mark_extraction_status( document_id=document_id, profile_manager=profile_manager, @@ -2669,15 +2681,6 @@ def dynamic_extractor( f"Extraction completed but status not saved." ) - # A fresh (non-cache-hit) extraction rewrote any persisted page - # images — stored VLM answers for this document are stale. - # No-op in OSS (cloud-only hook). - invalidate_vlm_answers_on_reextraction( - document_id=str(document_id), - profile_manager=profile_manager, - extract_file_path=extract_file_path, - ) - return extracted_text @staticmethod diff --git a/backend/prompt_studio/tests/test_vlm_utils.py b/backend/prompt_studio/tests/test_vlm_utils.py index d024770226..e0f6105fca 100644 --- a/backend/prompt_studio/tests/test_vlm_utils.py +++ b/backend/prompt_studio/tests/test_vlm_utils.py @@ -62,9 +62,7 @@ def test_deployment_validation_propagates(self, cloud_hooks: MagicMock) -> None: with pytest.raises(ValueError): vlm_utils.validate_workflow_for_deployment(SimpleNamespace()) - def test_invalidation_delegates_and_swallows_failure( - self, cloud_hooks: MagicMock - ) -> None: + def test_invalidation_delegates(self, cloud_hooks: MagicMock) -> None: profile = SimpleNamespace() vlm_utils.invalidate_vlm_answers_on_reextraction( document_id="d1", profile_manager=profile, extract_file_path="/e.txt" @@ -72,8 +70,15 @@ def test_invalidation_delegates_and_swallows_failure( cloud_hooks.invalidate_vlm_answers_on_reextraction.assert_called_once_with( document_id="d1", profile_manager=profile, extract_file_path="/e.txt" ) - # Invalidation failure must not fail the extraction it rides on. + + def test_invalidation_failure_propagates(self, cloud_hooks: MagicMock) -> None: + # The hook owns its error policy; a raised error must fail the + # re-extraction loudly (before the success marker commits) rather + # than silently leave stored answers stale against rewritten pages. cloud_hooks.invalidate_vlm_answers_on_reextraction.side_effect = RuntimeError - vlm_utils.invalidate_vlm_answers_on_reextraction( - document_id="d1", profile_manager=profile, extract_file_path="/e.txt" - ) # no raise + with pytest.raises(RuntimeError): + vlm_utils.invalidate_vlm_answers_on_reextraction( + document_id="d1", + profile_manager=SimpleNamespace(), + extract_file_path="/e.txt", + ) diff --git a/backend/prompt_studio/vlm_utils.py b/backend/prompt_studio/vlm_utils.py index 22ac149b49..18f9f44dea 100644 --- a/backend/prompt_studio/vlm_utils.py +++ b/backend/prompt_studio/vlm_utils.py @@ -75,17 +75,22 @@ def invalidate_vlm_answers_on_reextraction( ) -> None: """Invalidate stored VLM answers after a re-extraction rewrote pages/. - Called from the extraction choke point right after a successful - (non-cache-hit) extraction. Never raises — invalidation failure must - not fail the extraction itself; the cloud hook logs and degrades. + Called from the extraction choke point after a successful + (non-cache-hit) extraction, BEFORE the extraction-success marker is + committed. Exceptions propagate — the hook owns its error policy. + The current cloud hook is purely informational (Prompt Studio has no + read-side answer cache; output rows are overwritten per run, the + same semantics text-mode re-extraction has always had) and never + raises. A future hook that performs real invalidation must either + handle its own failures or let them fail the re-extraction loudly: + because the failure lands before the marker commits, a retry re-runs + extraction and invalidation instead of cache-hitting past a + stale-answer state. """ if not VLM_IMAGE_ANSWER_AVAILABLE: return - try: - _hooks.invalidate_vlm_answers_on_reextraction( - document_id=document_id, - profile_manager=profile_manager, - extract_file_path=extract_file_path, - ) - except Exception: - logger.exception("VLM answer invalidation failed after re-extraction") + _hooks.invalidate_vlm_answers_on_reextraction( + document_id=document_id, + profile_manager=profile_manager, + extract_file_path=extract_file_path, + ) From 92d85e56942d336105c95051734581139dc25160 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Thu, 6 Aug 2026 15:48:50 +0530 Subject: [PATCH 20/24] [MISC] Retrigger CI: review-thread confirmations added for the two flagged P1 threads No code change. The two threads Greptile flagged as resolved-without- explanation now carry developer confirmations (fixing commit b49587a0f on the invalidation thread; intentional-deferral rationale on the concurrent-reads thread). This empty commit re-fires the required Greptile Review check so the score reflects that state. Co-Authored-By: Claude Fable 5 From c4573311c63c0c47006dc676cf41675f7bffac37 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Thu, 6 Aug 2026 15:55:50 +0530 Subject: [PATCH 21/24] UN-2646 [FIX] Verify the pages dir is empty after reset, not just that rm returned MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Greptile caught a real gap in the directory reset: FileStorage.rm's S3-compatibility fallback (MissingContentMD5 → per-object deletes) only WARNS on individual deletion failures, so rm can return successfully while old page files survive. A survivor numbered within the new page range would silently join the replacement set — exactly the stale-page defect the reset exists to prevent. The writer now verifies the prefix is actually gone after rm (through a fresh listing, invalidating fsspec's dircache first) and raises ExtractorError on survivors instead of writing a contaminated set. Regression test drives an rm that silently leaves a file behind; the in-memory fixture's rm now also drops directory entries, matching real backends. Co-Authored-By: Claude Fable 5 --- .../x2text/llm_whisperer_v2/src/helper.py | 16 +++++++++++++++ unstract/sdk1/tests/llmw_image_fixtures.py | 5 +++++ unstract/sdk1/tests/test_llmw_image_helper.py | 20 +++++++++++++++++++ 3 files changed, 41 insertions(+) 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 c03019745c..b7e5e1a732 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 @@ -785,6 +785,22 @@ def persist_page_images( status_code=500, actual_err=e, ) from e + # ``fs.rm`` is not enough on its own: its S3-compatibility fallback + # (MissingContentMD5 → per-object deletes) only WARNS on individual + # failures, so a "successful" rm can leave survivors behind — which + # would silently join the new set as stale trailing pages. Verify the + # prefix is actually gone (through a fresh listing, not fsspec's + # dircache) and fail loudly otherwise. + invalidate = getattr(getattr(fs, "fs", None), "invalidate_cache", None) + if callable(invalidate): + invalidate(page_store_dir) + if fs.exists(page_store_dir): + raise ExtractorError( + "Previous page images survived the pre-write cleanup (partial " + f"delete): dir={page_store_dir}, provider={fs.provider.value}. " + "Refusing to write a new page set on top of stale pages.", + status_code=500, + ) fs.mkdir(create_parents=True, path=page_store_dir) write_with_retry = retry_with_exponential_backoff( diff --git a/unstract/sdk1/tests/llmw_image_fixtures.py b/unstract/sdk1/tests/llmw_image_fixtures.py index 6194e78602..e10d8d4343 100644 --- a/unstract/sdk1/tests/llmw_image_fixtures.py +++ b/unstract/sdk1/tests/llmw_image_fixtures.py @@ -92,6 +92,11 @@ def rm(self, path: str, recursive: bool = True) -> None: for key in list(self._files): if key == str(path) or key.startswith(prefix): del self._files[key] + # Real backends remove the directory itself too (local: rmtree; + # S3: the prefix stops existing once its objects are gone). + for d in list(self._dirs): + if d == str(path) or d.startswith(prefix): + self._dirs.discard(d) def mkdir(self, path: str, create_parents: bool = True) -> None: self._dirs.add(str(path)) diff --git a/unstract/sdk1/tests/test_llmw_image_helper.py b/unstract/sdk1/tests/test_llmw_image_helper.py index 36cd978578..83352556a2 100644 --- a/unstract/sdk1/tests/test_llmw_image_helper.py +++ b/unstract/sdk1/tests/test_llmw_image_helper.py @@ -206,6 +206,26 @@ def _rm_fails(path: str, recursive: bool = True) -> None: with pytest.raises(ExtractorError, match="clear previous page images"): H.persist_page_images(fs, "doc/pages", [(1, b"x")]) + def test_silently_partial_delete_raises_instead_of_serving_stale_pages( + self, + ) -> None: + # FileStorage.rm's S3-compatibility fallback deletes objects one at a + # time and only WARNS on per-object failures — so rm can "succeed" + # while files survive. The writer must detect survivors and refuse to + # write, else the stale page joins the new set as a trailing page. + fs = InMemoryFileStorage(provider=FileStorageProvider.S3) + H.persist_page_images(fs, "doc/pages", [(1, b"a"), (2, b"b"), (3, b"c")]) + + real_rm = type(fs).rm + + def _rm_leaves_survivor(path: str, recursive: bool = True) -> None: + real_rm(fs, path, recursive) + fs._files["doc/pages/page_003.png"] = b"c" # survived the delete + + fs.rm = _rm_leaves_survivor # type: ignore[method-assign] + with pytest.raises(ExtractorError, match="survived the pre-write cleanup"): + H.persist_page_images(fs, "doc/pages", [(1, b"x"), (2, b"y")]) + 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( From 6df67c81f9e231161ce9a3755c1f816cb84a6d12 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Thu, 6 Aug 2026 16:09:27 +0530 Subject: [PATCH 22/24] UN-2646 [DOCS] Document the accepted concurrency contract at the writer and reader MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The overlap between a same-document re-extraction and an in-flight prompt execution is a designed, accepted limitation: the reader fails with a typed, retryable error (IMAGE_OUTPUT_MISSING) rather than ever mixing old and new pages into one answer. This was agreed in review (atomic directory replacement does not exist on object storage; a generation-versioned scheme was considered and rejected as out of scope) but the code never said so — record the contract in the writer and reader docstrings so it reads as intent, not oversight. Co-Authored-By: Claude Fable 5 --- .../x2text/llm_whisperer_v2/src/helper.py | 15 +++++++++++++++ .../sdk1/adapters/x2text/page_image_loader.py | 9 +++++++++ 2 files changed, 24 insertions(+) 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 b7e5e1a732..074ad980a4 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 @@ -774,6 +774,21 @@ def persist_page_images( trailing images would otherwise be read back as part of the new set. A failed reset raises rather than risk serving another document's pages to the vision LLM. + + Concurrency — a DESIGNED, ACCEPTED limitation, not an oversight: a + prompt execution that reads this directory while a re-extraction of + the same document is rewriting it observes a missing or incomplete + set and fails with a *typed, retryable* error + (``PageImagesNotFoundError`` / ``PageImageSetIncompleteError``, both + surfaced to the user as ``IMAGE_OUTPUT_MISSING``). This loud + transient failure is deliberately preferred over the alternatives: + in-place overwrites can silently mix old and new pages into one + answer, and atomic directory replacement does not exist on the + object-storage backends ``FileStorage`` targets (S3 has no atomic + rename). A generation-versioned directory scheme with a pointer + object was considered and rejected as out of scope (PR #2210 + review); revisit only if same-document re-extract-while-answering + becomes a real workflow. """ try: if fs.exists(page_store_dir): diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/page_image_loader.py b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/page_image_loader.py index e293ef5df2..39fa60e4d8 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/page_image_loader.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/page_image_loader.py @@ -21,6 +21,15 @@ (post-write loss). Distinct from "not found" so remediation can differ. - ``PageCapExceededError`` — document larger than the page cap; callers must fail explicitly rather than silently truncate. + +Concurrency contract (designed, accepted): the writer resets and rewrites +the stable pages directory on re-extraction, so a read that overlaps a +same-document re-extraction may observe a missing or incomplete set. That +surfaces as the typed errors above — a loud, retryable failure — by +deliberate choice: silently mixing old and new pages into one answer is +the worse outcome, and atomic directory replacement does not exist on +object storage. See ``LLMWhispererHelper.persist_page_images`` for the +full rationale. """ import base64 From 178866fb0bb55e9ca80c86b028e50967cdccee56 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Thu, 6 Aug 2026 16:47:00 +0530 Subject: [PATCH 23/24] [MISC] Retrigger review: flagged threads reopened with developer responses in place No code change. The two threads previously flagged as resolved-without-response are now unresolved and carry the developer confirmations (fixing commit b49587a0f on the invalidation thread; intentional-deferral rationale on the concurrent-reads thread) so the review run reads them in their answered, open state. Co-Authored-By: Claude Fable 5 From 5ac756511101172c17248f0c1b7669e1c64591c3 Mon Sep 17 00:00:00 2001 From: pk-zipstack Date: Thu, 6 Aug 2026 17:46:33 +0530 Subject: [PATCH 24/24] UN-2646 [TEST] Pin the no-op-rm variant of the post-reset survivor guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The guard itself (verify the pages prefix is gone after rm, raise ExtractorError on survivors) landed in c4573311c together with a partial-delete regression test. This adds the degenerate variant the review asked for verbatim — rm as a complete no-op — asserting the write is rejected and the old set is left untouched. Co-Authored-By: Claude Fable 5 --- unstract/sdk1/tests/test_llmw_image_helper.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/unstract/sdk1/tests/test_llmw_image_helper.py b/unstract/sdk1/tests/test_llmw_image_helper.py index 83352556a2..6293578c24 100644 --- a/unstract/sdk1/tests/test_llmw_image_helper.py +++ b/unstract/sdk1/tests/test_llmw_image_helper.py @@ -226,6 +226,18 @@ def _rm_leaves_survivor(path: str, recursive: bool = True) -> None: with pytest.raises(ExtractorError, match="survived the pre-write cleanup"): H.persist_page_images(fs, "doc/pages", [(1, b"x"), (2, b"y")]) + def test_noop_rm_is_detected_and_write_rejected(self) -> None: + # Degenerate variant of the above: rm succeeds as a complete no-op + # (nothing deleted at all). The post-reset verification must reject + # the write outright — nothing may be written over the old set. + fs = InMemoryFileStorage(provider=FileStorageProvider.S3) + H.persist_page_images(fs, "doc/pages", [(1, b"a"), (2, b"b")]) + + fs.rm = lambda path, recursive=True: None # type: ignore[method-assign] + with pytest.raises(ExtractorError, match="survived the pre-write cleanup"): + H.persist_page_images(fs, "doc/pages", [(1, b"x")]) + assert fs.read(path="doc/pages/page_001.png", mode="rb") == b"a" # untouched + 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(