diff --git a/.gitignore b/.gitignore index f6837ce079..c655f4acff 100644 --- a/.gitignore +++ b/.gitignore @@ -653,7 +653,11 @@ docker/*.env !docker/sample*.env docker/public_tools.json docker/proxy_overrides.yaml -docker/compose.override.yaml +docker/compose.*.yaml +# ``docker/compose.debug.yaml`` is checked-in tooling — keep it out of the +# broader ``compose.*.yaml`` ignore so a delete + recreate doesn't make it +# look untracked, and so teammates can spot it. +!docker/compose.debug.yaml docker/workflow_data/ # Tool development @@ -696,6 +700,12 @@ CLAUDE.md CONTRIBUTION_GUIDE.md .mcp.json +# Codex +AGENTS.md + +# Pi +.pi/ + # Windsurf .qodo .windsurfrules diff --git a/backend/prompt_studio/lookup_utils.py b/backend/prompt_studio/lookup_utils.py new file mode 100644 index 0000000000..8cfb0a7286 --- /dev/null +++ b/backend/prompt_studio/lookup_utils.py @@ -0,0 +1,171 @@ +"""Shared utility for lookup operations. No-ops in OSS. + +Only the absence of ``pluggable_apps.lookups`` itself is treated as +"cloud not installed"; an ImportError from a transitive dependency +re-raises so we don't silently degrade to a no-op on a real bug. +""" + +import logging +from typing import Any + +logger = logging.getLogger(__name__) + +_CLOUD_LOOKUP_MODULES = { + # OSS images lack the parent ``pluggable_apps`` package, so include it. + "pluggable_apps", + "pluggable_apps.lookups", + "pluggable_apps.lookups.execution", + "pluggable_apps.lookups.output_enrichment", + "pluggable_apps.lookups.staleness", + "pluggable_apps.lookups.validation", + "pluggable_apps.lookups.models", +} + +try: + from pluggable_apps.lookups import execution as _execution + from pluggable_apps.lookups import output_enrichment as _output_enrichment + from pluggable_apps.lookups import staleness as _staleness + from pluggable_apps.lookups import validation as _validation + from pluggable_apps.lookups.models import LookupOutputResult as _LookupOutputResult + + LOOKUPS_AVAILABLE = True +except ImportError as e: + if e.name not in _CLOUD_LOOKUP_MODULES: + raise + LOOKUPS_AVAILABLE = False + + +def get_lookup_config(prompt) -> dict | None: + """Return lookup config for a prompt, or None if lookups are unavailable.""" + if not LOOKUPS_AVAILABLE: + return None + return _execution.build_lookup_config_for_prompt(prompt) + + +def get_lookup_configs_for_tool(tool, prompts=None) -> list[dict] | None: + """Return lookup configs for a tool (single pass), or None in OSS. + + ``prompts`` scopes validation to the run's prompts so unrelated + incomplete assignments on the tool don't block it. + """ + if not LOOKUPS_AVAILABLE: + return None + return _execution.build_lookup_configs_for_tool(tool, prompts=prompts) + + +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 so a run is only blocked when the + multi-var lookup is actually used by it. + """ + if not LOOKUPS_AVAILABLE: + return [] + _, names = _execution.has_multi_var_lookups(tool, prompt_ids=prompt_ids) + return names + + +def persist_lookup_output(prompt_output, prompt_lookup: dict) -> None: + """Persist lookup enrichment result. No-op in OSS.""" + if not LOOKUPS_AVAILABLE: + return + lookup_meta = prompt_lookup.get("meta", {}) + lookup_id = lookup_meta.get("lookup_id") + if not lookup_id: + return + 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=defaults, + ) + + +def enrich_prompt_output(prompt_output, data: dict) -> dict: + """Let cloud plugins enrich serialized prompt output with lookup data. + + No-op in OSS. + """ + if not LOOKUPS_AVAILABLE: + return data + return _output_enrichment.enrich_with_lookup_output(prompt_output, data) + + +def validate_lookups_for_export(prompts) -> tuple[dict, str | None]: + """Validate lookup assignments before export. Returns ({}, None) in OSS.""" + if not LOOKUPS_AVAILABLE: + return {}, None + return _validation.validate_lookups_for_export(prompts) + + +def get_latest_lookup_mutation_for_tool(tool): + """Max ``modified_at`` across lookup-related records linked to the tool + (version, reference file, assignment) — feeds the staleness banner. + None if unavailable or nothing linked. + """ + if not LOOKUPS_AVAILABLE: + return None + return _staleness.get_latest_lookup_mutation_for_tool(tool) + + +def get_original_value_if_enriched( + metadata: dict, prompt_key: str +) -> tuple[Any, dict] | None: + """Return ``(original_value, prompt_lookup_dict)`` if ``prompt_key`` was + enriched, or ``None`` otherwise. + + Pure metadata-shape check — safe to call even when LOOKUPS_AVAILABLE + is False (returns None because the shape won't match). + """ + if not isinstance(metadata, dict): + return None + lookup_outputs = metadata.get("lookup_outputs") or {} + prompt_lookup = lookup_outputs.get(prompt_key) + if isinstance(prompt_lookup, dict) and "original" in prompt_lookup: + return prompt_lookup.get("original"), prompt_lookup + return None + + +def attach_combined_output_enrichment(result: dict, enriched_by_key: dict) -> None: + """Stamp the combined-output payload with enriched-output metadata. + + Key name stays cloud-side so the FE-plugin shape can evolve without + coordinating with OSS. + """ + if not LOOKUPS_AVAILABLE: + return + _output_enrichment.attach_combined_output_enrichment(result, enriched_by_key) + + +def extract_prompt_output_enrichment(item) -> dict | None: + """Pick enriched-output data off a serialized prompt-output row. + + Returns a plugin-opaque dict (FE-only) or None when no enrichment + is present / plugin missing. + """ + if not LOOKUPS_AVAILABLE: + return None + return _output_enrichment.extract_prompt_output_enrichment(item) + + +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. + """ + if not LOOKUPS_AVAILABLE: + return { + "ok": True, + "draft_lookups": [], + "multi_var_lookups": [], + "incomplete_lookups": [], + "single_pass_enabled": bool( + getattr(tool, "single_pass_extraction_mode", False) + ), + } + return _validation.get_lookup_validation_for_tool(tool) 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 1e1802b776..cd2a12dac8 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/prompt_studio_helper.py b/backend/prompt_studio/prompt_studio_core_v2/prompt_studio_helper.py index ff2b345505..62cbf2d9a2 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 @@ -14,12 +14,17 @@ from django.db import transaction from django.db.models.manager import BaseManager from plugins import get_plugin +from rest_framework.exceptions import APIException from rest_framework.request import Request from utils.file_storage.constants import FileStorageKeys from utils.file_storage.helpers.prompt_studio_file_helper import PromptStudioFileHelper from utils.local_context import StateStore from backend.celery_service import app as celery_app +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, @@ -387,6 +392,9 @@ def _build_prompt_output( if webhook_enabled: output[TSPKeys.POSTPROCESSING_WEBHOOK_URL] = webhook_url + 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 output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_MONITOR_LLM] = [monitor_llm] @@ -798,6 +806,9 @@ def build_fetch_response_payload( if webhook_enabled: output[TSPKeys.POSTPROCESSING_WEBHOOK_URL] = webhook_url + 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 output[TSPKeys.EVAL_SETTINGS][TSPKeys.EVAL_SETTINGS_MONITOR_LLM] = [monitor_llm] @@ -1166,6 +1177,10 @@ def build_single_pass_payload( TSPKeys.SIMILARITY_TOP_K: default_profile.similarity_top_k, } + lookup_configs = get_lookup_configs_for_tool(tool, prompts=prompts) + if lookup_configs: + tool_settings["lookup_configs"] = lookup_configs + for p in prompts: if not p.prompt: raise EmptyPromptError() @@ -1607,6 +1622,9 @@ def _execute_single_prompt( is_single_pass=False, profile_manager_id=profile_manager_id, ) + except APIException: + # Validation responses are user-facing; DRF renders them as-is. + raise except Exception as e: logger.error( f"[{tool.tool_id}] Error while fetching response for " @@ -1672,6 +1690,9 @@ def _execute_prompts_in_single_pass( document_id=document_id, is_single_pass=True, ) + except APIException: + # Validation responses are user-facing; DRF renders them as-is. + raise except Exception as e: logger.error( f"[{tool.tool_id}] Error while fetching single pass response: {e}" @@ -1911,6 +1932,8 @@ def _fetch_response( output[TSPKeys.ENABLE_POSTPROCESSING_WEBHOOK] = webhook_enabled if webhook_enabled: output[TSPKeys.POSTPROCESSING_WEBHOOK_URL] = webhook_url + 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_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 99edefd0b4..30ea7045e2 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 @@ -32,6 +33,11 @@ from workflow_manager.endpoint_v2.models import WorkflowEndpoint 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, +) from prompt_studio.prompt_profile_manager_v2.constants import ( ProfileManagerErrors, ProfileManagerKeys, @@ -90,6 +96,29 @@ 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. + + Multi-var lookups only resolve correctly under single-pass; ``prompt_ids`` + scopes the gate so an unrelated multi-var lookup doesn't block runs that + don't actually use it. Caller must skip this on the SP path. + """ + names = get_multi_var_lookups_for_tool(custom_tool, prompt_ids=prompt_ids) + if not names: + return None + return Response( + { + "detail": ( + "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.""" @@ -477,6 +506,14 @@ 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) + # Must precede the lookup gate so missing prompt_id returns a clear 400. + if not prompt_id: + return Response( + {"error": "prompt id is required."}, + status=status.HTTP_400_BAD_REQUEST, + ) + if err := _multi_var_lookup_block_response(custom_tool, prompt_ids=[prompt_id]): + return err profile_manager_id: str = request.data.get( ToolStudioPromptKeys.PROFILE_MANAGER_ID ) @@ -485,13 +522,6 @@ def fetch_response(self, request: HttpRequest, pk: Any = None) -> Response: org_id = UserSessionUtils.get_organization_id(request) user_id = custom_tool.created_by.user_id - - # Resolve prompt — guard against missing / stale prompt_id - if not prompt_id: - return Response( - {"error": "prompt id is required."}, - status=status.HTTP_400_BAD_REQUEST, - ) try: prompt = ToolStudioPrompt.objects.get(pk=prompt_id) except ToolStudioPrompt.DoesNotExist: @@ -510,9 +540,8 @@ def fetch_response(self, request: HttpRequest, pk: Any = None) -> Response: document: DocumentManager = DocumentManager.objects.get(pk=document_id) doc_path = str(Path(doc_path) / document.document_name) - # Agentic table prompts have a separate executor worker. Build the - # payload via the cloud payload_modifier plugin and dispatch directly - # so the legacy answer_prompt path is bypassed. + # Agentic table prompts have their own executor — build payload via + # the cloud plugin and dispatch directly, bypassing answer_prompt. if prompt.enforce_type == ToolStudioPromptKeys.AGENTIC_TABLE: payload_modifier_plugin = get_plugin("payload_modifier") if not payload_modifier_plugin: @@ -633,6 +662,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( @@ -881,9 +912,17 @@ 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[ + prompt_studio_tool = serializer.validated_data.get( ProfileManagerKeys.PROMPT_STUDIO_TOOL - ] + ) + if not prompt_studio_tool: + # Write back into validated_data so perform_create() doesn't + # persist NULL and orphan the profile from every + # ``filter(prompt_studio_tool=...)`` query. + prompt_studio_tool = self.get_object() + serializer.validated_data[ProfileManagerKeys.PROMPT_STUDIO_TOOL] = ( + prompt_studio_tool + ) profile_count = ProfileManager.objects.filter( prompt_studio_tool=prompt_studio_tool ).count() @@ -1103,6 +1142,10 @@ def export_tool(self, request: Request, pk: Any = None) -> Response: force_export=force_export, ) + # Anchor for staleness checks (e.g. lookup-change banner). + 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, @@ -1116,6 +1159,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() @@ -1294,10 +1346,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) + # NULL last_exported_at → treat as clean to avoid false alarms + # on 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: 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 64932eb3f9..9c963cf3f6 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 @@ -2,8 +2,15 @@ import logging from typing import Any -from django.core.exceptions import ObjectDoesNotExist - +from django.core.exceptions import ValidationError +from django.db import IntegrityError + +from prompt_studio.lookup_utils import ( + attach_combined_output_enrichment, + extract_prompt_output_enrichment, + get_original_value_if_enriched, + persist_lookup_output, +) from prompt_studio.prompt_profile_manager_v2.models import ProfileManager from prompt_studio.prompt_studio_core_v2.exceptions import ( AnswerFetchError, @@ -170,6 +177,15 @@ def update_or_create_prompt_output( # TODO: use enums here output = outputs.get(prompt.prompt_key) + + # On enrichment, store the raw LLM output here; the enriched + # value is persisted separately via persist_lookup_output. + enrichment = get_original_value_if_enriched(metadata, prompt.prompt_key) + if enrichment is not None: + output, prompt_lookup = enrichment + else: + prompt_lookup = None + if prompt.enforce_type in { "json", "table", @@ -195,6 +211,18 @@ def update_or_create_prompt_output( word_confidence_data=prompt_word_confidence_data, ) + # Narrow except so plugin contract drift surfaces as a real + # error instead of being masked as a successful save. + if prompt_lookup: + try: + persist_lookup_output(prompt_output, prompt_lookup) + except (IntegrityError, ValidationError): + logger.error( + "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) @@ -242,41 +270,76 @@ def fetch_default_output_response( Returns: dict[str, Any]: Formatted JSON response for combined output. + Cloud plugins may attach an opaque enrichment payload. """ - # Initialize the result dictionary + from prompt_studio.lookup_utils import enrich_prompt_output + + # Memoise default-profile resolution per tool to avoid N+1 on this + # hot path (panel-switch latency). + default_profile_cache: dict[str, str | None] = {} + + def _resolve(tool_prompt: ToolStudioPrompt) -> str | None: + profile_manager_id = tool_prompt.profile_manager_id + if profile_manager_id: + return profile_manager_id + if not use_default_profile: + return None + tool_id = tool_prompt.tool_id_id + if tool_id not in default_profile_cache: + try: + default_profile_cache[tool_id] = ( + ProfileManager.get_default_llm_profile( + tool_prompt.tool_id + ).profile_id + ) + except DefaultProfileError: + default_profile_cache[tool_id] = None + return default_profile_cache[tool_id] + + prompts_to_query: list[tuple[ToolStudioPrompt, str]] = [] result: dict[str, Any] = {} - # Iterate over ToolStudioPrompt records 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: + profile_manager_id = _resolve(tool_prompt) + if profile_manager_id is None: result[tool_prompt.prompt_key] = "" continue - - if not profile_manager_id: - default_profile = ProfileManager.get_default_llm_profile( - tool_prompt.tool_id - ) - profile_manager_id = default_profile.profile_id - - try: - queryset = PromptStudioOutputManager.objects.filter( - prompt_id=prompt_id, - profile_manager=profile_manager_id, + prompts_to_query.append((tool_prompt, profile_manager_id)) + + # ``DISTINCT ON`` (Postgres) yields the latest row per + # (prompt_id, profile_manager_id) at the SQL layer. + outputs_index: dict[tuple[str, str], PromptStudioOutputManager] = {} + if prompts_to_query: + prompt_ids = [str(p.prompt_id) for p, _ in prompts_to_query] + profile_ids = list({pmid for _, pmid in prompts_to_query}) + outputs = ( + PromptStudioOutputManager.objects.filter( + prompt_id__in=prompt_ids, + profile_manager_id__in=profile_ids, is_single_pass_extract=False, document_manager_id=document_manager_id, ) - - if not queryset.exists(): - result[tool_prompt.prompt_key] = "" - continue - - for output in queryset: - result[tool_prompt.prompt_key] = output.output - except ObjectDoesNotExist: + .order_by("prompt_id", "profile_manager_id", "-modified_at") + .distinct("prompt_id", "profile_manager_id") + ) + outputs_index = { + (str(o.prompt_id), str(o.profile_manager_id)): o for o in outputs + } + + enrichment_by_key: dict[str, Any] = {} + for tool_prompt, profile_manager_id in prompts_to_query: + output = outputs_index.get( + (str(tool_prompt.prompt_id), str(profile_manager_id)) + ) + if output is None: result[tool_prompt.prompt_key] = "" + continue + result[tool_prompt.prompt_key] = output.output + enriched = enrich_prompt_output(output, {}) + bundle = extract_prompt_output_enrichment(enriched) + if bundle is not None: + enrichment_by_key[tool_prompt.prompt_key] = bundle + + attach_combined_output_enrichment(result, enrichment_by_key) return result 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..1c56e2323d 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,16 @@ def to_representation(self, instance): " | Process continued" ) data["coverage"] = {} + # log+continue: enrichment failure shouldn't 500 the list endpoint. + try: + data = enrich_prompt_output(instance, data) + except Exception as e: + logger.error( + "Error occurred while enriching prompt output for " + f"prompt_id {instance.prompt_id} (run_id={instance.run_id}): {e}" + " | Process continued" + ) + # Convert string to list try: context = data["context"] 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..44111dc744 100644 --- a/backend/prompt_studio/prompt_studio_output_manager_v2/views.py +++ b/backend/prompt_studio/prompt_studio_output_manager_v2/views.py @@ -5,11 +5,12 @@ from django.db.models import QuerySet from django.http import HttpRequest from rest_framework import status, viewsets -from rest_framework.exceptions import APIException +from rest_framework.exceptions import ValidationError from rest_framework.response import Response from rest_framework.versioning import URLPathVersioning from utils.common_utils import CommonUtils from utils.filtering import FilterHelper +from utils.user_context import UserContext from prompt_studio.prompt_studio_output_manager_v2.constants import ( PromptOutputManagerErrorMessage, @@ -61,6 +62,57 @@ 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. + + Backs the lookup Test panel's "Use Latest Outputs" button. Returns + raw extraction (not enriched) so the lookup can be tested fresh. + """ + tool_id = request.GET.get("tool_id") + keys_param = request.GET.get("prompt_keys", "") + if not tool_id: + # APIException(code=400) returns 500; ValidationError returns 400. + raise ValidationError(detail=PromptOutputManagerErrorMessage.TOOL_VALIDATION) + + prompt_keys = [k.strip() for k in keys_param.split(",") if k.strip()] + if not prompt_keys: + return Response({}, status=status.HTTP_200_OK) + + # Custom actions skip filter_queryset(), so OrganizationFilterBackend + # never runs — scope explicitly to prevent cross-tenant reads. + organization = UserContext.get_organization() + prompt_id_to_key = dict( + ToolStudioPrompt.objects.filter( + tool_id=tool_id, + tool_id__organization=organization, + prompt_key__in=prompt_keys, + ).values_list("prompt_id", "prompt_key") + ) + if not prompt_id_to_key: + return Response({}, status=status.HTTP_200_OK) + + # ``DISTINCT ON("prompt_id")`` keeps the latest row per prompt at + # the SQL layer to avoid materialising every doc × run combo. + outputs = ( + PromptStudioOutputManager.objects.filter( + prompt_id__in=prompt_id_to_key.keys(), + tool_id__organization=organization, + ) + .exclude(output__isnull=True) + .exclude(output__exact="") + .order_by("prompt_id", "-modified_at") + .distinct("prompt_id") + .values("prompt_id", "output") + ) + + result: dict[str, str] = {} + for row in outputs: + key = prompt_id_to_key.get(row["prompt_id"]) + if key: + result[key] = row["output"] + + 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 @@ -69,7 +121,7 @@ def get_output_for_tool_default(self, request: HttpRequest) -> Response: tool_validation_message = PromptOutputManagerErrorMessage.TOOL_VALIDATION tool_not_found = PromptOutputManagerErrorMessage.TOOL_NOT_FOUND if not tool_id: - raise APIException(detail=tool_validation_message, code=400) + raise ValidationError(detail=tool_validation_message) try: # Fetch ToolStudioPrompt records based on tool_id @@ -77,7 +129,7 @@ def get_output_for_tool_default(self, request: HttpRequest) -> Response: tool_id=tool_id ).order_by("sequence_number") except ObjectDoesNotExist: - raise APIException(detail=tool_not_found, code=400) + raise ValidationError(detail=tool_not_found) # Invoke helper method to frame and fetch default response. result: dict[str, Any] = OutputManagerHelper.fetch_default_output_response( 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 08d2c27baa..4fee8c10bc 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 @@ -297,6 +298,15 @@ def frame_export_json( settings, JsonSchemaKey.WORD_CONFIDENCE_POSTAMBLE.upper(), "" ) + # Validate only what will actually be exported — NOTES / inactive + # prompts never run, so incomplete lookups on them shouldn't fail export. + exportable_prompts = [ + p for p in prompts if p.prompt_type != JsonSchemaKey.NOTES and p.active + ] + lookup_configs, lookup_error = validate_lookups_for_export(exportable_prompts) + if lookup_error: + raise InValidCustomToolError(lookup_error) + for prompt in prompts: if prompt.prompt_type == JsonSchemaKey.NOTES or not prompt.active: continue @@ -355,6 +365,9 @@ def frame_export_json( output[JsonSchemaKey.POSTPROCESSING_WEBHOOK_URL] = ( prompt.postprocessing_webhook_url ) + 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 ( 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, diff --git a/backend/usage_v2/hooks.py b/backend/usage_v2/hooks.py new file mode 100644 index 0000000000..31d576e693 --- /dev/null +++ b/backend/usage_v2/hooks.py @@ -0,0 +1,27 @@ +"""Post-write hooks for ``UsageBatchCreateView``. + +Hooks fire inside the view's transaction; a failure rolls the batch back +so Usage rows and any side-table writes stay consistent. Records carry +an opaque ``cloud_extras`` dict that OSS forwards verbatim — plugins +read only the keys they own. +""" + +from collections.abc import Callable + +from .models import Usage + +PostWriteHook = Callable[[list[dict], list[Usage]], None] + +_post_write_hooks: list[PostWriteHook] = [] + + +def register_post_write_hook(fn: PostWriteHook) -> PostWriteHook: + # Idempotent: ready() can re-fire under test reloads / dev autoreload. + if fn not in _post_write_hooks: + _post_write_hooks.append(fn) + return fn + + +def run_post_write_hooks(records: list[dict], usage_objects: list[Usage]) -> None: + for hook in _post_write_hooks: + hook(records, usage_objects) 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..85c20c16e4 100644 --- a/backend/usage_v2/internal_views.py +++ b/backend/usage_v2/internal_views.py @@ -1,19 +1,35 @@ -"""Internal API views for Usage access by workers.""" +"""Internal API views for Usage access by workers. + +Mounted under ``/internal/`` and gated by ``InternalAPIAuthMiddleware``. +""" import logging +from django.db import transaction from django.http import JsonResponse from rest_framework import status +from rest_framework.exceptions import APIException, ValidationError from rest_framework.request import Request +from rest_framework.response import Response from rest_framework.views import APIView +from utils.user_context import UserContext from unstract.core.data_models import UsageResponseData from .helper import UsageHelper +from .hooks import run_post_write_hooks +from .models import Usage +from .serializers import UsageBatchCreateSerializer logger = logging.getLogger(__name__) +class UsagePersistError(APIException): + status_code = status.HTTP_500_INTERNAL_SERVER_ERROR + default_detail = "Failed to persist usage records." + default_code = "usage_persist_failed" + + class UsageInternalView(APIView): """Internal API view for workers to access usage data. @@ -133,3 +149,66 @@ 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) -> Response: + input_serializer = UsageBatchCreateSerializer(data=request.data) + input_serializer.is_valid(raise_exception=True) + records = input_serializer.validated_data["records"] + if not records: + return Response({"created": 0}, status=status.HTTP_200_OK) + + organization = UserContext.get_organization() + if organization is None: + logger.error( + "UsageBatchCreateView received %d records with no organization; " + "refusing to write rows that would be invisible to tenant dashboards", + len(records), + ) + raise ValidationError( + "Organization context missing. Worker must send X-Organization-ID." + ) + + usage_objects = [ + 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"), + # Coerce "" to None so the cross-field CheckConstraint passes. + llm_usage_reason=r.get("llm_usage_reason") or None, + 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), + project_id=r.get("project_id"), + prompt_id=r.get("prompt_id"), + execution_time_ms=r.get("execution_time_ms"), + status=r.get("status"), + error_message=r.get("error_message"), + ) + for r in records + ] + + try: + # Atomic with hooks: orphan Usage rows are worse than retrying. + with transaction.atomic(): + created = Usage.objects.bulk_create(usage_objects, batch_size=500) + run_post_write_hooks(records, created) + except Exception as e: + logger.error( + "bulk_create failed for %d usage records (org=%s): %s", + len(usage_objects), + organization.organization_id, + e, + exc_info=True, + ) + raise UsagePersistError() from e + return Response({"created": len(created)}, status=status.HTTP_201_CREATED) 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..191a31dd27 --- /dev/null +++ b/backend/usage_v2/migrations/0004_usage_metrics_fields.py @@ -0,0 +1,89 @@ +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"), + ("lookup", "Lookup"), + ], + db_comment="Reason for LLM usage. Empty if usage_type is 'embedding'. ", + max_length=255, + null=True, + ), + ), + migrations.AddField( + model_name="usage", + name="project_id", + field=models.UUIDField( + blank=True, + db_comment=( + "Prompt Studio project (tool) the call belongs to (no FK; " + "survives tool deletion). NULL for embeddings and historical " + "rows." + ), + null=True, + ), + ), + migrations.AddField( + model_name="usage", + name="prompt_id", + field=models.UUIDField( + blank=True, + db_comment=( + "Prompt key UUID that triggered the call (no FK; survives " + "prompt deletion). NULL for single-pass / embeddings / " + "historical rows." + ), + 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, + choices=[ + ("SUCCESS", "Success"), + ("ERROR", "Error"), + ("SKIPPED", "Skipped"), + ], + 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, + ), + ), + # Indexes on project_id / prompt_id moved to 0005 so they can run + # CONCURRENTLY — usage is billing-critical and a plain AddIndex takes + # a share-update lock for the duration of the build on large tables. + ] diff --git a/backend/usage_v2/migrations/0005_usage_reason_ref_created_idx.py b/backend/usage_v2/migrations/0005_usage_reason_ref_created_idx.py new file mode 100644 index 0000000000..c93d4de605 --- /dev/null +++ b/backend/usage_v2/migrations/0005_usage_reason_ref_created_idx.py @@ -0,0 +1,76 @@ +from django.db import migrations, models +from django.db.models import Q + + +class Migration(migrations.Migration): + """Build the project_id / prompt_id dashboard indexes without locking. + + CONCURRENTLY requires the migration to run outside a transaction, hence + atomic = False. RunSQL with IF NOT EXISTS makes a partial-apply + (process killed between SQL success and django_migrations insert) + recoverable on retry without manual --fake intervention. + """ + + atomic = False + + dependencies = [ + ("usage_v2", "0004_usage_metrics_fields"), + ] + + operations = [ + migrations.RunSQL( + sql=( + "CREATE INDEX CONCURRENTLY IF NOT EXISTS " + "idx_usage_project_created " + 'ON "usage" (project_id, created_at DESC);' + ), + reverse_sql=("DROP INDEX CONCURRENTLY IF EXISTS idx_usage_project_created;"), + state_operations=[ + migrations.AddIndex( + model_name="usage", + index=models.Index( + fields=["project_id", "-created_at"], + name="idx_usage_project_created", + ), + ), + ], + ), + migrations.RunSQL( + sql=( + "CREATE INDEX CONCURRENTLY IF NOT EXISTS " + "idx_usage_prompt_created " + 'ON "usage" (prompt_id, created_at DESC);' + ), + reverse_sql=("DROP INDEX CONCURRENTLY IF EXISTS idx_usage_prompt_created;"), + state_operations=[ + migrations.AddIndex( + model_name="usage", + index=models.Index( + fields=["prompt_id", "-created_at"], + name="idx_usage_prompt_created", + ), + ), + ], + ), + # Partial — only lookup-reason rows. Avoids heap-scanning all + # Usage rows when the dashboard groups by (run × prompt). + migrations.RunSQL( + sql=( + "CREATE INDEX CONCURRENTLY IF NOT EXISTS " + "idx_usage_lookup_recent " + 'ON "usage" (organization_id, created_at DESC) ' + "WHERE llm_usage_reason = 'lookup';" + ), + reverse_sql=("DROP INDEX CONCURRENTLY IF EXISTS idx_usage_lookup_recent;"), + state_operations=[ + migrations.AddIndex( + model_name="usage", + index=models.Index( + fields=["organization", "-created_at"], + name="idx_usage_lookup_recent", + condition=Q(llm_usage_reason="lookup"), + ), + ), + ], + ), + ] diff --git a/backend/usage_v2/models.py b/backend/usage_v2/models.py index 57ae3d143d..0267d4a68f 100644 --- a/backend/usage_v2/models.py +++ b/backend/usage_v2/models.py @@ -1,6 +1,7 @@ import uuid from django.db import models +from django.db.models import Q from utils.models.base_model import BaseModel, BaseModelManager from utils.models.organization_mixin import ( DefaultOrganizationManagerMixin, @@ -13,10 +14,21 @@ class UsageType(models.TextChoices): EMBEDDING = "embedding", "Embedding Usage" -class LLMUsageReason(models.TextChoices): - EXTRACTION = "extraction", "Extraction" - CHALLENGE = "challenge", "Challenge" - SUMMARIZE = "summarize", "Summarize" +class UsageStatus(models.TextChoices): + SUCCESS = "SUCCESS", "Success" + ERROR = "ERROR", "Error" + SKIPPED = "SKIPPED", "Skipped" + + +# Static union of OSS + cloud values — keeps OSS model state aligned with +# migration state so ``makemigrations --check`` doesn't drift in CI. + +LLM_USAGE_REASON_CHOICES: list[tuple[str, str]] = [ + ("extraction", "Extraction"), + ("challenge", "Challenge"), + ("summarize", "Summarize"), + ("lookup", "Lookup"), +] class UsageModelManager(DefaultOrganizationManagerMixin, BaseModelManager): @@ -52,7 +64,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'. ", @@ -67,6 +79,39 @@ 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") + project_id = models.UUIDField( + null=True, + blank=True, + db_comment=( + "Prompt Studio project (tool) the call belongs to (no FK; survives " + "tool deletion). NULL for embeddings and historical rows." + ), + ) + prompt_id = models.UUIDField( + null=True, + blank=True, + db_comment=( + "Prompt key UUID that triggered the call (no FK; survives prompt " + "deletion). NULL for single-pass / embeddings / historical rows." + ), + ) + 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, + choices=UsageStatus.choices, + 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() @@ -78,4 +123,17 @@ class Meta: indexes = [ models.Index(fields=["run_id"]), models.Index(fields=["execution_id"]), + models.Index( + fields=["project_id", "-created_at"], + name="idx_usage_project_created", + ), + models.Index( + fields=["prompt_id", "-created_at"], + name="idx_usage_prompt_created", + ), + models.Index( + fields=["organization", "-created_at"], + name="idx_usage_lookup_recent", + condition=Q(llm_usage_reason="lookup"), + ), ] diff --git a/backend/usage_v2/serializers.py b/backend/usage_v2/serializers.py index a823297083..7b7f8aaed5 100644 --- a/backend/usage_v2/serializers.py +++ b/backend/usage_v2/serializers.py @@ -21,3 +21,38 @@ class UsageSerializer(serializers.ModelSerializer): class Meta: model = Usage fields = "__all__" + + +class UsageRecordCreateSerializer(serializers.Serializer): + """Worker-emitted usage record. Required fields anchor billing-critical attribution.""" + + adapter_instance_id = serializers.CharField(required=True, allow_blank=False) + model_name = serializers.CharField(required=True, allow_blank=False) + usage_type = serializers.CharField(required=True, allow_blank=False) + + workflow_id = serializers.CharField(required=False, allow_blank=True, default="") + execution_id = serializers.CharField(required=False, allow_blank=True, default="") + run_id = serializers.UUIDField(required=False, allow_null=True, default=None) + llm_usage_reason = serializers.CharField( + required=False, allow_null=True, allow_blank=True, default=None + ) + embedding_tokens = serializers.IntegerField(required=False, default=0) + prompt_tokens = serializers.IntegerField(required=False, default=0) + completion_tokens = serializers.IntegerField(required=False, default=0) + total_tokens = serializers.IntegerField(required=False, default=0) + cost_in_dollars = serializers.FloatField(required=False, default=0.0) + project_id = serializers.UUIDField(required=False, allow_null=True, default=None) + prompt_id = serializers.UUIDField(required=False, allow_null=True, default=None) + execution_time_ms = serializers.IntegerField( + required=False, allow_null=True, default=None + ) + status = serializers.CharField(required=False, allow_null=True, default=None) + error_message = serializers.CharField( + required=False, allow_null=True, allow_blank=True, default=None + ) + # Opaque carrier forwarded to post-write hooks; OSS never reads it. + cloud_extras = serializers.DictField(required=False, allow_null=True, default=None) + + +class UsageBatchCreateSerializer(serializers.Serializer): + records = UsageRecordCreateSerializer(many=True, allow_empty=True) diff --git a/frontend/src/components/custom-tools/combined-output/CombinedOutput.jsx b/frontend/src/components/custom-tools/combined-output/CombinedOutput.jsx index 3036e7b9d6..7e30a99339 100644 --- a/frontend/src/components/custom-tools/combined-output/CombinedOutput.jsx +++ b/frontend/src/components/custom-tools/combined-output/CombinedOutput.jsx @@ -47,6 +47,67 @@ try { // The component will remain null if it is not available } +// OSS falls back to passthrough helpers — no enrichment. +let splitCombinedData = (data) => ({ combined: data, bundle: null }); +let buildEnrichedFromBundle = (_output, _bundle, _formatter) => ({}); +let getEnrichmentFromItem = (_item) => null; +try { + const mod = await import("../../../plugins/lookup-enriched-toggle/helpers"); + splitCombinedData = mod.splitCombinedData; + buildEnrichedFromBundle = mod.buildEnrichedFromBundle; + getEnrichmentFromItem = mod.getEnrichmentFromItem; +} catch {} + +const buildDefaultProfileOutputs = (data) => { + const { combined: payload, bundle } = splitCombinedData(data); + const output = Object.entries(payload).reduce((acc, [key, value]) => { + acc[key] = displayPromptResult(value, false); + return acc; + }, {}); + const enriched = buildEnrichedFromBundle(output, bundle, displayPromptResult); + return { + output, + enriched, + hasEnriched: bundle != null && Object.keys(enriched).length > 0, + }; +}; + +const buildPerPromptOutput = (item, data, selectedProfile) => { + const profileManager = selectedProfile || item?.profile_manager; + const outputDetails = data.find( + (outputValue) => + outputValue?.prompt_id === item?.prompt_id && + outputValue?.profile_manager === profileManager, + ); + const value = + outputDetails?.output?.length > 0 + ? displayPromptResult(outputDetails?.output, false) + : ""; + const enrichment = getEnrichmentFromItem(outputDetails); + const enrichedValue = enrichment?.output + ? displayPromptResult(enrichment.output, false) + : value; + return { value, enrichedValue, hasEnriched: !!enrichment?.output }; +}; + +const buildSelectedProfileOutputs = (data, prompts, selectedProfile) => { + const output = {}; + const enriched = {}; + let hasEnriched = false; + for (const item of prompts) { + if (item?.prompt_type === promptType.notes) continue; + const { + value, + enrichedValue, + hasEnriched: enrichedHit, + } = buildPerPromptOutput(item, data, selectedProfile); + output[item?.prompt_key] = value; + enriched[item?.prompt_key] = enrichedValue; + hasEnriched = hasEnriched || enrichedHit; + } + return { output, enriched: hasEnriched ? enriched : {}, hasEnriched }; +}; + function CombinedOutput({ docId, setFilledFields, selectedPrompts }) { const { details, @@ -59,6 +120,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( @@ -115,32 +177,12 @@ function CombinedOutput({ docId, setFilledFields, selectedPrompts }) { const res = await handleOutputApiRequest(); const data = res?.data || []; const prompts = details?.prompts || []; - - if (activeKey === "0" && !isSimplePromptStudio) { - const output = Object.entries(data).reduce((acc, [key, value]) => { - acc[key] = displayPromptResult(value, false); - return acc; - }, {}); - setCombinedOutput(output); - } 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, - ); - - acc[item?.prompt_key] = - outputDetails?.output?.length > 0 - ? displayPromptResult(outputDetails?.output, false) - : ""; - } - return acc; - }, {}); - setCombinedOutput(output); - } + const useDefaultProfile = activeKey === "0" && !isSimplePromptStudio; + const { output, enriched } = useDefaultProfile + ? buildDefaultProfileOutputs(data) + : buildSelectedProfileOutputs(data, prompts, selectedProfile); + setCombinedOutput(output); + setEnrichedOutput(enriched); } catch (err) { setAlertDetails( handleException(err, "Failed to generate combined output"), @@ -229,6 +271,7 @@ function CombinedOutput({ docId, setFilledFields, selectedPrompts }) { return ( { Prism.highlightAll(); - }, [combinedOutput]); + }, [combinedOutput, enrichedOutput, activeView]); + + 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 +60,24 @@ function JsonView({ /> ))} -
+
+ {EnrichedOutputToggle && ( + 0) + } + /> + )} +
@@ -51,6 +87,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/document-parser/DocumentParser.jsx b/frontend/src/components/custom-tools/document-parser/DocumentParser.jsx index 3d5c891e13..63b1a037c5 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"; @@ -32,6 +33,10 @@ try { // The component will remain null of it is not available } +// Module-scoped to avoid per-render recompilation. +const UUID_RE = + /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/i; + function DocumentParser({ addPromptInstance, scrollToBottom, @@ -42,6 +47,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 +114,28 @@ function DocumentParser({ } }, [scrollToBottom]); + // Cross-link from Lookup Studio: scroll to a specific prompt. + useEffect(() => { + const scrollToPromptId = searchParams.get("scrollTo"); + if ( + !scrollToPromptId || + !UUID_RE.test(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); + } + + searchParams.delete("scrollTo"); + setSearchParams(searchParams, { replace: true }); + }, [details?.prompts, searchParams]); + const promptUrl = (urlPath) => { return `/api/v1/unstract/${sessionDetails?.orgId}/prompt-studio/prompt/${urlPath}`; }; @@ -242,7 +270,7 @@ function DocumentParser({
{details?.prompts?.map((item) => { return ( -
+
Promise.resolve(true); + function Header({ setOpenSettings, handleUpdateTool, setOpenShareModal, setOpenCloneModal, + checkLookups = noopCheckLookups, }) { const [isExportLoading, setIsExportLoading] = useState(false); const { details, isPublicSource, markChangesAsExported } = @@ -129,7 +127,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 +137,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 +256,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 +267,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}`; @@ -527,6 +531,7 @@ Header.propTypes = { handleUpdateTool: PropTypes.func.isRequired, setOpenCloneModal: PropTypes.func.isRequired, setOpenShareModal: PropTypes.func.isRequired, + checkLookups: PropTypes.func, }; export { Header }; 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/custom-tools/profile-info-bar/ProfileInfoBar.css b/frontend/src/components/custom-tools/profile-info-bar/ProfileInfoBar.css index 8346c26d71..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,3 +1,10 @@ .profile-info-bar { + display: flex; + flex-wrap: wrap; + 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/components/custom-tools/prompt-card/Header.jsx b/frontend/src/components/custom-tools/prompt-card/Header.jsx index cea5e2d6d3..30fe75fe54 100644 --- a/frontend/src/components/custom-tools/prompt-card/Header.jsx +++ b/frontend/src/components/custom-tools/prompt-card/Header.jsx @@ -41,6 +41,38 @@ try { // The component will remain 'undefined' it is not available } +let LookupMenuItem; +try { + const mod = await import( + "../../../plugins/lookup-studio/prompt-card/LookupMenuItem" + ); + LookupMenuItem = mod.LookupMenuItem; +} catch {} + +let usePromptRunGatePlugin; +let lookupPluginLoadError; +try { + const mod = await import( + "../../../plugins/lookup-studio/hooks/usePromptRunGate" + ); + usePromptRunGatePlugin = mod.usePromptRunGate; +} catch (err) { + lookupPluginLoadError = err; +} + +// Sibling plugin loaded but this hook didn't — surface so the no-op fallback +// doesn't silently disable the run gate. +if (lookupPluginLoadError && LookupMenuItem) { + // eslint-disable-next-line no-console + console.warn( + "[Header] lookup-studio plugin loaded but usePromptRunGate failed to import", + lookupPluginLoadError, + ); +} + +// Stable identity — avoid conditional hook call. +const usePromptRunGate = usePromptRunGatePlugin || (() => null); + function Header({ promptDetails, promptKey, @@ -70,6 +102,7 @@ function Header({ isSimplePromptStudio, details, } = useCustomToolStore(); + const runGate = usePromptRunGate(promptDetails); const [items, setItems] = useState([]); const [isDisablePrompt, setIsDisablePrompt] = useState(null); @@ -268,6 +301,20 @@ function Header({ isPublicSource, }, ]; + if (LookupMenuItem && !isSimplePromptStudio) { + dropdownItems.splice( + -1, + 0, + { + type: "divider", + }, + { + label: , + key: "lookup", + }, + ); + } + if (isSimplePromptStudio) { dropdownItems.splice(0, 1); } @@ -346,7 +393,9 @@ function Header({ )} {!singlePassExtractMode && !isSimplePromptStudio && ( <> - + - + + {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 38f7d924a6..b8aaa4b393 100644 --- a/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx +++ b/frontend/src/components/custom-tools/prompt-card/PromptOutput.jsx @@ -51,6 +51,49 @@ 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 { + // Cloud-only plugin — absent in OSS builds; LookupOutputTabs stays null. +} + +let getEnrichedCopyText; +try { + const mod = await import( + "../../../plugins/lookup-studio/prompt-card/getEnrichedCopyText" + ); + getEnrichedCopyText = mod.getEnrichedCopyText; +} catch { + // Cloud-only plugin — absent in OSS builds; falls back to raw copy text. +} + +// Fallback to raw text — plugin throw on malformed enrichment shouldn't break Copy. +const resolveCopyText = (promptOutputId, fallbackText) => { + if (!getEnrichedCopyText) { + return fallbackText; + } + try { + const enriched = getEnrichedCopyText(promptOutputId); + return typeof enriched === "string" && enriched ? enriched : fallbackText; + } catch (err) { + // eslint-disable-next-line no-console + console.warn("[PromptOutput] getEnrichedCopyText threw:", err); + return fallbackText; + } +}; + +// Cloud wraps children in LookupOutputTabs; OSS passes through. +const renderWithLookupWrapper = (lookupProps, children) => + LookupOutputTabs ? ( + {children} + ) : ( + children + ); + function PromptOutput({ promptDetails, handleRun, @@ -194,32 +237,43 @@ function PromptOutput({ "highlighted-prompt-cell" }`} > - + {renderWithLookupWrapper( + { + promptId, + profileManagerId: defaultLlmProfile, + defaultLlmProfile, + promptOutputId: promptOutputData?.promptOutputId, + }, + , + )}
+ copyToClipboard={() => { copyOutputToClipboard( - displayPromptResult( - promptOutput, - true, - promptDetails?.enable_highlight, + resolveCopyText( + promptOutputData?.promptOutputId, + displayPromptResult( + promptOutput, + true, + promptDetails?.enable_highlight, + ), ), - ) - } + ); + }} /> - + {renderWithLookupWrapper( + { + promptId, + profileManagerId: profileId, + defaultLlmProfile, + promptOutputId: promptOutputData?.promptOutputId, + }, + , + )}
+ copyToClipboard={() => { copyOutputToClipboard( - displayPromptResult( - promptOutputData?.output, - true, + resolveCopyText( + promptOutputData?.promptOutputId, + displayPromptResult( + promptOutputData?.output, + true, + ), ), - ) - } + ); + }} />
diff --git a/frontend/src/components/custom-tools/token-usage/TokenUsage.jsx b/frontend/src/components/custom-tools/token-usage/TokenUsage.jsx index 846c11f0a2..3c52a40058 100644 --- a/frontend/src/components/custom-tools/token-usage/TokenUsage.jsx +++ b/frontend/src/components/custom-tools/token-usage/TokenUsage.jsx @@ -23,8 +23,13 @@ function TokenUsage({ tokenUsageId, isLoading }) { setTokens(tokenUsage[tokenUsageId]); // Update tokens state with the token usage data for the given tokenUsageId }, [tokenUsage, tokenUsageId]); - // If no tokens data is available, render nothing - if (!tokens || !Object.keys(tokens)?.length || isLoading) { + // BE returns all-null fields when no Usage rows match; treat as empty. + if ( + !tokens || + !Object.keys(tokens)?.length || + isLoading || + tokens?.total_tokens == null + ) { return "NA"; } diff --git a/frontend/src/components/custom-tools/tool-ide/ToolIde.jsx b/frontend/src/components/custom-tools/tool-ide/ToolIde.jsx index 93c8198593..1f7e21c259 100644 --- a/frontend/src/components/custom-tools/tool-ide/ToolIde.jsx +++ b/frontend/src/components/custom-tools/tool-ide/ToolIde.jsx @@ -44,6 +44,28 @@ try { } catch { // Do nothing if plugins are not loaded. } + +// Cloud-only — OSS stub. +let useLookupDirtySeed = () => {}; +try { + const mod = await import( + "../../../plugins/lookup-studio/hooks/useLookupDirtySeed.js" + ); + useLookupDirtySeed = mod.useLookupDirtySeed; +} catch {} + +// Cloud-only — OSS stub resolves true to skip the gate. +let useLookupExportGate = () => ({ + checkLookups: () => Promise.resolve(true), + modalEl: null, +}); +try { + const mod = await import( + "../../../plugins/lookup-studio/hooks/useLookupExportGate" + ); + useLookupExportGate = mod.useLookupExportGate; +} catch {} + function ToolIde() { const [openSettings, setOpenSettings] = useState(false); const customToolStore = useCustomToolStore(); @@ -76,6 +98,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 +201,9 @@ function ToolIde() { } }, [details?.tool_id]); + // Surfaces re-export banner when /lookups page edits made the tool stale. + useLookupDirtySeed(details?.tool_id); + // Cleanup abort controller on unmount useEffect(() => { return () => { @@ -189,6 +215,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 +267,7 @@ function ToolIde() { handleException, markChangesAsExported, setPostHogCustomEvent, + checkLookups, ]); const generateIndex = async (doc) => { @@ -340,11 +371,13 @@ function ToolIde() { isExporting={isExporting} /> )} + {lookupGateModalEl}
); case "link": { - // Protocol-relative URLs (`//evil.com/...`) also start with `/` - // so exclude them from the internal-route branch. + // Exclude protocol-relative `//evil.com` from internal-route branch. const isInternal = url?.startsWith("/") && !url.startsWith("//"); if (isInternal) { const resolvedUrl = orgName ? `/${orgName}${url}` : url; diff --git a/frontend/src/components/helpers/custom-tools/CustomToolsHelper.js b/frontend/src/components/helpers/custom-tools/CustomToolsHelper.js index 543ad9b14b..35c2854cb0 100644 --- a/frontend/src/components/helpers/custom-tools/CustomToolsHelper.js +++ b/frontend/src/components/helpers/custom-tools/CustomToolsHelper.js @@ -19,6 +19,18 @@ try { } catch { // Do nothing, Not-found Page will be triggered. } + +let fetchLookupAssignments; +let fetchLookupOutputs; +let resetLookupAssignments; +try { + const mod = await import( + "../../../plugins/lookup-studio/store/useFetchLookupAssignments" + ); + fetchLookupAssignments = mod.fetchLookupAssignments; + fetchLookupOutputs = mod.fetchLookupOutputs; + resetLookupAssignments = mod.resetLookupAssignments; +} catch {} function CustomToolsHelper() { const [isLoading, setIsLoading] = useState(true); const { id } = useParams(); @@ -115,6 +127,14 @@ function CustomToolsHelper() { .then((res) => { const data = res?.data; updatedCusTool["adapters"] = data; + + 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 +151,9 @@ function CustomToolsHelper() { setDefaultCustomTool(); emptyCusToolMessages(); resetTokenUsage(); + if (resetLookupAssignments) { + resetLookupAssignments(); + } }; }, []); diff --git a/frontend/src/components/logging/execution-logs/ExecutionLogs.jsx b/frontend/src/components/logging/execution-logs/ExecutionLogs.jsx index a9c72a0efb..8e44287e2d 100644 --- a/frontend/src/components/logging/execution-logs/ExecutionLogs.jsx +++ b/frontend/src/components/logging/execution-logs/ExecutionLogs.jsx @@ -54,14 +54,14 @@ function ExecutionLogs() { ? location.state?.from || `/${sessionDetails?.orgName}/logs` : null; - // State to pass back for scroll restoration + // Scroll-restoration wins; otherwise preserve caller's upstream UI state. const backRouteState = id && location.state?.scrollToCardId ? { scrollToCardId: location.state.scrollToCardId, cardExpanded: location.state.cardExpanded, } - : null; + : location.state?.backRouteState || null; const items = [ { diff --git a/frontend/src/components/navigations/side-nav-bar/SideNavBar.jsx b/frontend/src/components/navigations/side-nav-bar/SideNavBar.jsx index bb328b1a98..164fdab145 100644 --- a/frontend/src/components/navigations/side-nav-bar/SideNavBar.jsx +++ b/frontend/src/components/navigations/side-nav-bar/SideNavBar.jsx @@ -89,6 +89,12 @@ try { // Plugin unavailable } +let lookupStudioEnabled = false; +try { + await import("../../../plugins/lookup-studio"); + lookupStudioEnabled = true; +} catch {} + let manualReviewSettingsEnabled = false; try { await import("../../../plugins/manual-review/settings/Settings.jsx"); @@ -506,6 +512,19 @@ const SideNavBar = ({ collapsed, setCollapsed }) => { }); } + // Keep Prompt Studio highlighted on /lookups. Replace, don't mutate — + // `data` may alias the `menu` prop. + if (lookupStudioEnabled && isUnstract && data[0]?.subMenu) { + const onLookupPath = globalThis.location.pathname.startsWith( + `/${orgName}/lookups`, + ); + if (onLookupPath) { + data[0].subMenu = data[0].subMenu.map((el) => + el.id === 1.1 ? { ...el, active: true } : el, + ); + } + } + // Add HITL Review section if plugin is available and user has HITL role const isHITLRole = [ "unstract_reviewer", diff --git a/frontend/src/components/navigations/tool-nav-bar/ToolNavBar.jsx b/frontend/src/components/navigations/tool-nav-bar/ToolNavBar.jsx index 999630d596..e39c50ff1e 100644 --- a/frontend/src/components/navigations/tool-nav-bar/ToolNavBar.jsx +++ b/frontend/src/components/navigations/tool-nav-bar/ToolNavBar.jsx @@ -9,6 +9,7 @@ import "./ToolNavBar.css"; function ToolNavBar({ title, + titleAdornment, subtitle, onEditTitle, enableSearch, @@ -19,6 +20,7 @@ function ToolNavBar({ onNavigateBack, segmentFilter, segmentOptions, + segmentValue, onSearch, searchKey, }) { @@ -53,6 +55,7 @@ function ToolNavBar({ {title} + {titleAdornment} {onEditTitle && (