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..6c2870071f --- /dev/null +++ b/backend/adapter_processor_v2/image_output_gating.py @@ -0,0 +1,122 @@ +"""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: + """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 + + +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 + 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/adapter_processor_v2/serializers.py b/backend/adapter_processor_v2/serializers.py index c8545223bd..b63d891093 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..863b7d033e --- /dev/null +++ b/backend/adapter_processor_v2/tests/test_image_output_gating.py @@ -0,0 +1,151 @@ +"""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) + + +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/api_v2/serializers.py b/backend/api_v2/serializers.py index 3db7f53db6..e62a82b3db 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 6d8763e05a..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 @@ -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, @@ -76,7 +77,9 @@ 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 from unstract.sdk1.exceptions import IndexingError, SdkError from unstract.sdk1.execution.context import ExecutionContext @@ -443,6 +446,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() @@ -845,6 +849,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 @@ -891,6 +898,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() @@ -952,6 +960,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, @@ -1068,6 +1077,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 @@ -1145,6 +1157,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, @@ -1279,6 +1292,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: @@ -1396,6 +1413,84 @@ 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 + ) -> 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 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. + + 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: + return + adapter_id = getattr(x2text, "adapter_id", "") or "" + if not adapter_id.startswith("llmwhisperer|"): + return + metadata = x2text.metadata or {} + if ( + metadata.get(ImageOutputConstants.OUTPUT_MODE) + != 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, + status_code=400, + ) + @staticmethod def index_document( tool_id: str, @@ -2046,6 +2141,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() @@ -2470,6 +2566,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)) @@ -2553,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, 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 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..a0f5fea413 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_core_v2/tests/test_validate_image_output_pdf_only.py @@ -0,0 +1,126 @@ +"""Unit tests for the image-output PDF-only guard. + +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. +""" + +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.constants import ImageOutputConstants + +PromptStudioHelper = _psh_mod.PromptStudioHelper + +_LLMW_ADAPTER_ID = "llmwhisperer|a5e6b8af-3e1f-4a80-b006-d017e8e67f93" + + +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: + """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: + 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) == ImageOutputConstants.PDF_ONLY_ERROR + + @pytest.mark.parametrize("file_name", ["statement.pdf", "STATEMENT.PDF"]) + def test_pdf_passes_case_insensitively(self, file_name: str) -> None: + PromptStudioHelper._validate_image_output_pdf_only( + _profile({"output_mode": "image"}), file_name + ) + + +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"}, {}, 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 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).""" + + 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/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..e0f6105fca --- /dev/null +++ b/backend/prompt_studio/tests/test_vlm_utils.py @@ -0,0 +1,84 @@ +"""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_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 + + 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(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" + ) + + 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 + 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 new file mode 100644 index 0000000000..18f9f44dea --- /dev/null +++ b/backend/prompt_studio/vlm_utils.py @@ -0,0 +1,96 @@ +"""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 + 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). 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: + VLM_HOOKS_BROKEN = False + else: + VLM_HOOKS_BROKEN = True + logger.error( + "plugins.vlm_image_answer is present but backend_hooks failed " + "to import — image output mode is disabled (new saves rejected, " + "image-mode extraction blocked) until the install is repaired" + ) + + +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 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 + _hooks.invalidate_vlm_answers_on_reextraction( + document_id=document_id, + profile_manager=profile_manager, + extract_file_path=extract_file_path, + ) 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 c76e073e1b..0c4127cbca 100644 --- a/frontend/src/components/custom-tools/add-llm-profile/AddLlmProfile.jsx +++ b/frontend/src/components/custom-tools/add-llm-profile/AddLlmProfile.jsx @@ -352,10 +352,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 bd703e6538..feb89d8f6f 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,77 @@ 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." + ) + + # --- 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/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..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 @@ -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,6 +14,7 @@ class Modes(Enum): class OutputModes(Enum): LAYOUT_PRESERVING = "layout_preserving" TEXT = "text" + IMAGE = ImageOutputConstants.IMAGE_MODE class HTTPMethod(Enum): @@ -31,6 +34,12 @@ class WhispererEndpoint: STATUS = "whisper-status" RETRIEVE = "whisper-retrieve" HIGHLIGHTS = "highlights" + # 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" class WhispererEnv: @@ -47,6 +56,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" @@ -55,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" @@ -114,3 +133,69 @@ 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. + + 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 (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" + MESSAGE = "message" + + # 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_INTERMEDIATE = frozenset({"accepted", "processing", "queued"}) + + # --- Submit query params --- + IMAGE_FORMAT_PARAM = "format" + DEFAULT_IMAGE_FORMAT = "png" + FILE_NAME_PARAM = "file_name" + + # --- Per-page image naming / storage layout --- + # 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 + 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 ade89f7cba..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 @@ -1,5 +1,9 @@ import json import logging +import re +import time +import zipfile +import zlib from io import BytesIO from pathlib import Path from typing import Any @@ -13,12 +17,20 @@ ) 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, Modes, OutputModes, WhispererConfig, WhispererDefaults, + WhispererEndpoint, WhispererHeader, WhisperStatus, ) @@ -26,7 +38,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 +59,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 = WhispererDefaults.IMAGE_REQUEST_TIMEOUT, + 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. + 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 +115,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 +438,534 @@ 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 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), + # so no separate `0*` prefix is needed — keeping the pattern linear. + _PAGE_IMAGE_RE = re.compile(ImageOutputConfig.ZIP_PAGE_MEMBER_REGEX, re.IGNORECASE) + + @staticmethod + def _safe_json(response: Response) -> dict[str, Any]: + """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 + + @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. ``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 + 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) + 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 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} returned an " + f"unexpected status '{status or ''}': {msg}", + status_code=502, + ) + 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. + """ + # 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, + ) + buffer = BytesIO() + # 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 + + @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]] = [] + 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))) + 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 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. + 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]], expected_page_count: int | None + ) -> None: + """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 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 != expected_page_count: + raise ExtractorError( + "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, + ) + + # 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: + 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 _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, + page_store_dir: str, + pages: list[tuple[int, bytes]], + ) -> list[PageImageReference]: + """Write every page image to FileStorage with per-page retry. + + 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``. + + 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. + + 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): + 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.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( + 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: + 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}, " + 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, + ) -> tuple[str, list[PageImageReference]]: + """End-to-end image output flow (orchestrator). + + submit -> poll -> download+extract (ONCE) -> verify page count -> + 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_bytes = fs.read(path=input_file_path, mode="rb") + whisper_hash = LLMWhispererHelper.submit_pdf_to_images( + config, + BytesIO(input_bytes), + tag=tag, + file_name=Path(input_file_path).name, + ) + LLMWhispererHelper.poll_pdf_to_images_status(config, whisper_hash) + + pages = LLMWhispererHelper._download_and_extract( + config=config, whisper_hash=whisper_hash + ) + + # 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, + ) + references = LLMWhispererHelper.persist_page_images(fs, page_store_dir, pages) + logger.info( + "Image mode: completed job=%s pages=%d", whisper_hash, len(references) + ) + return whisper_hash, 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`` — + 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." + ) + + @staticmethod + def write_image_output( + fs: FileStorage, + output_file_path: str, + summary: str, + ) -> None: + """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( + path=str(output_file_path), + mode="w", + data=summary, + encoding="utf-8", + ) + except Exception as 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 3a48a57647..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 @@ -4,12 +4,19 @@ import os from typing import TYPE_CHECKING, Any -from unstract.sdk1.adapters.x2text.constants import X2TextConstants +from unstract.sdk1.adapters.exceptions import ExtractorError +from unstract.sdk1.adapters.x2text.constants import ( + ImageOutputConstants, + 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 +68,84 @@ def test_connection(self) -> bool: ) return True + @staticmethod + 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 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, + 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 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`` — 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, fs=fs) + whisper_hash, 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, + ) + summary = LLMWhispererHelper.build_image_output_summary(page_images) + # 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, + ) + logger.info( + "Image mode: returning %d page image reference(s) for %s", + len(page_images), + input_file_path, + ) + return TextExtractionResult( + extracted_text=summary, + extraction_metadata=TextExtractionMetadata( + whisper_hash=whisper_hash, + page_images=page_images, + ), + ) + def process( self, input_file_path: str, @@ -81,7 +166,30 @@ 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 + ) 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, + fs, + tag=kwargs.get(X2TextConstants.TAGS), + ) + logger.info( "HIGHLIGHT_DEBUG LLMWhispererV2.process: enable_highlight=%s", enable_highlight, 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..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 @@ -44,10 +44,15 @@ "title": "Output Mode", "enum": [ "layout_preserving", - "text" + "text", + "image" ], - "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)" + "enumNames": [ + "Layout Preserving", + "Text", + "Image (PDF only)" + ], + "default": "layout_preserving" }, "line_splitter_tolerance": { "type": "number", @@ -58,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": { @@ -111,35 +116,65 @@ "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": { + "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)." + } + } } - }, - "required": [ - "median_filter_size", - "gaussian_blur_radius" - ] - } + } + ] } 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..39fa60e4d8 --- /dev/null +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/page_image_loader.py @@ -0,0 +1,319 @@ +"""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. + +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 +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 new file mode 100644 index 0000000000..e10d8d4343 --- /dev/null +++ b/unstract/sdk1/tests/llmw_image_fixtures.py @@ -0,0 +1,205 @@ +"""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 + 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] + # 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)) + + 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", + length: int = -1, + **_: object, + ) -> bytes | str: + 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) + 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]: + 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, + fail_substrings: tuple[str, ...] = (), + **kwargs: object, + ) -> None: + """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( + 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 + 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) + + 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..57d5a9fb45 --- /dev/null +++ b/unstract/sdk1/tests/test_llm_whisperer_v2_constants.py @@ -0,0 +1,34 @@ +"""Unit tests for LLMWhisperer v2 adapter constants. + +Covers the image OutputModes value and the env-var-backed page-store retry +budget. +""" + +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) -> 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 new file mode 100644 index 0000000000..6293578c24 --- /dev/null +++ b/unstract/sdk1/tests/test_llmw_image_helper.py @@ -0,0 +1,406 @@ +"""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 +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 helper as helper_mod +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, + FlakyFileStorage, + InMemoryFileStorage, + make_page_zip, + minimal_png, +) + +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: + corrupt = io.BytesIO(CORRUPT_ZIP) + 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: + pages = H.extract_page_images_from_zip(io.BytesIO(make_page_zip(2))) + 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, 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, 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, expected_page_count=None) # no raise + + +class TestFolderKeyAndNaming: + 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 "doc-a" in dir_a and "doc-b" in dir_b + assert dir_a.endswith("pages") + + 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_when_no_output(self) -> None: + assert H.build_page_store_dir(None, "/docs/in.pdf") == "/docs/in/pages" + + @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: + # 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(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_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_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_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( + output_file_path=str(tmp_path / "out.txt"), + input_file_path=str(tmp_path / "in.pdf"), + ) + 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" + + +_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 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 + + +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: + """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" + ), + ] + 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) + + # 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 + + 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 new file mode 100644 index 0000000000..9f058b34d2 --- /dev/null +++ b/unstract/sdk1/tests/test_llmw_v2_process_image.py @@ -0,0 +1,200 @@ +"""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) + adapter = _adapter() + 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, + ) -> tuple[str, list[PageImageReference]]: + captured.update(kwargs) + return "run-1|doc-hash", 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"), + ) + # 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") + + # 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 + # 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 + assert write_calls["output_file_path"] == "out.txt" + assert write_calls["summary"] == expected_summary + + def test_empty_page_list_is_safe(self, monkeypatch: MonkeyPatch) -> None: + 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 == [] + assert result.extracted_text == LLMWhispererHelper.build_image_output_summary([]) + + def test_tag_forwarded_to_helper(self, monkeypatch: MonkeyPatch) -> None: + captured: dict[str, object] = {} + + def _capture(**kwargs: object) -> tuple[str, list]: + captured.update(kwargs) + 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 **_: ("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: + image_called = {"hit": False} + monkeypatch.setattr( + LLMWhispererHelper, + "get_page_images", + lambda **_: image_called.__setitem__("hit", True) or ("wh", []), + ) + adapter = _adapter(output_mode="image") + with pytest.raises(ExtractorError, match="PDF input only"): + adapter.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") + + 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..73f9650f62 --- /dev/null +++ b/unstract/sdk1/tests/test_page_image_loader.py @@ -0,0 +1,314 @@ +"""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 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 + +from tests.llmw_image_fixtures import InMemoryFileStorage, minimal_png + +_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_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 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..1245aaa221 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,19 @@ 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, + execution_source=str(execution_source or ""), + 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 +1784,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 +1891,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 +2334,37 @@ 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: + 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), + 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 "" + ), + stamped_mode=stamped_mode, + execution_source=str(params.get(PSKeys.EXECUTION_SOURCE) 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..08706b8db6 --- /dev/null +++ b/workers/executor/executors/vlm_image_answer.py @@ -0,0 +1,263 @@ +"""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 + + +_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 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. + """ + 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 + + 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). + 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 = "", + 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. Same detection semantics as + ``detect_image_mode_config``: stamp first, platform resolution only + for non-IDE sources. + """ + if not adapter_instance_id: + return + 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 " + "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..adc0fb4e25 --- /dev/null +++ b/workers/executor/tests/test_vlm_image_answer_bridge.py @@ -0,0 +1,353 @@ +"""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): + # 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): + 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 + + 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/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, 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()