diff --git a/frontend/src/components/custom-tools/add-llm-profile/AddLlmProfile.jsx b/frontend/src/components/custom-tools/add-llm-profile/AddLlmProfile.jsx index a020540b26..bac0e5325f 100644 --- a/frontend/src/components/custom-tools/add-llm-profile/AddLlmProfile.jsx +++ b/frontend/src/components/custom-tools/add-llm-profile/AddLlmProfile.jsx @@ -64,10 +64,6 @@ function AddLlmProfile({ useEffect(() => { setAdaptorProfilesDropdown(); - - return () => { - setEditLlmProfileId(null); - }; }, []); // Load retrieval strategies when tool_id is available (only once) @@ -163,7 +159,7 @@ function AddLlmProfile({ useEffect(() => { if (resetForm) { - form.resetFields(); + form.setFieldsValue(formDetails); setResetForm(false); } }, [formDetails]); diff --git a/frontend/src/components/custom-tools/manage-llm-profiles/ManageLlmProfiles.jsx b/frontend/src/components/custom-tools/manage-llm-profiles/ManageLlmProfiles.jsx index 3bb80d6cb2..15d2263f60 100644 --- a/frontend/src/components/custom-tools/manage-llm-profiles/ManageLlmProfiles.jsx +++ b/frontend/src/components/custom-tools/manage-llm-profiles/ManageLlmProfiles.jsx @@ -166,6 +166,7 @@ function ManageLlmProfiles() { }, [llmProfiles, defaultLlmProfile]); const handleAddNewLlmProfileBtnClick = () => { + setEditLlmProfileId(null); setIsAddLlm(true); try { diff --git a/frontend/src/components/settings/settings/Settings.css b/frontend/src/components/settings/settings/Settings.css index 5a08144c99..d89b9291b5 100644 --- a/frontend/src/components/settings/settings/Settings.css +++ b/frontend/src/components/settings/settings/Settings.css @@ -34,7 +34,6 @@ } .settings-menu-item:first-child { - font-weight: 600; margin-bottom: 4px; } diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index 2560756ed0..6f9b04ebce 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -1494,7 +1494,24 @@ def _execute_single_prompt( return if output.get(PSKeys.TYPE) == PSKeys.LINE_ITEM: - raise LegacyExecutorError(message="LINE_ITEM extraction is not supported.") + self._run_line_item_extraction( + output=output, + context=context, + structured_output=structured_output, + metadata=metadata, + metrics=metrics, + run_id=run_id, + execution_id=execution_id, + execution_source=execution_source, + platform_api_key=platform_api_key, + tool_id=tool_id, + doc_name=doc_name, + prompt_name=prompt_name, + file_path=file_path, + tool_settings=tool_settings, + shim=shim, + ) + return usage_kwargs = {"run_id": run_id, "execution_id": execution_id} try: @@ -1700,6 +1717,83 @@ def _run_table_extraction( ) shim.stream_log(f"Completed prompt: {prompt_name}") + def _run_line_item_extraction( + self, + output: dict[str, Any], + context: ExecutionContext, + structured_output: dict[str, Any], + metadata: dict[str, Any], + metrics: dict[str, Any], + run_id: str, + execution_id: str, + execution_source: str, + platform_api_key: str, + tool_id: str, + doc_name: str, + prompt_name: str, + file_path: str, + tool_settings: dict[str, Any], + shim: Any, + ) -> None: + """Delegate LINE_ITEM prompt to the line_item executor plugin.""" + from executor.executors.constants import PromptServiceConstants as PSKeys + + try: + line_item_executor = ExecutorRegistry.get("line_item") + except KeyError: + raise LegacyExecutorError( + message=( + "LINE_ITEM extraction requires the line_item executor " + "plugin. Install the line_item_extractor plugin." + ) + ) + line_item_ctx = ExecutionContext( + executor_name="line_item", + operation="line_item_extract", + run_id=run_id, + execution_source=execution_source, + organization_id=context.organization_id, + request_id=context.request_id, + executor_params={ + "llm_adapter_instance_id": output.get(PSKeys.LLM, ""), + "tool_settings": tool_settings, + "output": output, + "prompt": output.get(PSKeys.PROMPTX, ""), + "file_path": file_path, + "PLATFORM_SERVICE_API_KEY": platform_api_key, + "execution_id": execution_id, + "tool_id": tool_id, + "file_name": doc_name, + "prompt_name": prompt_name, + }, + ) + line_item_ctx._log_component = self._log_component + line_item_ctx.log_events_id = self._log_events_id + + shim.stream_log(f"Running line-item extraction for: {prompt_name}") + line_item_result = line_item_executor.execute(line_item_ctx) + + if line_item_result.success: + data = line_item_result.data or {} + structured_output[prompt_name] = data.get("output", "") + line_item_metrics = data.get("metadata", {}).get("metrics", {}) + metrics.setdefault(prompt_name, {}).update( + {"line_item_extraction": line_item_metrics} + ) + context_list = data.get("context") + if context_list: + metadata[PSKeys.CONTEXT][prompt_name] = context_list + shim.stream_log(f"Line-item extraction completed for: {prompt_name}") + logger.info("LINE_ITEM extraction completed: prompt=%s", prompt_name) + else: + structured_output[prompt_name] = "" + logger.error( + "LINE_ITEM extraction failed for prompt=%s: %s", + prompt_name, + line_item_result.error, + ) + shim.stream_log(f"Completed prompt: {prompt_name}") + @staticmethod def _apply_type_conversion( output: dict[str, Any], @@ -1732,16 +1826,19 @@ def _apply_type_conversion( ) elif output_type == PSKeys.EMAIL: - email_prompt = ( - f"Extract the email from the following text:\n{answer}" - f"\n\nOutput just the email. " - f"The email should be directly assignable to a string " - f"variable. No explanation is required. If you cannot " - f'extract the email, output "NA".' - ) - structured_output[prompt_name] = LegacyExecutor._convert_scalar_answer( - answer, llm, answer_prompt_svc, email_prompt - ) + if answer.lower() == "na": + structured_output[prompt_name] = answer + else: + email_prompt = ( + f"Extract the email from the following text:\n{answer}" + f"\n\nOutput just the email. " + f"The email should be directly assignable to a string " + f"variable. No explanation is required. If you cannot " + f'extract the email, output "NA".' + ) + structured_output[prompt_name] = answer_prompt_svc.run_completion( + llm=llm, prompt=email_prompt + ) elif output_type == PSKeys.DATE: date_prompt = ( diff --git a/workers/tests/test_answer_prompt.py b/workers/tests/test_answer_prompt.py index 77784b2465..d3b4505fb3 100644 --- a/workers/tests/test_answer_prompt.py +++ b/workers/tests/test_answer_prompt.py @@ -562,8 +562,8 @@ def test_invalid_strategy_skips_retrieval( ) result = executor._handle_answer_prompt(ctx) - # Answer stays "NA" which gets sanitized to None - assert result.data[PSKeys.OUTPUT]["field_a"] is None + # Answer stays "NA" (top-level NA is preserved, not sanitized) + assert result.data[PSKeys.OUTPUT]["field_a"] == "NA" class TestHandleAnswerPromptMultiPrompt: @@ -687,21 +687,21 @@ def test_vectordb_closed(self, mock_shim_cls, mock_deps): class TestNullSanitization: """Tests for _sanitize_null_values.""" - def test_na_string_becomes_none(self): - """Top-level 'NA' string → None.""" + def test_na_string_preserved(self): + """Top-level 'NA' string is preserved (not sanitized to None).""" from executor.executors.legacy_executor import LegacyExecutor output = {"field": "NA"} result = LegacyExecutor._sanitize_null_values(output) - assert result["field"] is None + assert result["field"] == "NA" - def test_na_case_insensitive(self): - """'na' (lowercase) → None.""" + def test_na_case_insensitive_preserved(self): + """Top-level 'na' (lowercase) is preserved (not sanitized to None).""" from executor.executors.legacy_executor import LegacyExecutor output = {"field": "na"} result = LegacyExecutor._sanitize_null_values(output) - assert result["field"] is None + assert result["field"] == "na" def test_nested_list_na(self): """NA in nested list items → None.""" diff --git a/workers/tests/test_line_item_extraction.py b/workers/tests/test_line_item_extraction.py new file mode 100644 index 0000000000..5a46b47b1f --- /dev/null +++ b/workers/tests/test_line_item_extraction.py @@ -0,0 +1,459 @@ +"""Tests for LegacyExecutor._run_line_item_extraction (LINE_ITEM type). + +Mirrors the structure of TABLE delegation tests in test_sanity_phase6d.py. + +Verifies: +1. Plugin missing → LegacyExecutorError with install hint. +2. Plugin success → output written + metrics merged + context propagated. +3. Plugin failure → empty output + error logged. +4. End-to-end through _execute_single_prompt with a LINE_ITEM prompt + in the structure-tool path (eager Celery + fake plugin). +""" + +import logging +from unittest.mock import MagicMock, patch + +import pytest +from executor.executors.answer_prompt import AnswerPromptService +from executor.executors.constants import PromptServiceConstants as PSKeys +from executor.executors.exceptions import LegacyExecutorError +from unstract.sdk1.execution.context import ExecutionContext +from unstract.sdk1.execution.executor import BaseExecutor +from unstract.sdk1.execution.registry import ExecutorRegistry +from unstract.sdk1.execution.result import ExecutionResult + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +@pytest.fixture(autouse=True) +def _clean_registry(): + """Ensure a clean executor registry for every test.""" + ExecutorRegistry.clear() + yield + ExecutorRegistry.clear() + + +def _get_legacy_executor(): + """Register and fetch the LegacyExecutor instance.""" + from executor.executors.legacy_executor import LegacyExecutor + + if "legacy" not in ExecutorRegistry.list_executors(): + ExecutorRegistry.register(LegacyExecutor) + return ExecutorRegistry.get("legacy") + + +def _make_line_item_prompt(): + """Build a LINE_ITEM prompt config dict (mirrors _execute_single_prompt + expectations). + """ + return { + PSKeys.NAME: "line_items", + PSKeys.PROMPT: "Extract all invoice line items.", + PSKeys.PROMPTX: "Extract all invoice line items.", + PSKeys.TYPE: PSKeys.LINE_ITEM, + PSKeys.CHUNK_SIZE: 0, + PSKeys.CHUNK_OVERLAP: 0, + PSKeys.LLM: "llm-1", + PSKeys.EMBEDDING: "emb-1", + PSKeys.VECTOR_DB: "vdb-1", + PSKeys.X2TEXT_ADAPTER: "x2t-1", + PSKeys.RETRIEVAL_STRATEGY: "simple", + } + + +def _make_context(): + """Build a minimal ExecutionContext for the answer_prompt path.""" + tool_settings = { + PSKeys.PREAMBLE: "", + PSKeys.POSTAMBLE: "", + PSKeys.GRAMMAR: [], + PSKeys.ENABLE_HIGHLIGHT: False, + } + return ExecutionContext( + executor_name="legacy", + operation="answer_prompt", + run_id="run-line-item-001", + execution_source="tool", + organization_id="org-test", + request_id="req-line-item-001", + executor_params={ + PSKeys.TOOL_SETTINGS: tool_settings, + PSKeys.OUTPUTS: [_make_line_item_prompt()], + PSKeys.TOOL_ID: "tool-1", + PSKeys.FILE_HASH: "hash123", + PSKeys.FILE_PATH: "/data/invoice.txt", + PSKeys.FILE_NAME: "invoice.txt", + PSKeys.PLATFORM_SERVICE_API_KEY: "pk-test", + }, + ) + + +def _standard_patches(executor): + """Common patches needed to drive _handle_answer_prompt → _execute_single_prompt + until it reaches the LINE_ITEM branch. + """ + llm = MagicMock(name="llm") + llm.get_metrics.return_value = {} + mock_llm_cls = MagicMock(return_value=llm) + return { + "_get_prompt_deps": patch.object( + executor, + "_get_prompt_deps", + return_value=( + AnswerPromptService, + MagicMock( + retrieve_complete_context=MagicMock( + return_value=["context chunk"] + ) + ), + MagicMock( + is_variables_present=MagicMock(return_value=False) + ), + None, # Index + mock_llm_cls, + MagicMock(), # EmbeddingCompat + MagicMock(), # VectorDB + ), + ), + "shim": patch( + "executor.executors.legacy_executor.ExecutorToolShim", + return_value=MagicMock(), + ), + "index_key": patch( + "unstract.sdk1.utils.indexing.IndexingUtils.generate_index_key", + return_value="doc-id-1", + ), + } + + +# --------------------------------------------------------------------------- +# Fake LineItemExecutor plugins +# --------------------------------------------------------------------------- + + +def _make_success_plugin( + output_value=None, + metrics=None, + context_list=None, +): + """Build a fake plugin class that returns a success ExecutionResult.""" + payload = {"output": output_value or {"items": [{"sku": "A1", "qty": 2}]}} + if metrics is not None: + payload["metadata"] = {"metrics": metrics} + if context_list is not None: + payload["context"] = context_list + + class _SuccessPlugin(BaseExecutor): + @property + def name(self) -> str: + return "line_item" + + def execute(self, context: ExecutionContext) -> ExecutionResult: + self.received_context = context + return ExecutionResult(success=True, data=payload) + + return _SuccessPlugin + + +def _make_failure_plugin(error="extraction blew up"): + class _FailurePlugin(BaseExecutor): + @property + def name(self) -> str: + return "line_item" + + def execute(self, context: ExecutionContext) -> ExecutionResult: + return ExecutionResult.failure(error=error) + + return _FailurePlugin + + +# --------------------------------------------------------------------------- +# 1. Plugin missing → clear error +# --------------------------------------------------------------------------- + + +class TestLineItemPluginMissing: + @patch("executor.executors.legacy_executor.ExecutorToolShim") + @patch( + "unstract.sdk1.utils.indexing.IndexingUtils.generate_index_key", + return_value="doc-id-1", + ) + def test_line_item_raises_when_plugin_missing( + self, _mock_key, _mock_shim_cls + ): + """LINE_ITEM prompt raises LegacyExecutorError with install hint.""" + executor = _get_legacy_executor() + ctx = _make_context() + patches = _standard_patches(executor) + + with patches["_get_prompt_deps"], patches["shim"], patches["index_key"]: + with pytest.raises( + LegacyExecutorError, + match="line_item_extractor plugin", + ): + executor._handle_answer_prompt(ctx) + + +# --------------------------------------------------------------------------- +# 2. Plugin success → output written + metrics + context +# --------------------------------------------------------------------------- + + +class TestLineItemPluginSuccess: + @patch("executor.executors.legacy_executor.ExecutorToolShim") + @patch( + "unstract.sdk1.utils.indexing.IndexingUtils.generate_index_key", + return_value="doc-id-1", + ) + def test_success_writes_output_and_merges_metrics( + self, _mock_key, _mock_shim_cls + ): + executor = _get_legacy_executor() + ctx = _make_context() + patches = _standard_patches(executor) + + plugin_cls = _make_success_plugin( + output_value={"items": [{"sku": "A1", "qty": 2}]}, + metrics={"llm_calls": 3}, + context_list=["full file body"], + ) + # Register so ExecutorRegistry.get("line_item") finds it + ExecutorRegistry.register(plugin_cls) + + with patches["_get_prompt_deps"], patches["shim"], patches["index_key"]: + result = executor._handle_answer_prompt(ctx) + + assert result.success is True + # structured_output[prompt_name] holds the plugin output + assert result.data["output"]["line_items"] == { + "items": [{"sku": "A1", "qty": 2}] + } + # Metrics are merged under the line_item_extraction sub-key + prompt_metrics = result.data["metrics"]["line_items"] + assert prompt_metrics["line_item_extraction"] == {"llm_calls": 3} + # Context list is propagated to metadata.context + assert result.data["metadata"][PSKeys.CONTEXT]["line_items"] == [ + "full file body" + ] + + @patch("executor.executors.legacy_executor.ExecutorToolShim") + @patch( + "unstract.sdk1.utils.indexing.IndexingUtils.generate_index_key", + return_value="doc-id-1", + ) + def test_success_passes_correct_executor_params( + self, _mock_key, _mock_shim_cls + ): + """Verify the sub-context built for the plugin has all expected + keys with the right values. + """ + executor = _get_legacy_executor() + ctx = _make_context() + patches = _standard_patches(executor) + + captured: dict = {} + + class _CapturePlugin(BaseExecutor): + @property + def name(self) -> str: + return "line_item" + + def execute(self, context: ExecutionContext) -> ExecutionResult: + captured["ctx"] = context + return ExecutionResult(success=True, data={"output": {}}) + + ExecutorRegistry.register(_CapturePlugin) + + with patches["_get_prompt_deps"], patches["shim"], patches["index_key"]: + executor._handle_answer_prompt(ctx) + + sub_ctx = captured["ctx"] + assert sub_ctx.executor_name == "line_item" + assert sub_ctx.operation == "line_item_extract" + assert sub_ctx.run_id == "run-line-item-001" + assert sub_ctx.organization_id == "org-test" + + params = sub_ctx.executor_params + assert params["llm_adapter_instance_id"] == "llm-1" + assert params["PLATFORM_SERVICE_API_KEY"] == "pk-test" + assert params["file_path"] == "/data/invoice.txt" + assert params["file_name"] == "invoice.txt" + assert params["tool_id"] == "tool-1" + assert params["prompt_name"] == "line_items" + assert params["prompt"] == "Extract all invoice line items." + # output dict and tool_settings are passed through + assert params["output"][PSKeys.NAME] == "line_items" + assert params["output"][PSKeys.TYPE] == PSKeys.LINE_ITEM + assert PSKeys.PREAMBLE in params["tool_settings"] + + +# --------------------------------------------------------------------------- +# 3. Plugin failure → empty output + error logged +# --------------------------------------------------------------------------- + + +class TestLineItemPluginFailure: + @patch("executor.executors.legacy_executor.ExecutorToolShim") + @patch( + "unstract.sdk1.utils.indexing.IndexingUtils.generate_index_key", + return_value="doc-id-1", + ) + def test_failure_writes_empty_output_and_logs( + self, _mock_key, _mock_shim_cls, caplog + ): + executor = _get_legacy_executor() + ctx = _make_context() + patches = _standard_patches(executor) + + ExecutorRegistry.register(_make_failure_plugin("plugin error")) + + with patches["_get_prompt_deps"], patches["shim"], patches["index_key"]: + with caplog.at_level( + logging.ERROR, + logger="executor.executors.legacy_executor", + ): + result = executor._handle_answer_prompt(ctx) + + assert result.success is True # answer_prompt itself does not raise + assert result.data["output"]["line_items"] == "" + # Failure logged + assert any( + "LINE_ITEM extraction failed" in rec.message + and "plugin error" in rec.message + for rec in caplog.records + ) + + +# --------------------------------------------------------------------------- +# 4. End-to-end through Celery eager mode (structure-tool path) +# --------------------------------------------------------------------------- + + +@pytest.fixture +def eager_app(): + """Configure executor Celery app for eager-mode testing.""" + from executor.worker import app + + original = { + "task_always_eager": app.conf.task_always_eager, + "task_eager_propagates": app.conf.task_eager_propagates, + "result_backend": app.conf.result_backend, + } + app.conf.update( + task_always_eager=True, + task_eager_propagates=False, + result_backend="cache+memory://", + ) + yield app + app.conf.update(original) + + +def _structure_tool_ctx(): + """Build an answer_prompt context with a single LINE_ITEM prompt for + the structure-tool path (execution_source='tool'). + """ + tool_settings = { + PSKeys.PREAMBLE: "Extract carefully.", + PSKeys.POSTAMBLE: "No commentary.", + PSKeys.GRAMMAR: [], + PSKeys.ENABLE_HIGHLIGHT: False, + PSKeys.ENABLE_CHALLENGE: False, + } + return ExecutionContext( + executor_name="legacy", + operation="answer_prompt", + run_id="run-line-item-e2e", + execution_source="tool", + organization_id="org-e2e", + request_id="req-e2e", + executor_params={ + PSKeys.TOOL_SETTINGS: tool_settings, + PSKeys.OUTPUTS: [_make_line_item_prompt()], + PSKeys.TOOL_ID: "tool-e2e", + PSKeys.FILE_HASH: "hash-e2e", + PSKeys.FILE_PATH: "/data/rent_roll.txt", + PSKeys.FILE_NAME: "rent_roll.txt", + PSKeys.PLATFORM_SERVICE_API_KEY: "pk-e2e", + }, + ) + + +class TestLineItemEndToEnd: + @patch("executor.executors.legacy_executor.ExecutorToolShim") + @patch( + "unstract.sdk1.utils.indexing.IndexingUtils.generate_index_key", + return_value="doc-id-e2e", + ) + @patch( + "executor.executors.legacy_executor.LegacyExecutor._get_prompt_deps" + ) + @patch( + "executor.executors.plugins.loader.ExecutorPluginLoader.get", + return_value=None, + ) + def test_celery_eager_chain_with_line_item_plugin( + self, + _mock_plugin_loader, + mock_deps, + _mock_index_utils, + _mock_shim_cls, + eager_app, + ): + """Push a LINE_ITEM payload through the full Celery eager chain + with a fake line_item plugin registered. + """ + # Re-register LegacyExecutor since the autouse fixture cleared it + from executor.executors.legacy_executor import LegacyExecutor + + ExecutorRegistry.register(LegacyExecutor) + + # Register fake line_item plugin + plugin_cls = _make_success_plugin( + output_value={ + "items": [ + {"unit": "1A", "rent": 1500}, + {"unit": "1B", "rent": 1700}, + ] + }, + metrics={"llm_calls": 2, "tokens": 1234}, + context_list=["rent roll body"], + ) + ExecutorRegistry.register(plugin_cls) + + # Mock the prompt deps so _execute_single_prompt can run far + # enough to hit the LINE_ITEM branch + llm = MagicMock(name="llm") + llm.get_metrics.return_value = {} + mock_llm_cls = MagicMock(return_value=llm) + mock_deps.return_value = ( + AnswerPromptService, + MagicMock( + retrieve_complete_context=MagicMock(return_value=["chunk"]) + ), + MagicMock(is_variables_present=MagicMock(return_value=False)), + None, + mock_llm_cls, + MagicMock(), + MagicMock(), + ) + + ctx = _structure_tool_ctx() + task = eager_app.tasks["execute_extraction"] + async_result = task.apply(args=[ctx.to_dict()]) + result_dict = async_result.get() + result = ExecutionResult.from_dict(result_dict) + + assert result.success is True + assert result.data["output"]["line_items"] == { + "items": [ + {"unit": "1A", "rent": 1500}, + {"unit": "1B", "rent": 1700}, + ] + } + assert ( + result.data["metrics"]["line_items"]["line_item_extraction"] + == {"llm_calls": 2, "tokens": 1234} + ) diff --git a/workers/tests/test_sanity_phase6d.py b/workers/tests/test_sanity_phase6d.py index cd40c1b685..46d2693465 100644 --- a/workers/tests/test_sanity_phase6d.py +++ b/workers/tests/test_sanity_phase6d.py @@ -201,9 +201,12 @@ def test_table_type_raises_when_plugin_missing( @patch("executor.executors.legacy_executor.ExecutorToolShim") @patch("unstract.sdk1.utils.indexing.IndexingUtils.generate_index_key", return_value="doc-id-1") - def test_line_item_type_raises_not_supported( + def test_line_item_type_raises_when_plugin_missing( self, mock_key, mock_shim_cls ): + """LINE_ITEM prompts raise install error when line_item plugin + is not registered (mirrors TABLE missing-plugin behavior). + """ mock_shim_cls.return_value = MagicMock() executor = _get_executor() ctx = _make_context(output_type=PSKeys.LINE_ITEM) # "line-item" @@ -211,8 +214,16 @@ def test_line_item_type_raises_not_supported( patches = _standard_patches(executor, llm) with patches["_get_prompt_deps"], patches["shim"], patches["index_key"]: - with pytest.raises(LegacyExecutorError, match="not supported"): - executor._handle_answer_prompt(ctx) + with patch( + "unstract.sdk1.execution.registry.ExecutorRegistry.get", + side_effect=KeyError( + "No executor registered with name 'line_item'" + ), + ): + with pytest.raises( + LegacyExecutorError, match="line_item_extractor plugin" + ): + executor._handle_answer_prompt(ctx) # ---------------------------------------------------------------------------