From 36463bb9b68314995ac451db33d1f345855416a6 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Thu, 2 Apr 2026 01:57:01 +0530 Subject: [PATCH 01/57] UN-2946 [FEAT] Add lightweight list serializer for Prompt Studio and prompt list endpoint - Add CustomToolListSerializer for the list action to avoid N+1 queries (profile lookups, prompt fetching, coverage calculation per tool) - Add ToolStudioPromptListSerializer with only prompt_id, prompt_key, enforce_type, sequence_number - Add GET /prompt-studio/prompt/?tool_id={uuid} list endpoint - List action uses select_related and Subquery annotation for prompt_count - Detail endpoint unchanged (still uses full CustomToolSerializer) Co-Authored-By: Claude Opus 4.6 (1M context) --- .../prompt_studio_v2/serializers.py | 17 +++++++++++++++++ backend/prompt_studio/prompt_studio_v2/urls.py | 13 ++++++++++--- backend/prompt_studio/prompt_studio_v2/views.py | 11 ++++++++++- 3 files changed, 37 insertions(+), 4 deletions(-) diff --git a/backend/prompt_studio/prompt_studio_v2/serializers.py b/backend/prompt_studio/prompt_studio_v2/serializers.py index e1adddc33c..6a4d28032d 100644 --- a/backend/prompt_studio/prompt_studio_v2/serializers.py +++ b/backend/prompt_studio/prompt_studio_v2/serializers.py @@ -5,6 +5,23 @@ from .models import ToolStudioPrompt +class ToolStudioPromptListSerializer(serializers.ModelSerializer): + """Lightweight serializer for listing prompts by tool. + + Returns only the fields needed for linking/display without + output data or coverage calculation. + """ + + class Meta: + model = ToolStudioPrompt + fields = [ + "prompt_id", + "prompt_key", + "enforce_type", + "sequence_number", + ] + + class ToolStudioPromptSerializer(AuditSerializer): class Meta: model = ToolStudioPrompt diff --git a/backend/prompt_studio/prompt_studio_v2/urls.py b/backend/prompt_studio/prompt_studio_v2/urls.py index 23e5f02438..0ce3d03545 100644 --- a/backend/prompt_studio/prompt_studio_v2/urls.py +++ b/backend/prompt_studio/prompt_studio_v2/urls.py @@ -3,6 +3,8 @@ from .views import ToolStudioPromptView +prompt_studio_prompt_list = ToolStudioPromptView.as_view({"get": "list"}) + prompt_studio_prompt_detail = ToolStudioPromptView.as_view( { "get": "retrieve", @@ -16,15 +18,20 @@ urlpatterns = format_suffix_patterns( [ + path( + "prompt/reorder/", + reorder_prompts, + name="reorder_prompts", + ), path( "prompt//", prompt_studio_prompt_detail, name="tool-studio-prompt-detail", ), path( - "prompt/reorder/", - reorder_prompts, - name="reorder_prompts", + "prompt/", + prompt_studio_prompt_list, + name="tool-studio-prompt-list", ), ] ) diff --git a/backend/prompt_studio/prompt_studio_v2/views.py b/backend/prompt_studio/prompt_studio_v2/views.py index f120baf15f..a540480274 100644 --- a/backend/prompt_studio/prompt_studio_v2/views.py +++ b/backend/prompt_studio/prompt_studio_v2/views.py @@ -10,7 +10,10 @@ from prompt_studio.prompt_studio_v2.constants import ToolStudioPromptKeys from prompt_studio.prompt_studio_v2.controller import PromptStudioController from prompt_studio.prompt_studio_v2.models import ToolStudioPrompt -from prompt_studio.prompt_studio_v2.serializers import ToolStudioPromptSerializer +from prompt_studio.prompt_studio_v2.serializers import ( + ToolStudioPromptListSerializer, + ToolStudioPromptSerializer, +) class ToolStudioPromptView(viewsets.ModelViewSet): @@ -28,8 +31,14 @@ class ToolStudioPromptView(viewsets.ModelViewSet): versioning_class = URLPathVersioning serializer_class = ToolStudioPromptSerializer + permission_classes: list[type[PromptAcesssToUser]] = [PromptAcesssToUser] + def get_serializer_class(self): + if self.action == "list": + return ToolStudioPromptListSerializer + return ToolStudioPromptSerializer + def get_queryset(self) -> QuerySet | None: filter_args = FilterHelper.build_filter_args( self.request, From 9300a4d67db47bf78818fd5d06f0541bb0ba9238 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Thu, 2 Apr 2026 01:57:29 +0530 Subject: [PATCH 02/57] UN-2946 [FEAT] Add Look-Ups plugin integration in sidebar nav and routes - Add lookup-studio plugin detection with dynamic import - Add PromptStudioPopoverContent for hover submenu (Projects / Look-Ups) following the same Popover pattern as HITL and Platform Settings - Register lookups/* route in useMainAppRoutes.js Co-Authored-By: Claude Opus 4.6 (1M context) --- .../navigations/side-nav-bar/SideNavBar.jsx | 118 ++++++++++++++++++ frontend/src/routes/useMainAppRoutes.js | 9 ++ 2 files changed, 127 insertions(+) diff --git a/frontend/src/components/navigations/side-nav-bar/SideNavBar.jsx b/frontend/src/components/navigations/side-nav-bar/SideNavBar.jsx index bb328b1a98..3f6f43733f 100644 --- a/frontend/src/components/navigations/side-nav-bar/SideNavBar.jsx +++ b/frontend/src/components/navigations/side-nav-bar/SideNavBar.jsx @@ -89,6 +89,14 @@ try { // Plugin unavailable } +let lookupStudioEnabled = false; +try { + await import("../../../plugins/lookup-studio"); + lookupStudioEnabled = true; +} catch { + // Plugin unavailable +} + let manualReviewSettingsEnabled = false; try { await import("../../../plugins/manual-review/settings/Settings.jsx"); @@ -258,6 +266,45 @@ HITLPopoverContent.propTypes = { navigate: PropTypes.func.isRequired, }; +const PROMPT_STUDIO_MENU_ITEMS = [ + { key: "projects", label: "Projects", subPath: "/tools" }, + { key: "lookups", label: "Look-Ups", subPath: "/lookups" }, +]; + +const getActivePromptStudioKey = (orgName) => { + const currentPath = globalThis.location.pathname; + if (currentPath.startsWith(`/${orgName}/lookups`)) { + return "lookups"; + } + return "projects"; +}; + +const PromptStudioPopoverContent = ({ orgName, navigate }) => { + const activeKey = getActivePromptStudioKey(orgName); + + return ( + + ); +}; + +PromptStudioPopoverContent.propTypes = { + orgName: PropTypes.string.isRequired, + navigate: PropTypes.func.isRequired, +}; + const SideNavBar = ({ collapsed, setCollapsed }) => { const navigate = useNavigate(); const { sessionDetails } = useSessionStore(); @@ -506,6 +553,14 @@ const SideNavBar = ({ collapsed, setCollapsed }) => { }); } + // Mark Prompt Studio item for popover rendering when lookups plugin is available + if (lookupStudioEnabled && isUnstract) { + const psItem = data[0]?.subMenu?.find((el) => el.id === 1.1); + if (psItem) { + psItem.hasLookupPopover = true; + } + } + // Add HITL Review section if plugin is available and user has HITL role const isHITLRole = [ "unstract_reviewer", @@ -700,6 +755,69 @@ const SideNavBar = ({ collapsed, setCollapsed }) => { ); } + // Prompt Studio with Look-Ups popover + if (el.hasLookupPopover) { + const psContent = ( + + { + if (!el.disable) { + navigate(el.path); + } + }} + data-testid="sidebar-prompt-studio" + > + side_icon + {!collapsed && ( +
+ + {el.title} + + + {el.description} + +
+ )} + + + ); + + if (el.disable) { + return
{psContent}
; + } + + return ( + + } + trigger="hover" + placement="rightTop" + arrow={false} + overlayClassName="settings-popover-overlay" + > + {psContent} + + ); + } + return ( } /> )} + {LookupStudio && } />} } /> } /> Date: Thu, 2 Apr 2026 14:11:02 +0530 Subject: [PATCH 03/57] UN-2946 [FIX] Use .get() fallback for prompt_studio_tool in create_profile_manager Co-Authored-By: Claude Opus 4.6 (1M context) --- backend/prompt_studio/prompt_studio_core_v2/views.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/backend/prompt_studio/prompt_studio_core_v2/views.py b/backend/prompt_studio/prompt_studio_core_v2/views.py index 22c1a378bb..e936e3dcae 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/views.py @@ -824,9 +824,10 @@ def create_profile_manager(self, request: HttpRequest, pk: Any = None) -> Respon serializer = ProfileManagerSerializer(data=request.data, context=context) serializer.is_valid(raise_exception=True) # Check for the maximum number of profiles constraint - prompt_studio_tool = serializer.validated_data[ - ProfileManagerKeys.PROMPT_STUDIO_TOOL - ] + prompt_studio_tool = ( + serializer.validated_data.get(ProfileManagerKeys.PROMPT_STUDIO_TOOL) + or self.get_object() + ) profile_count = ProfileManager.objects.filter( prompt_studio_tool=prompt_studio_tool ).count() From 4765e99bc5a3d1bd17360795b7e17d468e00a5b9 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Sat, 4 Apr 2026 03:40:14 +0530 Subject: [PATCH 04/57] UN-2946 [FEAT] Add Lookups V2 OSS integration hooks for post-extraction enrichment MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Add lookup_config export in prompt_studio_helper and registry_helper via cloud plugin guard (try/except ImportError) - Store raw output in PromptStudioOutputManager, enriched in cloud LookupOutputResult — preserving both for UI tab display - Add LookupEnrichmentProtocol and plugin call in post-extraction pipeline using ExecutorPluginLoader (no-op in OSS) - Track lookup LLM usage via standard metrics pipeline (usage_kwargs with run_id/execution_id, capture_metrics) - Move webhook postprocessing from answer_prompt to pipeline - Frontend: dynamic plugin imports for LookupMenuItem, LookupIndicator, LookupOutputTabs in prompt cards; fetch lookup outputs on page load - Add scroll-to-prompt support via query param in DocumentParser Co-Authored-By: Claude Opus 4.6 (1M context) --- .gitignore | 1 + .../prompt_studio_helper.py | 35 ++++++ .../output_manager_helper.py | 35 ++++++ .../prompt_studio_registry_helper.py | 11 ++ .../document-parser/DocumentParser.jsx | 23 +++- .../custom-tools/prompt-card/Header.jsx | 24 ++++ .../prompt-card/PromptCardItems.jsx | 13 ++ .../custom-tools/prompt-card/PromptOutput.jsx | 116 +++++++++++++----- .../helpers/custom-tools/CustomToolsHelper.js | 26 ++++ frontend/src/hooks/usePromptOutput.js | 14 +++ workers/executor/executors/answer_prompt.py | 25 +--- workers/executor/executors/legacy_executor.py | 89 ++++++++++++++ .../executor/executors/plugins/protocols.py | 7 ++ 13 files changed, 366 insertions(+), 53 deletions(-) diff --git a/.gitignore b/.gitignore index f6837ce079..cd4084ecdd 100644 --- a/.gitignore +++ b/.gitignore @@ -703,3 +703,4 @@ CONTRIBUTION_GUIDE.md # MCP servers .serena +.gstack/ diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index d0ffef3114..9b76a86782 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -387,6 +387,18 @@ def _build_prompt_output( if webhook_enabled: output[TSPKeys.POSTPROCESSING_WEBHOOK_URL] = webhook_url + # Lookup config (cloud plugin hook) + try: + from pluggable_apps.lookup_v1.execution import ( + build_lookup_config_for_prompt, + ) + + lookup_config = build_lookup_config_for_prompt(prompt) + if lookup_config: + output["lookup_config"] = lookup_config + except ImportError: + pass + output[TSPKeys.EVAL_SETTINGS] = {} output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_EVALUATE] = prompt.evaluate output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_MONITOR_LLM] = [monitor_llm] @@ -798,6 +810,18 @@ def build_fetch_response_payload( if webhook_enabled: output[TSPKeys.POSTPROCESSING_WEBHOOK_URL] = webhook_url + # Lookup config (cloud plugin hook) + try: + from pluggable_apps.lookup_v1.execution import ( + build_lookup_config_for_prompt, + ) + + lookup_config = build_lookup_config_for_prompt(prompt) + if lookup_config: + output["lookup_config"] = lookup_config + except ImportError: + pass + output[TSPKeys.EVAL_SETTINGS] = {} output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_EVALUATE] = prompt.evaluate output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_MONITOR_LLM] = [monitor_llm] @@ -1893,6 +1917,17 @@ def _fetch_response( output[TSPKeys.ENABLE_POSTPROCESSING_WEBHOOK] = webhook_enabled if webhook_enabled: output[TSPKeys.POSTPROCESSING_WEBHOOK_URL] = webhook_url + # Lookup config (cloud plugin hook) + try: + from pluggable_apps.lookup_v1.execution import ( + build_lookup_config_for_prompt, + ) + + lookup_config = build_lookup_config_for_prompt(prompt) + if lookup_config: + output["lookup_config"] = lookup_config + except ImportError: + pass # Eval settings for the prompt output[TSPKeys.EVAL_SETTINGS] = {} output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_EVALUATE] = prompt.evaluate diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py index 405b91e00f..d0197e128f 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py @@ -170,6 +170,15 @@ def update_or_create_prompt_output( # TODO: use enums here output = outputs.get(prompt.prompt_key) + + # If lookup enrichment ran, structured_output contains the enriched + # value. Restore the original raw LLM output for the prompt output + # table — the enriched value lives in LookupOutputResult instead. + lookup_outputs = metadata.get("lookup_outputs", {}) + prompt_lookup = lookup_outputs.get(prompt.prompt_key) + if prompt_lookup and "original" in prompt_lookup: + output = prompt_lookup["original"] + if prompt.enforce_type in {"json", "table", "record", "line-item"}: output = json.dumps(output) eval_metrics = outputs.get(f"{prompt.prompt_key}__evaluation", []) @@ -189,6 +198,32 @@ def update_or_create_prompt_output( word_confidence_data=prompt_word_confidence_data, ) + # Persist lookup outputs if present (cloud plugin) + if prompt_lookup: + try: + from pluggable_apps.lookup_v1.models import ( + LookupOutputResult, + ) + + lookup_meta = prompt_lookup.get("meta", {}) + lookup_id = lookup_meta.get("lookup_id") + if lookup_id: + LookupOutputResult.objects.update_or_create( + prompt_output=prompt_output, + defaults={ + "lookup_definition_id": lookup_id, + "output": prompt_lookup.get("enriched", ""), + }, + ) + except ImportError: + pass + except Exception: + logger.warning( + "Failed to persist lookup output for prompt %s", + prompt.prompt_key, + exc_info=True, + ) + # Serialize the instance serializer = PromptStudioOutputSerializer(prompt_output) serialized_data.append(serializer.data) diff --git a/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py b/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py index 6ce1f72095..87d4fa0def 100644 --- a/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py +++ b/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py @@ -355,6 +355,17 @@ def frame_export_json( output[JsonSchemaKey.POSTPROCESSING_WEBHOOK_URL] = ( prompt.postprocessing_webhook_url ) + # Lookup config (cloud plugin hook) + try: + from pluggable_apps.lookup_v1.execution import ( + build_lookup_config_for_prompt, + ) + + lookup_config = build_lookup_config_for_prompt(prompt) + if lookup_config: + output["lookup_config"] = lookup_config + except ImportError: + pass # Retaining the old fields in condition # for backward compatibility. To be removed in future. if ( diff --git a/frontend/src/components/custom-tools/document-parser/DocumentParser.jsx b/frontend/src/components/custom-tools/document-parser/DocumentParser.jsx index 3d5c891e13..c5e60a2c93 100644 --- a/frontend/src/components/custom-tools/document-parser/DocumentParser.jsx +++ b/frontend/src/components/custom-tools/document-parser/DocumentParser.jsx @@ -1,5 +1,6 @@ import PropTypes from "prop-types"; import { useEffect, useRef, useState } from "react"; +import { useSearchParams } from "react-router-dom"; import "./DocumentParser.css"; import { promptType } from "../../../helpers/GetStaticData"; @@ -42,6 +43,7 @@ function DocumentParser({ const [isChallenge, setIsChallenge] = useState(false); const [allTableSettings, setAllTableSettings] = useState([]); const bottomRef = useRef(null); + const [searchParams, setSearchParams] = useSearchParams(); const { details, isSimplePromptStudio, @@ -108,6 +110,25 @@ function DocumentParser({ } }, [scrollToBottom]); + // Handle scrollTo query param for cross-linking from Lookup Studio + useEffect(() => { + const scrollToPromptId = searchParams.get("scrollTo"); + if (!scrollToPromptId || !details?.prompts?.length) { + return; + } + + const el = document.querySelector(`[data-prompt-id="${scrollToPromptId}"]`); + if (el) { + el.scrollIntoView({ behavior: "smooth", block: "center" }); + el.classList.add("highlighted-prompt"); + setTimeout(() => el.classList.remove("highlighted-prompt"), 2000); + } + + // Clear the param so it doesn't re-trigger + searchParams.delete("scrollTo"); + setSearchParams(searchParams, { replace: true }); + }, [details?.prompts]); + const promptUrl = (urlPath) => { return `/api/v1/unstract/${sessionDetails?.orgId}/prompt-studio/prompt/${urlPath}`; }; @@ -242,7 +263,7 @@ function DocumentParser({
{details?.prompts?.map((item) => { return ( -
+
, + key: "lookup", + }, + ); + } + if (isSimplePromptStudio) { dropdownItems.splice(0, 1); } diff --git a/frontend/src/components/custom-tools/prompt-card/PromptCardItems.jsx b/frontend/src/components/custom-tools/prompt-card/PromptCardItems.jsx index 90f4936a3a..40439bce55 100644 --- a/frontend/src/components/custom-tools/prompt-card/PromptCardItems.jsx +++ b/frontend/src/components/custom-tools/prompt-card/PromptCardItems.jsx @@ -30,6 +30,16 @@ try { // The component will remain null of it is not available } +let LookupIndicator; +try { + const mod = await import( + "../../../plugins/lookup-studio/prompt-card/LookupIndicator" + ); + LookupIndicator = mod.LookupIndicator; +} catch { + // Not available in OSS +} + function PromptCardItems({ promptDetails, enforceTypeList, @@ -260,6 +270,9 @@ function PromptCardItems({ + {LookupIndicator && ( + + )} {details?.enable_highlight && diff --git a/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx b/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx index 88df329744..02e552ae44 100644 --- a/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx +++ b/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx @@ -51,6 +51,16 @@ try { // The component will remain null of it is not available } +let LookupOutputTabs; +try { + const mod = await import( + "../../../plugins/lookup-studio/prompt-card/LookupOutputTabs" + ); + LookupOutputTabs = mod.LookupOutputTabs; +} catch { + // Not available in OSS +} + function PromptOutput({ promptDetails, handleRun, @@ -193,20 +203,46 @@ function PromptOutput({ "highlighted-prompt-cell" }`} > - + {LookupOutputTabs ? ( + + + + ) : ( + + )}
- + {LookupOutputTabs ? ( + + + + ) : ( + + )}
{ const data = res?.data; updatedCusTool["adapters"] = data; + + // Fetch lookup data (cloud only, fire-and-forget) + if (fetchLookupAssignments) { + const toolId = updatedCusTool["details"]?.tool_id; + fetchLookupAssignments(axiosPrivate, sessionDetails?.orgId, toolId); + if (fetchLookupOutputs) { + fetchLookupOutputs(axiosPrivate, sessionDetails?.orgId, toolId); + } + } }) .catch((err) => { setAlertDetails(handleException(err, "Failed to load the custom tool")); @@ -131,6 +154,9 @@ function CustomToolsHelper() { setDefaultCustomTool(); emptyCusToolMessages(); resetTokenUsage(); + if (resetLookupAssignments) { + resetLookupAssignments(); + } }; }, []); diff --git a/frontend/src/hooks/usePromptOutput.js b/frontend/src/hooks/usePromptOutput.js index 530cdf2bc7..0e06070b7b 100644 --- a/frontend/src/hooks/usePromptOutput.js +++ b/frontend/src/hooks/usePromptOutput.js @@ -23,6 +23,16 @@ try { // The component will remain null of it is not available } +let handleLookupOutput; +try { + const mod = await import( + "../plugins/lookup-studio/prompt-card/handleLookupOutput" + ); + handleLookupOutput = mod.handleLookupOutput; +} catch { + // Not available in OSS +} + const usePromptOutput = () => { const { sessionDetails } = useSessionStore(); const { setTokenUsage, updateTokenUsage } = useTokenUsageStore(); @@ -125,6 +135,10 @@ const usePromptOutput = () => { wordConfidenceData: item?.word_confidence_data, }; + if (handleLookupOutput && item?.lookup_outputs) { + handleLookupOutput(item.prompt_output_id, item.lookup_outputs); + } + if (item?.is_single_pass_extract && isTokenUsageForSinglePassAdded) { return; } diff --git a/workers/executor/executors/answer_prompt.py b/workers/executor/executors/answer_prompt.py index c22c8aaafd..d1eef5b3be 100644 --- a/workers/executor/executors/answer_prompt.py +++ b/workers/executor/executors/answer_prompt.py @@ -338,27 +338,4 @@ def handle_json( structured_output[prompt_key] = {} return - highlight_data = None - if enable_highlight and metadata and PSKeys.HIGHLIGHT_DATA in metadata: - highlight_data = metadata[PSKeys.HIGHLIGHT_DATA].get(prompt_key) - - processed_data = parsed_data - updated_highlight_data = None - - webhook_enabled = output.get(PSKeys.ENABLE_POSTPROCESSING_WEBHOOK, False) - if webhook_enabled: - webhook_url = output.get(PSKeys.POSTPROCESSING_WEBHOOK_URL) - processed_data, updated_highlight_data = ( - AnswerPromptService._run_webhook_postprocess( - parsed_data=parsed_data, - webhook_url=webhook_url, - highlight_data=highlight_data, - ) - ) - - structured_output[prompt_key] = processed_data - - if enable_highlight and metadata and updated_highlight_data is not None: - metadata.setdefault(PSKeys.HIGHLIGHT_DATA, {})[prompt_key] = ( - updated_highlight_data - ) + structured_output[prompt_key] = parsed_data diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index 18e1f47749..0b06b24bc2 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -1654,6 +1654,15 @@ def _execute_single_prompt( ) shim.stream_log(f"Applied type conversion for: {prompt_name}") + self._run_post_extraction_pipeline( + output=output, + structured_output=structured_output, + metadata=metadata, + metrics=metrics, + shim=shim, + usage_kwargs=usage_kwargs, + ) + self._run_challenge_if_enabled( tool_settings=tool_settings, output=output, @@ -1844,6 +1853,86 @@ def _run_line_item_extraction( level=LogLevel.ERROR, ) + def _run_post_extraction_pipeline( + self, + output: dict[str, Any], + structured_output: dict[str, Any], + metadata: dict[str, Any], + metrics: dict[str, Any], + shim: Any, + usage_kwargs: dict[str, Any] | None = None, + ) -> None: + """Post-extraction pipeline: lookup enrichment, webhook postprocessing. + + Runs after type conversion, before challenge/evaluation. + """ + from executor.executors.answer_prompt import AnswerPromptService + from executor.executors.constants import PromptServiceConstants as PSKeys + from executor.executors.plugins import ExecutorPluginLoader + + prompt_name = output[PSKeys.NAME] + current_value = structured_output.get(prompt_name) + + # Step 1: Lookup enrichment (cloud plugin) + lookup_config = output.get("lookup_config") + lookup_cls = ExecutorPluginLoader.get("lookup-enrichment") + if lookup_config and current_value is not None and lookup_cls: + _, _, _, _, llm_cls, _, _ = self._get_prompt_deps() + llm_adapter_id = lookup_config.get("llm_adapter_id", "") + llm = llm_cls( + adapter_instance_id=llm_adapter_id, + tool=shim, + usage_kwargs={ + **(usage_kwargs or {}), + PSKeys.LLM_USAGE_REASON: "lookup", + }, + capture_metrics=True, + ) + + enricher = lookup_cls( + current_value=current_value, + lookup_config=lookup_config, + structured_output=structured_output, + llm=llm, + shim=shim, + ) + lookup_result = enricher.run() + + if lookup_result is not None: + metadata.setdefault("lookup_outputs", {})[prompt_name] = { + "original": str(current_value), + "enriched": lookup_result, + "meta": { + "lookup_id": lookup_config.get("lookup_id", ""), + "lookup_name": lookup_config.get("lookup_name", ""), + }, + } + structured_output[prompt_name] = lookup_result + shim.stream_log(f"Lookup enrichment complete for: {prompt_name}") + + metrics.setdefault(prompt_name, {})[f"{llm.get_usage_reason()}_llm"] = ( + llm.get_metrics() + ) + + # Step 2: Webhook postprocessing (JSON only, moved from handle_json) + output_type = output.get(PSKeys.TYPE, "") + webhook_enabled = output.get(PSKeys.ENABLE_POSTPROCESSING_WEBHOOK, False) + if webhook_enabled and output_type == PSKeys.JSON: + webhook_url = output.get(PSKeys.POSTPROCESSING_WEBHOOK_URL) + highlight_data = None + if metadata and PSKeys.HIGHLIGHT_DATA in metadata: + highlight_data = metadata.get(PSKeys.HIGHLIGHT_DATA, {}).get(prompt_name) + processed, updated_highlights = AnswerPromptService._run_webhook_postprocess( + parsed_data=structured_output.get(prompt_name), + webhook_url=webhook_url, + highlight_data=highlight_data, + ) + structured_output[prompt_name] = processed + if updated_highlights is not None and metadata: + metadata.setdefault(PSKeys.HIGHLIGHT_DATA, {})[prompt_name] = ( + updated_highlights + ) + @staticmethod def _apply_type_conversion( output: dict[str, Any], diff --git a/workers/executor/executors/plugins/protocols.py b/workers/executor/executors/plugins/protocols.py index fb4d676b37..a5ed148a5b 100644 --- a/workers/executor/executors/plugins/protocols.py +++ b/workers/executor/executors/plugins/protocols.py @@ -49,3 +49,10 @@ class EvaluationProtocol(Protocol): """Legacy executor: prompt evaluation.""" def run(self, **kwargs: Any) -> dict: ... + + +@runtime_checkable +class LookupEnrichmentProtocol(Protocol): + """Legacy executor: post-extraction lookup enrichment.""" + + def run(self) -> str | None: ... From 68e394ba6b0a455e6ea9919d7a571eb1abd6fac7 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Sat, 4 Apr 2026 03:51:31 +0530 Subject: [PATCH 05/57] UN-2946 Removed unnecessary gitignore --- .gitignore | 1 - 1 file changed, 1 deletion(-) diff --git a/.gitignore b/.gitignore index cd4084ecdd..f6837ce079 100644 --- a/.gitignore +++ b/.gitignore @@ -703,4 +703,3 @@ CONTRIBUTION_GUIDE.md # MCP servers .serena -.gstack/ From 059ceed1ff028be85cc0a7a86f263685e33f2915 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Sat, 4 Apr 2026 04:37:04 +0530 Subject: [PATCH 06/57] UN-2946 [REFACTOR] Deduplicate lookup config helper and add lookup usage reason - Extract get_lookup_config() to prompt_studio/lookup_utils.py, replacing 4 identical try/except ImportError blocks across prompt_studio_helper and prompt_studio_registry_helper - Add LOOKUP to LLMUsageReason choices (was missing, causing invalid choice on usage records from lookup enrichment LLM calls) - Migration: usage_v2/0004_add_lookup_usage_reason Co-Authored-By: Claude Opus 4.6 (1M context) --- backend/prompt_studio/lookup_utils.py | 26 ++++++++++++ .../prompt_studio_helper.py | 40 ++++--------------- .../prompt_studio_registry_helper.py | 14 ++----- .../0004_add_lookup_usage_reason.py | 28 +++++++++++++ backend/usage_v2/models.py | 1 + 5 files changed, 65 insertions(+), 44 deletions(-) create mode 100644 backend/prompt_studio/lookup_utils.py create mode 100644 backend/usage_v2/migrations/0004_add_lookup_usage_reason.py diff --git a/backend/prompt_studio/lookup_utils.py b/backend/prompt_studio/lookup_utils.py new file mode 100644 index 0000000000..f2e889f50f --- /dev/null +++ b/backend/prompt_studio/lookup_utils.py @@ -0,0 +1,26 @@ +"""Shared utility for lookup config resolution. + +Wraps the cloud-only build_lookup_config_for_prompt call so that +OSS callers don't repeat the try/except ImportError guard. +""" + +import logging + +logger = logging.getLogger(__name__) + + +def get_lookup_config(prompt) -> dict | None: + """Return lookup config for a prompt, or None if lookups are unavailable. + + This is a thin wrapper around the cloud plugin's + build_lookup_config_for_prompt. In OSS deployments where the plugin + is absent, it returns None silently. + """ + try: + from pluggable_apps.lookup_v1.execution import ( + build_lookup_config_for_prompt, + ) + + return build_lookup_config_for_prompt(prompt) + except ImportError: + return None diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index 9b76a86782..6744171ed4 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -20,6 +20,7 @@ from utils.local_context import StateStore from backend.celery_service import app as celery_app +from prompt_studio.lookup_utils import get_lookup_config from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_profile_manager_v2.profile_manager_helper import ( ProfileManagerHelper, @@ -387,17 +388,8 @@ def _build_prompt_output( if webhook_enabled: output[TSPKeys.POSTPROCESSING_WEBHOOK_URL] = webhook_url - # Lookup config (cloud plugin hook) - try: - from pluggable_apps.lookup_v1.execution import ( - build_lookup_config_for_prompt, - ) - - lookup_config = build_lookup_config_for_prompt(prompt) - if lookup_config: - output["lookup_config"] = lookup_config - except ImportError: - pass + if lookup_config := get_lookup_config(prompt): + output["lookup_config"] = lookup_config output[TSPKeys.EVAL_SETTINGS] = {} output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_EVALUATE] = prompt.evaluate @@ -810,17 +802,8 @@ def build_fetch_response_payload( if webhook_enabled: output[TSPKeys.POSTPROCESSING_WEBHOOK_URL] = webhook_url - # Lookup config (cloud plugin hook) - try: - from pluggable_apps.lookup_v1.execution import ( - build_lookup_config_for_prompt, - ) - - lookup_config = build_lookup_config_for_prompt(prompt) - if lookup_config: - output["lookup_config"] = lookup_config - except ImportError: - pass + if lookup_config := get_lookup_config(prompt): + output["lookup_config"] = lookup_config output[TSPKeys.EVAL_SETTINGS] = {} output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_EVALUATE] = prompt.evaluate @@ -1917,17 +1900,8 @@ def _fetch_response( output[TSPKeys.ENABLE_POSTPROCESSING_WEBHOOK] = webhook_enabled if webhook_enabled: output[TSPKeys.POSTPROCESSING_WEBHOOK_URL] = webhook_url - # Lookup config (cloud plugin hook) - try: - from pluggable_apps.lookup_v1.execution import ( - build_lookup_config_for_prompt, - ) - - lookup_config = build_lookup_config_for_prompt(prompt) - if lookup_config: - output["lookup_config"] = lookup_config - except ImportError: - pass + if lookup_config := get_lookup_config(prompt): + output["lookup_config"] = lookup_config # Eval settings for the prompt output[TSPKeys.EVAL_SETTINGS] = {} output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_EVALUATE] = prompt.evaluate diff --git a/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py b/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py index 87d4fa0def..c6cb80bb07 100644 --- a/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py +++ b/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py @@ -7,6 +7,7 @@ from django.db import IntegrityError from plugins import get_plugin +from prompt_studio.lookup_utils import get_lookup_config from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.models import CustomTool from prompt_studio.prompt_studio_core_v2.prompt_studio_helper import PromptStudioHelper @@ -355,17 +356,8 @@ def frame_export_json( output[JsonSchemaKey.POSTPROCESSING_WEBHOOK_URL] = ( prompt.postprocessing_webhook_url ) - # Lookup config (cloud plugin hook) - try: - from pluggable_apps.lookup_v1.execution import ( - build_lookup_config_for_prompt, - ) - - lookup_config = build_lookup_config_for_prompt(prompt) - if lookup_config: - output["lookup_config"] = lookup_config - except ImportError: - pass + if lookup_config := get_lookup_config(prompt): + output["lookup_config"] = lookup_config # Retaining the old fields in condition # for backward compatibility. To be removed in future. if ( diff --git a/backend/usage_v2/migrations/0004_add_lookup_usage_reason.py b/backend/usage_v2/migrations/0004_add_lookup_usage_reason.py new file mode 100644 index 0000000000..76c628d9ee --- /dev/null +++ b/backend/usage_v2/migrations/0004_add_lookup_usage_reason.py @@ -0,0 +1,28 @@ +# Generated by Django 4.2.1 on 2026-04-03 22:44 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("usage_v2", "0003_usage_usage_executi_4deb35_idx"), + ] + + operations = [ + migrations.AlterField( + model_name="usage", + name="llm_usage_reason", + field=models.CharField( + blank=True, + choices=[ + ("extraction", "Extraction"), + ("challenge", "Challenge"), + ("summarize", "Summarize"), + ("lookup", "Lookup"), + ], + db_comment="Reason for LLM usage. Empty if usage_type is 'embedding'. ", + max_length=255, + null=True, + ), + ), + ] diff --git a/backend/usage_v2/models.py b/backend/usage_v2/models.py index 8da3d751ba..dc983c56b5 100644 --- a/backend/usage_v2/models.py +++ b/backend/usage_v2/models.py @@ -17,6 +17,7 @@ class LLMUsageReason(models.TextChoices): EXTRACTION = "extraction", "Extraction" CHALLENGE = "challenge", "Challenge" SUMMARIZE = "summarize", "Summarize" + LOOKUP = "lookup", "Lookup" class UsageModelManager(DefaultOrganizationManagerMixin, models.Manager): From 7699441e45d6b9fb362a2668d7aeca41a43ffe52 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Sun, 5 Apr 2026 22:37:00 +0530 Subject: [PATCH 07/57] UN-2946 [FEAT] Add get_last_usage() to SDK1 LLM for token tracking Store prompt/completion/total token counts from the most recent complete() call on the LLM object itself, making usage data queryable without relying on the Audit pipeline roundtrip. Co-Authored-By: Claude Opus 4.6 (1M context) --- unstract/sdk1/src/unstract/sdk1/llm.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index c1730e6613..f751e7f87c 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -211,6 +211,7 @@ def __init__( # noqa: C901 if capture_metrics_from_platform is not None: self._capture_metrics = capture_metrics_from_platform self._metrics: dict[str, object] = {} + self._last_usage: Mapping[str, int] = {} def _get_adapter_info(self) -> str: """Build a display string identifying this adapter for errors.""" @@ -552,6 +553,10 @@ def get_model_name(self) -> str: def get_metrics(self) -> dict[str, object]: return self._metrics + def get_last_usage(self) -> Mapping[str, int]: + """Token usage from the most recent complete() call.""" + return self._last_usage + def get_usage_reason(self) -> object: return self.platform_kwargs.get("llm_usage_reason") @@ -573,6 +578,12 @@ def _record_usage( logger.info(f"[sdk1][LLM][{model}][{llm_api}] Prompt Tokens: {prompt_tokens}") logger.info(f"[sdk1][LLM][{model}][{llm_api}] LLM Usage: {all_tokens}") + self._last_usage = { + "prompt_tokens": all_tokens.prompt_llm_token_count, + "completion_tokens": all_tokens.completion_llm_token_count, + "total_tokens": all_tokens.total_llm_token_count, + } + Audit().push_usage_data( platform_api_key=self._platform_api_key, token_counter=all_tokens, @@ -963,6 +974,10 @@ def get_metrics(self) -> dict[str, object]: """Get captured metrics.""" return self._llm_instance.get_metrics() + def get_last_usage(self) -> Mapping[str, int]: + """Token usage from the most recent complete() call.""" + return self._llm_instance.get_last_usage() + def get_usage_reason(self) -> object: """Get usage reason from platform kwargs.""" return self._llm_instance.get_usage_reason() From 6e368961a5ddf821b1c4fc3bc3bdfcfad6cf8cb1 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Mon, 6 Apr 2026 00:16:19 +0530 Subject: [PATCH 08/57] UN-2946 [REFACTOR] Split post-extraction pipeline into lookup and webhook methods Move lookup result-application logic to the cloud plugin, matching the challenge plugin pattern where the plugin owns metadata mutation. Co-Authored-By: Claude Opus 4.6 (1M context) --- workers/executor/executors/legacy_executor.py | 120 +++++++++--------- 1 file changed, 62 insertions(+), 58 deletions(-) diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index 0b06b24bc2..7ccdcb656d 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -1654,7 +1654,7 @@ def _execute_single_prompt( ) shim.stream_log(f"Applied type conversion for: {prompt_name}") - self._run_post_extraction_pipeline( + self._run_lookup_enrichment( output=output, structured_output=structured_output, metadata=metadata, @@ -1662,6 +1662,11 @@ def _execute_single_prompt( shim=shim, usage_kwargs=usage_kwargs, ) + self._run_webhook_postprocessing( + output=output, + structured_output=structured_output, + metadata=metadata, + ) self._run_challenge_if_enabled( tool_settings=tool_settings, @@ -1853,7 +1858,7 @@ def _run_line_item_extraction( level=LogLevel.ERROR, ) - def _run_post_extraction_pipeline( + def _run_lookup_enrichment( self, output: dict[str, Any], structured_output: dict[str, Any], @@ -1862,76 +1867,75 @@ def _run_post_extraction_pipeline( shim: Any, usage_kwargs: dict[str, Any] | None = None, ) -> None: - """Post-extraction pipeline: lookup enrichment, webhook postprocessing. - - Runs after type conversion, before challenge/evaluation. - """ - from executor.executors.answer_prompt import AnswerPromptService + """Run lookup enrichment plugin if enabled and available.""" from executor.executors.constants import PromptServiceConstants as PSKeys from executor.executors.plugins import ExecutorPluginLoader prompt_name = output[PSKeys.NAME] current_value = structured_output.get(prompt_name) - # Step 1: Lookup enrichment (cloud plugin) lookup_config = output.get("lookup_config") lookup_cls = ExecutorPluginLoader.get("lookup-enrichment") - if lookup_config and current_value is not None and lookup_cls: - _, _, _, _, llm_cls, _, _ = self._get_prompt_deps() - llm_adapter_id = lookup_config.get("llm_adapter_id", "") - llm = llm_cls( - adapter_instance_id=llm_adapter_id, - tool=shim, - usage_kwargs={ - **(usage_kwargs or {}), - PSKeys.LLM_USAGE_REASON: "lookup", - }, - capture_metrics=True, - ) + if not (lookup_config and current_value is not None and lookup_cls): + return - enricher = lookup_cls( - current_value=current_value, - lookup_config=lookup_config, - structured_output=structured_output, - llm=llm, - shim=shim, - ) - lookup_result = enricher.run() - - if lookup_result is not None: - metadata.setdefault("lookup_outputs", {})[prompt_name] = { - "original": str(current_value), - "enriched": lookup_result, - "meta": { - "lookup_id": lookup_config.get("lookup_id", ""), - "lookup_name": lookup_config.get("lookup_name", ""), - }, - } - structured_output[prompt_name] = lookup_result - shim.stream_log(f"Lookup enrichment complete for: {prompt_name}") + _, _, _, _, llm_cls, _, _ = self._get_prompt_deps() + llm_adapter_id = lookup_config.get("llm_adapter_id", "") + llm = llm_cls( + adapter_instance_id=llm_adapter_id, + tool=shim, + usage_kwargs={ + **(usage_kwargs or {}), + PSKeys.LLM_USAGE_REASON: "lookup", + }, + capture_metrics=True, + ) - metrics.setdefault(prompt_name, {})[f"{llm.get_usage_reason()}_llm"] = ( - llm.get_metrics() - ) + enricher = lookup_cls( + current_value=current_value, + lookup_config=lookup_config, + structured_output=structured_output, + llm=llm, + shim=shim, + metadata=metadata, + prompt_name=prompt_name, + ) + enricher.run() + + metrics.setdefault(prompt_name, {})[f"{llm.get_usage_reason()}_llm"] = ( + llm.get_metrics() + ) - # Step 2: Webhook postprocessing (JSON only, moved from handle_json) + @staticmethod + def _run_webhook_postprocessing( + output: dict[str, Any], + structured_output: dict[str, Any], + metadata: dict[str, Any], + ) -> None: + """Run webhook postprocessing if enabled (JSON outputs only).""" + from executor.executors.answer_prompt import AnswerPromptService + from executor.executors.constants import PromptServiceConstants as PSKeys + + prompt_name = output[PSKeys.NAME] output_type = output.get(PSKeys.TYPE, "") webhook_enabled = output.get(PSKeys.ENABLE_POSTPROCESSING_WEBHOOK, False) - if webhook_enabled and output_type == PSKeys.JSON: - webhook_url = output.get(PSKeys.POSTPROCESSING_WEBHOOK_URL) - highlight_data = None - if metadata and PSKeys.HIGHLIGHT_DATA in metadata: - highlight_data = metadata.get(PSKeys.HIGHLIGHT_DATA, {}).get(prompt_name) - processed, updated_highlights = AnswerPromptService._run_webhook_postprocess( - parsed_data=structured_output.get(prompt_name), - webhook_url=webhook_url, - highlight_data=highlight_data, + if not (webhook_enabled and output_type == PSKeys.JSON): + return + + webhook_url = output.get(PSKeys.POSTPROCESSING_WEBHOOK_URL) + highlight_data = None + if metadata and PSKeys.HIGHLIGHT_DATA in metadata: + highlight_data = metadata.get(PSKeys.HIGHLIGHT_DATA, {}).get(prompt_name) + processed, updated_highlights = AnswerPromptService._run_webhook_postprocess( + parsed_data=structured_output.get(prompt_name), + webhook_url=webhook_url, + highlight_data=highlight_data, + ) + structured_output[prompt_name] = processed + if updated_highlights is not None and metadata: + metadata.setdefault(PSKeys.HIGHLIGHT_DATA, {})[prompt_name] = ( + updated_highlights ) - structured_output[prompt_name] = processed - if updated_highlights is not None and metadata: - metadata.setdefault(PSKeys.HIGHLIGHT_DATA, {})[prompt_name] = ( - updated_highlights - ) @staticmethod def _apply_type_conversion( From 5b8c06dfa6db8f737fea7ae44c460e7673f4e4e9 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Mon, 6 Apr 2026 02:46:32 +0530 Subject: [PATCH 09/57] UN-2946 [FEAT] Generic async extraction callbacks and WebSocket transport fallback Add reusable extraction_complete/extraction_error callback tasks to the ide_callback worker, replacing the need for Django-based celery workers for text extraction. Add ExtractionAPIClient for internal API calls. Add polling fallback to WebSocket transport for local dev reliability. Co-Authored-By: Claude Opus 4.6 (1M context) --- backend/utils/websocket_views.py | 2 +- frontend/src/helpers/SocketContext.js | 2 +- workers/ide_callback/tasks.py | 159 ++++++++++++++++++++ workers/shared/clients/extraction_client.py | 63 ++++++++ 4 files changed, 224 insertions(+), 2 deletions(-) create mode 100644 workers/shared/clients/extraction_client.py diff --git a/backend/utils/websocket_views.py b/backend/utils/websocket_views.py index a57521ca43..51998399c4 100644 --- a/backend/utils/websocket_views.py +++ b/backend/utils/websocket_views.py @@ -63,7 +63,7 @@ def emit_websocket(request): # Emit the WebSocket event _emit_websocket_event(room=room, event=event, data=message_data) - logger.debug(f"WebSocket event emitted: room={room}, event={event}") + logger.info(f"WebSocket event emitted: room={room}, event={event}") return JsonResponse( { diff --git a/frontend/src/helpers/SocketContext.js b/frontend/src/helpers/SocketContext.js index 6e6ace9a63..ec6a274e9e 100644 --- a/frontend/src/helpers/SocketContext.js +++ b/frontend/src/helpers/SocketContext.js @@ -16,7 +16,7 @@ const SocketProvider = ({ children }) => { // This ensures session cookies are sent (same-origin) and avoids // cross-origin WebSocket issues. const newSocket = io(getBaseUrl(), { - transports: ["websocket"], + transports: ["websocket", "polling"], path: "/api/v1/socket", }); setSocket(newSocket); diff --git a/workers/ide_callback/tasks.py b/workers/ide_callback/tasks.py index cb610678b1..e2b3b4dcf6 100644 --- a/workers/ide_callback/tasks.py +++ b/workers/ide_callback/tasks.py @@ -504,3 +504,162 @@ def ide_prompt_error( ) except Exception: logger.exception("ide_prompt_error callback failed") + + +# ------------------------------------------------------------------ +# Generic Text Extraction Callbacks +# +# Reusable extraction callbacks that route based on ``source`` in +# callback_kwargs (e.g. "lookup", future "prompt_studio"). +# ------------------------------------------------------------------ + + +def _get_extraction_client(): + from shared.clients.extraction_client import ExtractionAPIClient + + return ExtractionAPIClient() + + +@app.task(name="extraction_complete") +def extraction_complete( + result_dict: dict[str, Any], + callback_kwargs: dict[str, Any] | None = None, +) -> dict[str, Any]: + """Celery link callback after successful text extraction. + + Computes token count from extracted text, persists result via + internal API, and emits a WebSocket event. + """ + cb = callback_kwargs or {} + source = cb.get("source", "") + file_id = cb.get("file_id", "") + org_id = cb.get("org_id", "") + extracted_text_path = cb.get("extracted_text_path", "") + ws_room = cb.get("ws_room", "") + ws_event = cb.get("ws_event", "") + + api = _get_extraction_client() + ps_api = _get_api_client() + + try: + # Check executor-level failure + if not result_dict.get("success", False): + error_msg = result_dict.get("error", "Unknown executor error") + logger.error( + "extraction executor reported failure: source=%s file=%s error=%s", + source, + file_id, + error_msg, + ) + api.mark_extraction_error( + source=source, + file_id=file_id, + error=error_msg, + organization_id=org_id, + ) + if ws_room and ws_event: + _emit_websocket( + ps_api, + room=ws_room, + event=ws_event, + data={ + "file_id": file_id, + "status": "ERROR", + "error": error_msg[:500], + }, + ) + return {"status": "failed", "error": error_msg} + + extracted_text = result_dict.get("data", {}).get("extracted_text", "") + token_count = len(extracted_text) // 4 + + api.mark_extraction_complete( + source=source, + file_id=file_id, + token_count=token_count, + extracted_text_path=extracted_text_path, + organization_id=org_id, + ) + + if ws_room and ws_event: + _emit_websocket( + ps_api, + room=ws_room, + event=ws_event, + data={ + "file_id": file_id, + "status": "COMPLETED", + "token_count": token_count, + }, + ) + + logger.info( + "Extraction completed: source=%s file=%s tokens=%d", + source, + file_id, + token_count, + ) + return {"status": "completed", "file_id": file_id, "token_count": token_count} + + except Exception as e: + logger.exception( + "extraction_complete callback failed: source=%s file=%s", source, file_id + ) + if ws_room and ws_event: + try: + _emit_websocket( + ps_api, + room=ws_room, + event=ws_event, + data={ + "file_id": file_id, + "status": "ERROR", + "error": str(e)[:500], + }, + ) + except Exception: + pass + raise + + +@app.task(name="extraction_error") +def extraction_error( + failed_task_id: str, + callback_kwargs: dict[str, Any] | None = None, +) -> None: + """Celery link_error callback when an extraction task fails.""" + cb = callback_kwargs or {} + source = cb.get("source", "") + file_id = cb.get("file_id", "") + org_id = cb.get("org_id", "") + ws_room = cb.get("ws_room", "") + ws_event = cb.get("ws_event", "") + + api = _get_extraction_client() + ps_api = _get_api_client() + + try: + error_msg = _get_task_error(failed_task_id, default="Text extraction failed") + + api.mark_extraction_error( + source=source, + file_id=file_id, + error=error_msg, + organization_id=org_id, + ) + + if ws_room and ws_event: + _emit_websocket( + ps_api, + room=ws_room, + event=ws_event, + data={ + "file_id": file_id, + "status": "ERROR", + "error": error_msg[:500], + }, + ) + except Exception: + logger.exception( + "extraction_error callback failed: source=%s file=%s", source, file_id + ) diff --git a/workers/shared/clients/extraction_client.py b/workers/shared/clients/extraction_client.py new file mode 100644 index 0000000000..db6781be62 --- /dev/null +++ b/workers/shared/clients/extraction_client.py @@ -0,0 +1,63 @@ +"""Extraction API Client for text extraction callbacks. + +Used by the ide_callback worker to persist extraction results +through the backend's internal API endpoints. +""" + +import logging +from typing import Any + +from .base_client import BaseAPIClient + +logger = logging.getLogger(__name__) + +_EXTRACTION_COMPLETE_ENDPOINT = "v1/extraction/extraction-complete/" +_EXTRACTION_ERROR_ENDPOINT = "v1/extraction/extraction-error/" + + +class ExtractionAPIClient(BaseAPIClient): + """API client for generic text extraction callback endpoints.""" + + def mark_extraction_complete( + self, + source: str, + file_id: str, + token_count: int, + extracted_text_path: str, + organization_id: str | None = None, + **extra: Any, + ) -> dict[str, Any]: + """Notify backend that extraction succeeded.""" + payload: dict[str, Any] = { + "source": source, + "file_id": file_id, + "token_count": token_count, + "extracted_text_path": extracted_text_path, + **extra, + } + return self.post( + _EXTRACTION_COMPLETE_ENDPOINT, + data=payload, + organization_id=organization_id, + ) + + def mark_extraction_error( + self, + source: str, + file_id: str, + error: str, + organization_id: str | None = None, + **extra: Any, + ) -> dict[str, Any]: + """Notify backend that extraction failed.""" + payload: dict[str, Any] = { + "source": source, + "file_id": file_id, + "error": error, + **extra, + } + return self.post( + _EXTRACTION_ERROR_ENDPOINT, + data=payload, + organization_id=organization_id, + ) From df4956945ced453861bb540d330f28aa26d96ff4 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Mon, 6 Apr 2026 03:24:17 +0530 Subject: [PATCH 10/57] Reduce success notification duration from 2s to 1s for less intrusive UX --- frontend/src/store/alert-store.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/src/store/alert-store.js b/frontend/src/store/alert-store.js index 24549e09ce..5c5cf5c16d 100644 --- a/frontend/src/store/alert-store.js +++ b/frontend/src/store/alert-store.js @@ -4,7 +4,7 @@ import { create } from "zustand"; import { isNonNegativeNumber } from "../helpers/GetStaticData"; const DEFAULT_DURATION = 6; -const SUCCESS_DURATION = 2; +const SUCCESS_DURATION = 1; const STORE_VARIABLES = { alertDetails: { From 5e39f704ba2d13fa33a93814e75d358d1a65763b Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Mon, 6 Apr 2026 03:27:33 +0530 Subject: [PATCH 11/57] Revert "Reduce success notification duration from 2s to 1s for less intrusive UX" This reverts commit d6e136deefc3a8a5f90dc2007a4ba5c4d7effdf9. --- frontend/src/store/alert-store.js | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/src/store/alert-store.js b/frontend/src/store/alert-store.js index 5c5cf5c16d..24549e09ce 100644 --- a/frontend/src/store/alert-store.js +++ b/frontend/src/store/alert-store.js @@ -4,7 +4,7 @@ import { create } from "zustand"; import { isNonNegativeNumber } from "../helpers/GetStaticData"; const DEFAULT_DURATION = 6; -const SUCCESS_DURATION = 1; +const SUCCESS_DURATION = 2; const STORE_VARIABLES = { alertDetails: { From e4c023eb5f9b1c7c96fb8387dc9783163b52c709 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Tue, 7 Apr 2026 04:52:27 +0530 Subject: [PATCH 12/57] UN-2946 [REFACTOR] Pluggable lookup export validation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace inline DRAFT lookup check with pluggable cloud-only hook. Uses try/except ImportError pattern — zero lookup code in OSS. Collects all DRAFT lookups in one pass with markdown-linked error messages. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../prompt_studio_registry_helper.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py b/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py index c6cb80bb07..bb7841c108 100644 --- a/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py +++ b/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py @@ -7,7 +7,6 @@ from django.db import IntegrityError from plugins import get_plugin -from prompt_studio.lookup_utils import get_lookup_config from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.models import CustomTool from prompt_studio.prompt_studio_core_v2.prompt_studio_helper import PromptStudioHelper @@ -298,6 +297,19 @@ def frame_export_json( settings, JsonSchemaKey.WORD_CONFIDENCE_POSTAMBLE.upper(), "" ) + # Validate lookup assignments (cloud-only, no-op in OSS) + lookup_configs = {} + try: + from pluggable_apps.lookup_v1.validation import ( + validate_lookups_for_export, + ) + + lookup_configs, lookup_error = validate_lookups_for_export(prompts) + if lookup_error: + raise InValidCustomToolError(lookup_error) + except ImportError: + pass + for prompt in prompts: if prompt.prompt_type == JsonSchemaKey.NOTES or not prompt.active: continue @@ -356,8 +368,9 @@ def frame_export_json( output[JsonSchemaKey.POSTPROCESSING_WEBHOOK_URL] = ( prompt.postprocessing_webhook_url ) - if lookup_config := get_lookup_config(prompt): - output["lookup_config"] = lookup_config + prompt_id_str = str(prompt.prompt_id) + if prompt_id_str in lookup_configs: + output["lookup_config"] = lookup_configs[prompt_id_str] # Retaining the old fields in condition # for backward compatibility. To be removed in future. if ( From 9bf072e8e7fb7918cf25ea11a9be19509b4c412c Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Tue, 7 Apr 2026 16:52:32 +0530 Subject: [PATCH 13/57] UN-2946 [REFACTOR] Lookups V2 review cleanup - Consolidate cloud imports into lookup_utils.py with persist_lookup_output() and validate_lookups_for_export() wrappers - Fix LookupEnrichmentProtocol.run() return type to None matching challenge/evaluation pattern - Revert logger.info to logger.debug in websocket_views.py - Eliminate duplicated LookupOutputTabs ternary with renderWithLookupWrapper helper - Move lookups menu constants from SideNavBar.jsx to cloud plugin - Harden DocumentParser.jsx scrollTo with UUID validation and fix useEffect dependency - Revert SocketContext transport to ["websocket"] Co-Authored-By: Claude Opus 4.6 (1M context) --- backend/prompt_studio/lookup_utils.py | 44 ++++++++--- .../output_manager_helper.py | 20 +---- .../prompt_studio_registry_helper.py | 15 +--- backend/utils/websocket_views.py | 2 +- .../document-parser/DocumentParser.jsx | 10 ++- .../custom-tools/prompt-card/PromptOutput.jsx | 75 +++++++------------ .../navigations/side-nav-bar/SideNavBar.jsx | 43 +---------- frontend/src/helpers/SocketContext.js | 2 +- .../executor/executors/plugins/protocols.py | 2 +- 9 files changed, 81 insertions(+), 132 deletions(-) diff --git a/backend/prompt_studio/lookup_utils.py b/backend/prompt_studio/lookup_utils.py index f2e889f50f..823d9fc015 100644 --- a/backend/prompt_studio/lookup_utils.py +++ b/backend/prompt_studio/lookup_utils.py @@ -1,7 +1,7 @@ -"""Shared utility for lookup config resolution. +"""Shared utility for lookup operations. -Wraps the cloud-only build_lookup_config_for_prompt call so that -OSS callers don't repeat the try/except ImportError guard. +Wraps cloud-only lookup calls so that OSS callers don't repeat +the try/except ImportError guard. All functions are no-ops in OSS. """ import logging @@ -10,12 +10,7 @@ def get_lookup_config(prompt) -> dict | None: - """Return lookup config for a prompt, or None if lookups are unavailable. - - This is a thin wrapper around the cloud plugin's - build_lookup_config_for_prompt. In OSS deployments where the plugin - is absent, it returns None silently. - """ + """Return lookup config for a prompt, or None if lookups are unavailable.""" try: from pluggable_apps.lookup_v1.execution import ( build_lookup_config_for_prompt, @@ -24,3 +19,34 @@ def get_lookup_config(prompt) -> dict | None: return build_lookup_config_for_prompt(prompt) except ImportError: return None + + +def persist_lookup_output(prompt_output, prompt_lookup: dict) -> None: + """Persist lookup enrichment result. No-op in OSS.""" + try: + from pluggable_apps.lookup_v1.models import LookupOutputResult + + lookup_meta = prompt_lookup.get("meta", {}) + lookup_id = lookup_meta.get("lookup_id") + if lookup_id: + LookupOutputResult.objects.update_or_create( + prompt_output=prompt_output, + defaults={ + "lookup_definition_id": lookup_id, + "output": prompt_lookup.get("enriched", ""), + }, + ) + except ImportError: + pass + + +def validate_lookups_for_export(prompts) -> tuple[dict, str | None]: + """Validate lookup assignments before export. Returns ({}, None) in OSS.""" + try: + from pluggable_apps.lookup_v1.validation import ( + validate_lookups_for_export as _validate, + ) + + return _validate(prompts) + except ImportError: + return {}, None diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py index d0197e128f..0fb92c126d 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py @@ -4,6 +4,7 @@ from django.core.exceptions import ObjectDoesNotExist +from prompt_studio.lookup_utils import persist_lookup_output from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.exceptions import ( AnswerFetchError, @@ -198,25 +199,10 @@ def update_or_create_prompt_output( word_confidence_data=prompt_word_confidence_data, ) - # Persist lookup outputs if present (cloud plugin) + # Persist lookup outputs if present (cloud plugin, no-op in OSS) if prompt_lookup: try: - from pluggable_apps.lookup_v1.models import ( - LookupOutputResult, - ) - - lookup_meta = prompt_lookup.get("meta", {}) - lookup_id = lookup_meta.get("lookup_id") - if lookup_id: - LookupOutputResult.objects.update_or_create( - prompt_output=prompt_output, - defaults={ - "lookup_definition_id": lookup_id, - "output": prompt_lookup.get("enriched", ""), - }, - ) - except ImportError: - pass + persist_lookup_output(prompt_output, prompt_lookup) except Exception: logger.warning( "Failed to persist lookup output for prompt %s", diff --git a/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py b/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py index bb7841c108..92cfb6e160 100644 --- a/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py +++ b/backend/prompt_studio/prompt_studio_registry_v2/prompt_studio_registry_helper.py @@ -7,6 +7,7 @@ from django.db import IntegrityError from plugins import get_plugin +from prompt_studio.lookup_utils import validate_lookups_for_export from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.models import CustomTool from prompt_studio.prompt_studio_core_v2.prompt_studio_helper import PromptStudioHelper @@ -298,17 +299,9 @@ def frame_export_json( ) # Validate lookup assignments (cloud-only, no-op in OSS) - lookup_configs = {} - try: - from pluggable_apps.lookup_v1.validation import ( - validate_lookups_for_export, - ) - - lookup_configs, lookup_error = validate_lookups_for_export(prompts) - if lookup_error: - raise InValidCustomToolError(lookup_error) - except ImportError: - pass + lookup_configs, lookup_error = validate_lookups_for_export(prompts) + if lookup_error: + raise InValidCustomToolError(lookup_error) for prompt in prompts: if prompt.prompt_type == JsonSchemaKey.NOTES or not prompt.active: diff --git a/backend/utils/websocket_views.py b/backend/utils/websocket_views.py index 51998399c4..a57521ca43 100644 --- a/backend/utils/websocket_views.py +++ b/backend/utils/websocket_views.py @@ -63,7 +63,7 @@ def emit_websocket(request): # Emit the WebSocket event _emit_websocket_event(room=room, event=event, data=message_data) - logger.info(f"WebSocket event emitted: room={room}, event={event}") + logger.debug(f"WebSocket event emitted: room={room}, event={event}") return JsonResponse( { diff --git a/frontend/src/components/custom-tools/document-parser/DocumentParser.jsx b/frontend/src/components/custom-tools/document-parser/DocumentParser.jsx index c5e60a2c93..2b55e26807 100644 --- a/frontend/src/components/custom-tools/document-parser/DocumentParser.jsx +++ b/frontend/src/components/custom-tools/document-parser/DocumentParser.jsx @@ -111,9 +111,15 @@ function DocumentParser({ }, [scrollToBottom]); // Handle scrollTo query param for cross-linking from Lookup Studio + const UUID_RE = + /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i; useEffect(() => { const scrollToPromptId = searchParams.get("scrollTo"); - if (!scrollToPromptId || !details?.prompts?.length) { + if ( + !scrollToPromptId || + !UUID_RE.test(scrollToPromptId) || + !details?.prompts?.length + ) { return; } @@ -127,7 +133,7 @@ function DocumentParser({ // Clear the param so it doesn't re-trigger searchParams.delete("scrollTo"); setSearchParams(searchParams, { replace: true }); - }, [details?.prompts]); + }, [details?.prompts, searchParams]); const promptUrl = (urlPath) => { return `/api/v1/unstract/${sessionDetails?.orgId}/prompt-studio/prompt/${urlPath}`; diff --git a/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx b/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx index 02e552ae44..872c657b2d 100644 --- a/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx +++ b/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx @@ -61,6 +61,15 @@ try { // Not available in OSS } +// Wraps children in LookupOutputTabs when available (cloud), +// passes through children directly in OSS. +const renderWithLookupWrapper = (lookupProps, children) => + LookupOutputTabs ? ( + {children} + ) : ( + children + ); + function PromptOutput({ promptDetails, handleRun, @@ -203,31 +212,13 @@ function PromptOutput({ "highlighted-prompt-cell" }`} > - {LookupOutputTabs ? ( - - - - ) : ( + {renderWithLookupWrapper( + { + promptId, + profileManagerId: defaultLlmProfile, + defaultLlmProfile, + promptOutputId: promptOutputData?.promptOutputId, + }, + />, )}
- {LookupOutputTabs ? ( - - - - ) : ( + {renderWithLookupWrapper( + { + promptId, + profileManagerId: profileId, + defaultLlmProfile, + promptOutputId: promptOutputData?.promptOutputId, + }, + />, )}
{ - const currentPath = globalThis.location.pathname; - if (currentPath.startsWith(`/${orgName}/lookups`)) { - return "lookups"; - } - return "projects"; -}; - -const PromptStudioPopoverContent = ({ orgName, navigate }) => { - const activeKey = getActivePromptStudioKey(orgName); - - return ( - - ); -}; - -PromptStudioPopoverContent.propTypes = { - orgName: PropTypes.string.isRequired, - navigate: PropTypes.func.isRequired, -}; - const SideNavBar = ({ collapsed, setCollapsed }) => { const navigate = useNavigate(); const { sessionDetails } = useSessionStore(); diff --git a/frontend/src/helpers/SocketContext.js b/frontend/src/helpers/SocketContext.js index ec6a274e9e..6e6ace9a63 100644 --- a/frontend/src/helpers/SocketContext.js +++ b/frontend/src/helpers/SocketContext.js @@ -16,7 +16,7 @@ const SocketProvider = ({ children }) => { // This ensures session cookies are sent (same-origin) and avoids // cross-origin WebSocket issues. const newSocket = io(getBaseUrl(), { - transports: ["websocket", "polling"], + transports: ["websocket"], path: "/api/v1/socket", }); setSocket(newSocket); diff --git a/workers/executor/executors/plugins/protocols.py b/workers/executor/executors/plugins/protocols.py index a5ed148a5b..9ce1c5eb9e 100644 --- a/workers/executor/executors/plugins/protocols.py +++ b/workers/executor/executors/plugins/protocols.py @@ -55,4 +55,4 @@ def run(self, **kwargs: Any) -> dict: ... class LookupEnrichmentProtocol(Protocol): """Legacy executor: post-extraction lookup enrichment.""" - def run(self) -> str | None: ... + def run(self) -> None: ... From ae4ba0a130d1aee7914750d7f7a7773590a7e8fe Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Tue, 7 Apr 2026 19:45:01 +0530 Subject: [PATCH 14/57] UN-2946 [UI] Replace sidebar popover with in-page tabs for Lookups Move Prompt Studio / Look-Ups navigation from a hover popover on the sidebar into a Segmented control within the ToolNavBar. CustomTools dynamically imports LookupList from the plugin and renders tabs when available, falling back to projects-only view in OSS mode. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../list-of-tools/ListOfTools.jsx | 13 +++- .../navigations/side-nav-bar/SideNavBar.jsx | 73 ++----------------- .../navigations/tool-nav-bar/ToolNavBar.jsx | 3 + frontend/src/pages/CustomTools.jsx | 30 +++++++- 4 files changed, 48 insertions(+), 71 deletions(-) diff --git a/frontend/src/components/custom-tools/list-of-tools/ListOfTools.jsx b/frontend/src/components/custom-tools/list-of-tools/ListOfTools.jsx index dd1e55853b..86ab15ab16 100644 --- a/frontend/src/components/custom-tools/list-of-tools/ListOfTools.jsx +++ b/frontend/src/components/custom-tools/list-of-tools/ListOfTools.jsx @@ -48,7 +48,7 @@ DefaultCustomButtons.propTypes = { handleNewProjectBtnClick: PropTypes.func.isRequired, }; -function ListOfTools() { +function ListOfTools({ segmentOptions, segmentValue, onSegmentChange }) { const [isListLoading, setIsListLoading] = useState(false); const [openAddTool, setOpenAddTool] = useState(false); const [openImportTool, setOpenImportTool] = useState(false); @@ -373,12 +373,15 @@ function ListOfTools() { return ( <>
{defaultContent}
@@ -412,4 +415,10 @@ function ListOfTools() { ); } +ListOfTools.propTypes = { + segmentOptions: PropTypes.arrayOf(PropTypes.string), + segmentValue: PropTypes.string, + onSegmentChange: PropTypes.func, +}; + export { ListOfTools }; diff --git a/frontend/src/components/navigations/side-nav-bar/SideNavBar.jsx b/frontend/src/components/navigations/side-nav-bar/SideNavBar.jsx index b3b1d0a014..7cecc20333 100644 --- a/frontend/src/components/navigations/side-nav-bar/SideNavBar.jsx +++ b/frontend/src/components/navigations/side-nav-bar/SideNavBar.jsx @@ -90,11 +90,9 @@ try { } let lookupStudioEnabled = false; -let PromptStudioPopoverContent = null; try { - const mod = await import("../../../plugins/lookup-studio"); + await import("../../../plugins/lookup-studio"); lookupStudioEnabled = true; - PromptStudioPopoverContent = mod.PromptStudioPopoverContent; } catch { // Plugin unavailable } @@ -516,11 +514,13 @@ const SideNavBar = ({ collapsed, setCollapsed }) => { }); } - // Mark Prompt Studio item for popover rendering when lookups plugin is available + // Extend Prompt Studio active state to include /lookups paths if (lookupStudioEnabled && isUnstract) { const psItem = data[0]?.subMenu?.find((el) => el.id === 1.1); if (psItem) { - psItem.hasLookupPopover = true; + psItem.active = + psItem.active || + globalThis.location.pathname.startsWith(`/${orgName}/lookups`); } } @@ -718,69 +718,6 @@ const SideNavBar = ({ collapsed, setCollapsed }) => { ); } - // Prompt Studio with Look-Ups popover - if (el.hasLookupPopover) { - const psContent = ( - - { - if (!el.disable) { - navigate(el.path); - } - }} - data-testid="sidebar-prompt-studio" - > - side_icon - {!collapsed && ( -
- - {el.title} - - - {el.description} - -
- )} - - - ); - - if (el.disable) { - return
{psContent}
; - } - - return ( - - } - trigger="hover" - placement="rightTop" - arrow={false} - overlayClassName="settings-popover-overlay" - > - {psContent} - - ); - } - return ( @@ -111,6 +113,7 @@ ToolNavBar.propTypes = { previousRouteState: PropTypes.object, onNavigateBack: PropTypes.func, segmentOptions: PropTypes.array, + segmentValue: PropTypes.string, segmentFilter: PropTypes.func, onSearch: PropTypes.func, searchKey: PropTypes.string, diff --git a/frontend/src/pages/CustomTools.jsx b/frontend/src/pages/CustomTools.jsx index bd177301c7..1f6c2d3934 100644 --- a/frontend/src/pages/CustomTools.jsx +++ b/frontend/src/pages/CustomTools.jsx @@ -1,7 +1,35 @@ +import { useEffect, useState } from "react"; + import { ListOfTools } from "../components/custom-tools/list-of-tools/ListOfTools"; +const TAB_OPTIONS = ["Projects", "Look-Ups"]; + function CustomTools() { - return ; + const [LookupListComp, setLookupListComp] = useState(null); + const [activeTab, setActiveTab] = useState("Projects"); + + useEffect(() => { + import("../plugins/lookup-studio") + .then((mod) => setLookupListComp(() => mod.LookupList)) + .catch(() => {}); + }, []); + + // No lookup plugin = just render projects list (OSS mode) + if (!LookupListComp) { + return ; + } + + const tabProps = { + segmentOptions: TAB_OPTIONS, + segmentValue: activeTab, + onSegmentChange: setActiveTab, + }; + + return activeTab === "Projects" ? ( + + ) : ( + + ); } export { CustomTools }; From 8aa675884b51ca610f3e2f64acd416f1848e4724 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Wed, 8 Apr 2026 16:42:08 +0530 Subject: [PATCH 15/57] UN-2946 [FEAT] Deferred batch usage tracking with operation metrics Switch from eager per-call Audit HTTP push to a deferred batch write pattern for adapter usage. LLM/embedding calls stash records in-memory; the executor flushes them into ExecutionResult metadata; the Celery task batch-writes via a new internal endpoint. Adds 5 nullable columns to Usage (reference_id, reference_type, execution_time_ms, status, error_message) and a composite index for lookup dashboard queries. Extensible choice lists allow cloud plugins to register additional usage reasons and reference types. Co-Authored-By: Claude Opus 4.6 (1M context) --- backend/usage_v2/internal_urls.py | 5 + backend/usage_v2/internal_views.py | 41 +++++++++ .../0004_add_lookup_usage_reason.py | 28 ------ .../migrations/0004_usage_metrics_fields.py | 90 ++++++++++++++++++ backend/usage_v2/models.py | 86 ++++++++++++++++-- unstract/sdk1/src/unstract/sdk1/llm.py | 91 ++++++++++++++----- .../sdk1/src/unstract/sdk1/usage_handler.py | 46 ++++++++-- .../sdk1/src/unstract/sdk1/utils/common.py | 7 ++ workers/executor/executors/legacy_executor.py | 18 ++++ workers/executor/executors/usage.py | 81 ----------------- workers/executor/tasks.py | 21 +++++ workers/shared/clients/usage_client.py | 25 +++++ 12 files changed, 390 insertions(+), 149 deletions(-) delete mode 100644 backend/usage_v2/migrations/0004_add_lookup_usage_reason.py create mode 100644 backend/usage_v2/migrations/0004_usage_metrics_fields.py delete mode 100644 workers/executor/executors/usage.py diff --git a/backend/usage_v2/internal_urls.py b/backend/usage_v2/internal_urls.py index b5a8675554..e7e0082e99 100644 --- a/backend/usage_v2/internal_urls.py +++ b/backend/usage_v2/internal_urls.py @@ -17,4 +17,9 @@ internal_views.PagesProcessedInternalView.as_view(), name="aggregated-pages-processed", ), + path( + "batch/", + internal_views.UsageBatchCreateView.as_view(), + name="usage-batch-create", + ), ] diff --git a/backend/usage_v2/internal_views.py b/backend/usage_v2/internal_views.py index 3c28779c56..ca6642b443 100644 --- a/backend/usage_v2/internal_views.py +++ b/backend/usage_v2/internal_views.py @@ -6,10 +6,12 @@ from rest_framework import status from rest_framework.request import Request from rest_framework.views import APIView +from utils.user_context import UserContext from unstract.core.data_models import UsageResponseData from .helper import UsageHelper +from .models import Usage logger = logging.getLogger(__name__) @@ -133,3 +135,42 @@ def get(self, request: Request, file_execution_id: str) -> JsonResponse: }, status=status.HTTP_500_INTERNAL_SERVER_ERROR, ) + + +class UsageBatchCreateView(APIView): + """Bulk create usage records from worker finalization.""" + + def post(self, request: Request) -> JsonResponse: + records = request.data.get("records", []) + if not records: + return JsonResponse({"created": 0}, status=200) + + # Resolved by InternalAPIAuthMiddleware via StateStore + organization = UserContext.get_organization() + + usage_objects = [] + for r in records: + usage_objects.append( + Usage( + organization=organization, + workflow_id=r.get("workflow_id", ""), + execution_id=r.get("execution_id", ""), + adapter_instance_id=r.get("adapter_instance_id", ""), + run_id=r.get("run_id"), + usage_type=r.get("usage_type", "llm"), + llm_usage_reason=r.get("llm_usage_reason", ""), + model_name=r.get("model_name", ""), + embedding_tokens=r.get("embedding_tokens", 0), + prompt_tokens=r.get("prompt_tokens", 0), + completion_tokens=r.get("completion_tokens", 0), + total_tokens=r.get("total_tokens", 0), + cost_in_dollars=r.get("cost_in_dollars", 0.0), + reference_id=r.get("reference_id"), + reference_type=r.get("reference_type"), + execution_time_ms=r.get("execution_time_ms"), + status=r.get("status"), + error_message=r.get("error_message"), + ) + ) + created = Usage.objects.bulk_create(usage_objects) + return JsonResponse({"created": len(created)}, status=201) diff --git a/backend/usage_v2/migrations/0004_add_lookup_usage_reason.py b/backend/usage_v2/migrations/0004_add_lookup_usage_reason.py deleted file mode 100644 index 76c628d9ee..0000000000 --- a/backend/usage_v2/migrations/0004_add_lookup_usage_reason.py +++ /dev/null @@ -1,28 +0,0 @@ -# Generated by Django 4.2.1 on 2026-04-03 22:44 - -from django.db import migrations, models - - -class Migration(migrations.Migration): - dependencies = [ - ("usage_v2", "0003_usage_usage_executi_4deb35_idx"), - ] - - operations = [ - migrations.AlterField( - model_name="usage", - name="llm_usage_reason", - field=models.CharField( - blank=True, - choices=[ - ("extraction", "Extraction"), - ("challenge", "Challenge"), - ("summarize", "Summarize"), - ("lookup", "Lookup"), - ], - db_comment="Reason for LLM usage. Empty if usage_type is 'embedding'. ", - max_length=255, - null=True, - ), - ), - ] diff --git a/backend/usage_v2/migrations/0004_usage_metrics_fields.py b/backend/usage_v2/migrations/0004_usage_metrics_fields.py new file mode 100644 index 0000000000..03174f0677 --- /dev/null +++ b/backend/usage_v2/migrations/0004_usage_metrics_fields.py @@ -0,0 +1,90 @@ +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("usage_v2", "0003_usage_usage_executi_4deb35_idx"), + ] + + operations = [ + # Extend llm_usage_reason choices (cloud plugins append at runtime) + migrations.AlterField( + model_name="usage", + name="llm_usage_reason", + field=models.CharField( + blank=True, + choices=[ + ("extraction", "Extraction"), + ("challenge", "Challenge"), + ("summarize", "Summarize"), + ], + db_comment="Reason for LLM usage. Empty if usage_type is 'embedding'. ", + max_length=255, + null=True, + ), + ), + migrations.AddField( + model_name="usage", + name="reference_id", + field=models.UUIDField( + blank=True, + db_comment=( + "Polymorphic correlation ID (no FK constraint) linking to the " + "entity that triggered this usage. Interpret via reference_type. " + "OSS values: prompt_key UUID. " + "NULL for most operations; survives entity deletion." + ), + null=True, + ), + ), + migrations.AddField( + model_name="usage", + name="reference_type", + field=models.CharField( + blank=True, + choices=[("prompt_key", "Prompt Key")], + db_comment=( + "Discriminator for reference_id. " + "OSS values: 'prompt_key'. " + "NULL when reference_id is NULL." + ), + max_length=64, + null=True, + ), + ), + migrations.AddField( + model_name="usage", + name="execution_time_ms", + field=models.IntegerField( + blank=True, + db_comment="Wall-clock time for the operation in milliseconds", + null=True, + ), + ), + migrations.AddField( + model_name="usage", + name="status", + field=models.CharField( + blank=True, + db_comment="Operation outcome: SUCCESS, ERROR, or SKIPPED", + max_length=16, + null=True, + ), + ), + migrations.AddField( + model_name="usage", + name="error_message", + field=models.TextField( + blank=True, + db_comment="Error details when status is ERROR", + null=True, + ), + ), + migrations.AddIndex( + model_name="usage", + index=models.Index( + fields=["llm_usage_reason", "reference_id", "-created_at"], + name="idx_usage_reason_ref_created", + ), + ), + ] diff --git a/backend/usage_v2/models.py b/backend/usage_v2/models.py index dc983c56b5..365fdc956b 100644 --- a/backend/usage_v2/models.py +++ b/backend/usage_v2/models.py @@ -1,3 +1,4 @@ +import logging import uuid from django.db import models @@ -7,17 +8,44 @@ DefaultOrganizationMixin, ) +logger = logging.getLogger(__name__) + class UsageType(models.TextChoices): LLM = "llm", "LLM Usage" EMBEDDING = "embedding", "Embedding Usage" -class LLMUsageReason(models.TextChoices): - EXTRACTION = "extraction", "Extraction" - CHALLENGE = "challenge", "Challenge" - SUMMARIZE = "summarize", "Summarize" - LOOKUP = "lookup", "Lookup" +# ── Extensible choice lists ───────────────────────────────────────── +# OSS defines base values. Cloud plugins append via try-import so that +# Django validation accepts cloud-specific values when the plugin is +# installed, without leaking cloud details into OSS code. + +_LLM_USAGE_REASON_CHOICES: list[tuple[str, str]] = [ + ("extraction", "Extraction"), + ("challenge", "Challenge"), + ("summarize", "Summarize"), +] + +_REFERENCE_TYPE_CHOICES: list[tuple[str, str]] = [ + ("prompt_key", "Prompt Key"), +] + +try: + from pluggable_apps.lookup_v1.constants import ( + CLOUD_LLM_USAGE_REASON_CHOICES, + CLOUD_REFERENCE_TYPE_CHOICES, + ) + + _LLM_USAGE_REASON_CHOICES.extend(CLOUD_LLM_USAGE_REASON_CHOICES) + _REFERENCE_TYPE_CHOICES.extend(CLOUD_REFERENCE_TYPE_CHOICES) +except ImportError: + pass +except Exception: + logger.warning("Failed to load cloud usage choices", exc_info=True) + +LLM_USAGE_REASON_CHOICES = _LLM_USAGE_REASON_CHOICES +REFERENCE_TYPE_CHOICES = _REFERENCE_TYPE_CHOICES class UsageModelManager(DefaultOrganizationManagerMixin, models.Manager): @@ -25,6 +53,11 @@ class UsageModelManager(DefaultOrganizationManagerMixin, models.Manager): class Usage(DefaultOrganizationMixin, BaseModel): + # reference_type → reference_id mapping (no FK constraint): + # "prompt_key" → ToolStudioPrompt UUID (OSS) + # Cloud plugins register additional types via CLOUD_REFERENCE_TYPE_CHOICES. + # Usage records survive entity deletion. + id = models.UUIDField( primary_key=True, default=uuid.uuid4, @@ -53,7 +86,7 @@ class Usage(DefaultOrganizationMixin, BaseModel): ) llm_usage_reason = models.CharField( max_length=255, - choices=LLMUsageReason.choices, + choices=LLM_USAGE_REASON_CHOICES, null=True, blank=True, db_comment="Reason for LLM usage. Empty if usage_type is 'embedding'. ", @@ -68,6 +101,43 @@ class Usage(DefaultOrganizationMixin, BaseModel): ) total_tokens = models.IntegerField(db_comment="Total number of tokens used") cost_in_dollars = models.FloatField(db_comment="Total number of tokens used") + reference_id = models.UUIDField( + null=True, + blank=True, + db_comment=( + "Polymorphic correlation ID (no FK constraint) linking to the " + "entity that triggered this usage. Interpret via reference_type. " + "OSS values: prompt_key UUID. " + "NULL for most operations; survives entity deletion." + ), + ) + reference_type = models.CharField( + max_length=64, + choices=REFERENCE_TYPE_CHOICES, + null=True, + blank=True, + db_comment=( + "Discriminator for reference_id. " + "OSS values: 'prompt_key'. " + "NULL when reference_id is NULL." + ), + ) + execution_time_ms = models.IntegerField( + null=True, + blank=True, + db_comment="Wall-clock time for the operation in milliseconds", + ) + status = models.CharField( + max_length=16, + null=True, + blank=True, + db_comment="Operation outcome: SUCCESS, ERROR, or SKIPPED", + ) + error_message = models.TextField( + null=True, + blank=True, + db_comment="Error details when status is ERROR", + ) # Manager objects = UsageModelManager() @@ -79,4 +149,8 @@ class Meta: indexes = [ models.Index(fields=["run_id"]), models.Index(fields=["execution_id"]), + models.Index( + fields=["llm_usage_reason", "reference_id", "-created_at"], + name="idx_usage_reason_ref_created", + ), ] diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index f751e7f87c..56b27dc163 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -9,11 +9,10 @@ import litellm # from litellm import get_supported_openai_params -from litellm import get_max_tokens, token_counter +from litellm import get_max_tokens from pydantic import ValidationError from unstract.sdk1.adapters.constants import Common from unstract.sdk1.adapters.llm1 import adapters -from unstract.sdk1.audit import Audit from unstract.sdk1.constants import Common as SdkCommon from unstract.sdk1.constants import ToolEnv from unstract.sdk1.exceptions import LLMError, SdkError, strip_litellm_prefix @@ -21,7 +20,6 @@ from unstract.sdk1.tool.base import BaseTool from unstract.sdk1.utils.common import ( LLMResponseCompat, - TokenCounterCompat, capture_metrics, ) from unstract.sdk1.utils.retry_utils import ( @@ -211,7 +209,7 @@ def __init__( # noqa: C901 if capture_metrics_from_platform is not None: self._capture_metrics = capture_metrics_from_platform self._metrics: dict[str, object] = {} - self._last_usage: Mapping[str, int] = {} + self._pending_usage: list[dict] = [] def _get_adapter_info(self) -> str: """Build a display string identifying this adapter for errors.""" @@ -555,11 +553,27 @@ def get_metrics(self) -> dict[str, object]: def get_last_usage(self) -> Mapping[str, int]: """Token usage from the most recent complete() call.""" - return self._last_usage + if not self._pending_usage: + return {} + last = self._pending_usage[-1] + return { + "prompt_tokens": last["prompt_tokens"], + "completion_tokens": last["completion_tokens"], + "total_tokens": last["total_tokens"], + } def get_usage_reason(self) -> object: return self.platform_kwargs.get("llm_usage_reason") + def flush_pending_usage(self) -> list[dict]: + """Return and clear all pending usage records. + + Called by the executor at finalization to collect records for batch write. + """ + records = self._pending_usage + self._pending_usage = [] + return records + def _record_usage( self, model: str, @@ -567,29 +581,52 @@ def _record_usage( usage: Mapping[str, int] | None, llm_api: str, ) -> None: - prompt_tokens = token_counter(model=model, messages=messages) usage_data: Mapping[str, int] = usage or {} - all_tokens = TokenCounterCompat( - prompt_tokens=usage_data.get("prompt_tokens", 0), - completion_tokens=usage_data.get("completion_tokens", 0), - total_tokens=usage_data.get("total_tokens", 0), + prompt_tokens = usage_data.get("prompt_tokens", 0) + completion_tokens = usage_data.get("completion_tokens", 0) + total_tokens = usage_data.get("total_tokens", 0) + + logger.info( + "[sdk1][LLM][%s][%s] Usage: prompt=%d completion=%d total=%d", + model, + llm_api, + prompt_tokens, + completion_tokens, + total_tokens, ) - logger.info(f"[sdk1][LLM][{model}][{llm_api}] Prompt Tokens: {prompt_tokens}") - logger.info(f"[sdk1][LLM][{model}][{llm_api}] LLM Usage: {all_tokens}") - - self._last_usage = { - "prompt_tokens": all_tokens.prompt_llm_token_count, - "completion_tokens": all_tokens.completion_llm_token_count, - "total_tokens": all_tokens.total_llm_token_count, - } - - Audit().push_usage_data( - platform_api_key=self._platform_api_key, - token_counter=all_tokens, - event_type="llm", - model_name=model, - kwargs={"provider": self.adapter.get_provider(), **self.platform_kwargs}, + try: + prompt_cost, compl_cost = litellm.cost_per_token( + model=model, + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + ) + cost = prompt_cost + compl_cost + except Exception: + cost = 0.0 + + # Strip provider prefix (e.g. "azure/gpt-4o" → "gpt-4o") for storage, + # matching the old Audit.push_usage_data() behavior. + display_model = model.split("/", 1)[-1] if model else model + + self._pending_usage.append( + { + "usage_type": "llm", + "model_name": display_model, + "provider": self.adapter.get_provider(), + "adapter_instance_id": self.platform_kwargs.get( + "adapter_instance_id", "" + ), + "run_id": self.platform_kwargs.get("run_id", ""), + "execution_id": self.platform_kwargs.get("execution_id", ""), + "llm_usage_reason": self.platform_kwargs.get("llm_usage_reason", ""), + "prompt_tokens": prompt_tokens, + "completion_tokens": completion_tokens, + "total_tokens": total_tokens, + "embedding_tokens": 0, + "cost_in_dollars": cost, + "status": "SUCCESS", + } ) # Finish reasons indicating a safety/policy refusal across providers: @@ -982,6 +1019,10 @@ def get_usage_reason(self) -> object: """Get usage reason from platform kwargs.""" return self._llm_instance.get_usage_reason() + def flush_pending_usage(self) -> list[dict]: + """Return and clear all pending usage records.""" + return self._llm_instance.flush_pending_usage() + def test_connection(self) -> bool: """Test connection to the LLM provider.""" return self._llm_instance.test_connection() diff --git a/unstract/sdk1/src/unstract/sdk1/usage_handler.py b/unstract/sdk1/src/unstract/sdk1/usage_handler.py index 0ffa7c43e2..44a819b44c 100644 --- a/unstract/sdk1/src/unstract/sdk1/usage_handler.py +++ b/unstract/sdk1/src/unstract/sdk1/usage_handler.py @@ -1,9 +1,9 @@ from typing import Any +import litellm from llama_index.core.callbacks import CBEventType, TokenCountingHandler from llama_index.core.callbacks.base_handler import BaseCallbackHandler from llama_index.core.embeddings import BaseEmbedding -from unstract.sdk1.audit import Audit from unstract.sdk1.constants import LogLevel from unstract.sdk1.tool.stream import StreamMixin @@ -57,6 +57,7 @@ def __init__( self._verbose = verbose self.token_counter = token_counter self.embed_model = embed_model + self._pending_usage: list[dict] = [] self.platform_api_key = platform_api_key super().__init__( log_level=log_level, # StreamMixin's args @@ -102,16 +103,43 @@ def on_event_end( and payload is not None ): model_name = self.embed_model.model_name - # Need to push the data to via platform service + embedding_tokens = self.token_counter.total_embedding_token_count self.stream_log( - log=f"Pushing embedding usage for model {model_name}", + log=f"Recording embedding usage for model {model_name}", level=LogLevel.DEBUG, ) - Audit(log_level=self.log_level).push_usage_data( - platform_api_key=self.platform_api_key, - token_counter=self.token_counter, - event_type=event_type, - model_name=self.embed_model.model_name, - kwargs=self.kwargs, + + try: + prompt_cost, _ = litellm.cost_per_token( + model=model_name, + prompt_tokens=embedding_tokens, + completion_tokens=0, + ) + cost = prompt_cost + except Exception: + cost = 0.0 + + display_model = model_name.split("/", 1)[-1] if model_name else model_name + + self._pending_usage.append( + { + "usage_type": "embedding", + "model_name": display_model, + "adapter_instance_id": self.kwargs.get("adapter_instance_id", ""), + "run_id": self.kwargs.get("run_id", ""), + "execution_id": self.kwargs.get("execution_id", ""), + "embedding_tokens": embedding_tokens, + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "cost_in_dollars": cost, + "status": "SUCCESS", + } ) self.token_counter.reset_counts() + + def flush_pending_usage(self) -> list[dict]: + """Return and clear all pending usage records.""" + records = self._pending_usage + self._pending_usage = [] + return records diff --git a/unstract/sdk1/src/unstract/sdk1/utils/common.py b/unstract/sdk1/src/unstract/sdk1/utils/common.py index 86b0f173b6..c2c3fdbca5 100644 --- a/unstract/sdk1/src/unstract/sdk1/utils/common.py +++ b/unstract/sdk1/src/unstract/sdk1/utils/common.py @@ -250,6 +250,13 @@ def wrapper(self: object, *args: object, **kwargs: object) -> object: # If the key isn't in self._metrics, set it to new_metrics self._metrics = new_metrics + # Stamp timing onto the most recent pending usage record + pending = getattr(self, "_pending_usage", []) + if pending: + time_taken = new_metrics.get(time_taken_key) + if time_taken is not None: + pending[-1]["execution_time_ms"] = int(time_taken * 1000) + return result return wrapper diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index 7ccdcb656d..7a0a41a424 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -79,6 +79,7 @@ def execute(self, context: ExecutionContext) -> ExecutionResult: # Extract log streaming info (set by tasks.py for IDE sessions). self._log_events_id: str = context.log_events_id or "" self._log_component: dict[str, str] = getattr(context, "_log_component", {}) + self._usage_records: list[dict[str, Any]] = [] handler_name = self._OPERATION_MAP.get(context.operation) if handler_name is None: @@ -107,6 +108,11 @@ def execute(self, context: ExecutionContext) -> ExecutionResult: context.run_id, result.success, ) + # Attach collected usage records to the result metadata + if self._usage_records: + result.metadata.setdefault("usage_records", []).extend( + self._usage_records + ) return result except LegacyExecutorError as exc: elapsed = time.monotonic() - start @@ -1390,6 +1396,7 @@ def _run_challenge_if_enabled( metadata=metadata, ) challenger.run() + self._usage_records.extend(challenge_llm.flush_pending_usage()) shim.stream_log(f"Challenge verification completed for: {prompt_name}") logger.info("Challenge completed: prompt=%s", prompt_name) @@ -1702,6 +1709,15 @@ def _execute_single_prompt( f"{llm.get_usage_reason()}_llm": llm.get_metrics(), } ) + self._usage_records.extend(llm.flush_pending_usage()) + # Flush embedding usage from callback handlers + if chunk_size > 0: + try: + for handler in embedding.callback_manager.handlers: + if hasattr(handler, "flush_pending_usage"): + self._usage_records.extend(handler.flush_pending_usage()) + except Exception: + pass if vector_db: vector_db.close() @@ -1901,6 +1917,7 @@ def _run_lookup_enrichment( prompt_name=prompt_name, ) enricher.run() + self._usage_records.extend(llm.flush_pending_usage()) metrics.setdefault(prompt_name, {})[f"{llm.get_usage_reason()}_llm"] = ( llm.get_metrics() @@ -2150,6 +2167,7 @@ def _handle_summarize(self, context: ExecutionContext) -> ExecutionResult: shim.stream_log("Running document summarization...") summary = answer_prompt_svc.run_completion(llm=llm, prompt=prompt) + self._usage_records.extend(llm.flush_pending_usage()) logger.info("Summarization completed: run_id=%s", context.run_id) shim.stream_log("Summarization completed") return ExecutionResult( diff --git a/workers/executor/executors/usage.py b/workers/executor/executors/usage.py deleted file mode 100644 index ab6296eaeb..0000000000 --- a/workers/executor/executors/usage.py +++ /dev/null @@ -1,81 +0,0 @@ -"""Usage tracking helper for the executor worker. - -Ported from prompt-service/.../helpers/usage.py. -Flask/DB dependencies removed — usage data is pushed via the SDK1 -``Audit`` class (HTTP to platform API) and returned directly in -``ExecutionResult.metadata`` instead of querying the DB. - -Note: The SDK1 adapters (LLM, EmbeddingCompat) already call -``Audit().push_usage_data()`` internally. This helper is for -explicit push calls outside of adapter operations (e.g. rent rolls). -""" - -import logging -from typing import Any - -logger = logging.getLogger(__name__) - - -class UsageHelper: - @staticmethod - def push_usage_data( - event_type: str, - kwargs: dict[str, Any], - platform_api_key: str, - token_counter: Any = None, - model_name: str = "", - ) -> bool: - """Push usage data to the audit service. - - Wraps ``Audit().push_usage_data()`` with validation and - error handling. - - Args: - event_type: Type of usage event (e.g. "llm", "embedding"). - kwargs: Context dict (run_id, execution_id, etc.). - platform_api_key: API key for platform service auth. - token_counter: Token counter with usage metrics. - model_name: Name of the model used. - - Returns: - True if successful, False otherwise. - """ - if not kwargs or not isinstance(kwargs, dict): - logger.error("Invalid kwargs provided to push_usage_data") - return False - - if not platform_api_key or not isinstance(platform_api_key, str): - logger.error("Invalid platform_api_key provided to push_usage_data") - return False - - try: - from unstract.sdk1.audit import Audit - - logger.debug( - "Pushing usage data for event_type=%s model=%s", - event_type, - model_name, - ) - - Audit().push_usage_data( - platform_api_key=platform_api_key, - token_counter=token_counter, - model_name=model_name, - event_type=event_type, - kwargs=kwargs, - ) - - logger.info("Successfully pushed usage data for %s", model_name) - return True - except Exception: - logger.exception("Error pushing usage data") - return False - - @staticmethod - def format_float_positional(value: float, precision: int = 10) -> str: - """Format a float without scientific notation. - - Removes trailing zeros for clean display of cost values. - """ - formatted: str = f"{value:.{precision}f}" - return formatted.rstrip("0").rstrip(".") if "." in formatted else formatted diff --git a/workers/executor/tasks.py b/workers/executor/tasks.py index 77d5ecaebd..ea02b7c353 100644 --- a/workers/executor/tasks.py +++ b/workers/executor/tasks.py @@ -6,7 +6,9 @@ """ from celery import shared_task +from shared.clients import UsageAPIClient from shared.enums.task_enums import TaskName +from shared.infrastructure.config import WorkerConfig from shared.infrastructure.logging import WorkerLogger from unstract.sdk1.execution.context import ExecutionContext @@ -97,6 +99,25 @@ def execute_extraction(self, execution_context_dict: dict) -> dict: orchestrator = ExecutionOrchestrator() result = orchestrator.execute(context) + # Batch write usage records collected during execution + usage_records = result.metadata.get("usage_records", []) + if usage_records: + try: + config = WorkerConfig() + with UsageAPIClient(config) as usage_client: + usage_client.set_organization_context(context.organization_id) + usage_client.bulk_create_usage( + usage_records, + organization_id=context.organization_id, + ) + except Exception: + logger.warning( + "Failed to flush %d usage records for run_id=%s", + len(usage_records), + context.run_id, + exc_info=True, + ) + logger.info( "execute_extraction complete: celery_task_id=%s request_id=%s success=%s", self.request.id, diff --git a/workers/shared/clients/usage_client.py b/workers/shared/clients/usage_client.py index a7ba8c7013..f61f65e27a 100644 --- a/workers/shared/clients/usage_client.py +++ b/workers/shared/clients/usage_client.py @@ -171,6 +171,31 @@ def get_aggregated_token_count( message="Failed to retrieve usage data", ) + def bulk_create_usage( + self, records: list[dict], organization_id: str | None = None + ) -> bool: + """Bulk create usage records at execution finalization. + + Args: + records: List of usage record dicts to create. + organization_id: Optional organization ID override. + + Returns: + True if records were created successfully. + """ + if not records: + return True + try: + response = self.post( + "v1/usage/batch/", + data={"records": records}, + organization_id=organization_id, + ) + return response.get("success", False) or "created" in response + except Exception as e: + logger.error("Failed to bulk create usage records: %s", e) + return False + def get_aggregated_pages_processed( self, file_execution_id: str | uuid.UUID, organization_id: str | None = None ) -> int | None: From 0d29c9d631594c1acf19dec3f248264fd34cc80f Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Fri, 10 Apr 2026 14:15:22 +0530 Subject: [PATCH 16/57] UN-2946 [FEAT] Add plugin hook for lookup output enrichment in serializer Bridge function in lookup_utils.py lets cloud plugins enrich PromptStudioOutputSerializer with lookup data (enriched output, lookup name). Enables real-time lookup results via WebSocket without page refresh. Co-Authored-By: Claude Opus 4.6 (1M context) --- backend/prompt_studio/lookup_utils.py | 27 ++++++++++++++++--- .../serializers.py | 3 +++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/backend/prompt_studio/lookup_utils.py b/backend/prompt_studio/lookup_utils.py index 823d9fc015..10783e405c 100644 --- a/backend/prompt_studio/lookup_utils.py +++ b/backend/prompt_studio/lookup_utils.py @@ -29,17 +29,36 @@ def persist_lookup_output(prompt_output, prompt_lookup: dict) -> None: lookup_meta = prompt_lookup.get("meta", {}) lookup_id = lookup_meta.get("lookup_id") if lookup_id: + defaults = { + "lookup_definition_id": lookup_id, + "output": prompt_lookup.get("enriched", ""), + } + version_id = lookup_meta.get("version_id") + if version_id: + defaults["version_id"] = version_id LookupOutputResult.objects.update_or_create( prompt_output=prompt_output, - defaults={ - "lookup_definition_id": lookup_id, - "output": prompt_lookup.get("enriched", ""), - }, + defaults=defaults, ) except ImportError: pass +def enrich_prompt_output(prompt_output, data: dict) -> dict: + """Let cloud plugins enrich serialized prompt output with lookup data. + + No-op in OSS. + """ + try: + from pluggable_apps.lookup_v1.output_enrichment import ( + enrich_with_lookup_output, + ) + + return enrich_with_lookup_output(prompt_output, data) + except ImportError: + return data + + def validate_lookups_for_export(prompts) -> tuple[dict, str | None]: """Validate lookup assignments before export. Returns ({}, None) in OSS.""" try: diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/serializers.py b/backend/prompt_studio/prompt_studio_output_manager_v2/serializers.py index 275e4a0956..e4276dbb23 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/serializers.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/serializers.py @@ -4,6 +4,7 @@ from usage_v2.helper import UsageHelper from backend.serializers import AuditSerializer +from prompt_studio.lookup_utils import enrich_prompt_output from .models import PromptStudioOutputManager from .output_manager_util import OutputManagerUtils @@ -47,6 +48,8 @@ def to_representation(self, instance): " | Process continued" ) data["coverage"] = {} + data = enrich_prompt_output(instance, data) + # Convert string to list try: context = data["context"] From 17ffe3802f3c4bacef9e51e861d6ff8e0e53ed79 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Sun, 12 Apr 2026 11:22:19 +0530 Subject: [PATCH 17/57] UN-2946 [FEAT] Add lookup usage observability with error handling and metadata passthrough - Wire usage_kwargs_extra from lookup config into LLM usage_kwargs for execution observability - Add error handling around enricher.run() with explicit ERROR usage records - Generic passthrough of _usage_kwargs into usage records for arbitrary metadata (e.g. reference_id) Co-Authored-By: Claude Opus 4.6 (1M context) --- unstract/sdk1/src/unstract/sdk1/llm.py | 1 + workers/executor/executors/legacy_executor.py | 31 +++++++++++++++---- 2 files changed, 26 insertions(+), 6 deletions(-) diff --git a/unstract/sdk1/src/unstract/sdk1/llm.py b/unstract/sdk1/src/unstract/sdk1/llm.py index 56b27dc163..eb31c52ad6 100644 --- a/unstract/sdk1/src/unstract/sdk1/llm.py +++ b/unstract/sdk1/src/unstract/sdk1/llm.py @@ -626,6 +626,7 @@ def _record_usage( "embedding_tokens": 0, "cost_in_dollars": cost, "status": "SUCCESS", + **self._usage_kwargs, } ) diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index 7a0a41a424..42210210ef 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -1903,6 +1903,7 @@ def _run_lookup_enrichment( usage_kwargs={ **(usage_kwargs or {}), PSKeys.LLM_USAGE_REASON: "lookup", + **lookup_config.get("usage_kwargs_extra", {}), }, capture_metrics=True, ) @@ -1916,12 +1917,30 @@ def _run_lookup_enrichment( metadata=metadata, prompt_name=prompt_name, ) - enricher.run() - self._usage_records.extend(llm.flush_pending_usage()) - - metrics.setdefault(prompt_name, {})[f"{llm.get_usage_reason()}_llm"] = ( - llm.get_metrics() - ) + try: + enricher.run() + except Exception as e: + logger.warning("Lookup enrichment failed for %s: %s", prompt_name, e) + error_record = { + "usage_type": "llm", + "llm_usage_reason": "lookup", + "model_name": lookup_config.get("llm_adapter_id", "unknown"), + "status": "ERROR", + "error_message": str(e)[:2000], + "prompt_tokens": 0, + "completion_tokens": 0, + "total_tokens": 0, + "embedding_tokens": 0, + "cost_in_dollars": 0.0, + **(usage_kwargs or {}), + **lookup_config.get("usage_kwargs_extra", {}), + } + self._usage_records.append(error_record) + finally: + self._usage_records.extend(llm.flush_pending_usage()) + metrics.setdefault(prompt_name, {})[f"{llm.get_usage_reason()}_llm"] = ( + llm.get_metrics() + ) @staticmethod def _run_webhook_postprocessing( From 6aad216f5c57feb87adf0090d1559adcee7674cf Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Wed, 15 Apr 2026 15:22:22 +0530 Subject: [PATCH 18/57] UN-2946 [FEAT] Support enriched output copy and lookup drawer plugin hooks Add dynamic import of getEnrichedCopyText so the copy button copies enriched lookup output when the Enriched tab is active. Applied to both single-pass and multi-profile output paths. Co-Authored-By: Claude Opus 4.6 (1M context) --- .../custom-tools/prompt-card/PromptOutput.jsx | 48 +++++++++++++------ 1 file changed, 33 insertions(+), 15 deletions(-) diff --git a/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx b/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx index 872c657b2d..a1e492bb29 100644 --- a/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx +++ b/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx @@ -61,6 +61,16 @@ try { // Not available in OSS } +let getEnrichedCopyText; +try { + const mod = await import( + "../../../plugins/lookup-studio/prompt-card/getEnrichedCopyText" + ); + getEnrichedCopyText = mod.getEnrichedCopyText; +} catch { + // Not available in OSS +} + // Wraps children in LookupOutputTabs when available (cloud), // passes through children directly in OSS. const renderWithLookupWrapper = (lookupProps, children) => @@ -237,15 +247,19 @@ function PromptOutput({
+ copyToClipboard={() => { + const enrichedText = getEnrichedCopyText?.( + promptOutputData?.promptOutputId, + ); copyOutputToClipboard( - displayPromptResult( - promptOutput, - true, - promptDetails?.enable_highlight, - ), - ) - } + enrichedText || + displayPromptResult( + promptOutput, + true, + promptDetails?.enable_highlight, + ), + ); + }} /> + copyToClipboard={() => { + const enrichedText = getEnrichedCopyText?.( + promptOutputData?.promptOutputId, + ); copyOutputToClipboard( - displayPromptResult( - promptOutputData?.output, - true, - ), - ) - } + enrichedText || + displayPromptResult( + promptOutputData?.output, + true, + ), + ); + }} />
From dafefa9d82cb72cde00ad9334b210f960af4a745 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Tue, 21 Apr 2026 23:44:10 +0530 Subject: [PATCH 19/57] UN-2946 [FEAT] Lookup export validation gate, raw-latest helper, and modified_at fix - Add /prompt-studio//lookup-validation/ endpoint backing the FE Export/Deploy gate; multi-var block check accepts prompt_ids so a single prompt run isn't blocked by an unrelated multi-var lookup. - Add /prompt-output/latest-by-keys/ endpoint that returns the most recent raw output per prompt_key for the test panel's "Use Latest Outputs" helper. - Fix prompt output modified_at not refreshing on re-runs (QuerySet.update bypasses auto_now); set timezone.now() explicitly in the update args. - lookup_utils: bridge get_lookup_validation_for_tool and get_multi_var_lookups_for_tool with prompt_ids scoping. - Header wires useLookupExportGate via try-import (no-op stub in OSS). - TokenUsage treats all-null Usage rows as empty. - CombinedOutput / JsonView build enriched dict from metadata.lookup_outputs to back the Raw|Enriched output toggle. - .gitignore: widen docker/compose.*.yaml. Co-Authored-By: Claude Opus 4.7 (1M context) --- .gitignore | 2 +- backend/prompt_studio/lookup_utils.py | 50 +++++++++++++++ .../prompt_studio_helper.py | 7 ++- .../prompt_studio_core_v2/urls.py | 9 +++ .../prompt_studio_core_v2/views.py | 47 ++++++++++++++ .../output_manager_helper.py | 21 ++++++- .../prompt_studio_output_manager_v2/urls.py | 6 ++ .../prompt_studio_output_manager_v2/views.py | 49 +++++++++++++++ .../combined-output/CombinedOutput.jsx | 61 ++++++++++++++----- .../custom-tools/combined-output/JsonView.jsx | 46 ++++++++++++-- .../components/custom-tools/header/Header.jsx | 24 +++++++- .../profile-info-bar/ProfileInfoBar.css | 2 + .../custom-tools/token-usage/TokenUsage.jsx | 9 ++- 13 files changed, 306 insertions(+), 27 deletions(-) diff --git a/.gitignore b/.gitignore index f6837ce079..f0c0cb8438 100644 --- a/.gitignore +++ b/.gitignore @@ -653,7 +653,7 @@ docker/*.env !docker/sample*.env docker/public_tools.json docker/proxy_overrides.yaml -docker/compose.override.yaml +docker/compose.*.yaml docker/workflow_data/ # Tool development diff --git a/backend/prompt_studio/lookup_utils.py b/backend/prompt_studio/lookup_utils.py index 10783e405c..5625c33878 100644 --- a/backend/prompt_studio/lookup_utils.py +++ b/backend/prompt_studio/lookup_utils.py @@ -21,6 +21,34 @@ def get_lookup_config(prompt) -> dict | None: return None +def get_lookup_configs_for_tool(tool) -> list[dict] | None: + """Return lookup configs for a tool (single pass), or None in OSS.""" + try: + from pluggable_apps.lookup_v1.execution import ( + build_lookup_configs_for_tool, + ) + + return build_lookup_configs_for_tool(tool) + except ImportError: + return None + + +def get_multi_var_lookups_for_tool(tool, prompt_ids=None) -> list[str]: + """Return names of multi-variable lookups linked to the tool, [] in OSS. + + ``prompt_ids`` scopes the check to a specific subset of linked prompts + so single / bulk runs only block when a lookup the run actually uses + is multi-variable. + """ + try: + from pluggable_apps.lookup_v1.execution import has_multi_var_lookups + + _, names = has_multi_var_lookups(tool, prompt_ids=prompt_ids) + return names + except ImportError: + return [] + + def persist_lookup_output(prompt_output, prompt_lookup: dict) -> None: """Persist lookup enrichment result. No-op in OSS.""" try: @@ -69,3 +97,25 @@ def validate_lookups_for_export(prompts) -> tuple[dict, str | None]: return _validate(prompts) except ImportError: return {}, None + + +def get_lookup_validation_for_tool(tool) -> dict: + """Pre-emptive lookup validation for FE Export / Deploy gating. + + Returns an "always ok" payload in OSS so the FE gate is a no-op. + """ + try: + from pluggable_apps.lookup_v1.validation import ( + get_lookup_validation_for_tool as _validate, + ) + + return _validate(tool) + except ImportError: + return { + "ok": True, + "draft_lookups": [], + "multi_var_lookups": [], + "single_pass_enabled": bool( + getattr(tool, "single_pass_extraction_mode", False) + ), + } diff --git a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index 6744171ed4..ee8f1a4c8b 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py +++ b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py @@ -20,7 +20,7 @@ from utils.local_context import StateStore from backend.celery_service import app as celery_app -from prompt_studio.lookup_utils import get_lookup_config +from prompt_studio.lookup_utils import get_lookup_config, get_lookup_configs_for_tool from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_profile_manager_v2.profile_manager_helper import ( ProfileManagerHelper, @@ -1173,6 +1173,11 @@ def build_single_pass_payload( TSPKeys.SIMILARITY_TOP_K: default_profile.similarity_top_k, } + # Inject lookup configs for single pass enrichment + lookup_configs = get_lookup_configs_for_tool(tool) + if lookup_configs: + tool_settings["lookup_configs"] = lookup_configs + for p in prompts: if not p.prompt: raise EmptyPromptError() diff --git a/backend/prompt_studio/prompt_studio_core_v2/urls.py b/backend/prompt_studio/prompt_studio_core_v2/urls.py index 86cbb97dd3..9163e8736e 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/urls.py +++ b/backend/prompt_studio/prompt_studio_core_v2/urls.py @@ -66,6 +66,10 @@ prompt_studio_task_status = PromptStudioCoreView.as_view({"get": "task_status"}) +prompt_studio_lookup_validation = PromptStudioCoreView.as_view( + {"get": "lookup_validation"} +) + urlpatterns = format_suffix_patterns( [ @@ -165,5 +169,10 @@ prompt_studio_task_status, name="prompt-studio-task-status", ), + path( + "prompt-studio//lookup-validation/", + prompt_studio_lookup_validation, + name="prompt-studio-lookup-validation", + ), ] ) diff --git a/backend/prompt_studio/prompt_studio_core_v2/views.py b/backend/prompt_studio/prompt_studio_core_v2/views.py index e936e3dcae..96fd4c8702 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/views.py @@ -32,6 +32,10 @@ from workflow_manager.endpoint_v2.models import WorkflowEndpoint from backend.celery_service import app as celery_app +from prompt_studio.lookup_utils import ( + get_lookup_validation_for_tool, + get_multi_var_lookups_for_tool, +) from prompt_studio.prompt_profile_manager_v2.constants import ( ProfileManagerErrors, ProfileManagerKeys, @@ -89,6 +93,34 @@ logger = logging.getLogger(__name__) +def _multi_var_lookup_block_response(custom_tool, prompt_ids=None): + """Block non-SP runs when a linked lookup has >1 input variable. + + ``prompt_ids`` scopes the check to the prompt(s) being run so a + multi-var lookup attached to an unrelated prompt in the same project + doesn't block a single-var lookup's run. + + Returns a Response object (HTTP 400) when a block applies, or None + to let the caller proceed. + """ + if getattr(custom_tool, "single_pass_extraction_mode", False): + return None + names = get_multi_var_lookups_for_tool(custom_tool, prompt_ids=prompt_ids) + if not names: + return None + return Response( + { + "error": ( + "Multi-variable lookup(s) " + f"{', '.join(names)} are linked to prompts in this project. " + "These can only run in single pass extraction mode. " + "Enable single pass or unlink the lookup before running." + ) + }, + status=status.HTTP_400_BAD_REQUEST, + ) + + class PromptStudioCoreView(viewsets.ModelViewSet): """Viewset to handle all Custom tool related operations.""" @@ -476,6 +508,10 @@ def fetch_response(self, request: HttpRequest, pk: Any = None) -> Response: document_id: str = request.data.get(ToolStudioPromptKeys.DOCUMENT_ID) prompt_id: str = request.data.get(ToolStudioPromptKeys.ID) run_id: str = request.data.get(ToolStudioPromptKeys.RUN_ID) + if err := _multi_var_lookup_block_response( + custom_tool, prompt_ids=[prompt_id] if prompt_id else None + ): + return err profile_manager_id: str = request.data.get( ToolStudioPromptKeys.PROFILE_MANAGER_ID ) @@ -577,6 +613,8 @@ def bulk_fetch_response(self, request: HttpRequest, pk: Any = None) -> Response: {"error": "prompt_ids is required and must be non-empty."}, status=status.HTTP_400_BAD_REQUEST, ) + if err := _multi_var_lookup_block_response(custom_tool, prompt_ids=prompt_ids): + return err document_id: str = request.data.get(ToolStudioPromptKeys.DOCUMENT_ID) run_id: str = request.data.get(ToolStudioPromptKeys.RUN_ID) profile_manager_id: str = request.data.get( @@ -1060,6 +1098,15 @@ def export_tool(self, request: Request, pk: Any = None) -> Response: status=status.HTTP_200_OK, ) + @action(detail=True, methods=["get"], url_path="lookup-validation") + def lookup_validation(self, request: Request, pk: Any = None) -> Response: + """Pre-emptive lookup gating for Export / API Deployment buttons. + + Cloud-only check; OSS returns ``ok: True`` so the FE proceeds. + """ + custom_tool = self.get_object() + return Response(get_lookup_validation_for_tool(custom_tool)) + @action(detail=True, methods=["get"]) def export_tool_info(self, request: Request, pk: Any = None) -> Response: custom_tool = self.get_object() diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py index 0fb92c126d..056e7feb5d 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/output_manager_helper.py @@ -3,6 +3,7 @@ from typing import Any from django.core.exceptions import ObjectDoesNotExist +from django.utils import timezone from prompt_studio.lookup_utils import persist_lookup_output from prompt_studio.prompt_profile_manager_v2.models import ProfileManager @@ -107,6 +108,8 @@ def update_or_create_prompt_output( "highlight_data": highlight_data, "confidence_data": confidence_data, "word_confidence_data": word_confidence_data, + # QuerySet.update() bypasses auto_now on modified_at. + "modified_at": timezone.now(), } PromptStudioOutputManager.objects.filter( document_manager=document_manager, @@ -257,17 +260,20 @@ def fetch_default_output_response( Returns: dict[str, Any]: Formatted JSON response for combined output. + Includes a reserved ``_lookup_outputs`` key with per-prompt + enriched data when lookups are configured. """ - # Initialize the result dictionary + from prompt_studio.lookup_utils import enrich_prompt_output + result: dict[str, Any] = {} - # Iterate over ToolStudioPrompt records + lookup_outputs: dict[str, Any] = {} + for tool_prompt in tool_studio_prompts: if tool_prompt.prompt_type == PSOMKeys.NOTES: continue prompt_id = str(tool_prompt.prompt_id) profile_manager_id = tool_prompt.profile_manager_id - # If profile_manager is not set, skip this record if not profile_manager_id and not use_default_profile: result[tool_prompt.prompt_key] = "" continue @@ -292,6 +298,15 @@ def fetch_default_output_response( for output in queryset: result[tool_prompt.prompt_key] = output.output + # Check for lookup enrichment + enriched = enrich_prompt_output(output, {}) + if "lookup_outputs" in enriched: + lookup_outputs[tool_prompt.prompt_key] = enriched[ + "lookup_outputs" + ] except ObjectDoesNotExist: result[tool_prompt.prompt_key] = "" + + if lookup_outputs: + result["_lookup_outputs"] = lookup_outputs return result diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/urls.py b/backend/prompt_studio/prompt_studio_output_manager_v2/urls.py index 61ec8540fa..77270c22fe 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/urls.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/urls.py @@ -7,6 +7,7 @@ get_output_for_tool_default = PromptStudioOutputView.as_view( {"get": "get_output_for_tool_default"} ) +latest_outputs_by_keys = PromptStudioOutputView.as_view({"get": "latest_outputs_by_keys"}) urlpatterns = format_suffix_patterns( [ @@ -16,5 +17,10 @@ get_output_for_tool_default, name="prompt-default-profile-outputs", ), + path( + "prompt-output/latest-by-keys/", + latest_outputs_by_keys, + name="prompt-output-latest-by-keys", + ), ] ) diff --git a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py index c0c002d803..5a31c481dc 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py @@ -61,6 +61,55 @@ def get_queryset(self) -> QuerySet | None: return queryset + def latest_outputs_by_keys(self, request: HttpRequest) -> Response: + """Return the most recent raw output value per source prompt key. + + Used by the lookup Test panel's "Use Latest Outputs" button to + pre-fill {{input.X}} fields from prior prompt runs in the linked + tool. Always returns the raw extraction — the enriched value would + already include lookup post-processing, which would defeat the + purpose of testing the lookup with sample inputs. + """ + tool_id = request.GET.get("tool_id") + keys_param = request.GET.get("prompt_keys", "") + if not tool_id: + raise APIException( + detail=PromptOutputManagerErrorMessage.TOOL_VALIDATION, + code=400, + ) + + prompt_keys = [k.strip() for k in keys_param.split(",") if k.strip()] + if not prompt_keys: + return Response({}, status=status.HTTP_200_OK) + + prompt_id_to_key = dict( + ToolStudioPrompt.objects.filter( + tool_id=tool_id, prompt_key__in=prompt_keys + ).values_list("prompt_id", "prompt_key") + ) + if not prompt_id_to_key: + return Response({}, status=status.HTTP_200_OK) + + outputs = ( + PromptStudioOutputManager.objects.filter( + prompt_id__in=prompt_id_to_key.keys() + ) + .exclude(output__isnull=True) + .exclude(output__exact="") + .order_by("-modified_at") + .values("prompt_id", "output") + ) + + result: dict[str, str] = {} + for row in outputs: + key = prompt_id_to_key.get(row["prompt_id"]) + if key and key not in result: + result[key] = row["output"] + if len(result) == len(prompt_id_to_key): + break + + return Response(result, status=status.HTTP_200_OK) + def get_output_for_tool_default(self, request: HttpRequest) -> Response: # Get the tool_id from request parameters # TODO: Setup Serializer here diff --git a/frontend/src/components/custom-tools/combined-output/CombinedOutput.jsx b/frontend/src/components/custom-tools/combined-output/CombinedOutput.jsx index 3036e7b9d6..20993ffc4b 100644 --- a/frontend/src/components/custom-tools/combined-output/CombinedOutput.jsx +++ b/frontend/src/components/custom-tools/combined-output/CombinedOutput.jsx @@ -59,6 +59,7 @@ function CombinedOutput({ docId, setFilledFields, selectedPrompts }) { } = useCustomToolStore(); const [combinedOutput, setCombinedOutput] = useState({}); + const [enrichedOutput, setEnrichedOutput] = useState({}); const [isOutputLoading, setIsOutputLoading] = useState(false); const [adapterData, setAdapterData] = useState([]); const [activeKey, setActiveKey] = useState( @@ -117,29 +118,60 @@ function CombinedOutput({ docId, setFilledFields, selectedPrompts }) { const prompts = details?.prompts || []; if (activeKey === "0" && !isSimplePromptStudio) { + const lookupOutputs = data?._lookup_outputs || {}; const output = Object.entries(data).reduce((acc, [key, value]) => { + if (key === "_lookup_outputs") return acc; acc[key] = displayPromptResult(value, false); return acc; }, {}); setCombinedOutput(output); + + if (Object.keys(lookupOutputs).length > 0) { + const enriched = {}; + for (const [key, val] of Object.entries(output)) { + const lookupData = lookupOutputs[key]; + enriched[key] = lookupData?.output + ? displayPromptResult(lookupData.output, false) + : val; + } + setEnrichedOutput(enriched); + } else { + setEnrichedOutput({}); + } } else { - const output = prompts.reduce((acc, item) => { - if (item?.prompt_type !== promptType.notes) { - const profileManager = selectedProfile || item?.profile_manager; - const outputDetails = data.find( - (outputValue) => - outputValue?.prompt_id === item?.prompt_id && - outputValue?.profile_manager === profileManager, - ); + const output = {}; + const enriched = {}; + let hasEnriched = false; + + for (const item of prompts) { + if (item?.prompt_type === promptType.notes) continue; + const profileManager = selectedProfile || item?.profile_manager; + const outputDetails = data.find( + (outputValue) => + outputValue?.prompt_id === item?.prompt_id && + outputValue?.profile_manager === profileManager, + ); + + output[item?.prompt_key] = + outputDetails?.output?.length > 0 + ? displayPromptResult(outputDetails?.output, false) + : ""; - acc[item?.prompt_key] = - outputDetails?.output?.length > 0 - ? displayPromptResult(outputDetails?.output, false) - : ""; + // Build enriched output from lookup_outputs + const lookupData = outputDetails?.lookup_outputs; + if (lookupData?.output) { + enriched[item?.prompt_key] = displayPromptResult( + lookupData.output, + false, + ); + hasEnriched = true; + } else { + enriched[item?.prompt_key] = output[item?.prompt_key]; } - return acc; - }, {}); + } + setCombinedOutput(output); + setEnrichedOutput(hasEnriched ? enriched : {}); } } catch (err) { setAlertDetails( @@ -229,6 +261,7 @@ function CombinedOutput({ docId, setFilledFields, selectedPrompts }) { return ( { Prism.highlightAll(); - }, [combinedOutput]); + }, [combinedOutput, enrichedOutput, activeView]); + + // Reset to Raw when enriched data is not available + useEffect(() => { + if (!enrichedOutput || Object.keys(enrichedOutput).length === 0) { + setActiveView("Raw"); + } + }, [enrichedOutput]); + + const displayOutput = + activeView === "Enriched" && + enrichedOutput && + Object.keys(enrichedOutput).length > 0 + ? enrichedOutput + : combinedOutput; return (
@@ -34,14 +61,24 @@ function JsonView({ /> ))} -
+
+ {EnrichedOutputToggle && ( + 0) + } + /> + )} +
@@ -51,6 +88,7 @@ function JsonView({ JsonView.propTypes = { combinedOutput: PropTypes.object.isRequired, + enrichedOutput: PropTypes.object, handleTabChange: PropTypes.func, adapterData: PropTypes.array, selectedProfile: PropTypes.string, diff --git a/frontend/src/components/custom-tools/header/Header.jsx b/frontend/src/components/custom-tools/header/Header.jsx index 15e41c8a0a..5421f299b6 100644 --- a/frontend/src/components/custom-tools/header/Header.jsx +++ b/frontend/src/components/custom-tools/header/Header.jsx @@ -18,6 +18,10 @@ import "./Header.css"; let SinglePassToggleSwitch; let CloneButton; let PromptShareButton; +let useLookupExportGate = () => ({ + checkLookups: () => Promise.resolve(true), + modalEl: null, +}); try { const mod = await import( "../../../plugins/single-pass-toggle-switch/SinglePassToggleSwitch" @@ -26,6 +30,14 @@ try { } catch { // The variable will remain undefined if the component is not available. } +try { + const mod = await import( + "../../../plugins/lookup-studio/hooks/useLookupExportGate" + ); + useLookupExportGate = mod.useLookupExportGate; +} catch { + // OSS — gate stays a no-op resolving true. +} try { const mod = await import( "../../../plugins/prompt-studio-public-share/public-share-btn/PromptShareButton.jsx" @@ -72,6 +84,7 @@ function Header({ const [isApiDeploymentLoading, setIsApiDeploymentLoading] = useState(false); const [editModalOpen, setEditModalOpen] = useState(false); const [editForm] = Form.useForm(); + const { checkLookups, modalEl: lookupGateModalEl } = useLookupExportGate(); const handleExport = ( selectedUsers, @@ -129,7 +142,7 @@ function Header({ setConfirmModalVisible(false); }, [lastExportParams, handleExport]); - const handleShare = (isEdit) => { + const handleShare = async (isEdit) => { try { setPostHogCustomEvent("ps_exported_tool", { info: `Clicked on the 'Export' button`, @@ -139,6 +152,9 @@ function Header({ // If an error occurs while setting custom posthog event, ignore it and continue } + const ok = await checkLookups(details?.tool_id, "export"); + if (!ok) return; + const requestOptions = { method: "GET", url: `/api/v1/unstract/${sessionDetails?.orgId}/prompt-studio/export/${details?.tool_id}`, @@ -255,7 +271,7 @@ function Header({ }); }; - const handleCreateApiDeployment = () => { + const handleCreateApiDeployment = async () => { try { setPostHogCustomEvent("intent_create_api_deployment_from_prompt_studio", { info: "Clicked Create API Deployment in tool IDE", @@ -266,6 +282,9 @@ function Header({ // If an error occurs while setting custom posthog event, ignore it and continue } + const ok = await checkLookups(details?.tool_id, "API deployment"); + if (!ok) return; + // Check for existing API deployments before proceeding setIsApiDeploymentLoading(true); const path = `/api/v1/unstract/${sessionDetails.orgId}`; @@ -422,6 +441,7 @@ function Header({ return ( <> + {lookupGateModalEl} Date: Wed, 22 Apr 2026 01:19:15 +0530 Subject: [PATCH 20/57] UN-2946 [UI] Fix combined output pill overlap and preserve Look-Ups tab on back MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ProfileInfoBar: swap Row/Col for plain flex-wrap div — kills Ant Row negative-margin quirk that overlapped wrapped pills in combined output. - CustomTools: honor location.state.activeTab so back navigation from lookup detail lands on the Look-Ups tab instead of defaulting to Projects. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../profile-info-bar/ProfileInfoBar.css | 9 ++- .../profile-info-bar/ProfileInfoBar.jsx | 62 +++++++------------ frontend/src/pages/CustomTools.jsx | 13 +++- 3 files changed, 43 insertions(+), 41 deletions(-) diff --git a/frontend/src/components/custom-tools/profile-info-bar/ProfileInfoBar.css b/frontend/src/components/custom-tools/profile-info-bar/ProfileInfoBar.css index f9f89d03a9..da866a895f 100644 --- a/frontend/src/components/custom-tools/profile-info-bar/ProfileInfoBar.css +++ b/frontend/src/components/custom-tools/profile-info-bar/ProfileInfoBar.css @@ -1,5 +1,10 @@ .profile-info-bar { - margin-bottom: 10px; + display: flex; flex-wrap: wrap; - gap: 4px 0; + gap: 6px 8px; + margin-bottom: 10px; +} + +.profile-info-bar .ant-tag { + margin: 0; } diff --git a/frontend/src/components/custom-tools/profile-info-bar/ProfileInfoBar.jsx b/frontend/src/components/custom-tools/profile-info-bar/ProfileInfoBar.jsx index 8f52a68a49..9341753987 100644 --- a/frontend/src/components/custom-tools/profile-info-bar/ProfileInfoBar.jsx +++ b/frontend/src/components/custom-tools/profile-info-bar/ProfileInfoBar.jsx @@ -1,4 +1,4 @@ -import { Col, Row, Tag } from "antd"; +import { Tag } from "antd"; import PropTypes from "prop-types"; import "./ProfileInfoBar.css"; @@ -10,43 +10,29 @@ const ProfileInfoBar = ({ profiles, profileId }) => { } return ( - - - - Profile Name: {profile?.profile_name} - - - - - Chunk Size: {profile?.chunk_size} - - - - - Vector Store: {profile?.vector_store} - - - - - Embedding Model: {profile?.embedding_model} - - - - - LLM: {profile?.llm} - - - - - X2Text: {profile?.x2text} - - - - - Reindex: {profile?.reindex ? "Yes" : "No"} - - - +
+ + Profile Name: {profile?.profile_name} + + + Chunk Size: {profile?.chunk_size} + + + Vector Store: {profile?.vector_store} + + + Embedding Model: {profile?.embedding_model} + + + LLM: {profile?.llm} + + + X2Text: {profile?.x2text} + + + Reindex: {profile?.reindex ? "Yes" : "No"} + +
); }; diff --git a/frontend/src/pages/CustomTools.jsx b/frontend/src/pages/CustomTools.jsx index 1f6c2d3934..9543fb5c16 100644 --- a/frontend/src/pages/CustomTools.jsx +++ b/frontend/src/pages/CustomTools.jsx @@ -1,12 +1,16 @@ import { useEffect, useState } from "react"; +import { useLocation } from "react-router-dom"; import { ListOfTools } from "../components/custom-tools/list-of-tools/ListOfTools"; const TAB_OPTIONS = ["Projects", "Look-Ups"]; function CustomTools() { + const location = useLocation(); const [LookupListComp, setLookupListComp] = useState(null); - const [activeTab, setActiveTab] = useState("Projects"); + const [activeTab, setActiveTab] = useState( + location.state?.activeTab || "Projects", + ); useEffect(() => { import("../plugins/lookup-studio") @@ -14,6 +18,13 @@ function CustomTools() { .catch(() => {}); }, []); + // Honor tab from navigation state on subsequent entries + useEffect(() => { + if (location.state?.activeTab) { + setActiveTab(location.state.activeTab); + } + }, [location.state?.activeTab]); + // No lookup plugin = just render projects list (OSS mode) if (!LookupListComp) { return ; From d63767b514aff77142654203dfd5d3ef420aa5ea Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Wed, 22 Apr 2026 04:07:35 +0530 Subject: [PATCH 21/57] UN-2946 [FEAT] Add last_exported_at and wire lookup staleness bridge Introduces nullable last_exported_at on CustomTool (populated on first successful export) so staleness checks can compare against downstream mutations without a data backfill. NULL is treated as "unknown" and suppresses the lookup-dirty flag to avoid false alarms on pre-feature projects. Adds the get_latest_lookup_mutation_for_tool bridge in lookup_utils so OSS stays decoupled from the cloud plugin. Co-Authored-By: Claude Opus 4.7 (1M context) --- backend/prompt_studio/lookup_utils.py | 16 ++++++++++++++ .../0007_customtool_last_exported_at.py | 21 +++++++++++++++++++ .../prompt_studio_core_v2/models.py | 9 ++++++++ .../prompt_studio_core_v2/views.py | 17 +++++++++++++++ 4 files changed, 63 insertions(+) create mode 100644 backend/prompt_studio/prompt_studio_core_v2/migrations/0007_customtool_last_exported_at.py diff --git a/backend/prompt_studio/lookup_utils.py b/backend/prompt_studio/lookup_utils.py index 5625c33878..474f9a5d52 100644 --- a/backend/prompt_studio/lookup_utils.py +++ b/backend/prompt_studio/lookup_utils.py @@ -99,6 +99,22 @@ def validate_lookups_for_export(prompts) -> tuple[dict, str | None]: return {}, None +def get_latest_lookup_mutation_for_tool(tool): + """Return the max modified_at across all lookup-related records linked to + the tool (version, reference file, assignment). Used for banner staleness. + + Returns None if lookups are unavailable or nothing is linked. + """ + try: + from pluggable_apps.lookup_v1.staleness import ( + get_latest_lookup_mutation_for_tool as _get, + ) + + return _get(tool) + except ImportError: + return None + + def get_lookup_validation_for_tool(tool) -> dict: """Pre-emptive lookup validation for FE Export / Deploy gating. diff --git a/backend/prompt_studio/prompt_studio_core_v2/migrations/0007_customtool_last_exported_at.py b/backend/prompt_studio/prompt_studio_core_v2/migrations/0007_customtool_last_exported_at.py new file mode 100644 index 0000000000..91539c9cf9 --- /dev/null +++ b/backend/prompt_studio/prompt_studio_core_v2/migrations/0007_customtool_last_exported_at.py @@ -0,0 +1,21 @@ +# Generated by Django 4.2.1 on 2026-04-21 20:20 + +from django.db import migrations, models + + +class Migration(migrations.Migration): + dependencies = [ + ("prompt_studio_core_v2", "0006_add_custom_data_to_customtool"), + ] + + operations = [ + migrations.AddField( + model_name="customtool", + name="last_exported_at", + field=models.DateTimeField( + blank=True, + db_comment="Timestamp of the last successful export; NULL if never exported since the field was introduced.", + null=True, + ), + ), + ] diff --git a/backend/prompt_studio/prompt_studio_core_v2/models.py b/backend/prompt_studio/prompt_studio_core_v2/models.py index 406c157efc..f7d1d2b9e5 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/models.py +++ b/backend/prompt_studio/prompt_studio_core_v2/models.py @@ -161,6 +161,15 @@ class CustomTool(DefaultOrganizationMixin, BaseModel): db_comment="Flag to share this custom tool with all users in the organization", ) + # NULL on pre-feature tools; populated on first successful export. + # Drives staleness checks (e.g. lookup-change banner) without requiring + # a data backfill. + last_exported_at = models.DateTimeField( + null=True, + blank=True, + db_comment="Timestamp of the last successful export; NULL if never exported since the field was introduced.", + ) + objects = CustomToolModelManager() def delete(self, organization_id=None, *args, **kwargs): diff --git a/backend/prompt_studio/prompt_studio_core_v2/views.py b/backend/prompt_studio/prompt_studio_core_v2/views.py index 96fd4c8702..a41a40017f 100644 --- a/backend/prompt_studio/prompt_studio_core_v2/views.py +++ b/backend/prompt_studio/prompt_studio_core_v2/views.py @@ -14,6 +14,7 @@ from django.db import IntegrityError from django.db.models import Count, OuterRef, QuerySet, Subquery from django.http import HttpRequest, HttpResponse +from django.utils import timezone from file_management.constants import FileInformationKey as FileKey from file_management.exceptions import FileNotFound from permissions.permission import IsOwner, IsOwnerOrSharedUserOrSharedToOrg @@ -33,6 +34,7 @@ from backend.celery_service import app as celery_app from prompt_studio.lookup_utils import ( + get_latest_lookup_mutation_for_tool, get_lookup_validation_for_tool, get_multi_var_lookups_for_tool, ) @@ -1085,6 +1087,11 @@ def export_tool(self, request: Request, pk: Any = None) -> Response: force_export=force_export, ) + # Record export timestamp so staleness checks (e.g. lookup-change + # banner) can compare against mutations that happened afterwards. + custom_tool.last_exported_at = timezone.now() + custom_tool.save(update_fields=["last_exported_at"]) + # Notify HubSpot about first tool export notify_hubspot_event( user=request.user, @@ -1285,10 +1292,20 @@ def check_deployment_usage(self, request: Request, pk: Any = None) -> Response: instance: CustomTool = self.get_object() is_used, workflow_ids = self._check_tool_usage_in_workflows(instance) + # Lookup staleness: NULL last_exported_at means we can't compare, + # so treat as clean (don't false-alarm pre-feature projects). + is_lookup_dirty = False + if instance.last_exported_at is not None: + latest = get_latest_lookup_mutation_for_tool(instance) + is_lookup_dirty = ( + latest is not None and latest > instance.last_exported_at + ) + deployment_info: dict = { "is_used": is_used, "deployment_types": [], "message": "", + "is_lookup_dirty": is_lookup_dirty, } if is_used and workflow_ids: From eced338dd159148e1be54a52259204f828e8a875 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Wed, 22 Apr 2026 04:07:41 +0530 Subject: [PATCH 22/57] UN-2946 [FEAT] Stream lookup enrichment failures to workflow logs When enricher.run() raises, surface a user-visible ERROR log line in the workflow execution log alongside the existing usage record. Keeps lookup failures observable next to the other pre/post lookup lines we already emit. Co-Authored-By: Claude Opus 4.7 (1M context) --- workers/executor/executors/legacy_executor.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/workers/executor/executors/legacy_executor.py b/workers/executor/executors/legacy_executor.py index 42210210ef..f66e6d9fa0 100644 --- a/workers/executor/executors/legacy_executor.py +++ b/workers/executor/executors/legacy_executor.py @@ -1921,6 +1921,11 @@ def _run_lookup_enrichment( enricher.run() except Exception as e: logger.warning("Lookup enrichment failed for %s: %s", prompt_name, e) + lookup_label = lookup_config.get("lookup_name") or prompt_name + shim.stream_log( + f"Lookup `{lookup_label}` failed: {str(e)[:200]}", + level=LogLevel.ERROR, + ) error_record = { "usage_type": "llm", "llm_usage_reason": "lookup", From 58f2dbd162e25279e86ec48f0313c59a45b37a52 Mon Sep 17 00:00:00 2001 From: Chandrasekharan M Date: Wed, 22 Apr 2026 04:08:09 +0530 Subject: [PATCH 23/57] UN-2946 [UI] Wire lookup dirty-seed and export gate into ToolIde Loads useLookupDirtySeed (server-side is_lookup_dirty) and useLookupExportGate from the cloud plugin via dynamic imports so the reminder banner reflects lookup changes across page reloads and the banner's Export flow goes through the same validation modal as the main buttons. Also adds a titleAdornment slot on ToolNavBar for rendering the onboarding tooltip and relaxes EmptyState.text to accept nodes for the tagline + link composition. Co-Authored-By: Claude Opus 4.7 (1M context) --- .../custom-tools/tool-ide/ToolIde.jsx | 42 +++++++++++++++++++ .../navigations/tool-nav-bar/ToolNavBar.jsx | 3 ++ .../widgets/empty-state/EmptyState.jsx | 2 +- 3 files changed, 46 insertions(+), 1 deletion(-) diff --git a/frontend/src/components/custom-tools/tool-ide/ToolIde.jsx b/frontend/src/components/custom-tools/tool-ide/ToolIde.jsx index 93c8198593..ab87f6f19c 100644 --- a/frontend/src/components/custom-tools/tool-ide/ToolIde.jsx +++ b/frontend/src/components/custom-tools/tool-ide/ToolIde.jsx @@ -44,6 +44,36 @@ try { } catch { // Do nothing if plugins are not loaded. } + +// Cloud-only hook that seeds hasUnsavedChanges from server-side +// lookup-staleness. No-op stub in OSS. +let useLookupDirtySeed = () => { + // no-op +}; +try { + const mod = await import( + "../../../plugins/lookup-studio/hooks/useLookupDirtySeed.js" + ); + useLookupDirtySeed = mod.useLookupDirtySeed; +} catch { + // Do nothing if plugin is not loaded. +} + +// Cloud-only lookup export validation gate. OSS stub resolves true so +// the reminder bar's "Export" button proceeds directly. +let useLookupExportGate = () => ({ + checkLookups: () => Promise.resolve(true), + modalEl: null, +}); +try { + const mod = await import( + "../../../plugins/lookup-studio/hooks/useLookupExportGate" + ); + useLookupExportGate = mod.useLookupExportGate; +} catch { + // OSS — gate stays a no-op resolving true. +} + function ToolIde() { const [openSettings, setOpenSettings] = useState(false); const customToolStore = useCustomToolStore(); @@ -76,6 +106,7 @@ function ToolIde() { const isCheckingUsageRef = useRef(false); const hasCheckedForCurrentSessionRef = useRef(false); const abortControllerRef = useRef(null); + const { checkLookups, modalEl: lookupGateModalEl } = useLookupExportGate(); useEffect(() => { if (openShareModal) { @@ -178,6 +209,11 @@ function ToolIde() { } }, [details?.tool_id]); + // Cloud plugin seeds hasUnsavedChanges when a linked lookup has been + // edited since the tool's last export — surfaces the re-export banner + // for mutations made on the standalone /lookups page. No-op in OSS. + useLookupDirtySeed(details?.tool_id); + // Cleanup abort controller on unmount useEffect(() => { return () => { @@ -189,6 +225,10 @@ function ToolIde() { // Handle export from reminder bar const handleExportFromReminder = useCallback(async () => { + const ok = await checkLookups(details?.tool_id, "export"); + if (!ok) { + return; + } setIsExporting(true); try { const requestOptions = { @@ -237,6 +277,7 @@ function ToolIde() { handleException, markChangesAsExported, setPostHogCustomEvent, + checkLookups, ]); const generateIndex = async (doc) => { @@ -340,6 +381,7 @@ function ToolIde() { isExporting={isExporting} /> )} + {lookupGateModalEl}
{title} + {titleAdornment} {onEditTitle && ( - +