sea-dev 1.0.0

package/apps/workers/src/workflows/extraction.ts

@@ -0,0 +1,866 @@
+import { getWorkerConfig } from "@/lib/config";
+import { captureEvent } from "@/lib/posthog";
+import { Rag } from "@sea/ai/rag";
+import { GCPClient } from "@sea/clients/gcp";
+import { HatchetClientMgr } from "@sea/clients/hatchet";
+import {
+  AzureDocumentClassifier,
+  FormStructureExtractor,
+  LayoutExtractor,
+} from "@sea/core/data-extraction";
+import { extractSpecificPages } from "@sea/core/data-extraction/pdf-utils";
+import { db, schema } from "@sea/db";
+import type { DocumentExtractionRequest } from "@sea/schemas/core/data-extraction";
+import { nanoid } from "@sea/util/id";
+import { invariant } from "@sea/util/invariant";
+import { differenceInSeconds } from "date-fns";
+import { and, asc, desc, eq, inArray } from "drizzle-orm";
+
+const config = getWorkerConfig();
+const hatchet = new HatchetClientMgr(config.hatchet).getHatchet();
+
+interface BankDataType {
+  "Bank Statement": {
+    BankName: string;
+    AccountName?: string;
+    AccountHolderName?: string;
+    StatementStartDate?: string;
+    StatementEndDate?: string;
+    StatementPeriodStart?: string;
+    StatementPeriodEnd?: string;
+  };
+}
+
+/**
+ * Generates a human-readable name for a submission based on extracted bank data
+ * Falls back to "Submission {id}" if data is unavailable
+ *
+ * @param data - Extracted bank statement data containing bank name and dates
+ * @param id - Submission ID to use as fallback
+ * @returns Generated name like "BankName 2024-01-01 to 2024-01-31" or default
+ */
+async function generateSubmissionName(data: BankDataType, id: string): Promise<string | null> {
+  const defaultName = `Submission ${id.substring(0, 4)}`;
+  if (!data || typeof data !== "object") return defaultName;
+
+  // Try to use BankName and dates from Bank Statement if available
+  const bankStatement = data["Bank Statement"];
+  if (bankStatement && !Array.isArray(bankStatement)) {
+    const bankName = bankStatement.BankName;
+    const fromDate = bankStatement.StatementStartDate ?? bankStatement.StatementPeriodStart;
+    const toDate = bankStatement.StatementEndDate ?? bankStatement.StatementPeriodEnd;
+
+    if (bankName && typeof bankName === "string") {
+      let name = bankName;
+
+      // Add date range if available
+      if (fromDate && typeof fromDate === "string" && toDate && typeof toDate === "string") {
+        name += ` ${fromDate} to ${toDate}`;
+      }
+
+      return name;
+    }
+  }
+
+  // Fallback to "Submission" + first 4 characters of id
+  return defaultName;
+}
+
+/**
+ * Generates a sensible name for a deal based on submissions' extracted data
+ * Aggregates information from multiple submissions to create a deal-level name
+ */
+async function generateDealName(
+  submissions: Array<BankDataType>,
+  dealId: string,
+): Promise<string | null> {
+  const defaultName = `Deal ${dealId.substring(0, 4)}`;
+  if (submissions.length === 0) return defaultName;
+
+  const firstSubmission = submissions[0];
+  if (!firstSubmission || typeof firstSubmission !== "object") {
+    return defaultName;
+  }
+
+  // Try to use AccountName or AccountHolderName from Bank Statement if available
+  const bankStatement = firstSubmission["Bank Statement"];
+  if (bankStatement && !Array.isArray(bankStatement)) {
+    const accountName = bankStatement.AccountName;
+    if (accountName && typeof accountName === "string") {
+      return accountName;
+    }
+
+    const accountHolderName = bankStatement.AccountHolderName;
+    if (accountHolderName && typeof accountHolderName === "string") {
+      return accountHolderName;
+    }
+  }
+
+  // Fallback to "Deal" + first 4 characters of dealId
+  return defaultName;
+}
+
+/**
+ * Updates the progress of a document extraction by adding reasoning messages
+ * to the conversation. Creates or updates the latest assistant message with
+ * progress text.
+ *
+ * @param conversationId - ID of the conversation to update
+ * @param text - Progress message to display to the user
+ */
+async function updateProgress(conversationId: string, text: string) {
+  await db.transaction(async (tx) => {
+    const latestMessage = await tx.query.message.findFirst({
+      where: and(
+        eq(schema.message.conversationId, conversationId),
+        eq(schema.message.role, "assistant"),
+      ),
+      orderBy: desc(schema.message.createdAt),
+    });
+
+    if (latestMessage) {
+      await tx
+        .update(schema.message)
+        .set({
+          parts: [...latestMessage.parts, { type: "reasoning", text }],
+          updatedAt: new Date(),
+        })
+        .where(eq(schema.message.id, latestMessage.id));
+    } else {
+      await tx.insert(schema.message).values({
+        id: nanoid(),
+        conversationId,
+        role: "assistant",
+        parts: [{ type: "reasoning", text }],
+      });
+    }
+  });
+}
+
+// ============================================================================
+// Workflow
+// ============================================================================
+
+export const extractDataUnified = hatchet.workflow<DocumentExtractionRequest>({
+  name: "extract-data-unified",
+  onEvents: ["document:uploaded-unified"],
+});
+
+// ============================================================================
+// Types
+// ============================================================================
+
+interface SplitFileInfo {
+  documentExtractionId: string;
+  documentId: string;
+  blobKey: string;
+  name: string;
+  pageStart: number;
+  pageEnd: number;
+}
+
+interface CategorizationOutput {
+  splitFiles: SplitFileInfo[];
+  originalDocumentExtractionId: string;
+}
+
+/**
+ * Helper to convert page ranges from classification to page number arrays
+ */
+function getPageNumbers(
+  pages: Array<{ pageNumber: number; endPageNumber?: number }> | undefined,
+): number[] {
+  if (!pages || pages.length === 0) return [];
+
+  const pageNumbers: number[] = [];
+  for (const range of pages) {
+    const start = range.pageNumber;
+    const end = range.endPageNumber ?? start;
+    for (let i = start; i <= end; i++) {
+      pageNumbers.push(i);
+    }
+  }
+  return pageNumbers;
+}
+
+// ============================================================================
+// Tasks
+// ============================================================================
+
+/**
+ * Task 0: Categorization — split document and create child extractions
+ * Downloads the PDF, splits by page ranges, creates new document + documentExtraction
+ * records for each split. Non-PDF or single-chunk documents pass through unchanged.
+ * Only runs if the document is marked as composite.
+ */
+export const categorization = extractDataUnified.task({
+  name: "categorization",
+  timeout: "10m",
+  fn: async (data): Promise<CategorizationOutput> => {
+    const documentExtraction = await db.query.documentExtraction.findFirst({
+      where: eq(schema.documentExtraction.id, data.documentExtractionId),
+      with: {
+        document: true,
+        submission: {
+          with: {
+            conversation: true,
+          },
+        },
+      },
+    });
+
+    invariant(documentExtraction);
+    invariant(documentExtraction.document);
+    invariant(documentExtraction.submission);
+
+    // Update status to processing
+    await db
+      .update(schema.documentExtraction)
+      .set({ status: "processing" })
+      .where(eq(schema.documentExtraction.id, data.documentExtractionId));
+
+    // Only classify PDFs (skip non-PDF documents)
+    if (documentExtraction.document.contentType !== "application/pdf") {
+      console.warn("Skipping classification for non-PDF document");
+      return { splitFiles: [], originalDocumentExtractionId: data.documentExtractionId };
+    }
+
+    // Only classify documents stored on GCP (Azure download/upload not yet supported here)
+    if (documentExtraction.document.storageProvider !== "gcp") {
+      console.warn("Skipping classification for non-GCP document");
+      return { splitFiles: [], originalDocumentExtractionId: data.documentExtractionId };
+    }
+
+    await updateProgress(
+      documentExtraction.submission.conversationId,
+      "Analyzing document structure",
+    );
+
+    const gcpClient = new GCPClient(config.gcp);
+    const pdfBuffer = await gcpClient.downloadFile(documentExtraction.document.blobKey);
+
+    // Classify the document using Azure Document Intelligence
+    let classifierId: string | null = null;
+    switch (documentExtraction.submission.organizationId) {
+      case "rsvnsseb46qohsct2n6lh":
+        classifierId = "ilc-classification-model";
+        break;
+      case "f4ic8gv3a5z5ncir55s5b":
+        classifierId = "real-estate-classifier";
+        break;
+    }
+    if (!classifierId || !config.azure.documentProcessor) {
+      console.warn("No classifier model ID configured, skipping classification");
+      return {
+        splitFiles: [],
+        originalDocumentExtractionId: data.documentExtractionId,
+      };
+    }
+
+    const classifier = new AzureDocumentClassifier(config.azure.documentProcessor);
+    const classificationResult = await classifier.classifyDocumentAsync(
+      pdfBuffer.toString("base64"),
+      classifierId,
+      "application/pdf",
+      documentExtraction.document.composite ? "auto" : "none",
+    );
+    const documentSplits = classificationResult.documents.map((doc) => ({
+      name: doc.docType,
+      pages: getPageNumbers(doc.pages),
+      confidence: doc.confidence,
+    }));
+
+    // If document is not composite, just store the classification label and return
+    if (!documentExtraction.document.composite) {
+      const firstClassification = documentSplits[0];
+      if (firstClassification) {
+        await db
+          .update(schema.document)
+          .set({ classificationLabel: firstClassification.name })
+          .where(eq(schema.document.id, documentExtraction.documentId));
+        console.log(
+          `Classification: labeled document as "${firstClassification.name}" (confidence: ${firstClassification.confidence})`,
+        );
+      }
+      return {
+        splitFiles: [],
+        originalDocumentExtractionId: data.documentExtractionId,
+      };
+    }
+
+    // Document is composite - split it
+    await updateProgress(
+      documentExtraction.submission.conversationId,
+      `Splitting document into ${documentSplits.length} parts`,
+    );
+
+    const splitFiles: SplitFileInfo[] = [];
+    const baseBlobKey = documentExtraction.document.blobKey.replace(/\.pdf$/i, "");
+
+    for (const split of documentSplits) {
+      if (split.pages.length === 0) {
+        console.warn(`Skipping split "${split.name}" with no pages`);
+        continue;
+      }
+
+      const splitBuffer = await extractSpecificPages(pdfBuffer, split.pages);
+
+      const pagesLabel = split.pages.join("-");
+      const splitBlobKey = `${baseBlobKey}_pages_${pagesLabel}.pdf`;
+      await gcpClient.uploadFile(splitBuffer, splitBlobKey, "application/pdf");
+
+      // Get signed URL for the split document
+      const { url: blobUrl, expiry: blobUrlExpiry } = await gcpClient.getSignedUrl(splitBlobKey);
+
+      const [splitDocument] = await db
+        .insert(schema.document)
+        .values({
+          dealId: documentExtraction.document.dealId,
+          name: split.name,
+          contentType: "application/pdf",
+          storageProvider: "gcp",
+          blobKey: splitBlobKey,
+          blobUrl,
+          blobUrlExpiry,
+          uploaded: true,
+          parentId: documentExtraction.documentId,
+          classificationLabel: split.name,
+        })
+        .returning();
+      invariant(splitDocument);
+
+      const [splitExtraction] = await db
+        .insert(schema.documentExtraction)
+        .values({
+          documentId: splitDocument.id,
+          submissionId: documentExtraction.submissionId,
+          status: "processing",
+          extractionModel: documentExtraction.extractionModel,
+          extractionTrigger: documentExtraction.extractionTrigger,
+        })
+        .returning();
+      invariant(splitExtraction);
+
+      splitFiles.push({
+        documentExtractionId: splitExtraction.id,
+        documentId: splitDocument.id,
+        blobKey: splitBlobKey,
+        name: split.name,
+        pageStart: Math.min(...split.pages),
+        pageEnd: Math.max(...split.pages),
+      });
+    }
+
+    console.log(
+      `Categorization: split document into ${splitFiles.length} parts for extraction ${data.documentExtractionId}`,
+    );
+
+    return {
+      splitFiles,
+      originalDocumentExtractionId: data.documentExtractionId,
+    };
+  },
+});
+
+/**
+ * Helper: extract layout for a single documentExtractionId.
+ * Shared by extractLayout for both split and non-split paths.
+ */
+async function extractLayoutForExtraction(documentExtractionId: string) {
+  const documentExtraction = await db.query.documentExtraction.findFirst({
+    where: eq(schema.documentExtraction.id, documentExtractionId),
+    with: {
+      document: true,
+      submission: {
+        with: {
+          form: {
+            with: {
+              groups: {
+                orderBy: asc(schema.formGroup.order),
+              },
+            },
+          },
+        },
+      },
+    },
+  });
+
+  invariant(documentExtraction);
+  invariant(documentExtraction.submission);
+  invariant(documentExtraction.document);
+  invariant(documentExtraction.submission.form);
+
+  // Update status to processing
+  await db
+    .update(schema.documentExtraction)
+    .set({ status: "processing" })
+    .where(eq(schema.documentExtraction.id, documentExtractionId));
+
+  const extractor = new LayoutExtractor(config);
+  const hasPretrainedGroups = documentExtraction.submission.form.groups.some(
+    (g) => g.type === "bank_statement",
+  );
+  const extractionModel = hasPretrainedGroups ? "bank_statement" : "custom";
+  const provider = hasPretrainedGroups ? "gcp" : "azure";
+  const rawContentUri = await extractor.extractLayout(
+    documentExtraction.document.blobKey,
+    documentExtraction.document.contentType,
+    {
+      extractionModel,
+      provider,
+    },
+  );
+
+  await db
+    .update(schema.document)
+    .set({ rawContentUri })
+    .where(eq(schema.document.id, documentExtraction.documentId));
+
+  return { rawContentUri, conversationId: documentExtraction.submission.conversationId };
+}
+
+/**
+ * Task 1: Extract generic layout from document(s)
+ * Depends on categorization. If the document was split, processes each split in parallel.
+ */
+export const extractLayout = extractDataUnified.task({
+  name: "extract-layout",
+  timeout: "10m",
+  parents: [categorization],
+  fn: async (data, ctx) => {
+    try {
+      const catOutput = (await ctx.parentOutput(categorization)) as CategorizationOutput;
+      const isSplit = catOutput.splitFiles.length > 0;
+      const extractionIds = isSplit
+        ? catOutput.splitFiles.map((f) => f.documentExtractionId)
+        : [data.documentExtractionId];
+
+      // Update progress once using the original extraction's conversation
+      const originalExtraction = await db.query.documentExtraction.findFirst({
+        where: eq(schema.documentExtraction.id, data.documentExtractionId),
+        with: { submission: true },
+      });
+      invariant(originalExtraction);
+      await updateProgress(
+        originalExtraction.submission.conversationId,
+        isSplit
+          ? `Processing layouts for ${extractionIds.length} document parts`
+          : "Processing document layouts",
+      );
+
+      await Promise.all(extractionIds.map((id) => extractLayoutForExtraction(id)));
+
+      return { splitFiles: catOutput.splitFiles };
+    } catch (error) {
+      console.error(
+        `Failed to extract generic layout for document extraction ${data.documentExtractionId}:`,
+        error,
+      );
+      throw error;
+    }
+  },
+});
+
+/**
+ * Helper: index a single document extraction for RAG.
+ */
+async function indexSingleDocument(documentExtractionId: string) {
+  const documentExtraction = await db.query.documentExtraction.findFirst({
+    where: eq(schema.documentExtraction.id, documentExtractionId),
+    with: {
+      document: true,
+    },
+  });
+  invariant(documentExtraction);
+  invariant(documentExtraction.document);
+  invariant(documentExtraction.document.rawContentUri);
+
+  const extractor = new LayoutExtractor(config);
+  const layout = await extractor.getLayout(documentExtraction.document.rawContentUri);
+  const rag = new Rag([layout]);
+  await rag.vectorize();
+
+  if (rag.chunks?.length) {
+    const parentChunks = rag.chunks.filter((c) => c.type === "parent");
+    const childChunks = rag.chunks.filter((c) => c.type === "child");
+    const parentIdMap = new Map<string, string>();
+
+    if (parentChunks.length > 0) {
+      const insertedParents = await db
+        .insert(schema.documentChunks)
+        .values(
+          parentChunks.map((chunk) => ({
+            documentId: documentExtraction.documentId,
+            submissionId: documentExtraction.submissionId,
+            text: chunk.text,
+            chunkType: chunk.type,
+            embedding: chunk.embedding,
+            headingPath: chunk.metadata.headingPath || [],
+            nodeType: chunk.metadata.nodeType,
+            originalId: chunk.id,
+            childIds: [],
+            metadata: chunk.metadata,
+          })),
+        )
+        .returning({ id: schema.documentChunks.id });
+
+      insertedParents.forEach((row, i) => {
+        const ragId = parentChunks[i]!.id;
+        parentIdMap.set(ragId, row.id);
+      });
+    }
+
+    if (childChunks.length > 0) {
+      await db.insert(schema.documentChunks).values(
+        childChunks.map((chunk) => {
+          const ragParentId = chunk.parentId;
+          const dbParentId = ragParentId ? parentIdMap.get(ragParentId) : undefined;
+
+          return {
+            documentId: documentExtraction.documentId,
+            submissionId: documentExtraction.submissionId,
+            text: chunk.text,
+            chunkType: chunk.type,
+            embedding: chunk.embedding,
+            headingPath: chunk.metadata.headingPath || [],
+            nodeType: chunk.metadata.nodeType,
+            originalId: chunk.metadata.originalId,
+            parentId: dbParentId,
+            childIds: [],
+          };
+        }),
+      );
+    }
+  }
+
+  await db
+    .update(schema.document)
+    .set({ indexed: true })
+    .where(eq(schema.document.id, documentExtraction.documentId));
+}
+
+/**
+ * Task 1.5: Index document(s) for RAG (Retrieval-Augmented Generation)
+ * Depends on layout extraction completing. Fans out over split files if present.
+ */
+export const indexDocument = extractDataUnified.task({
+  name: "index-document",
+  timeout: "5m",
+  parents: [extractLayout],
+  fn: async (inputData, ctx) => {
+    try {
+      const catOutput = (await ctx.parentOutput(categorization)) as CategorizationOutput;
+      const extractionIds =
+        catOutput.splitFiles.length > 0
+          ? catOutput.splitFiles.map((f) => f.documentExtractionId)
+          : [inputData.documentExtractionId];
+
+      await Promise.all(extractionIds.map((id) => indexSingleDocument(id)));
+    } catch (error) {
+      console.error(
+        `Failed to index document for document extraction ${inputData.documentExtractionId}:`,
+        error,
+      );
+    }
+  },
+});
+
+/**
+ * Task 2: Structure and store form data
+ * Depends on layout extraction completing. Fans out over split files if present.
+ * Aggregates results from all splits into the parent extraction.
+ */
+export const structureAndStoreFormData = extractDataUnified.task({
+  name: "structure-and-store-form-data",
+  timeout: "20m",
+  parents: [extractLayout],
+  fn: async (inputData, ctx) => {
+    const catOutput = (await ctx.parentOutput(categorization)) as CategorizationOutput;
+    const isSplit = catOutput.splitFiles.length > 0;
+
+    const documentExtraction = await db.query.documentExtraction.findFirst({
+      where: eq(schema.documentExtraction.id, inputData.documentExtractionId),
+      with: {
+        document: true,
+        submission: {
+          with: {
+            form: { with: { groups: { with: { fields: true, derived: true } } } },
+            conversation: true,
+            deal: true,
+          },
+        },
+      },
+    });
+
+    invariant(documentExtraction);
+    invariant(documentExtraction.document);
+    invariant(documentExtraction.submission);
+    invariant(documentExtraction.submission.form);
+    invariant(documentExtraction.submission.conversation);
+
+    await updateProgress(
+      documentExtraction.submission.conversationId,
+      isSplit
+        ? `Structuring data from ${catOutput.splitFiles.length} document parts`
+        : "Structuring form group data",
+    );
+
+    // Build context from deal hint and form hint
+    const contextParts: string[] = [];
+    if (documentExtraction.submission.deal?.hint) {
+      contextParts.push(`Deal context: ${documentExtraction.submission.deal.hint}`);
+    }
+    if (documentExtraction.submission.form.hint) {
+      contextParts.push(`Form context: ${documentExtraction.submission.form.hint}`);
+    }
+    if (inputData.context) {
+      contextParts.push(inputData.context);
+    }
+    const context = contextParts.length > 0 ? contextParts.join("\n\n") : undefined;
+
+    // Determine which extraction IDs to process
+    const extractionIds = isSplit
+      ? catOutput.splitFiles.map((f) => f.documentExtractionId)
+      : [inputData.documentExtractionId];
+
+    // Run structureAndStore for each extraction in parallel
+    await Promise.all(
+      extractionIds.map(async (extractionId) => {
+        const extractor = new FormStructureExtractor({
+          layoutConfig: config,
+          extractionModel: inputData.options.extractionModel,
+          reasoningEffort: inputData.options.reasoningEffort,
+          form: documentExtraction.submission.form,
+          userPrompt: inputData.userPrompt,
+          context,
+        });
+        await extractor.structureAndStore(extractionId);
+      }),
+    );
+
+    // Gather citations from all extractions (split or original)
+    const citationIds = (
+      await db.query.documentCitation.findMany({
+        where: inArray(schema.documentCitation.documentExtractionId, extractionIds),
+        columns: { id: true },
+      })
+    ).map((c) => c.id);
+
+    const submissionGroups = await db.query.submissionGroup.findMany({
+      where: eq(schema.submissionGroup.submissionId, documentExtraction.submission.id),
+      with: {
+        formGroup: true,
+        items: {
+          with: {
+            formField: true,
+            versions: {
+              where:
+                citationIds.length > 0
+                  ? inArray(schema.submissionItemVersion.sourceId, citationIds)
+                  : undefined,
+              orderBy: [
+                desc(schema.submissionItemVersion.selected),
+                desc(schema.submissionItemVersion.updatedAt),
+              ],
+              limit: 1,
+            },
+          },
+        },
+      },
+      orderBy: [asc(schema.submissionGroup.formGroupId), asc(schema.submissionGroup.order)],
+    });
+
+    // Build aggregated data from all splits
+    const data: Record<
+      string,
+      | Record<string, number | string | boolean | null>
+      | Array<Record<string, number | string | boolean | null>>
+    > = {};
+
+    for (const group of submissionGroups) {
+      const groupData = group.items.reduce(
+        (acc, item) => ({
+          ...acc,
+          [item.formField.name]: item.versions[0]?.value ?? null,
+        }),
+        {} as Record<string, number | string | boolean | null>,
+      );
+
+      if (group.formGroup.multi) {
+        if (!Array.isArray(data[group.formGroup.name])) {
+          data[group.formGroup.name] = [];
+        }
+        (
+          data[group.formGroup.name] as Array<Record<string, number | string | boolean | null>>
+        ).push(groupData);
+      } else {
+        data[group.formGroup.name] = groupData;
+      }
+    }
+
+    // Mark parent extraction as done with aggregated data
+    await db
+      .update(schema.documentExtraction)
+      .set({ status: "done", data })
+      .where(eq(schema.documentExtraction.id, inputData.documentExtractionId));
+
+    // Also mark child extractions as done
+    if (isSplit) {
+      await db
+        .update(schema.documentExtraction)
+        .set({ status: "done" })
+        .where(inArray(schema.documentExtraction.id, extractionIds));
+    }
+
+    const createdAt = documentExtraction.createdAt ? new Date(documentExtraction.createdAt) : null;
+
+    const processingSeconds =
+      createdAt && !Number.isNaN(createdAt.getTime())
+        ? differenceInSeconds(new Date(), createdAt)
+        : undefined;
+    captureEvent({
+      event: "document_extraction_completed",
+      distinctId: documentExtraction.submission.organizationId,
+      properties: {
+        document_extraction_id: inputData.documentExtractionId,
+        document_id: documentExtraction.documentId,
+        submission_id: documentExtraction.submissionId,
+        form_id: documentExtraction.submission.formId,
+        deal_id: documentExtraction.submission.dealId,
+        model: documentExtraction.extractionModel,
+        organization_id: documentExtraction.submission.organizationId,
+        split_count: isSplit ? catOutput.splitFiles.length : 0,
+        ...(processingSeconds !== undefined && {
+          processing_seconds: processingSeconds,
+        }),
+      },
+      groups: { company: documentExtraction.submission.organizationId },
+    });
+  },
+});
+
+/**
+ * Task 3: Generate user-friendly names for submission and deal
+ * Depends on form data being extracted and stored
+ * Auto-generates names based on extracted bank statement data if names were auto-generated
+ */
+export const generateSubmissionAndDealNames = extractDataUnified.task({
+  name: "generate-submission-and-deal-names",
+  timeout: "1m",
+  parents: [structureAndStoreFormData],
+  fn: async (inputData) => {
+    const documentExtraction = await db.query.documentExtraction.findFirst({
+      where: eq(schema.documentExtraction.id, inputData.documentExtractionId),
+      with: {
+        document: true,
+        submission: {
+          with: { form: { with: { groups: { with: { fields: true } } } }, deal: true },
+        },
+      },
+    });
+    invariant(documentExtraction);
+
+    const tasksToRun: Promise<void>[] = [];
+    // auto-generate submission name if it's auto-generated and we have data
+    if (documentExtraction.submission.nameAutoGenerated) {
+      tasksToRun.push(
+        (async () => {
+          const generatedSubmissionName = await generateSubmissionName(
+            documentExtraction.data as BankDataType,
+            documentExtraction.submission.id,
+          );
+          if (generatedSubmissionName) {
+            await db
+              .update(schema.submission)
+              .set({ name: generatedSubmissionName, nameAutoGenerated: false })
+              .where(
+                and(
+                  eq(schema.submission.id, documentExtraction.submission.id),
+                  eq(schema.submission.nameAutoGenerated, true),
+                ),
+              );
+          }
+        })(),
+      );
+    }
+
+    // same for deal
+    if (documentExtraction.submission.deal.nameAutoGenerated) {
+      tasksToRun.push(
+        (async () => {
+          const generatedDealName = await generateDealName(
+            [documentExtraction.data as BankDataType],
+            documentExtraction.submission.dealId,
+          );
+          if (generatedDealName) {
+            await db
+              .update(schema.deal)
+              .set({ name: generatedDealName, nameAutoGenerated: false })
+              .where(
+                and(
+                  eq(schema.deal.id, documentExtraction.submission.dealId),
+                  eq(schema.deal.nameAutoGenerated, true),
+                ),
+              );
+          }
+        })(),
+      );
+    }
+
+    await Promise.allSettled(tasksToRun);
+  },
+});
+
+/**
+ * Cleanup task that runs on workflow failure
+ * Marks the parent document extraction as failed.
+ * Child split extractions (if any) are identified by matching the parent document's
+ * blobKey prefix to avoid accidentally failing unrelated extractions in the same submission.
+ */
+export const cleanup = extractDataUnified.onFailure({
+  name: "cleanup",
+  fn: async (data) => {
+    try {
+      await db
+        .update(schema.documentExtraction)
+        .set({ status: "failed" })
+        .where(eq(schema.documentExtraction.id, data.documentExtractionId));
+
+      // Find child split extractions by looking for documents whose blobKey
+      // is derived from the parent document's blobKey (e.g. "key_pages_1-3.pdf")
+      const parentExtraction = await db.query.documentExtraction.findFirst({
+        where: eq(schema.documentExtraction.id, data.documentExtractionId),
+        with: { document: true },
+      });
+      if (parentExtraction?.document) {
+        const baseBlobKey = parentExtraction.document.blobKey.replace(/\.pdf$/i, "");
+        // Find split documents by blob key prefix pattern
+        const splitDocuments = await db.query.document.findMany({
+          where: eq(schema.document.dealId, parentExtraction.document.dealId),
+          columns: { id: true, blobKey: true },
+        });
+        const splitDocIds = splitDocuments
+          .filter((d) => d.blobKey.startsWith(baseBlobKey + "_pages_"))
+          .map((d) => d.id);
+
+        if (splitDocIds.length > 0) {
+          await db
+            .update(schema.documentExtraction)
+            .set({ status: "failed" })
+            .where(
+              and(
+                inArray(schema.documentExtraction.documentId, splitDocIds),
+                eq(schema.documentExtraction.submissionId, parentExtraction.submissionId),
+              ),
+            );
+        }
+      }
+
+      console.error(
+        `Failed to extract structured data for document extraction ${data.documentExtractionId}`,
+      );
+    } catch (error) {
+      console.error(`Failed to update status during cleanup:`, error);
+    }
+  },
+});