UN-4044 [FEAT] Agent-KV extraction engine — public API, schema compiler and hardened code sandbox

backend/agent_kv/apps.py

@@ -0,0 +1,6 @@
+from django.apps import AppConfig
+
+
+class AgentKvConfig(AppConfig):
+    default_auto_field = "django.db.models.BigAutoField"
+    name = "agent_kv"

backend/agent_kv/constants.py

@@ -0,0 +1,52 @@
+#: v1 accepts exactly one extractor (`kv`), so the job row does not carry which
+#: one it ran. When fan-out lands this becomes per-job state rather than a
+#: constant -- the wire format (spec §7.0) is already shaped for that.
+V1_EXTRACTOR_NAME = "kv"
+
+STAGE_NAMES = [
+    "document_processing",
+    "extraction",
+    "qa",
+    "challenge",
+    "normalize",
+    "constraints",
+    "codegen",
+    "code_execution",
+]
+EXECUTOR_NAME = "agentic_kv"
+OPERATION_KV_EXTRACT = "kv_extract"
+EXECUTION_SOURCE = "agent_kv_api"
+
+#: The table extractor, served by the cloud `agentic_table` plugin's blind-API
+#: operation. Same executor as the IDE table path (and therefore the same,
+#: already-wired `celery_executor_agentic_table` queue) with a second operation
+#: -- a new executor name would derive a new queue needing wiring at five
+#: sites, and an unwired queue accepts work and drains nothing, silently.
+TABLE_EXTRACTOR_NAME = "table"
+
+TABLE_EXECUTOR_NAME = "agentic_table"
+OPERATION_TABLE_EXTRACT_API = "table_extract_api"
+
+#: Which executor and operation each extractor dispatches to. `dispatch_job`
+#: reads this rather than hardcoding one pair, so adding an extractor is one
+#: entry here plus its options serializer and stage list.
+EXTRACTOR_ROUTES = {
+    V1_EXTRACTOR_NAME: (EXECUTOR_NAME, OPERATION_KV_EXTRACT),
+    TABLE_EXTRACTOR_NAME: (TABLE_EXECUTOR_NAME, OPERATION_TABLE_EXTRACT_API),
+}
+
+#: The table engine reports one coarse stage: it has no node-level progress
+#: hooks (the IDE path gets `stream_log` only), so inventing finer stages here
+#: would describe progress the executor cannot actually report.
+TABLE_STAGE_NAMES = ["table_extraction"]
+
+#: Stage names ARE wire format -- they are returned to clients -- and they are
+#: extractor-specific (`qa`/`challenge`/`codegen` mean nothing to the table
+#: extractor). `_status_document` filters a job's recorded stages through the
+#: list for the extractor that ran: `StageReportView` persists whatever name
+#: the executor sends, so without a per-extractor list a table job's stages
+#: would be stored and then silently filtered out of every status response.
+STAGE_NAMES_BY_EXTRACTOR = {
+    V1_EXTRACTOR_NAME: STAGE_NAMES,
+    TABLE_EXTRACTOR_NAME: TABLE_STAGE_NAMES,
+}

backend/agent_kv/dispatch.py

@@ -0,0 +1,158 @@
+"""Executor dispatch glue (spec §5.3). One dispatch per job; UUID task_id."""
+
+import logging
+import uuid
+
+from celery import signature
+from django.conf import settings
+from django.utils import timezone
+
+from agent_kv.constants import EXECUTION_SOURCE, EXTRACTOR_ROUTES
+from agent_kv.models import AgentKVJob, JobStatus
+from unstract.sdk1.execution.context import ExecutionContext
+
+logger = logging.getLogger(__name__)
+
+CALLBACK_QUEUE = "agent_kv_callback"
+
+
+class DispatchError(Exception):
+    """Enqueue failed; the caller terminalizes the job (spec §5.3)."""
+
+
+def _dispatcher():
+    # No `celery_app`: UN-4046 removed that parameter when the routing
+    # dispatcher's Celery branch went with the pg_queue_enabled flag. Passing it
+    # raises TypeError, so every submit failed to dispatch -- and because it
+    # fails at the call rather than at import, nothing catches it until a real
+    # request is made.
+    from pg_queue.executor_rpc import get_executor_dispatcher
+
+    return get_executor_dispatcher()
+
+
+def _platform_api_key(job) -> str:
+    # Lazy import: avoids Django app registry init order (mirrors
+    # PromptStudioHelper._get_platform_api_key).
+    from platform_settings_v2.platform_auth_service import (
+        PlatformAuthenticationService,
+    )
+
+    # ``get_active_platform_key`` takes the org's public *slug*
+    # (``Organization.organization_id``, e.g. ``org_abc123``) and resolves it
+    # via ``get_organization_by_org_id`` -- NOT the row's UUID primary key that
+    # ``job.organization_id`` holds. Passing the PK here silently resolves to
+    # no organization and every dispatch fails with ``ActiveKeyNotFound``
+    # (caught live in the Task 13b integration run).
+    org_slug = job.organization.organization_id
+    platform_key = PlatformAuthenticationService.get_active_platform_key(org_slug)
+    if not platform_key:
+        raise DispatchError(f"No active platform key for org {org_slug}")
+    return str(platform_key.key)
+
+
+def dispatch_job(job, *, extractor: str, schema: dict, options: dict) -> None:
+    executor_name, operation = EXTRACTOR_ROUTES[extractor]
+    org_id = str(job.organization_id)
+    # Everything that can fail — platform-key lookup, context construction,
+    # and the enqueue call itself — lives inside this try so no internal
+    # failure (e.g. a transient DB error resolving the platform key) can
+    # escape as a raw, uncaught exception. Only the post-success bookkeeping
+    # below runs outside it.
+    try:
+        job.task_id = uuid.uuid4()
+        context = ExecutionContext(
+            executor_name=executor_name,
+            operation=operation,
+            run_id=str(job.id),
+            execution_source=EXECUTION_SOURCE,
+            organization_id=org_id,
+            executor_params={
+                "job_id": str(job.id),
+                "input_ref": job.input_ref,
+                "schema": schema,
+                "options": options,
+                "platform_api_key": _platform_api_key(job),
+                # The CAP the engine must enforce (spec §6.1/§6.6), not the
+                # measured count -- job.pages_total is None for Excel (no
+                # pre-OCR page concept), which would otherwise leave the
+                # engine with nothing to check the post-OCR virtual-page cap
+                # against. The measured count still rides along separately.
+                "max_pages": settings.AGENT_KV_MAX_PAGES,
+                "pages_total": job.pages_total,
+            },
+        )
+        cb_kwargs = {"callback_kwargs": {"job_id": str(job.id), "org_id": org_id}}
+        _dispatcher().dispatch_with_callback(
+            context,
+            on_success=signature(
+                "agent_kv_complete", kwargs=cb_kwargs, queue=CALLBACK_QUEUE
+            ),
+            on_error=signature("agent_kv_error", kwargs=cb_kwargs, queue=CALLBACK_QUEUE),
+            task_id=str(job.task_id),
+        )
+    except DispatchError:
+        raise
+    except Exception as e:
+        raise DispatchError(str(e)) from e
+    job.status = JobStatus.DISPATCHED
+    job.dispatched_at = timezone.now()
+    # Guarded queryset UPDATE, not job.save(): a plain save would blindly
+    # overwrite whatever status this job already raced to. Concretely: the
+    # executor can fail (or the job be cancelled) essentially instantly
+    # after enqueue, and its finalize callback can land -- marking the row
+    # FAILED/CANCELLED -- before this post-enqueue bookkeeping runs. An
+    # unconditional save() here would rewrite that terminal status back to
+    # DISPATCHED, un-terminalizing the job forever (nothing else ever
+    # revisits a DISPATCHED row). Only a still-PENDING row is advanced; a
+    # row this UPDATE doesn't match is left exactly as the winning writer
+    # left it. `modified_at` is stamped automatically by
+    # BaseModelQuerySet.update() (utils/models/base_model.py).
+    # Everything below is POST-ENQUEUE bookkeeping. The task is already on the
+    # queue, so a failure here must never be reported as a failed dispatch:
+    # `SubmitView` turns a DispatchError into a FAILED job, and the executor
+    # would then run, callback, and find a terminal row it cannot write to --
+    # the caller told nothing was billed for work that did run. Wrapped rather
+    # than left to propagate, which is what the single pre-review UPDATE did.
+    #
+    # Losing the bookkeeping entirely is recoverable: sweep phase 1 reaps a
+    # still-PENDING row with no `dispatched_at`, and phase 2's
+    # `dispatched_at IS NULL` arm covers the non-PENDING case.
+    try:
+        _record_dispatch(job)
+    except Exception:
+        logger.exception(
+            "agent-kv: dispatch bookkeeping failed for job %s after enqueue "
+            "(task is queued; the sweep will reconcile)",
+            job.id,
+        )
+
+
+def _record_dispatch(job) -> None:
+    """Persist task_id/status/dispatched_at against whatever the row raced to."""
+    advanced = AgentKVJob.objects.filter(id=job.id, status=JobStatus.PENDING).update(
+        task_id=job.task_id,
+        status=job.status,
+        dispatched_at=job.dispatched_at,
+    )
+    if not advanced:
+        # The row moved off PENDING between the enqueue above and this write.
+        # The benign case is a terminal status (the guard's whole purpose) --
+        # but there is a non-terminal one: StageReportView promotes
+        # PENDING -> RUNNING on the executor's FIRST stage report, which can
+        # easily land before this bookkeeping. The guard above then matches 0
+        # rows and `dispatched_at` stays NULL -- and a non-terminal row with a
+        # NULL `dispatched_at` is invisible to BOTH sweep phases: phase 1
+        # requires `status=PENDING`, phase 2 filters `dispatched_at__lt=cutoff`
+        # and SQL `NULL < x` is never true. The job reports `running` forever
+        # and `GET result` 409s for the life of the row, with nothing able to
+        # recover it.
+        #
+        # So stamp the dispatch bookkeeping for any still-non-terminal row,
+        # WITHOUT touching `status`: the row genuinely was dispatched, and
+        # moving RUNNING back to DISPATCHED would lose the executor's own
+        # progress. `dispatched_at__isnull=True` keeps this idempotent and
+        # stops a retry overwriting the original dispatch time.
+        AgentKVJob.objects.filter(id=job.id, dispatched_at__isnull=True).exclude(
+            status__in=list(AgentKVJob.TERMINAL)
+        ).update(task_id=job.task_id, dispatched_at=job.dispatched_at)

backend/agent_kv/exceptions.py

@@ -0,0 +1,16 @@
+from rest_framework.exceptions import APIException
+
+
+class EngineUnavailable(APIException):
+    status_code = 501
+    default_detail = "agent-kv engine not available on this deployment"
+
+
+class RateLimited(APIException):
+    status_code = 429
+    default_detail = "Too many requests"
+
+
+class JobNotFound(APIException):
+    status_code = 404
+    default_detail = "Job not found"