diff --git a/workers/api-deployment/tasks.py b/workers/api-deployment/tasks.py index fa7a218d1a..43afdbff27 100644 --- a/workers/api-deployment/tasks.py +++ b/workers/api-deployment/tasks.py @@ -22,6 +22,11 @@ from shared.infrastructure.logging.workflow_logger import WorkerWorkflowLogger from shared.patterns.retry.utils import retry from shared.processing.files import FileProcessingUtils +from shared.utils.manual_review_factory import ( + AUTO_ENABLE_HIGHLIGHT, + get_manual_review_workflow_util, + has_confidence_rules, +) from shared.workflow.execution import WorkerExecutionContext, WorkflowOrchestrationUtils from shared.workflow.execution.tool_validation import validate_workflow_tool_instances from worker import app @@ -665,6 +670,24 @@ def _run_workflow_api( total_files=total_files, ) + # The destination evaluates API rules only for a hitl_queue_name request; + # a hitl_packet_id skips rule evaluation, so highlight would go unused. + # ponytail: the API + use_file_history path evaluates DB rules, which + # this check does not cover, so a DB confidence rule there gets no + # highlight. use_file_history is hidden from the API schema and meant for + # demos, but a caller can still send it; add a "DB" check if it matters. + # Highlight adds confidence_data to the result metadata and fills + # line_numbers and whisper_hash, which include_metadata callers receive + # (highlight_data too, when the org's API highlight flag is on). The keys + # are kept on purpose, so callers can see why a document was routed. + auto_enable_highlight = ( + bool(kwargs.get("hitl_queue_name")) + and not kwargs.get("hitl_packet_id") + and has_confidence_rules( + get_manual_review_workflow_util(api_client, schema_name), workflow_id, "API" + ) + ) + # Create batch tasks following the exact Django pattern batch_tasks = [] execution_mode_str = ( @@ -687,6 +710,8 @@ def _run_workflow_api( total_files=total_files, **kwargs, ) + if auto_enable_highlight: + file_data.manual_review_config[AUTO_ENABLE_HIGHLIGHT] = True # Calculate manual review decisions for this specific batch if file_data.manual_review_config.get("review_required", False): diff --git a/workers/file_processing/structure_tool_task.py b/workers/file_processing/structure_tool_task.py index 36c875b7e6..ef069c2a5b 100644 --- a/workers/file_processing/structure_tool_task.py +++ b/workers/file_processing/structure_tool_task.py @@ -28,6 +28,10 @@ ) from shared.enums.task_enums import TaskName from shared.infrastructure.context import StateStore +from shared.utils.manual_review_factory import ( + AUTO_ENABLE_HIGHLIGHT, + HIGHLIGHT_AUTO_ENABLED_MESSAGE, +) from unstract.sdk1.constants import ToolEnv, UsageKwargs from unstract.sdk1.execution.context import ExecutionContext @@ -365,6 +369,20 @@ def _execute_structure_tool_impl(params: dict) -> dict: execution_data_dir=execution_data_dir, ) + # HITL confidence rules need highlight data. Override a copy so the + # user's saved setting is unchanged. Agentic projects are skipped above: + # workflow runs give them no highlight data, so confidence can't work. + # Known limit: with an extractor other than LLMWhisperer V2, no highlight + # plugin, or summarize-as-source, highlight still gives no confidence and + # rules still score 100; this line is logged anyway. + # Highlight also changes extraction: the prompt gets the line-number + # postamble and LLMWhisperer V2 adds line numbers, so answers may differ + # from a run with highlight off. This is accepted: confidence needs the + # line numbers the postamble asks for. + if params.get(AUTO_ENABLE_HIGHLIGHT): + tool_instance_metadata = {**tool_instance_metadata, _SK.ENABLE_HIGHLIGHT: True} + shim.stream_log(HIGHLIGHT_AUTO_ENABLED_MESSAGE) + # ---- Step 3: Profile overrides ---- _handle_profile_overrides(exec_metadata, platform_helper, tool_metadata) diff --git a/workers/general/tasks.py b/workers/general/tasks.py index 861ddc7ddc..bf80de2b04 100644 --- a/workers/general/tasks.py +++ b/workers/general/tasks.py @@ -42,7 +42,11 @@ from shared.patterns.retry.utils import circuit_breaker from shared.processing.files import FileProcessingUtils from shared.processing.types import FileDataValidator, TypeConverter -from shared.utils.manual_review_factory import get_manual_review_service +from shared.utils.manual_review_factory import ( + AUTO_ENABLE_HIGHLIGHT, + get_manual_review_service, + has_confidence_rules, +) from shared.workflow.connectors.source import WorkerSourceConnector from shared.workflow.execution import ( WorkerExecutionContext, @@ -978,6 +982,11 @@ def _orchestrate_file_processing_general( total_files=len(source_files), ) logger.debug("Global file data configured for manual review") + # Run-wide: rules are evaluated after extraction, so every file needs + # confidence data, including files not sent to review. Where highlight + # can't produce it, see the override in structure_tool_task.py. + if has_confidence_rules(workflow_util, workflow_id, "DB"): + global_file_data.manual_review_config[AUTO_ENABLE_HIGHLIGHT] = True # Pre-calculate file decisions for ALL files based on total count - not per batch! q_file_no_list = global_file_data.manual_review_config.get("q_file_no_list", []) logger.info( diff --git a/workers/shared/utils/manual_review_factory.py b/workers/shared/utils/manual_review_factory.py index 69c1011553..c53efb3ee0 100644 --- a/workers/shared/utils/manual_review_factory.py +++ b/workers/shared/utils/manual_review_factory.py @@ -17,6 +17,13 @@ MANUAL_REVIEW_PLUGIN_NAME = "manual_review" +# manual_review_config key: the workflow has HITL confidence rules, so the +# structure tool runs with highlight on to produce confidence scores. +AUTO_ENABLE_HIGHLIGHT = "auto_enable_highlight" +HIGHLIGHT_AUTO_ENABLED_MESSAGE = ( + "Highlight enabled automatically because HITL confidence rules are configured" +) + class ManualReviewServiceProtocol(Protocol): """Protocol defining the interface for manual review services.""" @@ -383,6 +390,16 @@ def get_manual_review_workflow_util(api_client: Any, organization_id: str) -> An return _create_null_workflow_util() +def has_confidence_rules(workflow_util: Any, workflow_id: str, rule_type: str) -> bool: + """Whether the workflow's HITL rules of ``rule_type`` ("DB" or "API") + contain a confidence condition. + + False in OSS, and with a plugin that lacks the method. + """ + check = getattr(workflow_util, "has_confidence_rules", None) + return bool(check and check(workflow_id, rule_type)) + + def has_manual_review_plugin(): return has_client_plugin(MANUAL_REVIEW_PLUGIN_NAME) diff --git a/workers/shared/workflow/execution/service.py b/workers/shared/workflow/execution/service.py index ebf8bce17c..cf22199d72 100644 --- a/workers/shared/workflow/execution/service.py +++ b/workers/shared/workflow/execution/service.py @@ -53,6 +53,7 @@ from ...api.internal_client import InternalAPIClient from ...infrastructure.logging import WorkerLogger from ...utils.error_utils import get_user_friendly_error_message +from ...utils.manual_review_factory import AUTO_ENABLE_HIGHLIGHT from ..destination_connector import ( DestinationConfig, WorkerDestinationConnector, @@ -847,7 +848,12 @@ def _execute_workflow_with_service( # Continue with normal execution if check fails # Step 3: Build and execute workflow - self._build_and_execute_workflow(execution_service, file_name) + review_config = file_processing_context.file_data.manual_review_config + self._build_and_execute_workflow( + execution_service, + file_name, + auto_enable_highlight=review_config.get(AUTO_ENABLE_HIGHLIGHT, False), + ) return True @@ -969,7 +975,10 @@ def _prepare_workflow_input_file( raise file_prep_error def _build_and_execute_workflow( - self, execution_service: WorkflowExecutionService, file_name: str + self, + execution_service: WorkflowExecutionService, + file_name: str, + auto_enable_highlight: bool = False, ) -> None: """Build and execute the workflow. @@ -977,7 +986,9 @@ def _build_and_execute_workflow( execute_structure_tool task instead of the Docker container flow. """ if self._is_structure_tool_workflow(execution_service): - self._execute_structure_tool_workflow(execution_service, file_name) + self._execute_structure_tool_workflow( + execution_service, file_name, auto_enable_highlight + ) else: # Original Docker-based flow (unchanged) execution_service.build_workflow() @@ -1015,7 +1026,10 @@ def _is_structure_tool_workflow( return False def _execute_structure_tool_workflow( - self, execution_service: WorkflowExecutionService, file_name: str + self, + execution_service: WorkflowExecutionService, + file_name: str, + auto_enable_highlight: bool = False, ) -> None: """Execute structure tool as a task instead of a Docker container. @@ -1061,6 +1075,7 @@ def _execute_structure_tool_workflow( "messaging_channel": getattr(execution_service, "messaging_channel", ""), "file_hash": metadata.get("source_hash", ""), "exec_metadata": metadata, + AUTO_ENABLE_HIGHLIGHT: auto_enable_highlight, } # Call synchronously (same process, in-band) diff --git a/workers/tests/test_barrier.py b/workers/tests/test_barrier.py index 7bf65d1825..0a276134c4 100644 --- a/workers/tests/test_barrier.py +++ b/workers/tests/test_barrier.py @@ -773,5 +773,135 @@ def test_manual_review_required_invokes_decision_helper_per_batch( ) +class TestAutoEnableHighlightOrchestration: + """HITL confidence rules set ``auto_enable_highlight``. + + The orchestrators check their own flow's rules once (API deployments the + API rules, only for a ``hitl_queue_name`` request without a + ``hitl_packet_id``; the rest the DB rules) and put the flag in + ``manual_review_config``, which the file workers read. + """ + + @staticmethod + def _util(confidence_rule_type): + """A workflow util with a confidence rule of one type, or none.""" + util = MagicMock(name="workflow_util") + util.has_confidence_rules.side_effect = ( + lambda workflow_id, rule_type: rule_type == confidence_rule_type + ) + return util + + @pytest.mark.parametrize( + "confidence_rule_type, hitl_kwargs, has_rules", + [ + ("API", {"hitl_queue_name": "q"}, True), + ("DB", {"hitl_queue_name": "q"}, False), + (None, {"hitl_queue_name": "q"}, False), + # The destination evaluates no API rules for these requests. + ("API", {}, False), + ("API", {"hitl_queue_name": "q", "hitl_packet_id": "p"}, False), + ], + ) + def test_api_deployment_sets_flag_on_every_batch( + self, monkeypatch, confidence_rule_type, hitl_kwargs, has_rules + ): + from shared.utils.manual_review_factory import AUTO_ENABLE_HIGHLIGHT + + mocks = _setup_workflow_api_mocks(monkeypatch, num_batches=2) + util = self._util(confidence_rule_type) + monkeypatch.setattr( + mocks.api_tasks, "get_manual_review_workflow_util", lambda *a: util + ) + configs = [] + monkeypatch.setattr( + mocks.api_tasks, + "_create_file_data", + lambda **kwargs: MagicMock(manual_review_config={}), + ) + monkeypatch.setattr( + mocks.api_tasks, + "_create_batch_data", + lambda **kwargs: configs.append( + dict(kwargs["file_data"].manual_review_config) + ), + ) + + mocks.api_tasks._run_workflow_api( + api_client=mocks.api_client, + schema_name="org_test", + workflow_id="wf-1", + execution_id="exec-1", + hash_values_of_files={"f1": MagicMock(name="file_1")}, + scheduled=False, + execution_mode=None, + pipeline_id="pipe-1", + use_file_history=False, + task_id="task-1", + **hitl_kwargs, + ) + + if hitl_kwargs.get("hitl_queue_name") and not hitl_kwargs.get("hitl_packet_id"): + util.has_confidence_rules.assert_called_once_with("wf-1", "API") + else: + util.has_confidence_rules.assert_not_called() + assert len(configs) == 2 + assert all(c.get(AUTO_ENABLE_HIGHLIGHT, False) is has_rules for c in configs) + + @pytest.mark.parametrize( + "confidence_rule_type, has_rules", [("DB", True), ("API", False), (None, False)] + ) + def test_general_sets_flag_on_global_file_data( + self, monkeypatch, confidence_rule_type, has_rules + ): + from general import tasks as general_tasks + from shared.utils.manual_review_factory import AUTO_ENABLE_HIGHLIGHT + + api_client = MagicMock(name="api_client") + util = self._util(confidence_rule_type) + util.create_workflow_file_data_with_manual_review.return_value = MagicMock( + manual_review_config={} + ) + service = MagicMock(name="manual_review_service") + service.get_workflow_util.return_value = util + monkeypatch.setattr( + general_tasks, "get_manual_review_service", lambda **kwargs: service + ) + monkeypatch.setattr( + general_tasks, + "_get_file_batches_general", + lambda **kwargs: [MagicMock(name="batch_1")], + ) + configs = [] + + def create_batch_data(**kwargs): + configs.append(dict(kwargs["file_data"].manual_review_config)) + return MagicMock(files=[]) + + monkeypatch.setattr( + general_tasks, "_create_batch_data_general", create_batch_data + ) + monkeypatch.setattr(general_tasks, "_log_batch_statistics_to_ui", MagicMock()) + monkeypatch.setattr( + general_tasks.WorkflowOrchestrationUtils, + "create_chord_execution", + MagicMock(return_value=MagicMock(id="chord-id")), + ) + + general_tasks._orchestrate_file_processing_general( + api_client=api_client, + workflow_id="wf-1", + execution_id="exec-1", + source_files={"f1": MagicMock(name="file_1")}, + pipeline_id="pipe-1", + scheduled=False, + execution_mode=None, + use_file_history=False, + organization_id="org_test", + ) + + util.has_confidence_rules.assert_called_once_with("wf-1", "DB") + assert configs[0].get(AUTO_ENABLE_HIGHLIGHT, False) is has_rules + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/workers/tests/test_structure_tool_pipeline.py b/workers/tests/test_structure_tool_pipeline.py index 4c2812bf66..80dcded85b 100644 --- a/workers/tests/test_structure_tool_pipeline.py +++ b/workers/tests/test_structure_tool_pipeline.py @@ -894,7 +894,10 @@ class TestStructureToolParamsPassthrough: "_is_structure_tool_workflow", return_value=True, ) - def test_structure_tool_params_passthrough(self, mock_is_struct, mock_exec_struct): + @pytest.mark.parametrize("auto_enable", [True, False]) + def test_structure_tool_params_passthrough( + self, mock_is_struct, mock_exec_struct, auto_enable + ): from shared.workflow.execution.service import ( WorkerWorkflowExecutionService, ) @@ -904,10 +907,68 @@ def test_structure_tool_params_passthrough(self, mock_is_struct, mock_exec_struc mock_exec_service = MagicMock() mock_exec_service.tool_instances = [MagicMock()] - service._build_and_execute_workflow(mock_exec_service, "test.pdf") + service._build_and_execute_workflow( + mock_exec_service, "test.pdf", auto_enable_highlight=auto_enable + ) # Verify _execute_structure_tool_workflow was called - mock_exec_struct.assert_called_once_with(mock_exec_service, "test.pdf") + mock_exec_struct.assert_called_once_with( + mock_exec_service, "test.pdf", auto_enable + ) + + @pytest.mark.parametrize( + "review_config, expected", + [({"auto_enable_highlight": True}, True), ({}, False)], + ) + @patch( + "shared.workflow.execution.service.WorkerWorkflowExecutionService." + "_build_and_execute_workflow" + ) + @patch( + "shared.workflow.execution.service.WorkerWorkflowExecutionService." + "_prepare_workflow_input_file", + return_value=None, + ) + @patch( + "shared.workflow.execution.service.WorkerWorkflowExecutionService." + "_compile_workflow", + return_value=True, + ) + def test_auto_enable_highlight_read_from_review_config( + self, mock_compile, mock_prepare, mock_build, review_config, expected + ): + from shared.workflow.execution.service import ( + WorkerWorkflowExecutionService, + ) + + context = MagicMock() + context.file_data.manual_review_config = review_config + + assert WorkerWorkflowExecutionService()._execute_workflow_with_service( + execution_service=MagicMock(), + file_processing_context=context, + file_name="invoice.pdf", + workflow_file_execution_id="", + execution_id="exec-1", + workflow_id="wf-1", + ) + assert mock_build.call_args.kwargs["auto_enable_highlight"] is expected + + @pytest.mark.parametrize( + "util, expected", + [ + (MagicMock(has_confidence_rules=MagicMock(return_value=True)), True), + (MagicMock(has_confidence_rules=MagicMock(return_value=False)), False), + # Older plugin without the method + (object(), False), + ], + ) + def test_has_confidence_rules_helper(self, util, expected): + from shared.utils.manual_review_factory import has_confidence_rules + + assert has_confidence_rules(util, "wf-1", "DB") is expected + if isinstance(util, MagicMock): + util.has_confidence_rules.assert_called_once_with("wf-1", "DB") @patch( "shared.workflow.execution.service.WorkerWorkflowExecutionService." @@ -963,6 +1024,154 @@ def test_source_file_name_uses_real_filename_not_sentinel( assert params["source_file_name"] == "invoice.pdf" assert params["source_file_name"] != "SOURCE" + @pytest.mark.parametrize("auto_enable", [True, False]) + @patch( + "shared.workflow.execution.service.WorkerWorkflowExecutionService." + "_get_platform_service_api_key", + return_value="sk-test", + ) + @patch("file_processing.structure_tool_task.execute_structure_tool") + def test_auto_enable_highlight_passed_in_params( + self, mock_execute_struct, mock_get_key, auto_enable + ): + """The service passes the flag on and leaves the metadata alone.""" + from shared.workflow.execution.service import ( + WorkerWorkflowExecutionService, + ) + + saved_metadata = {"enable_highlight": False} + mock_exec_service = MagicMock() + mock_exec_service.tool_instances = [MagicMock(metadata=saved_metadata)] + mock_exec_service.file_handler.get_workflow_metadata.return_value = {} + mock_execute_struct.return_value = {"success": True} + + WorkerWorkflowExecutionService()._execute_structure_tool_workflow( + mock_exec_service, "invoice.pdf", auto_enable_highlight=auto_enable + ) + + params = mock_execute_struct.call_args[0][0] + assert params["auto_enable_highlight"] is auto_enable + assert params["tool_instance_metadata"] == {"enable_highlight": False} + + +class TestAutoEnableHighlight: + """``auto_enable_highlight`` turns highlight on for the run.""" + + @pytest.mark.parametrize( + "user_highlight, auto_enable, expected", + [ + (False, True, True), # HITL confidence rule present + (True, False, True), # user setting respected + (False, False, False), + ], + ) + @patch(_PATCH_SHIM) + @patch(_PATCH_FILE_STORAGE) + @patch(_PATCH_PLATFORM_HELPER) + @patch(_PATCH_DISPATCHER) + def test_regular_overrides_a_copy( + self, + mock_get_executor_dispatcher, + mock_create_ph, + mock_get_fs, + mock_shim_cls, + base_params, + tool_metadata_regular, + mock_fs, + mock_platform_helper, + user_highlight, + auto_enable, + expected, + ): + """Word confidence and the saved metadata stay as the user set them.""" + from file_processing.structure_tool_task import ( + _execute_structure_tool_impl as execute_structure_tool, + ) + from shared.utils.manual_review_factory import ( + HIGHLIGHT_AUTO_ENABLED_MESSAGE, + ) + + mock_get_fs.return_value = mock_fs + mock_create_ph.return_value = mock_platform_helper + mock_platform_helper.get_prompt_studio_tool.return_value = { + "tool_metadata": tool_metadata_regular, + } + dispatcher_instance = create_autospec(PgExecutionDispatcher, instance=True) + mock_get_executor_dispatcher.return_value = dispatcher_instance + dispatcher_instance.dispatch.return_value = _make_pipeline_result() + + saved_metadata = base_params["tool_instance_metadata"] + saved_metadata["enable_highlight"] = user_highlight + saved_metadata["enable_word_confidence"] = False + base_params["auto_enable_highlight"] = auto_enable + + execute_structure_tool(base_params) + + ep = dispatcher_instance.dispatch.call_args[0][0].executor_params + assert ep["extract_params"]["enable_highlight"] is expected + assert ep["answer_params"]["tool_settings"]["enable_word_confidence"] is False + assert saved_metadata["enable_highlight"] is user_highlight + messages = [ + c.args[0] for c in mock_shim_cls.return_value.stream_log.call_args_list + ] + assert (HIGHLIGHT_AUTO_ENABLED_MESSAGE in messages) is auto_enable + + @patch("unstract.sdk1.x2txt.X2Text") + @patch(_PATCH_SHIM) + @patch(_PATCH_FILE_STORAGE) + @patch(_PATCH_PLATFORM_HELPER) + @patch(_PATCH_DISPATCHER) + def test_agentic_keeps_user_setting( + self, + mock_get_executor_dispatcher, + mock_create_ph, + mock_get_fs, + mock_shim_cls, + mock_x2text_cls, + base_params, + mock_fs, + mock_platform_helper, + ): + """Agentic runs get no highlight data, so the flag is ignored.""" + from file_processing.structure_tool_task import ( + _execute_structure_tool_impl as execute_structure_tool, + ) + from shared.utils.manual_review_factory import ( + HIGHLIGHT_AUTO_ENABLED_MESSAGE, + ) + + mock_get_fs.return_value = mock_fs + mock_create_ph.return_value = mock_platform_helper + mock_platform_helper.get_prompt_studio_tool.return_value = None + mock_platform_helper.get_agentic_studio_tool.return_value = { + "tool_metadata": { + "name": "Agentic Project", + "project_id": "ap-001", + "json_schema": {"field": "string"}, + }, + } + dispatcher_instance = create_autospec(PgExecutionDispatcher, instance=True) + mock_get_executor_dispatcher.return_value = dispatcher_instance + dispatcher_instance.dispatch.return_value = ExecutionResult( + success=True, data={"output": {"field": "value"}} + ) + mock_x2text_instance = MagicMock() + mock_x2text_instance.process.return_value = MagicMock( + extracted_text="extracted text" + ) + mock_x2text_cls.return_value = mock_x2text_instance + + base_params["tool_instance_metadata"]["enable_highlight"] = False + base_params["auto_enable_highlight"] = True + + execute_structure_tool(base_params) + + assert mock_x2text_instance.process.call_args.kwargs["enable_highlight"] is False + messages = [ + c.args[0] for c in mock_shim_cls.return_value.stream_log.call_args_list + ] + assert HIGHLIGHT_AUTO_ENABLED_MESSAGE not in messages + class TestHelperFunctions: """Test standalone helper functions."""