UN-3266 [FEAT] Async Executor Backend for Prompt Studio

backend/api_v2/api_deployment_views.py

@@ -3,7 +3,6 @@
 import uuid
 from typing import Any
 
-from configuration.models import Configuration
 from django.db.models import F, OuterRef, QuerySet, Subquery
 from django.http import HttpResponse
 from permissions.permission import IsOwner, IsOwnerOrSharedUserOrSharedToOrg
@@ -211,31 +210,15 @@ def get(
                 status=status.HTTP_422_UNPROCESSABLE_ENTITY,
             )
 
-        # Process completed execution
         response_status = status.HTTP_422_UNPROCESSABLE_ENTITY
         if execution_status_value == CeleryTaskState.COMPLETED.value:
             response_status = status.HTTP_200_OK
-            # Check if highlight data should be removed using configuration registry
-            api_deployment = deployment_execution_dto.api
-            organization = api_deployment.organization if api_deployment else None
-            enable_highlight = False  # Safe default if the key is unavailable (e.g., OSS)
-            # Check if the configuration key exists (Cloud deployment) or use settings (OSS)
-            from configuration.config_registry import ConfigurationRegistry
-
-            if ConfigurationRegistry.is_config_key_available(
-                "ENABLE_HIGHLIGHT_API_DEPLOYMENT"
-            ):
-                enable_highlight = Configuration.get_value_by_organization(
-                    config_key="ENABLE_HIGHLIGHT_API_DEPLOYMENT",
-                    organization=organization,
-                )
-            if not enable_highlight:
-                response.remove_result_metadata_keys(["highlight_data"])
-                response.remove_result_metadata_keys(["extracted_text"])
-            if not include_metadata:
-                response.remove_result_metadata_keys()
-            if not include_metrics:
-                response.remove_result_metrics()
+            DeploymentHelper.process_completed_execution(
+                response=response,
+                deployment_execution_dto=deployment_execution_dto,
+                include_metadata=include_metadata,
+                include_metrics=include_metrics,
+            )
         return Response(
             data={
                 "status": response.execution_status,

backend/api_v2/deployment_helper.py

@@ -258,8 +258,11 @@ def execute_workflow(
             result.status_api = DeploymentHelper.construct_status_endpoint(
                 api_endpoint=api.api_endpoint, execution_id=execution_id
             )
-            # Check if highlight data should be removed using configuration registry
+            # Ensure workflow identification keys are always in item metadata
             organization = api.organization if api else None
+            org_id = str(organization.organization_id) if organization else ""
+            cls._enrich_result_with_workflow_metadata(result, organization_id=org_id)
+            # Check if highlight data should be removed using configuration registry
             enable_highlight = False  # Safe default if the key is unavailable (e.g., OSS)
             from configuration.config_registry import ConfigurationRegistry
 
@@ -273,8 +276,10 @@ def execute_workflow(
             if not enable_highlight:
                 result.remove_result_metadata_keys(["highlight_data"])
                 result.remove_result_metadata_keys(["extracted_text"])
-            if not include_metadata:
-                result.remove_result_metadata_keys()
+            if include_metadata or include_metrics:
+                cls._enrich_result_with_usage_metadata(result)
+            if not include_metadata and not include_metrics:
+                result.remove_inner_result_metadata()
             if not include_metrics:
                 result.remove_result_metrics()
         except Exception as error:
@@ -293,6 +298,144 @@ def execute_workflow(
             )
         return APIExecutionResponseSerializer(result).data
 
+    @staticmethod
+    def _enrich_item_inner_metadata(
+        item: dict, file_exec_id: str, usage_helper: Any
+    ) -> None:
+        """Inject per-model usage breakdown into item['result']['metadata']."""
+        inner_result = item.get("result")
+        if not isinstance(inner_result, dict):
+            return
+        metadata = inner_result.get("metadata")
+        if not isinstance(metadata, dict):
+            return
+        usage_by_model = usage_helper.get_usage_by_model(file_exec_id)
+        if usage_by_model:
+            metadata.update(usage_by_model)
+
+    @staticmethod
+    def _enrich_item_top_metadata(
+        item: dict, file_exec_id: str, usage_helper: Any
+    ) -> None:
+        """Inject aggregated usage totals into item['metadata']['usage']."""
+        item_metadata = item.get("metadata")
+        if not isinstance(item_metadata, dict):
+            return
+        aggregated = usage_helper.get_aggregated_token_count(file_exec_id)
+        if aggregated:
+            aggregated["file_execution_id"] = file_exec_id
+            item_metadata["usage"] = aggregated
+
+    @staticmethod
+    def _enrich_result_with_usage_metadata(result: ExecutionResponse) -> None:
+        """Enrich each file result's metadata with usage data.
+
+        For each file_execution_id:
+        1. Injects per-model cost arrays (extraction_llm, challenge_llm,
+           embedding) into item["result"]["metadata"].
+        2. Injects aggregated usage totals into item["metadata"]["usage"],
+           matching the legacy response format.
+        """
+        if not isinstance(result.result, list):
+            return
+
+        from usage_v2.helper import UsageHelper
+
+        for item in result.result:
+            if not isinstance(item, dict):
+                continue
+            file_exec_id = item.get("file_execution_id")
+            if not file_exec_id:
+                continue
+            DeploymentHelper._enrich_item_inner_metadata(item, file_exec_id, UsageHelper)
+            DeploymentHelper._enrich_item_top_metadata(item, file_exec_id, UsageHelper)
+
+    @staticmethod
+    def _enrich_item_workflow_metadata(
+        item: dict,
+        file_exec_id: str,
+        fe_lookup: dict,
+        workflow_execution: Any,
+        organization_id: str,
+        tag_names: list[str],
+    ) -> None:
+        """Populate workflow identification keys into item['metadata']."""
+        if not isinstance(item.get("metadata"), dict):
+            item["metadata"] = {}
+        metadata = item["metadata"]
+        fe = fe_lookup.get(str(file_exec_id))
+        we = fe.workflow_execution if fe else workflow_execution
+        if fe:
+            metadata.setdefault("source_name", fe.file_name)
+            metadata.setdefault("source_hash", fe.file_hash or "")
+            metadata.setdefault("file_execution_id", str(fe.id))
+            metadata.setdefault("total_elapsed_time", fe.execution_time)
+        if we:
+            metadata.setdefault("workflow_id", str(we.workflow_id))
+            metadata.setdefault("execution_id", str(we.id))
+            metadata.setdefault(
+                "workflow_start_time",
+                we.created_at.timestamp() if we.created_at else None,
+            )
+        metadata.setdefault("organization_id", organization_id)
+        metadata.setdefault("tags", tag_names)
+
+    @staticmethod
+    def _enrich_result_with_workflow_metadata(
+        result: ExecutionResponse,
+        organization_id: str,
+    ) -> None:
+        """Ensure workflow identification keys are always present in item metadata.
+
+        Uses setdefault() — fills in MISSING keys only, never overwrites
+        values already present from the workers cache.
+        """
+        if not isinstance(result.result, list):
+            return
+
+        from workflow_manager.file_execution.models import WorkflowFileExecution
+
+        # 1. Collect file_execution_ids
+        file_exec_ids = [
+            item.get("file_execution_id")
+            for item in result.result
+            if isinstance(item, dict) and item.get("file_execution_id")
+        ]
+        if not file_exec_ids:
+            return
+
+        # 2. Batch query (single JOIN query for all file executions)
+        fe_lookup = {
+            str(fe.id): fe
+            for fe in WorkflowFileExecution.objects.filter(
+                id__in=file_exec_ids
+            ).select_related("workflow_execution")
+        }
+
+        # 3. Get execution-level data (tags) — one M2M query
+        workflow_execution = None
+        tag_names: list[str] = []
+        if fe_lookup:
+            first_fe = next(iter(fe_lookup.values()))
+            workflow_execution = first_fe.workflow_execution
+            tag_names = list(workflow_execution.tags.values_list("name", flat=True))
+
+        # 4. Enrich each item
+        for item in result.result:
+            if not isinstance(item, dict):
+                continue
+            file_exec_id = item.get("file_execution_id")
+            if not file_exec_id:
+                continue
+            DeploymentHelper._enrich_item_workflow_metadata(
+                item=item,
+                file_exec_id=file_exec_id,
+                fe_lookup=fe_lookup,
+                workflow_execution=workflow_execution,
+                organization_id=organization_id,
+                tag_names=tag_names,
+            )
+
     @staticmethod
     def get_execution_status(execution_id: str) -> ExecutionResponse:
         """Current status of api execution.
@@ -308,6 +451,42 @@ def get_execution_status(execution_id: str) -> ExecutionResponse:
         )
         return execution_response
 
+    @staticmethod
+    def process_completed_execution(
+        response: ExecutionResponse,
+        deployment_execution_dto: Any,
+        include_metadata: bool,
+        include_metrics: bool,
+    ) -> None:
+        """Enrich and clean up the response for a completed execution."""
+        from configuration.config_registry import ConfigurationRegistry
+
+        api_deployment = deployment_execution_dto.api
+        organization = api_deployment.organization if api_deployment else None
+        org_id = str(organization.organization_id) if organization else ""
+        DeploymentHelper._enrich_result_with_workflow_metadata(
+            response, organization_id=org_id
+        )
+        enable_highlight = False
+        if ConfigurationRegistry.is_config_key_available(
+            "ENABLE_HIGHLIGHT_API_DEPLOYMENT"
+        ):
+            from configuration.models import Configuration
+
+            enable_highlight = Configuration.get_value_by_organization(
+                config_key="ENABLE_HIGHLIGHT_API_DEPLOYMENT",
+                organization=organization,
+            )
+        if not enable_highlight:
+            response.remove_result_metadata_keys(["highlight_data"])
+            response.remove_result_metadata_keys(["extracted_text"])
+        if include_metadata or include_metrics:
+            DeploymentHelper._enrich_result_with_usage_metadata(response)
+        if not include_metadata and not include_metrics:
+            response.remove_inner_result_metadata()
+        if not include_metrics:
+            response.remove_result_metrics()
+
     @staticmethod
     def fetch_presigned_file(url: str) -> InMemoryUploadedFile:
         """Fetch a file from a presigned URL and convert it to an uploaded file.

backend/backend/worker_celery.py

@@ -0,0 +1,105 @@
+"""Lightweight Celery app for dispatching tasks to worker-v2 workers.
+
+The Django backend already has a Celery app for internal tasks (beat,
+periodic tasks, etc.) whose broker URL is set via CELERY_BROKER_URL.
+Workers use the same broker. This module provides a second Celery app
+instance that reuses the same broker URL (from Django settings) but
+bypasses Celery's env-var-takes-priority behaviour so it can coexist
+with the main Django Celery app in the same process.
+
+Problem: Celery reads the ``CELERY_BROKER_URL`` environment variable
+with highest priority — overriding constructor args, ``conf.update()``,
+and ``config_from_object()``.
+
+Solution: Subclass Celery and override ``connection_for_write`` /
+``connection_for_read`` so they always use our explicit broker URL,
+bypassing the config resolution chain entirely.
+"""
+
+import logging
+from urllib.parse import quote_plus
+
+from celery import Celery
+from django.conf import settings
+from kombu import Queue
+
+logger = logging.getLogger(__name__)
+
+_worker_app: Celery | None = None
+
+
+class _WorkerDispatchCelery(Celery):
+    """Celery subclass that forces an explicit broker URL.
+
+    Works around Celery's env-var-takes-priority behaviour where
+    ``CELERY_BROKER_URL`` always overrides per-app configuration.
+    The connection methods are the actual points where Celery opens
+    AMQP/Redis connections, so overriding them is both sufficient
+    and safe.
+    """
+
+    _explicit_broker: str | None = None
+
+    def connection_for_write(self, url=None, *args, **kwargs):
+        return super().connection_for_write(url or self._explicit_broker, *args, **kwargs)
+
+    def connection_for_read(self, url=None, *args, **kwargs):
+        return super().connection_for_read(url or self._explicit_broker, *args, **kwargs)
+
+
+def get_worker_celery_app() -> Celery:
+    """Get or create a Celery app for dispatching to worker-v2 workers.
+
+    The app uses:
+    - Same broker as the workers (built from CELERY_BROKER_BASE_URL,
+      CELERY_BROKER_USER, CELERY_BROKER_PASS via Django settings)
+    - Same PostgreSQL result backend as the Django Celery app
+
+    Returns:
+        Celery app configured for worker-v2 dispatch.
+    """
+    global _worker_app
+    if _worker_app is not None:
+        return _worker_app
+
+    # Reuse the broker URL already built by Django settings (base.py)
+    # from CELERY_BROKER_BASE_URL + CELERY_BROKER_USER + CELERY_BROKER_PASS
+    broker_url = settings.CELERY_BROKER_URL
+
+    # Reuse the same PostgreSQL result backend as Django's Celery app
+    result_backend = (
+        f"db+postgresql://{settings.DB_USER}:"
+        f"{quote_plus(settings.DB_PASSWORD)}"
+        f"@{settings.DB_HOST}:{settings.DB_PORT}/"
+        f"{settings.CELERY_BACKEND_DB_NAME}"
+    )
+
+    app = _WorkerDispatchCelery(
+        "worker-dispatch",
+        set_as_current=False,
+        fixups=[],
+    )
+    # Store the explicit broker URL for use in connection overrides
+    app._explicit_broker = broker_url
+
+    app.conf.update(
+        result_backend=result_backend,
+        task_queues=[Queue("executor")],
+        task_serializer="json",
+        accept_content=["json"],
+        result_serializer="json",
+        result_extended=True,
+    )
+
+    _worker_app = app
+    # Log broker host only (mask credentials)
+    safe_broker = broker_url.split("@")[-1] if "@" in broker_url else broker_url
+    safe_backend = (
+        result_backend.split("@")[-1] if "@" in result_backend else result_backend
+    )
+    logger.info(
+        "Created worker dispatch Celery app (broker=%s, result_backend=%s)",
+        safe_broker,
+        safe_backend,
+    )
+    return _worker_app