From 7b95a98c8b508ba14fe80acb32b60cec7e8defe2 Mon Sep 17 00:00:00 2001 From: Athul Date: Mon, 31 Aug 2026 19:45:37 +0530 Subject: [PATCH] UN-4042 [FIX] Bill the pages LLMWhisperer extracted, not every page in the file `push_usage_details` opened the input file and counted its pages, so an adapter configured with `pages_to_extract = "1-5"` billed all 199 pages of a 199-page PDF while LLMWhisperer billed the 5 it actually read. LLMWhisperer already returns the number it bills. `assemble_document` (`tasks_v2.py:1834`) computes `processed_page_count` once and uses that same value for both `whisper_metadata` (line 2032) and `page_count_total` (line 2095), and the whole `assembled_document` reaches the adapter through `/whisper-retrieve`. The V2 adapter was discarding it. - `TextExtractionResult` carries an optional `page_count` - `LLMWhispererHelper.get_processed_page_count` reads and validates it off the response - `push_usage_details` prefers it, falling back to counting the file when an adapter reports nothing Scoped to PDFs. LLMWhisperer also applies its 3000-words rule to sheets and text where Unstract hardcodes 1 page; adopting that raises those bills and is tracked separately in UN-4043. The billed count can now only go down or stay equal, never up. CoverForce, July 2026: 1,622 documents, 24,078 pages charged and never extracted. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_019zLkLjqzwoN5Lm8thPqV9p --- .../src/unstract/sdk1/adapters/x2text/dto.py | 12 + .../x2text/llm_whisperer_v2/src/helper.py | 24 ++ .../llm_whisperer_v2/src/llm_whisperer_v2.py | 1 + unstract/sdk1/src/unstract/sdk1/x2txt.py | 60 +++-- unstract/sdk1/tests/test_x2txt_page_usage.py | 225 ++++++++++++++++++ 5 files changed, 304 insertions(+), 18 deletions(-) create mode 100644 unstract/sdk1/tests/test_x2txt_page_usage.py diff --git a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/dto.py b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/dto.py index 95c60bbe8c..894d97f084 100644 --- a/unstract/sdk1/src/unstract/sdk1/adapters/x2text/dto.py +++ b/unstract/sdk1/src/unstract/sdk1/adapters/x2text/dto.py @@ -10,5 +10,17 @@ class TextExtractionMetadata: @dataclass class TextExtractionResult: + """Outcome of a text extraction. + + Attributes: + extracted_text: The text the extractor produced. + extraction_metadata: Extractor-specific metadata, if any. + page_count: Pages the extractor actually processed. Set by adapters that + report it, so that a document extracted with a page range is billed + for the pages read rather than every page in the file (UN-4042). + None when the adapter reports no count. + """ + extracted_text: str extraction_metadata: TextExtractionMetadata | None = None + page_count: int | None = None 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..f58937706f 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 @@ -313,6 +313,30 @@ def make_highlight_data_request( ) return retrieve_response + @staticmethod + def get_processed_page_count(response: dict[str, Any]) -> int | None: + """Pages LLMWhisperer actually processed for this extraction. + + This is the same count LLMWhisperer bills, so reading it keeps Unstract's + page usage in step with a `pages_to_extract` range instead of counting + every page in the uploaded file (UN-4042). + + Args: + response: Decoded `/whisper-retrieve` body returned by `make_request`. + + Returns: + The processed page count, or None if the response does not carry a + usable one — callers fall back to counting the input file. + """ + metadata = response.get("whisper_metadata") + if not isinstance(metadata, dict): + return None + page_count = metadata.get("processed_page_count") + # bool is a subclass of int, so it has to be rejected explicitly. + if isinstance(page_count, bool) or not isinstance(page_count, int): + return None + return page_count if page_count > 0 else None + @staticmethod def extract_text_from_response( output_file_path: str | None, 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..3fc6210d5a 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 @@ -108,4 +108,5 @@ def process( fs=fs, ), extraction_metadata=metadata, + page_count=LLMWhispererHelper.get_processed_page_count(response), ) diff --git a/unstract/sdk1/src/unstract/sdk1/x2txt.py b/unstract/sdk1/src/unstract/sdk1/x2txt.py index 2024f8cdbc..e229c63bbd 100644 --- a/unstract/sdk1/src/unstract/sdk1/x2txt.py +++ b/unstract/sdk1/src/unstract/sdk1/x2txt.py @@ -115,7 +115,12 @@ def process( input_file_path, output_file_path, fs, **kwargs ) # The will be executed each and every time text extraction takes place - self.push_usage_details(input_file_path, mime_type, fs=fs) + self.push_usage_details( + input_file_path, + mime_type, + fs=fs, + page_count=text_extraction_result.page_count, + ) return text_extraction_result def push_usage_details( @@ -123,31 +128,50 @@ def push_usage_details( input_file_path: str, mime_type: str, fs: FileStorage | None = None, + page_count: int | None = None, ) -> None: + """Report the pages consumed by a completed extraction. + + Args: + input_file_path: File that was extracted. + mime_type: MIME type of that file. + fs: File storage to read the file through. + page_count: Pages the extractor actually processed, when it reports + them. Preferred over counting the input file, because an adapter + configured with a `pages_to_extract` range reads fewer pages than + the document holds (UN-4042). None for adapters that report no + count, which keeps the earlier file-based behaviour. + """ if fs is None: fs = FileStorage(provider=FileStorageProvider.LOCAL) file_size = ToolUtils.get_file_size(input_file_path, fs) if mime_type == MimeType.PDF: - pdf_contents = io.BytesIO(fs.read(path=input_file_path, mode="rb")) - with pdfplumber.open(pdf_contents) as pdf: - # calculate the number of pages - page_count = len(pdf.pages) - Audit().push_page_usage_data( - platform_api_key=self._tool.get_env_or_die(ToolEnv.PLATFORM_API_KEY), - file_size=file_size, - file_type=mime_type, - page_count=page_count, - kwargs=self._usage_kwargs, + billable_pages = ( + page_count + if page_count is not None and page_count > 0 + else self._get_pdf_page_count(input_file_path, fs) ) else: # TODO: Calculate page usage for other file types (3000 words = 1 page) # We are allowing certain image types,and raw texts. We will consider them # as single page documents as there in no concept of page numbers. - Audit().push_page_usage_data( - platform_api_key=self._tool.get_env_or_die(ToolEnv.PLATFORM_API_KEY), - file_size=file_size, - file_type=mime_type, - page_count=1, - kwargs=self._usage_kwargs, - ) + # Deliberately not using the extractor's count here: LLMWhisperer + # applies the 3000-words rule to sheets and text, so adopting it would + # raise these bills. Tracked separately in UN-4043. + billable_pages = 1 + + Audit().push_page_usage_data( + platform_api_key=self._tool.get_env_or_die(ToolEnv.PLATFORM_API_KEY), + file_size=file_size, + file_type=mime_type, + page_count=billable_pages, + kwargs=self._usage_kwargs, + ) + + @staticmethod + def _get_pdf_page_count(input_file_path: str, fs: FileStorage) -> int: + """Number of pages held by the PDF at `input_file_path`.""" + pdf_contents = io.BytesIO(fs.read(path=input_file_path, mode="rb")) + with pdfplumber.open(pdf_contents) as pdf: + return len(pdf.pages) diff --git a/unstract/sdk1/tests/test_x2txt_page_usage.py b/unstract/sdk1/tests/test_x2txt_page_usage.py new file mode 100644 index 0000000000..1eef88b70c --- /dev/null +++ b/unstract/sdk1/tests/test_x2txt_page_usage.py @@ -0,0 +1,225 @@ +"""Tests for page usage accounting in X2Text. + +Regression cover for UN-4042: page usage was counted from the input file, so a +199-page PDF with ``pages_to_extract = "1-5"`` was billed 199 pages while +LLMWhisperer billed the 5 it actually extracted. +""" + +from typing import Any +from unittest.mock import MagicMock, patch + +import pytest +from unstract.sdk1.adapters.x2text.dto import TextExtractionResult +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, +) +from unstract.sdk1.constants import MimeType +from unstract.sdk1.x2txt import X2Text + + +@pytest.fixture +def x2text() -> X2Text: + """An X2Text with no adapter wired up — page usage needs neither.""" + tool = MagicMock() + tool.get_env_or_die.return_value = "test-platform-api-key" + return X2Text(tool=tool, usage_kwargs={"run_id": "test-run"}) + + +@pytest.fixture +def pdf_fs() -> MagicMock: + """A FileStorage standing in for a 199-page PDF.""" + fs = MagicMock() + fs.read.return_value = b"%PDF-1.4 fake" + return fs + + +def _pushed_page_count(audit: MagicMock) -> int: + """The page_count handed to the platform by a single push.""" + audit.return_value.push_page_usage_data.assert_called_once() + return audit.return_value.push_page_usage_data.call_args.kwargs["page_count"] + + +class TestPushUsageDetails: + """What X2Text.push_usage_details reports to the audit service.""" + + @patch("unstract.sdk1.x2txt.Audit") + @patch("unstract.sdk1.x2txt.pdfplumber") + @patch("unstract.sdk1.x2txt.ToolUtils.get_file_size", return_value=1234) + def test_pdf_uses_extractor_page_count( + self, + _size: MagicMock, + pdfplumber: MagicMock, + audit: MagicMock, + x2text: X2Text, + pdf_fs: MagicMock, + ) -> None: + """UN-4042: 5 pages extracted from a 199-page PDF bills 5, not 199.""" + pdfplumber.open.return_value.__enter__.return_value.pages = [None] * 199 + + x2text.push_usage_details("in.pdf", MimeType.PDF, fs=pdf_fs, page_count=5) + + assert _pushed_page_count(audit) == 5 + + @patch("unstract.sdk1.x2txt.Audit") + @patch("unstract.sdk1.x2txt.pdfplumber") + @patch("unstract.sdk1.x2txt.ToolUtils.get_file_size", return_value=1234) + def test_pdf_without_extractor_count_falls_back_to_file( + self, + _size: MagicMock, + pdfplumber: MagicMock, + audit: MagicMock, + x2text: X2Text, + pdf_fs: MagicMock, + ) -> None: + """Adapters that report nothing keep the pre-UN-4042 behaviour.""" + pdfplumber.open.return_value.__enter__.return_value.pages = [None] * 199 + + x2text.push_usage_details("in.pdf", MimeType.PDF, fs=pdf_fs, page_count=None) + + assert _pushed_page_count(audit) == 199 + + @pytest.mark.parametrize("bad_count", [0, -3]) + @patch("unstract.sdk1.x2txt.Audit") + @patch("unstract.sdk1.x2txt.pdfplumber") + @patch("unstract.sdk1.x2txt.ToolUtils.get_file_size", return_value=1234) + def test_pdf_rejects_non_positive_count( + self, + _size: MagicMock, + pdfplumber: MagicMock, + audit: MagicMock, + bad_count: int, + x2text: X2Text, + pdf_fs: MagicMock, + ) -> None: + """A zero or negative count must never bill zero pages.""" + pdfplumber.open.return_value.__enter__.return_value.pages = [None] * 199 + + x2text.push_usage_details("in.pdf", MimeType.PDF, fs=pdf_fs, page_count=bad_count) + + assert _pushed_page_count(audit) == 199 + + @patch("unstract.sdk1.x2txt.Audit") + @patch("unstract.sdk1.x2txt.ToolUtils.get_file_size", return_value=99) + def test_non_pdf_ignores_extractor_count( + self, + _size: MagicMock, + audit: MagicMock, + x2text: X2Text, + ) -> None: + """Non-PDF counting is UN-4043; this fix must not raise those bills.""" + x2text.push_usage_details("in.txt", MimeType.TEXT, fs=MagicMock(), page_count=20) + + assert _pushed_page_count(audit) == 1 + + +class TestGetProcessedPageCount: + """Reading processed_page_count out of an LLMWhisperer V2 response.""" + + @pytest.mark.parametrize( + "response,expected", + [ + ({"whisper_metadata": {"processed_page_count": 5}}, 5), + ({"whisper_metadata": {"processed_page_count": 1}}, 1), + ({}, None), + ({"whisper_metadata": None}, None), + ({"whisper_metadata": "nope"}, None), + ({"whisper_metadata": {}}, None), + ({"whisper_metadata": {"processed_page_count": 0}}, None), + ({"whisper_metadata": {"processed_page_count": -1}}, None), + ({"whisper_metadata": {"processed_page_count": "5"}}, None), + ({"whisper_metadata": {"processed_page_count": 5.0}}, None), + # isinstance(True, int) is True — must not bill 1 page for a bool. + ({"whisper_metadata": {"processed_page_count": True}}, None), + ], + ) + def test_validates_response( + self, response: dict[str, Any], expected: int | None + ) -> None: + assert LLMWhispererHelper.get_processed_page_count(response) == expected + + +class TestProcessForwardsPageCount: + """The wiring: what the adapter reports is what gets billed.""" + + @patch("unstract.sdk1.x2txt.Audit") + @patch("unstract.sdk1.x2txt.pdfplumber") + @patch("unstract.sdk1.x2txt.ToolUtils.get_file_size", return_value=1234) + def test_process_bills_the_count_the_adapter_reported( + self, + _size: MagicMock, + pdfplumber: MagicMock, + audit: MagicMock, + x2text: X2Text, + pdf_fs: MagicMock, + ) -> None: + """A 199-page PDF extracted with a 5-page range is billed 5.""" + pdfplumber.open.return_value.__enter__.return_value.pages = [None] * 199 + pdf_fs.mime_type.return_value = MimeType.PDF + x2text._x2text_instance = MagicMock() + x2text._x2text_instance.process.return_value = TextExtractionResult( + extracted_text="five pages of text", page_count=5 + ) + + result = x2text.process("in.pdf", fs=pdf_fs) + + assert _pushed_page_count(audit) == 5 + assert result.page_count == 5 + + @patch("unstract.sdk1.x2txt.Audit") + @patch("unstract.sdk1.x2txt.pdfplumber") + @patch("unstract.sdk1.x2txt.ToolUtils.get_file_size", return_value=1234) + def test_process_falls_back_for_adapters_reporting_nothing( + self, + _size: MagicMock, + pdfplumber: MagicMock, + audit: MagicMock, + x2text: X2Text, + pdf_fs: MagicMock, + ) -> None: + """LlamaParse, Unstructured and NoOp report no count — behaviour is unchanged.""" + pdfplumber.open.return_value.__enter__.return_value.pages = [None] * 199 + pdf_fs.mime_type.return_value = MimeType.PDF + x2text._x2text_instance = MagicMock() + x2text._x2text_instance.process.return_value = TextExtractionResult( + extracted_text="all of it" + ) + + x2text.process("in.pdf", fs=pdf_fs) + + assert _pushed_page_count(audit) == 199 + + +class TestLLMWhispererV2ReportsPageCount: + """The V2 adapter lifts processed_page_count off the whisper response.""" + + @patch.object(LLMWhispererHelper, "extract_text_from_response", return_value="text") + @patch.object(LLMWhispererHelper, "send_whisper_request") + def test_page_count_taken_from_whisper_metadata( + self, send: MagicMock, _extract: MagicMock + ) -> None: + send.return_value = { + "whisper_hash_v2": "hash-1", + "result_text": "text", + "whisper_metadata": { + "processed_page_count": 5, + "requested_page_count": 5, + "total_page_count": 199, + }, + } + + result = LLMWhispererV2({}).process("in.pdf", fs=MagicMock()) + + assert result.page_count == 5 + + @patch.object(LLMWhispererHelper, "extract_text_from_response", return_value="text") + @patch.object(LLMWhispererHelper, "send_whisper_request") + def test_page_count_is_none_when_metadata_absent( + self, send: MagicMock, _extract: MagicMock + ) -> None: + """Older LLMWhisperer deployments omit whisper_metadata entirely.""" + send.return_value = {"whisper_hash_v2": "hash-1", "result_text": "text"} + + result = LLMWhispererV2({}).process("in.pdf", fs=MagicMock()) + + assert result.page_count is None