diff --git a/workers/executor/executors/constants.py b/workers/executor/executors/constants.py index 9eddab8423..cd5bd2ec07 100644 --- a/workers/executor/executors/constants.py +++ b/workers/executor/executors/constants.py @@ -51,6 +51,15 @@ class PromptServiceConstants: CHALLENGE = "challenge" ENABLE_CHALLENGE = "enable_challenge" EXTRACTION = "extraction" + # Reserved namespace for file-level (whole-document) metrics so they cannot + # collide with user-defined output/prompt names at the top level of the + # metrics dict. Mirrors the "file" key of the API response. + FILE = "_file" + # Text extraction (X2Text) of the source document. Named to stay distinct + # from "extraction_llm", the extraction-purpose LLM call beside it. + TEXT_EXTRACTION = "text_extraction" + # Wire-format key every metrics producer reports its duration under. + TIME_TAKEN = "time_taken(s)" SUMMARIZE = "summarize" SINGLE_PASS_EXTRACTION = "single-pass-extraction" SIMPLE_PROMPT_STUDIO = "simple-prompt-studio" diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index b672d56ab6..7daaeda6ed 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -619,10 +619,12 @@ def _failure(child_result: ExecutionResult) -> ExecutionResult: ) step = 1 + extraction_metrics: dict = {} try: # ---- Step 1: Extract ---- if not skip_extraction: step += 1 + extraction_start = time.monotonic() extract_ctx = ExecutionContext( executor_name=context.executor_name, operation=Operation.EXTRACT.value, @@ -640,6 +642,15 @@ def _failure(child_result: ExecutionResult) -> ExecutionResult: return _failure(extract_result) _absorb(extract_result) extracted_text = extract_result.data.get(IKeys.EXTRACTED_TEXT, "") + extraction_time = time.monotonic() - extraction_start + # Nest under a reserved "_file" namespace so the metric never + # collides with a user-defined output named "text_extraction" + # during the top-level metrics merge (see _merge_pipeline_metrics). + extraction_metrics = { + PSKeys.FILE: { + PSKeys.TEXT_EXTRACTION: {PSKeys.TIME_TAKEN: extraction_time} + } + } # ---- Step 2: Summarize (if enabled) ---- if is_summarization: @@ -700,6 +711,7 @@ def _failure(child_result: ExecutionResult) -> ExecutionResult: source_file_name=source_file_name, extracted_text=extracted_text, index_metrics=index_metrics, + extraction_metrics=extraction_metrics, ) output_map = structured_output.get(PSKeys.OUTPUT, {}) or {} @@ -800,17 +812,29 @@ def _finalize_pipeline_result( source_file_name: str, extracted_text: str, index_metrics: dict, + extraction_metrics: dict[str, dict] | None = None, ) -> None: - """Populate metadata/metrics in structured_output after pipeline completion.""" + """Populate metadata/metrics in structured_output after pipeline completion. + + Args: + extraction_metrics: File-level text extraction timing, shaped as + ``{"_file": {"text_extraction": {"time_taken(s)": float}}}``. + Nested under the reserved ``_file`` namespace to avoid + colliding with user-defined output names during the merge. + ``None``/empty when extraction is skipped. + """ if "metadata" not in structured_output: structured_output["metadata"] = {} structured_output["metadata"]["file_name"] = source_file_name if extracted_text: structured_output["metadata"]["extracted_text"] = extracted_text - if index_metrics: + new_metrics = self._merge_pipeline_metrics( + index_metrics or {}, extraction_metrics or {} + ) + if new_metrics: existing_metrics = structured_output.get("metrics", {}) structured_output["metrics"] = self._merge_pipeline_metrics( - existing_metrics, index_metrics + existing_metrics, new_metrics ) def _run_pipeline_summarize( @@ -955,8 +979,6 @@ def _index_pipeline_output( index_records: list[dict], ) -> None: """Index a single structure-pipeline output entry in-place.""" - import datetime - chunk_size = output.get("chunk-size", 0) if chunk_size == 0: return @@ -977,7 +999,7 @@ def _index_pipeline_output( return seen_params.add(param_key) - indexing_start = datetime.datetime.now() + indexing_start = time.monotonic() logger.info( "Pipeline indexing: chunk_size=%s chunk_overlap=%s vector_db=%s", chunk_size, @@ -1029,9 +1051,9 @@ def _index_pipeline_output( if child_records: index_records.extend(child_records) - elapsed = (datetime.datetime.now() - indexing_start).total_seconds() + elapsed = time.monotonic() - indexing_start output_name = output.get("name", "") - index_metrics[output_name] = {"indexing": {"time_taken(s)": elapsed}} + index_metrics[output_name] = {"indexing": {PSKeys.TIME_TAKEN: elapsed}} @staticmethod def _merge_pipeline_metrics(metrics1: dict, metrics2: dict) -> dict: diff --git a/workers/file_processing/structure_tool_task.py b/workers/file_processing/structure_tool_task.py index 61334e63fc..4b2f02bfcc 100644 --- a/workers/file_processing/structure_tool_task.py +++ b/workers/file_processing/structure_tool_task.py @@ -196,6 +196,79 @@ def _should_skip_extraction_for_smart_table( return False +def _agentic_extraction_seconds(agentic_metrics: dict[str, Any]) -> float: + """Total X2Text seconds reported by the agentic_table executor. + + Each agentic prompt extracts the document itself, so the durations sum: + the figure is time *spent* extracting, not the wall-clock span of the + agentic step. Returns 0.0 when the executor reports no extraction timing, + which leaves the file-level bucket untouched rather than writing a zero. + """ + from executor.executors.constants import PromptServiceConstants as PSKeys + + total = 0.0 + for prompt_metrics in agentic_metrics.values(): + if not isinstance(prompt_metrics, dict): + continue + entry = prompt_metrics.get(PSKeys.TEXT_EXTRACTION) + if not isinstance(entry, dict): + continue + value = entry.get(PSKeys.TIME_TAKEN) + if isinstance(value, (int, float)) and not isinstance(value, bool): + total += float(value) + return total + + +def _merge_agentic_metrics( + structured_output: dict[str, Any], agentic_metrics: dict[str, Any] +) -> None: + """Fold the agentic executor's metrics into the tool's metrics dict. + + Two things happen here, and they are separate: + + * Per-prompt metrics land beside the legacy pipeline's, keyed by prompt + name under ``table_extraction`` — the same shape + ``LegacyExecutor._run_table_extraction`` uses for the non-agentic table + plugin, so both table routes read alike downstream. + * Any X2Text duration the executor reported is *added* to the file-level + ``_file.text_extraction`` bucket. The legacy pipeline times only its own + extraction, so without this an agentic-only run reports no extraction + time at all and a mixed run under-reports it (UN-2771 / UN-4084). + """ + from executor.executors.constants import PromptServiceConstants as PSKeys + + if not agentic_metrics: + return + + metrics = structured_output.setdefault("metrics", {}) + for prompt_name, prompt_metrics in agentic_metrics.items(): + key = prompt_name + if key == PSKeys.FILE: + # `_file` is reserved for whole-document metrics, so a prompt with + # that name gets a de-collided key. Merging it into the reserved + # bucket would conflate per-prompt and file-level figures; dropping + # it would lose every other count the prompt reported. The space + # cannot appear in the reserved key, so the two can never alias. + key = f"{prompt_name} (prompt)" + logger.warning( + "Agentic prompt named %r collides with the reserved file-level " + "metrics namespace; reporting its metrics under %r instead", + prompt_name, + key, + ) + metrics.setdefault(key, {}).update({"table_extraction": prompt_metrics}) + + extraction_seconds = _agentic_extraction_seconds(agentic_metrics) + if not extraction_seconds: + return + file_bucket = metrics.setdefault(PSKeys.FILE, {}).setdefault( + PSKeys.TEXT_EXTRACTION, {} + ) + file_bucket[PSKeys.TIME_TAKEN] = ( + file_bucket.get(PSKeys.TIME_TAKEN, 0.0) + extraction_seconds + ) + + # ----------------------------------------------------------------------- # Main Celery task # ----------------------------------------------------------------------- @@ -443,6 +516,7 @@ def _execute_structure_tool_impl(params: dict) -> dict: # PDF written alongside INFILE by the source connector. agentic_source_path = str(execution_run_data_folder / "SOURCE") agentic_results: dict[str, Any] = {} + agentic_metrics: dict[str, Any] = {} for at_output in agentic_table_outputs: at_settings = at_output.get("agentic_table_settings") or {} json_structure = at_settings.get("json_structure") @@ -490,6 +564,12 @@ def _execute_structure_tool_impl(params: dict) -> dict: return at_result.to_dict() at_output_data = at_result.data.get("output", {}) or {} agentic_results[at_output[_SK.NAME]] = at_output_data.get("tables", []) + # The executor runs its own X2Text, so its metrics carry an extraction + # duration the legacy pipeline's timer never sees. Same result shape the + # non-agentic table plugin uses (LegacyExecutor._run_table_extraction). + at_metrics = (at_result.data.get("metadata") or {}).get("metrics") or {} + if at_metrics: + agentic_metrics[at_output[_SK.NAME]] = at_metrics # ---- Step 6b: Dispatch legacy structure_pipeline ---- # Skipped entirely when every prompt is agentic_table — the legacy @@ -536,12 +616,16 @@ def _execute_structure_tool_impl(params: dict) -> dict: structured_output = pipeline_result.data if agentic_results: structured_output.setdefault("output", {}).update(agentic_results) + _merge_agentic_metrics(structured_output, agentic_metrics) else: - # All-agentic case: skip the legacy pipeline entirely. + # All-agentic case: skip the legacy pipeline entirely. The metrics + # dict is still populated, so an agentic-only deployment reports the + # same _file.text_extraction the legacy path does. structured_output = { "output": agentic_results, "metadata": {"agentic_only": True}, } + _merge_agentic_metrics(structured_output, agentic_metrics) pipeline_elapsed = 0.0 # ---- Step 7: Write output files ---- diff --git a/workers/tests/test_agentic_operations.py b/workers/tests/test_agentic_operations.py index f16e566d01..4d0fec4977 100644 --- a/workers/tests/test_agentic_operations.py +++ b/workers/tests/test_agentic_operations.py @@ -262,3 +262,140 @@ def test_agentic_ops_use_default_log_component(self, op): "table_extract", "smart_table_extract", ) + + +# --------------------------------------------------------------------------- +# Agentic extraction timing (UN-2771) +# --------------------------------------------------------------------------- + + +class TestAgenticExtractionMetrics: + """The agentic executor's X2Text duration reaches the file-level metric. + + The legacy pipeline times only its own extraction. Agentic prompts extract + the document inside their own executor, so without folding their reported + duration in, an agentic-only deployment reports no extraction time at all + and a mixed deployment under-reports it. + """ + + @staticmethod + def _agentic(seconds): + """One prompt's metrics as the executor reports them.""" + return {"invoices": {"text_extraction": {"time_taken(s)": seconds}}} + + def test_sums_across_prompts(self): + """Each agentic prompt extracts the file itself, so durations add.""" + from file_processing.structure_tool_task import _agentic_extraction_seconds + + metrics = { + "invoices": {"text_extraction": {"time_taken(s)": 1.5}}, + "receipts": {"text_extraction": {"time_taken(s)": 2.25}}, + } + assert _agentic_extraction_seconds(metrics) == pytest.approx(3.75) + + @pytest.mark.parametrize( + "metrics", + [ + {}, + {"invoices": {}}, + {"invoices": {"text_extraction": {}}}, + {"invoices": {"text_extraction": {"time_taken(s)": None}}}, + {"invoices": {"text_extraction": "not-a-dict"}}, + {"invoices": "not-a-dict"}, + # bool is an int subclass; True must not count as 1 second + {"invoices": {"text_extraction": {"time_taken(s)": True}}}, + ], + ) + def test_missing_or_malformed_timing_yields_zero(self, metrics): + """A plugin that reports no duration must not fabricate one.""" + from file_processing.structure_tool_task import _agentic_extraction_seconds + + assert _agentic_extraction_seconds(metrics) == 0.0 + + def test_all_agentic_run_reports_the_file_metric(self): + """The agentic-only path has no legacy pipeline, so it starts empty.""" + from file_processing.structure_tool_task import _merge_agentic_metrics + + structured_output = {"output": {}, "metadata": {"agentic_only": True}} + _merge_agentic_metrics(structured_output, self._agentic(4.0)) + + metrics = structured_output["metrics"] + assert metrics["_file"]["text_extraction"]["time_taken(s)"] == pytest.approx(4.0) + + def test_mixed_run_adds_to_the_legacy_duration(self): + """The under-reporting case: both extractors ran, both must count.""" + from file_processing.structure_tool_task import _merge_agentic_metrics + + structured_output = { + "metrics": {"_file": {"text_extraction": {"time_taken(s)": 1.0}}} + } + _merge_agentic_metrics(structured_output, self._agentic(2.5)) + + total = structured_output["metrics"]["_file"]["text_extraction"]["time_taken(s)"] + assert total == pytest.approx(3.5) + + def test_per_prompt_metrics_land_beside_the_legacy_ones(self): + """Same shape LegacyExecutor._run_table_extraction uses.""" + from file_processing.structure_tool_task import _merge_agentic_metrics + + structured_output = {"metrics": {"field_a": {"llm": {"time_taken(s)": 2.0}}}} + _merge_agentic_metrics(structured_output, self._agentic(1.0)) + + metrics = structured_output["metrics"] + assert metrics["field_a"] == {"llm": {"time_taken(s)": 2.0}} + assert metrics["invoices"]["table_extraction"] == { + "text_extraction": {"time_taken(s)": 1.0} + } + + def test_no_agentic_prompts_changes_nothing(self): + """A purely legacy run must not grow an empty metrics dict.""" + from file_processing.structure_tool_task import _merge_agentic_metrics + + structured_output = {"output": {}} + _merge_agentic_metrics(structured_output, {}) + + assert "metrics" not in structured_output + + def test_executor_reporting_no_timing_leaves_the_bucket_absent(self): + """Per-prompt metrics still surface; a zero duration is not written.""" + from file_processing.structure_tool_task import _merge_agentic_metrics + + structured_output = {} + _merge_agentic_metrics( + structured_output, {"invoices": {"table_rows": {"count": 12}}} + ) + + metrics = structured_output["metrics"] + assert metrics["invoices"]["table_extraction"] == {"table_rows": {"count": 12}} + assert "_file" not in metrics + + def test_prompt_named_file_gets_a_de_collided_key(self): + """`_file` is reserved, so a prompt with that name moves aside. + + Neither outcome the two namespaces could otherwise produce is + acceptable: merging into the reserved bucket conflates per-prompt with + file-level figures, and dropping the entry loses every other metric the + prompt reported. It is rehomed instead, and nothing is lost. + """ + from file_processing.structure_tool_task import _merge_agentic_metrics + + structured_output = {} + _merge_agentic_metrics( + structured_output, + {"_file": {"text_extraction": {"time_taken(s)": 3.0}, "rows": 12}}, + ) + + metrics = structured_output["metrics"] + # The reserved bucket carries file-level metrics only... + assert metrics["_file"] == {"text_extraction": {"time_taken(s)": 3.0}} + # ...and the prompt keeps everything it reported, under its own key. + assert metrics["_file (prompt)"]["table_extraction"] == { + "text_extraction": {"time_taken(s)": 3.0}, + "rows": 12, + } + + def test_de_collided_key_cannot_alias_the_reserved_one(self): + """The rehomed key is distinct from the reserved key by construction.""" + from executor.executors.constants import PromptServiceConstants as PSKeys + + assert " " not in PSKeys.FILE diff --git a/workers/tests/test_structure_pipeline.py b/workers/tests/test_structure_pipeline.py index 6ece5c23b2..80119faacf 100644 --- a/workers/tests/test_structure_pipeline.py +++ b/workers/tests/test_structure_pipeline.py @@ -8,6 +8,7 @@ from unittest.mock import MagicMock, patch import pytest + from unstract.sdk1.execution.context import ExecutionContext, Operation from unstract.sdk1.execution.result import ExecutionResult @@ -293,6 +294,168 @@ def test_index_metrics_merged(self, executor): assert "indexing" in metrics["field_a"] +# --------------------------------------------------------------------------- +# Tests — Text extraction timing metric +# --------------------------------------------------------------------------- + + +class TestExtractionMetrics: + """Extraction duration is reported under the reserved _file namespace. + + The metric is nested as ``metrics["_file"]["text_extraction"]`` rather than a + bare top-level key, because sibling keys in this dict are user-defined + output/prompt names — a prompt named "text_extraction" would otherwise + collide with it. + """ + + def _executor_with_extraction(self, executor, answer_metrics=None): + executor._handle_extract = MagicMock( + return_value=ExecutionResult(success=True, data={"extracted_text": "text"}) + ) + executor._handle_index = MagicMock( + return_value=ExecutionResult(success=True, data={"doc_id": "d1"}) + ) + executor._handle_answer_prompt = MagicMock( + return_value=ExecutionResult( + success=True, + data={"output": {}, "metrics": answer_metrics or {}}, + ) + ) + return executor + + def _run(self, executor): + ctx = _make_pipeline_context( + { + "extract_params": _base_extract_params(), + "index_template": _base_index_template(), + "answer_params": _base_answer_params(), + "pipeline_options": _base_pipeline_options(), + } + ) + return executor._handle_structure_pipeline(ctx) + + def test_extraction_metrics_recorded(self, executor): + """A normal run reports a non-negative float extraction duration.""" + self._executor_with_extraction(executor) + result = self._run(executor) + + assert result.success + time_taken = result.data["metrics"]["_file"]["text_extraction"]["time_taken(s)"] + assert isinstance(time_taken, float) + assert time_taken >= 0 + + def test_extraction_metric_does_not_collide_with_output_name(self, executor): + """A prompt named "text_extraction" coexists with the file metric.""" + self._executor_with_extraction( + executor, + answer_metrics={"text_extraction": {"llm": {"time_taken(s)": 2.0}}}, + ) + result = self._run(executor) + + metrics = result.data["metrics"] + # The user's prompt metrics are untouched... + assert metrics["text_extraction"] == {"llm": {"time_taken(s)": 2.0}} + # ...and the file-level timing lives in its own namespace. + assert "time_taken(s)" in metrics["_file"]["text_extraction"] + + def test_index_and_extraction_metrics_merged(self, executor): + """Per-output indexing and pipeline extraction metrics coexist.""" + self._executor_with_extraction( + executor, + answer_metrics={"field_a": {"llm": {"time_taken(s)": 2.0}}}, + ) + executor._run_pipeline_index = MagicMock( + return_value=({"field_a": {"indexing": {"time_taken(s)": 0.5}}}, []) + ) + result = self._run(executor) + + metrics = result.data["metrics"] + assert "llm" in metrics["field_a"] + assert "indexing" in metrics["field_a"] + assert "time_taken(s)" in metrics["_file"]["text_extraction"] + + def test_skip_extraction_records_no_extraction_metric(self, executor): + """Smart-table path skips extract, so no extraction metric is added.""" + executor._handle_extract = MagicMock() + executor._handle_index = MagicMock() + executor._handle_answer_prompt = MagicMock( + return_value=ExecutionResult(success=True, data={"output": {}}) + ) + opts = _base_pipeline_options() + opts["skip_extraction_and_indexing"] = True + + ctx = _make_pipeline_context( + { + "extract_params": _base_extract_params(), + "index_template": _base_index_template(), + "answer_params": _base_answer_params(), + "pipeline_options": opts, + } + ) + result = executor._handle_structure_pipeline(ctx) + + assert result.success + executor._handle_extract.assert_not_called() + assert "_file" not in result.data.get("metrics", {}) + + def test_single_pass_reports_same_file_metric(self, executor): + """Single pass gains the _file bucket beside its flat metrics. + + Single pass skips indexing but shares ``_finalize_pipeline_result``, so + it needs no special-casing: its metrics stay flat and simply gain the + same ``_file`` sibling the multi-prompt path gets. + """ + self._executor_with_extraction(executor) + executor._handle_single_pass_extraction = MagicMock( + return_value=ExecutionResult( + success=True, + data={ + "output": {}, + "metrics": { + "context_retrieval": {"time_taken(s)": 0.4}, + "extraction_llm": {"time_taken(s)": 2.0}, + }, + }, + ) + ) + opts = _base_pipeline_options() + opts["is_single_pass_enabled"] = True + + ctx = _make_pipeline_context( + { + "extract_params": _base_extract_params(), + "index_template": _base_index_template(), + "answer_params": _base_answer_params(), + "pipeline_options": opts, + } + ) + result = executor._handle_structure_pipeline(ctx) + + assert result.success + metrics = result.data["metrics"] + # Flat single-pass metrics survive untouched... + assert metrics["context_retrieval"] == {"time_taken(s)": 0.4} + assert metrics["extraction_llm"] == {"time_taken(s)": 2.0} + # ...beside the same _file bucket the multi-prompt path produces. + assert "time_taken(s)" in metrics["_file"]["text_extraction"] + + def test_extract_failure_records_no_extraction_metric(self, executor): + """Timing is taken after the failure early-return, so a failed extract + must not report a duration. + """ + executor._handle_extract = MagicMock( + return_value=ExecutionResult.failure(error="x2text error") + ) + executor._handle_index = MagicMock() + executor._handle_answer_prompt = MagicMock() + executor._finalize_pipeline_result = MagicMock() + + result = self._run(executor) + + assert not result.success + executor._finalize_pipeline_result.assert_not_called() + + # --------------------------------------------------------------------------- # Tests — Embedding usage record propagation # ---------------------------------------------------------------------------