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 ce7fbea0d1..4cbe93bb23 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/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 # ---------------------------------------------------------------------------