diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_variable_service.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_variable_service.py index 008ba078c0..2a4a7f7b3b 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_variable_service.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_variable_service.py @@ -64,6 +64,41 @@ def identify_variable_type(variable: str) -> VariableType: variable_type = VariableType.STATIC return variable_type + @staticmethod + def find_unresolvable_single_pass_variables(prompt: str) -> list[str]: + """Variables in ``prompt`` that cannot resolve under single-pass extraction. + + UN-2900. Single pass builds ONE combined prompt — every field is declared + up front in a single JSON schema and answered in one LLM call — so no + prompt's output exists to feed another prompt's variable. The runtime + reflects this: the enterprise ``single_pass_extraction`` plugin calls the + shared replacement service with ``structured_output={}``, and both + ``replace_static_variable`` and ``replace_dynamic_variable`` return the + prompt UNCHANGED when their lookup misses. The literal ``{{...}}`` is then + sent to the LLM, silently degrading the answer. + + CUSTOM_DATA is the one exception: it is resolved from the tool's own + ``custom_data`` and never consults the variable map, so it works + identically in both modes and is not reported here. + + Returns the offending variable strings, in prompt order, or an empty list. + + NOTE: classification pairs with ``VariableReplacementService`` in the + worker (``executor/executors/variable_replacement.py``), which keeps its + own copy of these regexes. If one side's patterns change, this validation + and the runtime behaviour will disagree. + """ + unresolvable: list[str] = [] + for variable in PromptStudioVariableService.extract_variables_from_prompt( + prompt=prompt + ): + variable_type = PromptStudioVariableService.identify_variable_type( + variable=variable + ) + if variable_type != VariableType.CUSTOM_DATA: + unresolvable.append(variable) + return unresolvable + @staticmethod def extract_variables_from_prompt(prompt: str) -> list[str]: variable: list[str] = [] diff --git a/backend/prompt_studio/prompt_studio_core_v2/serializers.py b/backend/prompt_studio/prompt_studio_core_v2/serializers.py index acb3a243d7..0dead92f4f 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/serializers.py +++ b/backend/prompt_studio/prompt_studio_core_v2/serializers.py @@ -186,9 +186,15 @@ def to_representation(self, instance): # type: ignore ) # Fetch prompt instances - prompt_instances: ToolStudioPrompt = ToolStudioPrompt.objects.filter( - tool_id=data.get(TSKeys.TOOL_ID) - ).order_by("sequence_number") + # select_related("tool_id"): ToolStudioPromptSerializer's + # single_pass_unresolvable_variables (UN-2900) reads the parent tool's + # single_pass_extraction_mode, which would otherwise be one query per + # prompt here. + prompt_instances: ToolStudioPrompt = ( + ToolStudioPrompt.objects.filter(tool_id=data.get(TSKeys.TOOL_ID)) + .select_related("tool_id") + .order_by("sequence_number") + ) data["created_by_email"] = ( instance.created_by.email if instance.created_by else "" diff --git a/backend/prompt_studio/prompt_studio_v2/serializers.py b/backend/prompt_studio/prompt_studio_v2/serializers.py index 4026f32b94..5463113539 100644 --- a/backend/prompt_studio/prompt_studio_v2/serializers.py +++ b/backend/prompt_studio/prompt_studio_v2/serializers.py @@ -23,6 +23,8 @@ class Meta: class ToolStudioPromptSerializer(AuditSerializer): + single_pass_unresolvable_variables = serializers.SerializerMethodField() + class Meta: model = ToolStudioPrompt fields = "__all__" @@ -30,6 +32,30 @@ class Meta: # the DRF auto-validator that 400s on re-save / PUT before the view runs. validators = [] + def get_single_pass_unresolvable_variables(self, obj) -> list[str]: + """UN-2900: variables in this prompt that single pass cannot resolve. + + Empty unless the parent tool has single-pass extraction enabled. Lives + on this serializer rather than the tool serializer so the same value is + returned by BOTH the tool detail fetch (which nests this serializer per + prompt) and the prompt save response, without duplicating the logic. + + ``tool_id`` is the FK to CustomTool; callers that serialize many prompts + should ``select_related("tool_id")`` to avoid a query per prompt. + """ + from prompt_studio.prompt_studio_core_v2.prompt_variable_service import ( + PromptStudioVariableService, + ) + + tool = getattr(obj, "tool_id", None) + if not tool or not getattr(tool, "single_pass_extraction_mode", False): + return [] + if not obj.prompt: + return [] + return PromptStudioVariableService.find_unresolvable_single_pass_variables( + prompt=obj.prompt + ) + class ToolStudioIndexSerializer(serializers.Serializer): file_name = serializers.CharField() diff --git a/backend/tool_instance_v2/tool_instance_helper.py b/backend/tool_instance_v2/tool_instance_helper.py index 138e48b5b2..d44a77c2a0 100644 --- a/backend/tool_instance_v2/tool_instance_helper.py +++ b/backend/tool_instance_v2/tool_instance_helper.py @@ -486,33 +486,57 @@ def validate_adapter_permissions( adapter_ids: set[str] = set() for llm in tool.properties.adapter.language_models: + adapter_id = None if llm.is_enabled and llm.adapter_id: - adapter_id = tool_meta[llm.adapter_id] + adapter_id = tool_meta.get(llm.adapter_id) elif llm.is_enabled: - adapter_id = tool_meta[AdapterPropertyKey.DEFAULT_LLM_ADAPTER_ID] + adapter_id = tool_meta.get(AdapterPropertyKey.DEFAULT_LLM_ADAPTER_ID) - adapter_ids.add(adapter_id) + # UN-2953: a tool instance may carry "" for an adapter id. + # Adding it made the schema validator compare "" against the + # UUID enum and raise, once per validation pass. + if adapter_id: + adapter_ids.add(adapter_id) for vdb in tool.properties.adapter.vector_stores: + adapter_id = None if vdb.is_enabled and vdb.adapter_id: - adapter_id = tool_meta[vdb.adapter_id] + adapter_id = tool_meta.get(vdb.adapter_id) elif vdb.is_enabled: - adapter_id = tool_meta[AdapterPropertyKey.DEFAULT_VECTOR_DB_ADAPTER_ID] + adapter_id = tool_meta.get( + AdapterPropertyKey.DEFAULT_VECTOR_DB_ADAPTER_ID + ) - adapter_ids.add(adapter_id) + # UN-2953: a tool instance may carry "" for an adapter id. + # Adding it made the schema validator compare "" against the + # UUID enum and raise, once per validation pass. + if adapter_id: + adapter_ids.add(adapter_id) for embedding in tool.properties.adapter.embedding_services: + adapter_id = None if embedding.is_enabled and embedding.adapter_id: - adapter_id = tool_meta[embedding.adapter_id] + adapter_id = tool_meta.get(embedding.adapter_id) elif embedding.is_enabled: - adapter_id = tool_meta[AdapterPropertyKey.DEFAULT_EMBEDDING_ADAPTER_ID] + adapter_id = tool_meta.get( + AdapterPropertyKey.DEFAULT_EMBEDDING_ADAPTER_ID + ) - adapter_ids.add(adapter_id) + # UN-2953: a tool instance may carry "" for an adapter id. + # Adding it made the schema validator compare "" against the + # UUID enum and raise, once per validation pass. + if adapter_id: + adapter_ids.add(adapter_id) for text_extractor in tool.properties.adapter.text_extractors: + adapter_id = None if text_extractor.is_enabled and text_extractor.adapter_id: - adapter_id = tool_meta[text_extractor.adapter_id] + adapter_id = tool_meta.get(text_extractor.adapter_id) elif text_extractor.is_enabled: - adapter_id = tool_meta[AdapterPropertyKey.DEFAULT_X2TEXT_ADAPTER_ID] + adapter_id = tool_meta.get(AdapterPropertyKey.DEFAULT_X2TEXT_ADAPTER_ID) - adapter_ids.add(adapter_id) + # UN-2953: a tool instance may carry "" for an adapter id. + # Adding it made the schema validator compare "" against the + # UUID enum and raise, once per validation pass. + if adapter_id: + adapter_ids.add(adapter_id) ToolInstanceHelper.validate_adapter_access(user=user, adapter_ids=adapter_ids) diff --git a/unstract/connectors/src/unstract/connectors/databases/bigquery/bigquery.py b/unstract/connectors/src/unstract/connectors/databases/bigquery/bigquery.py index d9e05a7055..75437216a9 100644 --- a/unstract/connectors/src/unstract/connectors/databases/bigquery/bigquery.py +++ b/unstract/connectors/src/unstract/connectors/databases/bigquery/bigquery.py @@ -13,6 +13,7 @@ from unstract.connectors.databases.exceptions import ( BigQueryForbiddenException, BigQueryNotFoundException, + BigQueryValueException, ColumnMissingException, ) from unstract.connectors.databases.sql_safety import ( @@ -106,13 +107,14 @@ def _sanitize_for_bigquery(data: Any) -> Any: if data == 0: return 0.0 - # Limit total significant figures to 15 for IEEE 754 compatibility - # BigQuery PARSE_JSON requires values that round-trip cleanly - # For large numbers (like Unix timestamps), this reduces decimal precision - # For small numbers (like costs), full precision is preserved - magnitude = math.floor(math.log10(abs(data))) + 1 - safe_decimals = max(0, 15 - magnitude) - return float(f"{data:.{safe_decimals}f}") + # Limit total significant figures to 15 for IEEE 754 compatibility. + # BigQuery PARSE_JSON requires values that round-trip cleanly. + # + # UN-3176: `g` asks for significant figures directly. The previous + # form asked for a DECIMAL-place count derived from the magnitude, + # which is not the same quantity and drifted from 15 sig figs at + # both ends of the range. + return float(f"{data:.15g}") elif isinstance(data, dict): return {k: BigQuery._sanitize_for_bigquery(v) for k, v in data.items()} @@ -318,6 +320,17 @@ def execute_query( detail=e.message, table_name=table_name ) from e except google.api_core.exceptions.BadRequest as e: + # UN-3176: BigQuery returns BadRequest for VALUE-level failures as + # well as for schema mismatches. Mapping them all to + # ColumnMissingException told users to check their columns when the + # columns were fine (e.g. a float that will not round-trip through + # PARSE_JSON), which misdirects the investigation. Discriminate + # before wrapping. + if BigQuery._is_value_error(e): + logger.error(f"Value rejected by BigQuery on insert: {str(e)}") + raise BigQueryValueException( + detail=e.message, table_name=table_name + ) from e logger.error(f"Column missing in inserting data: {str(e)}") db, schema, table = table_name.split(".") raise ColumnMissingException( @@ -327,6 +340,38 @@ def execute_query( table_name=table, ) from e + @staticmethod + def _is_value_error(e: Any) -> bool: + """True if a BigQuery BadRequest is about the DATA, not the schema. + + UN-3176. Matches known value-level signatures -- notably the PARSE_JSON + round-trip failure in this ticket -- against the exception's own text + and then against each entry of its structured ``errors`` payload. The + payload is checked separately because BigQuery's REST path fills it + with plain dicts, and ``GoogleAPICallError.__str__`` folds an entry + into the string only when it exposes ``.code``/``.message`` attributes. + Unknown shapes fall through to the existing column-missing behaviour, + so this only ever narrows a message that was already wrong for these + cases. + """ + value_error_markers = ( + "parse_json", + "round-trip through string representation", + "invalid json", + "cannot round-trip", + "failed to parse json", + ) + text = f"{getattr(e, 'message', '') or ''} {str(e)}".lower() + if any(marker in text for marker in value_error_markers): + return True + for error in getattr(e, "errors", None) or []: + if not isinstance(error, dict): + continue + message = str(error.get("message", "")).lower() + if any(marker in message for marker in value_error_markers): + return True + return False + def get_information_schema(self, table_name: str) -> dict[str, str]: """Function to generate information schema of the big query table. diff --git a/unstract/connectors/src/unstract/connectors/databases/exceptions.py b/unstract/connectors/src/unstract/connectors/databases/exceptions.py index b1a9b30853..c699ff4542 100644 --- a/unstract/connectors/src/unstract/connectors/databases/exceptions.py +++ b/unstract/connectors/src/unstract/connectors/databases/exceptions.py @@ -90,6 +90,29 @@ def __init__(self, detail: str, table_name: str) -> None: super().__init__(detail=final_detail) +class BigQueryValueException(UnstractDBConnectorException): + """A BigQuery BadRequest caused by the DATA, not the table schema. + + UN-3176: BigQuery returns BadRequest for value-level failures (a float that + will not round-trip through PARSE_JSON, a malformed JSON literal) as well as + for genuine schema mismatches. Mapping every BadRequest to + ColumnMissingException told users to "make sure all the columns exist" when + the columns were fine, which sent at least one investigation down the wrong + path. + """ + + def __init__(self, detail: Any, table_name: str) -> None: + default_detail = ( + f"Error writing to '{table_name}'. \n" + f"BigQuery rejected a value in the row being inserted -- the table " + f"schema is not the problem. This usually means a number could not " + f"be represented exactly, or a JSON column received text that is " + f"not valid JSON.\n" + ) + final_detail = _format_exception_detail(default_detail, detail) + super().__init__(detail=final_detail) + + class ColumnMissingException(UnstractDBConnectorException): def __init__( self, diff --git a/unstract/connectors/tests/databases/test_bigquery_db.py b/unstract/connectors/tests/databases/test_bigquery_db.py index 695d28a337..de6d157926 100644 --- a/unstract/connectors/tests/databases/test_bigquery_db.py +++ b/unstract/connectors/tests/databases/test_bigquery_db.py @@ -2,10 +2,13 @@ from unittest.mock import MagicMock, patch import google.api_core.exceptions + from unstract.connectors.databases.bigquery.bigquery import BigQuery from unstract.connectors.databases.exceptions import ( BigQueryForbiddenException, BigQueryNotFoundException, + BigQueryValueException, + ColumnMissingException, ) @@ -130,6 +133,111 @@ def test_exception_empty_detail(self): # When detail is empty, should not have "Details:" section self.assertNotIn("Details:", error_msg) + def test_bad_request_value_error_routes_to_value_exception(self): + """UN-3176: a BadRequest about the DATA is not a missing column.""" + value_error_msg = ( + "400 Invalid value: cannot round-trip through string representation; " + "error in PARSE_JSON expression" + ) + mock_error = google.api_core.exceptions.BadRequest(value_error_msg) + mock_error.message = value_error_msg + + context = self._execute_query_with_mock_error( + mock_error, BigQueryValueException + ) + + error_msg = str(context.exception.detail) + self.assertIn("BigQuery rejected a value", error_msg) + # The old message sent users to check a schema that was never wrong. + self.assertNotIn("make sure all the columns exist", error_msg) + self.assertIn("test.dataset.table", error_msg) + + def test_bad_request_value_error_detected_from_errors_payload_only(self): + """The marker lives ONLY in the structured payload, never in str(e). + + ``GoogleAPICallError.__str__`` folds an ``errors`` entry into the string + only when it exposes ``.code``/``.message`` ATTRIBUTES; BigQuery's REST + path fills ``errors`` with plain dicts, which fail that check. So this + case is reachable only by the payload loop in ``_is_value_error`` -- + delete that loop and this test fails while every message-based test + above still passes. + """ + mock_error = google.api_core.exceptions.BadRequest( + "400 Request failed", + errors=[ + { + "reason": "invalidQuery", + "message": "Cannot round-trip through string representation", + } + ], + ) + mock_error.message = "400 Request failed" + + # Guard the premise: if the marker ever leaks into str(e), this test + # would pass through the text branch and prove nothing. + self.assertNotIn("round-trip", str(mock_error).lower()) + + context = self._execute_query_with_mock_error( + mock_error, BigQueryValueException + ) + self.assertIn("BigQuery rejected a value", str(context.exception.detail)) + + def test_bad_request_schema_error_still_routes_to_column_missing(self): + """A genuine schema BadRequest must keep its existing behaviour.""" + schema_error_msg = "400 no such field: unknown_column" + mock_error = google.api_core.exceptions.BadRequest(schema_error_msg) + mock_error.message = schema_error_msg + + context = self._execute_query_with_mock_error( + mock_error, ColumnMissingException + ) + self.assertIn("column", str(context.exception.detail).lower()) + + +class TestBigQuerySanitizeForBigQuery(unittest.TestCase): + """UN-3176: the 15-significant-figure limit for PARSE_JSON compatibility. + + The previous implementation derived a DECIMAL-place count from the value's + magnitude, which is a different quantity from significant figures. The + values below are the two classes where the two forms actually disagree -- + an ordinary round number such as 3.14159 is preserved identically by both + and would pass against either implementation. + """ + + def test_large_magnitude_is_limited_to_15_significant_figures(self): + """Above 10^15 the old decimal count floored at 0 and kept every digit.""" + self.assertEqual( + BigQuery._sanitize_for_bigquery(1234567890123456.0), + 1234567890123460.0, + ) + + def test_value_just_below_a_power_of_ten_keeps_15_significant_figures(self): + """``log10`` returns an exact integer here, inflating the magnitude by one. + + The old form therefore asked for one decimal place too few and emitted + 14 significant figures, collapsing this value to 0.0001. + """ + value = 9.99999999999999e-05 + self.assertEqual(BigQuery._sanitize_for_bigquery(value), value) + + def test_special_values_are_dropped(self): + """NaN/Inf cannot be represented in JSON and must not reach BigQuery.""" + self.assertIsNone(BigQuery._sanitize_for_bigquery(float("nan"))) + self.assertIsNone(BigQuery._sanitize_for_bigquery(float("inf"))) + self.assertIsNone(BigQuery._sanitize_for_bigquery(float("-inf"))) + + def test_zero_and_modest_precision_are_preserved(self): + self.assertEqual(BigQuery._sanitize_for_bigquery(0.0), 0.0) + self.assertEqual(BigQuery._sanitize_for_bigquery(0.001228), 0.001228) + + def test_nested_structures_are_sanitized(self): + self.assertEqual( + BigQuery._sanitize_for_bigquery( + {"rows": [{"time": 1760509016.282637}], "name": "x"} + ), + {"rows": [{"time": 1760509016.28264}], "name": "x"}, + ) + if __name__ == "__main__": unittest.main() 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 04d16a3e7a..95644c7222 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 @@ -58,8 +58,12 @@ "line_splitter_strategy": { "type": "string", "title": "Line Splitter Strategy", - "enum": ["left-priority", "mid-priority", "right-priority"], - "default":"left-priority", + "enum": [ + "left-priority", + "mid-priority", + "right-priority" + ], + "default": "left-priority", "description": "An advanced option for customizing the line splitting process." }, "horizontal_stretch_factor": { @@ -93,6 +97,14 @@ "default": false, "description": "States whether to reproduce horizontal lines in the document. Note: This parameter is not applicable if `mode` chosen is `native_text` and will not work if `mark_vertical_lines` is set to `false`." }, + "word_confidence_threshold": { + "type": "number", + "title": "Word confidence threshold", + "default": 0.3, + "minimum": 0.0, + "maximum": 1.0, + "description": "Minimum OCR confidence a word must reach to be included in the extracted text. Lower this when words are dropped because the scan is faint or noisy. Note: This parameter is not applicable if `mode` chosen is `native_text`." + }, "tag": { "type": "string", "title": "Tag", diff --git a/unstract/sdk1/src/unstract/sdk1/exceptions.py b/unstract/sdk1/src/unstract/sdk1/exceptions.py index 40889af877..1531ca6dd9 100644 --- a/unstract/sdk1/src/unstract/sdk1/exceptions.py +++ b/unstract/sdk1/src/unstract/sdk1/exceptions.py @@ -2,6 +2,7 @@ import re import openai +from litellm import exceptions as litellm_exceptions logger = logging.getLogger(__name__) @@ -85,6 +86,23 @@ class RateLimitError(SdkError): DEFAULT_MESSAGE = "Running into rate limit errors, please try again later" +class ContextWindowExceededError(SdkError): + """The prompt (plus context) exceeded the model's context window. + + UN-3133: litellm raises ``ContextWindowExceededError``, which derives from + ``BadRequestError`` and so from ``openai.APIError``. Without a distinct + class it was wrapped as a plain ``SdkError`` carrying the provider's raw + 400 text, so the UI showed a generic "Error from " and the real + cause was only findable by reading worker logs. + """ + + DEFAULT_MESSAGE = ( + "The document and prompt together exceed the model's context window. " + "Reduce the chunk size, limit the pages extracted, or use a model with " + "a larger context window." + ) + + class FileStorageError(SdkError): DEFAULT_MESSAGE = ( "Error while connecting with the storage. " @@ -156,6 +174,26 @@ def parse_litellm_err(e: Exception, provider_name: str | None = None) -> SdkErro ) cleaned_message = strip_litellm_prefix(str(e)) + + # UN-3133: surface a context-window overflow as its own error type with an + # actionable message. It is a BadRequestError subclass, so without this it + # collapses into the generic wrap below and the user sees only the + # provider's raw 400. + if isinstance(e, litellm_exceptions.ContextWindowExceededError): + err = ContextWindowExceededError( + ContextWindowExceededError.DEFAULT_MESSAGE, + actual_err=e, + status_code=status_code, + ) + # Return early: the generic tail below would overwrite the actionable + # guidance with "Error from ." and leave only the raw 400. + # The provider text is kept underneath it for support/debugging. + err.message = ( + f"{ContextWindowExceededError.DEFAULT_MESSAGE}" + f"\n```\n{cleaned_message}\n```" + ) + return err + err = SdkError(cleaned_message, actual_err=e, status_code=status_code) if not provider_name: diff --git a/unstract/sdk1/src/unstract/sdk1/x2txt.py b/unstract/sdk1/src/unstract/sdk1/x2txt.py index 2024f8cdbc..2c769cf635 100644 --- a/unstract/sdk1/src/unstract/sdk1/x2txt.py +++ b/unstract/sdk1/src/unstract/sdk1/x2txt.py @@ -118,6 +118,55 @@ def process( self.push_usage_details(input_file_path, mime_type, fs=fs) return text_extraction_result + @staticmethod + def _parse_pages_to_extract(pages_to_extract: str, total_pages: int) -> int: + """Count the pages selected by an LLMWhisperer ``pages_to_extract`` spec. + + The spec is a comma separated list of single pages and ranges, where a + range may be open ended (``50-`` means "page 50 to the end"). Pages are + 1-indexed and may overlap, so they are collected into a set and clamped + to the document length. An empty spec means "all pages". + """ + selected: set[int] = set() + for part in pages_to_extract.split(","): + part = part.strip() + if not part: + continue + if "-" in part: + start_str, _, end_str = part.partition("-") + try: + start = int(start_str) + except ValueError: + continue + end = total_pages + if end_str: + try: + end = int(end_str) + except ValueError: + continue + selected.update(range(max(start, 1), min(end, total_pages) + 1)) + else: + try: + page = int(part) + except ValueError: + continue + if 1 <= page <= total_pages: + selected.add(page) + return len(selected) + + def _get_billable_page_count(self, page_count: int) -> int: + """Narrow ``page_count`` to the pages the adapter will actually extract. + + Falls back to the full count whenever the setting is absent, empty or + unparseable, so usage is never under-reported by a malformed value. + """ + config = getattr(self._x2text_instance, "config", None) or {} + pages_to_extract = str(config.get("pages_to_extract", "") or "").strip() + if not pages_to_extract: + return page_count + selected = self._parse_pages_to_extract(pages_to_extract, page_count) + return selected or page_count + def push_usage_details( self, input_file_path: str, @@ -133,6 +182,10 @@ def push_usage_details( with pdfplumber.open(pdf_contents) as pdf: # calculate the number of pages page_count = len(pdf.pages) + # UN-3038: when the adapter restricts extraction to a page range, + # only those pages are actually processed, so bill for them rather + # than for every page in the document. + page_count = self._get_billable_page_count(page_count) Audit().push_page_usage_data( platform_api_key=self._tool.get_env_or_die(ToolEnv.PLATFORM_API_KEY), file_size=file_size, diff --git a/unstract/sdk1/tests/test_llm_whisperer_v2_params.py b/unstract/sdk1/tests/test_llm_whisperer_v2_params.py index d8ca030e31..0016ce11a0 100644 --- a/unstract/sdk1/tests/test_llm_whisperer_v2_params.py +++ b/unstract/sdk1/tests/test_llm_whisperer_v2_params.py @@ -1,10 +1,16 @@ -"""Tests for the query params the LLMWhisperer V2 adapter sends.""" +"""Tests for the LLMWhisperer V2 adapter's query params and page-range billing. + +Both read the same ``pages_to_extract`` setting: the adapter sends it to the +service, and ``X2Text`` bills for the pages it selects. +""" import pytest from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.dto import ( WhispererRequestParams, ) from unstract.sdk1.adapters.x2text.llm_whisperer_v2.src.helper import LLMWhispererHelper +from unstract.sdk1.constants import MimeType +from unstract.sdk1.x2txt import X2Text def _params(config: dict) -> dict: @@ -47,3 +53,117 @@ def test_page_separator_read_under_legacy_config_key() -> None: assert params["page_separator"] == "<<< {{page_no}} >>>" assert "page_seperator" not in params + + +class _FakeAdapter: + """Stands in for the adapter instance, which only needs ``.config`` here.""" + + def __init__(self, config: dict | None) -> None: + self.config = config + + +def _billable(config: dict | None, page_count: int) -> int: + """Run ``_get_billable_page_count`` against a stub adapter. + + ``X2Text.__init__`` requires a live ``BaseTool``, and the method under test + reads nothing but ``self._x2text_instance.config``. + """ + x2text = X2Text.__new__(X2Text) + x2text._x2text_instance = _FakeAdapter(config) if config is not None else None + return x2text._get_billable_page_count(page_count) + + +@pytest.mark.parametrize( + ("spec", "total_pages", "expected"), + [ + ("1,3,5", 10, 3), + ("2-4", 10, 3), + ("50-", 60, 11), # open-ended range runs to the last page + ("1-3,2-4", 10, 4), # overlapping ranges are counted once + ("0-3", 10, 3), # pages are 1-indexed; page 0 does not exist + ("1-999", 10, 10), # clamped to the document length + ("99", 10, 0), # a single out-of-range page selects nothing + ("5-2", 10, 0), # an inverted range selects nothing + ], +) +def test_parse_pages_to_extract_counts_selected_pages( + spec: str, total_pages: int, expected: int +) -> None: + """UN-3038: billing follows the pages the adapter actually extracts.""" + assert X2Text._parse_pages_to_extract(spec, total_pages) == expected + + +def test_billable_page_count_narrows_to_the_selected_range() -> None: + """The whole point of UN-3038: 3 pages of a 100-page document bill as 3.""" + assert _billable({"pages_to_extract": "2-4"}, 100) == 3 + + +@pytest.mark.parametrize( + "config", + [ + {"pages_to_extract": ""}, # setting present but empty + {"pages_to_extract": " "}, # whitespace only + {"pages_to_extract": "not-a-range"}, # unparseable + {"pages_to_extract": "99"}, # parses, but selects nothing + {}, # setting absent + None, # no adapter instance at all + ], +) +def test_billable_page_count_falls_back_to_the_full_count(config: dict | None) -> None: + """Usage must never be UNDER-reported because a setting was malformed. + + Each of these makes the page selection unusable; billing then falls back to + every page in the document rather than to zero. + """ + assert _billable(config, 10) == 10 + + +def test_push_usage_details_reports_the_narrowed_page_count( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The narrowing must be wired into the value Audit actually receives. + + Asserting on ``_get_billable_page_count`` alone would still pass if the + call were dropped from ``push_usage_details``, which is the only place the + number becomes a bill. + """ + import unstract.sdk1.x2txt as x2txt_module + + recorded: dict[str, int] = {} + + class _FakeAudit: + def push_page_usage_data(self, **kwargs: object) -> None: + recorded["page_count"] = kwargs["page_count"] + + class _FakePdf: + pages = [object()] * 100 + + def __enter__(self) -> "_FakePdf": + return self + + def __exit__(self, *exc: object) -> None: + return None + + monkeypatch.setattr(x2txt_module, "Audit", _FakeAudit) + monkeypatch.setattr(x2txt_module.pdfplumber, "open", lambda _: _FakePdf()) + monkeypatch.setattr( + x2txt_module.ToolUtils, "get_file_size", staticmethod(lambda *a, **k: 1024) + ) + + class _FakeFs: + def read(self, **kwargs: object) -> bytes: + return b"%PDF-1.4" + + class _FakeTool: + def get_env_or_die(self, key: str) -> str: + return "test-key" + + x2text = X2Text.__new__(X2Text) + x2text._x2text_instance = _FakeAdapter({"pages_to_extract": "2-4"}) + x2text._tool = _FakeTool() + x2text._usage_kwargs = {} + + x2text.push_usage_details("doc.pdf", MimeType.PDF, fs=_FakeFs()) + + # 100-page document, 3 pages extracted -> 3 pages billed. + assert recorded["page_count"] == 3 diff --git a/workers/executor/executors/retrievers/keyword_table.py b/workers/executor/executors/retrievers/keyword_table.py index 5c7db9ab4e..e2909146e3 100644 --- a/workers/executor/executors/retrievers/keyword_table.py +++ b/workers/executor/executors/retrievers/keyword_table.py @@ -51,9 +51,13 @@ def retrieve(self) -> set[str]: llm=llm, # Use the provided LLM instead of defaulting to OpenAI ) - # Create retriever from keyword index + # Create retriever from keyword index. + # UN-2902: KeywordTableIndex's retriever caps results with + # `num_chunks_per_query` (default 10), not `similarity_top_k` -- + # that kwarg is accepted and ignored, so a profile asking for 3 + # chunks still got 10. keyword_retriever = keyword_index.as_retriever( - similarity_top_k=self.top_k, + num_chunks_per_query=self.top_k, ) # Retrieve nodes using keyword matching