sea-dev 1.0.0

package/libs/core/src/data-extraction/classification/azure.ts

@@ -0,0 +1,168 @@
+import type {
+  AnalyzeOperationOutput,
+  AnalyzeResultOutput,
+  DocumentIntelligenceClient,
+} from "@azure-rest/ai-document-intelligence";
+import createClient, {
+  getLongRunningPoller,
+} from "@azure-rest/ai-document-intelligence";
+import { AzureKeyCredential } from "@azure/core-auth";
+import type { AzureDocumentProcessorConfig } from "@sea/config";
+import type {
+  ClassificationResult,
+  DocumentClassificationResponse,
+} from "@sea/schemas/core/data-extraction";
+
+/**
+ * Azure Document Intelligence classifier for document classification
+ * Handles document classification using Azure's custom classification models
+ */
+export class AzureDocumentClassifier {
+  private documentIntelligenceClient: DocumentIntelligenceClient;
+
+  constructor(config: AzureDocumentProcessorConfig) {
+    this.documentIntelligenceClient = createClient(
+      config.documentIntelligenceEndpoint,
+      new AzureKeyCredential(config.documentIntelligenceKey),
+    );
+  }
+
+  /**
+   * Performs asynchronous document classification using a custom classifier model
+   * @param base64Data - Base64-encoded document data
+   * @param classifierId - The custom classifier model ID
+   * @param mimeType - MIME type of the document
+   * @param split - Document splitting mode ("auto" | "none" | "perPage"). Default is "auto" but "none" is recommended for v4.0 API
+   * @returns Promise<DocumentClassificationResponse> - The classification results
+   * @throws Error if the classification operation fails or Document Intelligence client is not configured
+   */
+  public async classifyDocumentAsync(
+    base64Data: string,
+    classifierId: string,
+    mimeType: "image/png" | "application/pdf" = "application/pdf",
+    split: "auto" | "none" | "perPage" = "auto",
+  ): Promise<DocumentClassificationResponse> {
+    console.info(
+      `Running classification for ${mimeType} with classifier ${classifierId}`,
+    );
+    if (!this.documentIntelligenceClient) {
+      throw new Error(
+        "Document Intelligence client is not configured. Please provide credentials in the config.",
+      );
+    }
+    try {
+      const initialResponse = await this.documentIntelligenceClient
+        .path("/documentClassifiers/{classifierId}:analyze", classifierId)
+        .post({
+          contentType: "application/json",
+          body: {
+            base64Source: base64Data,
+          },
+          queryParameters: {
+            // Using the latest API version
+            "api-version": "2024-11-30",
+            _overload: "classifyDocument",
+            ...(split && split !== "none" ? { split } : {}),
+          },
+        });
+
+      if (initialResponse.status !== "202") {
+        throw new Error(
+          `Classification failed with status: ${initialResponse.status}`,
+        );
+      }
+
+      const poller = getLongRunningPoller(
+        this.documentIntelligenceClient,
+        initialResponse,
+      );
+      const result = await poller.pollUntilDone();
+      if (result.status !== "200") {
+        throw new Error(
+          `Classification completed with error status: ${result.status}`,
+        );
+      }
+
+      const { analyzeResult } = result.body as AnalyzeOperationOutput;
+      if (!analyzeResult) {
+        throw new Error(
+          "No documents were classified or classification failed",
+        );
+      }
+
+      // Transform Azure response to our standardized format
+      return this.transformAzureClassificationResult(analyzeResult);
+    } catch (error) {
+      console.error("Error in classifyDocumentAsync:", error);
+      throw new Error(
+        `Failed to classify document using classifier ${classifierId}: ${error instanceof Error ? error.message : String(error)}`,
+      );
+    }
+  }
+
+  /**
+   * Classifies multiple document buffers with Azure Document Intelligence
+   * @param fileBuffers - Array of document buffers to classify
+   * @param classifierId - The classifier model ID to use
+   * @param mimeType - MIME type of the documents
+   * @param split - Document splitting mode
+   * @returns Promise<DocumentClassificationResponse[]> - Array of classification results
+   */
+  public async classifyDocuments(
+    fileBuffers: Buffer<ArrayBufferLike>[],
+    classifierId: string,
+    mimeType: string,
+    split: "auto" | "none" | "perPage" = "auto",
+  ): Promise<DocumentClassificationResponse[]> {
+    try {
+      const classifications = fileBuffers.map(async (b) => {
+        const base64File = b.toString("base64");
+        return await this.classifyDocumentAsync(
+          base64File,
+          classifierId,
+          mimeType as "image/png" | "application/pdf",
+          split,
+        );
+      });
+      return await Promise.all(classifications);
+    } catch (error) {
+      console.error("Error performing classification:", error);
+      throw new Error(
+        `Classification failed: ${error instanceof Error ? error.message : "Unknown error"}`,
+      );
+    }
+  }
+
+  /**
+   * Transforms Azure classification result to our standardized format
+   * @param azureResult - Azure classification output
+   * @returns Standardized DocumentClassificationResponse
+   * @private
+   */
+  private transformAzureClassificationResult(
+    azureResult: AnalyzeResultOutput,
+  ): DocumentClassificationResponse {
+    const documents: ClassificationResult[] = [];
+
+    if (azureResult.documents) {
+      for (const doc of azureResult.documents) {
+        const pages =
+          doc.boundingRegions?.map((region) => ({
+            pageNumber: region.pageNumber ?? 1,
+          })) ?? [];
+
+        documents.push({
+          docType: doc.docType,
+          confidence: doc.confidence ?? 0,
+          pages: pages.length > 0 ? pages : undefined,
+        });
+      }
+    }
+
+    return {
+      apiVersion: azureResult.apiVersion,
+      modelId: azureResult.modelId,
+      documents,
+    };
+  }
+}

package/libs/core/src/data-extraction/classification/index.ts

@@ -0,0 +1 @@
+export * from "./azure";

package/libs/core/src/data-extraction/dal.ts

@@ -0,0 +1,246 @@
+import type {
+  InsertDocumentCitation,
+  InsertSubmissionGroup,
+  InsertSubmissionItem,
+  InsertSubmissionItemVersion,
+  FormGroupHydrated,
+} from "@sea/db/types";
+import type {
+  StructuredGroup,
+  ValidationData,
+  BankStatementReconciliation,
+} from "./types";
+import { db, schema } from "@sea/db";
+import { asc, eq } from "drizzle-orm";
+import { invariant } from "@sea/util/invariant";
+import { nanoid } from "@sea/util/id";
+import { formatISODate } from "@sea/util/fmt";
+import { ValueWithReferenceSchema } from "./types";
+import type { SubmissionItemValue } from "@sea/db/schema";
+import type { ReferenceMap } from "@sea/schemas/core/layout";
+
+/**
+ * Stores extracted data to the submission by creating submission groups, items, and citations.
+ * Handles both single and multi-instance form groups, validating and linking all data to source references.
+ *
+ * @param params - Configuration object
+ * @param params.submissionId - Target submission ID
+ * @param params.documentExtractionId - Source document extraction ID
+ * @param params.formGroupId - Form group being processed
+ * @param params.data - Structured extracted data
+ * @param params.referenceMap - Maps reference IDs to document locations
+ * @param params.validationData - Optional validation results for marking invalid groups
+ */
+export async function storeExtractedGroupData({
+  submissionId,
+  documentExtractionId,
+  formGroupId,
+  data,
+  referenceMap,
+  validationData,
+  reconciliationData,
+}: {
+  submissionId: string;
+  documentExtractionId: string;
+  formGroupId: string;
+  data: StructuredGroup;
+  referenceMap: ReferenceMap;
+  validationData?: ValidationData;
+  reconciliationData?: BankStatementReconciliation;
+}) {
+  if (data === null) {
+    console.warn(`Deserialization aborted: data is null`);
+    return;
+  }
+
+  const submissionGroups: InsertSubmissionGroup[] = [];
+  const submissionItems: InsertSubmissionItem[] = [];
+  const itemVersions: InsertSubmissionItemVersion[] = [];
+  const documentCitations: InsertDocumentCitation[] = [];
+
+  const formGroup = await db.query.formGroup.findFirst({
+    where: eq(schema.formGroup.id, formGroupId),
+    with: {
+      fields: true,
+      derived: true,
+      submissionGroups: {
+        where: eq(schema.submissionGroup.submissionId, submissionId),
+        with: { items: true },
+        orderBy: asc(schema.submissionGroup.order),
+      },
+    },
+  });
+  invariant(formGroup, `No form group found for id ${formGroupId}`);
+
+  if (formGroup.multi) {
+    if (typeof data !== "object" || !Array.isArray(data)) {
+      console.warn(`Form group marked as multi, but data not an array`);
+      return;
+    }
+    const orderOffset =
+      formGroup.submissionGroups.length > 0
+        ? Math.max(...formGroup.submissionGroups.map((g) => g.order)) + 1
+        : 0;
+
+    for (const [order, groupElement] of data.entries()) {
+      const isValid = !validationData?.[formGroup.name]?.some(
+        (v) => v.index === order,
+      );
+      const newSubmissionGroup: InsertSubmissionGroup = {
+        id: nanoid(),
+        submissionId,
+        formGroupId: formGroup.id,
+        order: order + orderOffset,
+        isValid,
+      };
+      submissionGroups.push(newSubmissionGroup);
+
+      const newSubmissionItems = formGroup.fields.map((field) => ({
+        id: nanoid(),
+        submissionGroupId: newSubmissionGroup.id!,
+        formFieldId: field.id,
+      }));
+      submissionItems.push(...newSubmissionItems);
+
+      processGroupData(
+        documentExtractionId,
+        formGroup,
+        newSubmissionItems,
+        itemVersions,
+        documentCitations,
+        groupElement,
+        referenceMap,
+        order,
+      );
+    }
+  } else {
+    const submissionGroup = formGroup.submissionGroups[0];
+    invariant(submissionGroup);
+
+    processGroupData(
+      documentExtractionId,
+      formGroup,
+      submissionGroup.items,
+      itemVersions,
+      documentCitations,
+      data,
+      referenceMap,
+    );
+  }
+
+  await db.transaction(async (tx) => {
+    if (submissionGroups.length > 0)
+      await tx.insert(schema.submissionGroup).values(submissionGroups);
+    if (submissionItems.length > 0)
+      await tx.insert(schema.submissionItem).values(submissionItems);
+    if (documentCitations.length > 0)
+      await tx.insert(schema.documentCitation).values(documentCitations);
+    if (itemVersions.length > 0)
+      await tx.insert(schema.submissionItemVersion).values(itemVersions);
+
+    // Store reconciliation data in DocumentExtraction with breakdown if validation data exists
+    if (reconciliationData) {
+      const transactions = validationData?.Transactions;
+      const reconciliationWithBreakdown = {
+        ...reconciliationData,
+        breakdown: transactions && transactions.length > 0 ? transactions : [],
+      };
+
+      await tx
+        .update(schema.documentExtraction)
+        .set({ reconciliation: reconciliationWithBreakdown })
+        .where(eq(schema.documentExtraction.id, documentExtractionId));
+    }
+  });
+}
+
+/**
+ * Processes extracted data for a single form group and creates submission items with citations.
+ * Links field values to their source locations in the document via the reference map.
+ *
+ * @param documentExtractionId - Unique identifier for the document extraction
+ * @param formGroup - Form group definition with field metadata
+ * @param submissionItems - Array to populate with submission items
+ * @param itemVersions - Array to populate with submission item versions
+ * @param documentCitations - Array to populate with document citations
+ * @param groupData - Extracted data for this group
+ * @param referenceMap - Maps reference IDs to their source locations and text
+ * @param order - Optional order index for multi-instance groups
+ */
+export async function processGroupData(
+  documentExtractionId: string,
+  formGroup: FormGroupHydrated,
+  submissionItems: InsertSubmissionItem[],
+  itemVersions: InsertSubmissionItemVersion[],
+  documentCitations: InsertDocumentCitation[],
+  groupData: StructuredGroup,
+  referenceMap: ReferenceMap,
+  order: number | null = null,
+) {
+  if (!groupData || typeof groupData !== "object" || Array.isArray(groupData)) {
+    console.warn(
+      `Single group element ${formGroup.name} is missing or data is not an object, skipping...`,
+    );
+    return;
+  }
+
+  for (const field of formGroup.fields) {
+    const submissionItem = submissionItems.find(
+      (item) => item.formFieldId === field.id,
+    );
+    const valueWithRef = ValueWithReferenceSchema.parse(
+      groupData[field.name] ?? null,
+    );
+    if (!valueWithRef || !submissionItem) continue;
+
+    const { value, refId } = valueWithRef;
+    const sourceRefs = refId.split(",");
+
+    const newCitation: InsertDocumentCitation = {
+      id: nanoid(),
+      documentExtractionId,
+      path: `${formGroup.name}${order !== null ? `[${order}]` : ""}.${field.name}`,
+      text: "",
+      sourceRef: refId,
+      boundingRegions: {},
+    };
+    const newItemVersion: InsertSubmissionItemVersion = {
+      submissionItemId: submissionItem.id!,
+      sourceType: "document" as const,
+      selected: true,
+      sourceId: newCitation.id,
+      value: parseItemValue(value),
+    };
+
+    for (const sourceRef of sourceRefs) {
+      const citationDataForBlock = referenceMap[sourceRef];
+      if (!citationDataForBlock) continue;
+      newCitation.boundingRegions![sourceRef] =
+        citationDataForBlock.boundingRegions;
+      newCitation.text += citationDataForBlock.text;
+    }
+
+    documentCitations.push(newCitation);
+    itemVersions.push(newItemVersion);
+  }
+}
+
+/**
+ * Parses and normalizes field values into the submission item value format.
+ * Dates are converted to ISO format strings, booleans to numbers (0/1), and other values are passed through.
+ *
+ * @param value - Raw value from extraction (string, number, boolean, or Date)
+ * @returns Normalized value suitable for storing in submission items
+ */
+function parseItemValue(
+  value: string | number | boolean | Date,
+): SubmissionItemValue {
+  if (value instanceof Date) {
+    return formatISODate(value);
+  }
+  if (typeof value === "boolean") {
+    return Number(value);
+  }
+
+  return value;
+}