sea-dev 1.0.0

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

@@ -0,0 +1,730 @@
+import type {
+  AnalyzedDocumentOutput,
+  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 {
+  AzureDocumentResult,
+  AzureFieldOutput,
+  AzureParagraph,
+  AzureSection,
+  AzureTable,
+  AzureTableCell,
+} from "@sea/clients/azure";
+import type { Polygon, LayoutItemType } from "@sea/schemas/core/layout";
+import {
+  type PretrainedEntity,
+  type EntityType,
+  type ReferenceMap,
+  Cell,
+  DocumentLayout,
+  Footer,
+  Header,
+  Paragraph,
+  Row,
+  Section,
+  Table,
+  Heading,
+} from "@sea/schemas/core/layout";
+import { MAX_PAGES_PER_CHUNK } from "../pdf-utils";
+
+/**
+ * Azure Document Intelligence adapter for document extraction
+ * Handles document processing using Azure's Document Intelligence API
+ */
+export class AzureDocumentProcessor {
+  private documentIntelligenceClient: DocumentIntelligenceClient;
+
+  constructor(config: AzureDocumentProcessorConfig) {
+    this.documentIntelligenceClient = createClient(
+      config.documentIntelligenceEndpoint,
+      new AzureKeyCredential(config.documentIntelligenceKey),
+    );
+  }
+
+  /**
+   * Performs asynchronous document processing using a prebuilt model (e.g., bank statements, invoices)
+   * @param base64Data - Base64-encoded document data
+   * @param modelId - The prebuilt model ID (default: "prebuilt-document")
+   * @param mimeType - MIME type of the document
+   * @returns Promise<AnalyzeResultOutput> - The extracted data from the document
+   * @throws Error if the processing operation fails or Document Intelligence client is not configured
+   */
+  public async runExtractionAsync(
+    base64Data: string,
+    modelId: string,
+    mimeType: "image/png" | "application/pdf" = "application/pdf",
+  ): Promise<AnalyzeResultOutput> {
+    console.info(`Running extraction for ${mimeType} with model ${modelId}`);
+    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("/documentModels/{modelId}:analyze", modelId)
+        .post({
+          contentType: "application/json",
+          body: { base64Source: base64Data },
+        });
+
+      if (initialResponse.status !== "202") {
+        throw new Error(
+          `Analysis failed with status: ${initialResponse.status}`,
+        );
+      }
+      const poller = getLongRunningPoller(
+        this.documentIntelligenceClient,
+        initialResponse,
+      );
+      const result = await poller.pollUntilDone();
+      if (result.status !== "200") {
+        throw new Error(
+          `Analysis completed with error status: ${result.status}`,
+        );
+      }
+      const analysisOutput = result.body as AnalyzeOperationOutput;
+      const analyzeResult = analysisOutput.analyzeResult;
+      if (!analyzeResult) {
+        throw new Error("No documents were analyzed or extraction failed");
+      }
+
+      // Return the extracted data
+      return analyzeResult;
+    } catch (error) {
+      console.error("Error in runExtractionAsync:", error);
+      throw new Error(
+        `Failed to extract data using prebuilt model ${modelId}: ${error instanceof Error ? error.message : String(error)}`,
+      );
+    }
+  }
+
+  /**
+   * Processes multiple document buffers with Azure Document Intelligence
+   * @param fileBuffers - Array of document buffers to process
+   * @param modelId - The model ID to use for extraction
+   * @param mimeType - MIME type of the documents
+   * @returns Promise<AnalyzeResultOutput[]> - Array of extraction results
+   */
+  public async runExtraction(
+    fileBuffers: Buffer<ArrayBufferLike>[],
+    modelId: string,
+    mimeType: string,
+  ): Promise<AnalyzeResultOutput[]> {
+    try {
+      const extractions = fileBuffers.map(async (b) => {
+        const base64File = b.toString("base64");
+        return await this.runExtractionAsync(
+          base64File,
+          modelId,
+          mimeType as "image/png" | "application/pdf",
+        );
+      });
+      return await Promise.all(extractions);
+    } catch (error) {
+      console.error("Error performing extraction:", error);
+      throw new Error(
+        `Extraction failed: ${error instanceof Error ? error.message : "Unknown error"}`,
+      );
+    }
+  }
+}
+
+type PageInfo = {
+  pageNumber?: number;
+  width?: number;
+  height?: number;
+};
+
+/**
+ * Merges multiple Azure document chunks into a single document with corrected page offsets.
+ *
+ * Each chunk's pages are offset by their position multiplied by MAX_PAGES_PER_CHUNK
+ * to maintain correct page numbering across document chunks.
+ *
+ * @param layouts - Array of Azure document chunks to merge
+ * @returns Single merged Azure document with corrected page references
+ */
+export function mergeChunkedLayoutsAzure(
+  layouts: AzureDocumentResult[],
+): AzureDocumentResult {
+  const [mergedDocument] = layouts;
+  if (!mergedDocument) {
+    throw new Error("Cannot merge empty layouts array");
+  }
+  for (const [i, layout] of layouts.slice(1).entries()) {
+    const pageOffset = (i + 1) * MAX_PAGES_PER_CHUNK;
+    const sectionOffset = mergedDocument.sections?.length ?? 0;
+    const paragraphOffset = mergedDocument.paragraphs?.length ?? 0;
+    const tableOffset = mergedDocument.tables?.length ?? 0;
+    const figureOffset = mergedDocument.figures?.length ?? 0;
+
+    // merge text content
+    if (layout.content) {
+      mergedDocument.content += layout.content;
+    }
+
+    // offset pages
+    if (layout.pages) {
+      layout.pages.forEach((page) => {
+        if (page.pageNumber !== undefined) {
+          page.pageNumber += pageOffset;
+        }
+      });
+      mergedDocument.pages?.push(...layout.pages);
+    }
+
+    // merge section elements - need to offset the indices
+    for (const section of layout.sections ?? []) {
+      section.elements = section.elements?.map((e) => {
+        const { kind, index } = parseAzureElementPath(e);
+        if (!kind) {
+          throw new Error(`Invalid element path: ${e}`);
+        }
+        switch (kind) {
+          case "sections": {
+            return `/${kind}/${index + sectionOffset}`;
+          }
+          case "paragraphs": {
+            return `/${kind}/${index + paragraphOffset}`;
+          }
+          case "tables": {
+            return `/${kind}/${index + tableOffset}`;
+          }
+          case "figures": {
+            return `/${kind}/${index + figureOffset}`;
+          }
+          default: {
+            throw new Error(`Unrecognized element kind: ${kind}`);
+          }
+        }
+      });
+      mergedDocument.sections?.push(section);
+    }
+
+    // offset paragraph page numbers
+    if (layout.paragraphs) {
+      layout.paragraphs.forEach((paragraph) =>
+        offsetBoundingRegions(paragraph, pageOffset),
+      );
+      mergedDocument.paragraphs?.push(...layout.paragraphs);
+    }
+
+    // offset table page numbers
+    if (layout.tables) {
+      layout.tables.forEach((table) => offsetTablePages(table, pageOffset));
+      mergedDocument.tables?.push(...layout.tables);
+    }
+
+    // offset table page numbers
+    if (layout.figures) {
+      layout.figures.forEach((figure) =>
+        offsetBoundingRegions(figure, pageOffset),
+      );
+      mergedDocument.figures?.push(...layout.figures);
+    }
+
+    // TODO: do proper offset of documents, for now not used
+    // if (layout.documents) {
+    //   layout.documents.forEach((doc) =>
+    //     offsetDocumentPages(doc, pageOffset),
+    //   );
+    //   mergedDocument.documents?.push(...layout.documents);
+    // }
+  }
+
+  return mergedDocument;
+}
+
+/**
+ * Offsets all pageNumbers in boundingRegions by the given offset.
+ */
+function offsetBoundingRegions(
+  item: { boundingRegions?: Array<{ pageNumber?: number }> },
+  pageOffset: number,
+) {
+  item.boundingRegions?.forEach((region) => {
+    if (region.pageNumber !== undefined) {
+      region.pageNumber += pageOffset;
+    }
+  });
+}
+
+/**
+ * Offsets pageNumbers in a table and all its cells.
+ */
+function offsetTablePages(table: AzureTable, pageOffset: number) {
+  offsetBoundingRegions(table, pageOffset);
+  table.cells?.forEach((cell) => offsetBoundingRegions(cell, pageOffset));
+}
+
+/**
+ * Transforms Azure document results into a hierarchical DocumentLayout structure.
+ *
+ * Azure documents can have explicit sections or just paragraphs/tables. This method
+ * handles both cases, maintaining document order and properly nesting content within
+ * sections when available. All elements are registered in the reference map.
+ *
+ * @param result - Azure document extraction result
+ * @returns DocumentLayout with sections, paragraphs, tables, headers, and footers
+ */
+export function serializeLayoutAzure(
+  result: AzureDocumentResult,
+): DocumentLayout {
+  const currentLayout = new DocumentLayout([], {}, {});
+  const parsedSections: Set<string> = new Set();
+  for (const [index, section] of result.sections?.entries() ?? []) {
+    if (parsedSections.has(`/sections/${index}`)) continue;
+    const parsedSection = parseAzureSection(
+      section,
+      currentLayout.documentLayout,
+      result,
+      parsedSections,
+      currentLayout.referenceMap,
+      result.pages,
+    );
+    currentLayout.documentLayout.push(parsedSection);
+  }
+
+  // TODO: parse elements that are not in a section
+
+  return currentLayout;
+}
+
+/**
+ * Converts an Azure section into a structured Section object.
+ *
+ * Azure sections reference their content via element paths (e.g., "/paragraphs/0").
+ * This method resolves those references, identifies the section title (first element
+ * with title/heading role), and parses all remaining paragraphs and tables into the
+ * section's content.
+ *
+ * @param section - Azure section object with element references
+ * @param sectionIndex - Index of this section in the document
+ * @param doc - Full Azure document result for resolving references
+ * @param referenceMap - Map to register element locations
+ * @param pages - Page information for normalizing coordinates
+ * @returns Structured Section object, or null if the section has no valid elements
+ *
+ * @private
+ */
+function parseAzureSection(
+  section: AzureSection,
+  container: Section | LayoutItemType[],
+  doc: AzureDocumentResult,
+  parsedSections: Set<string>,
+  referenceMap: ReferenceMap,
+  pages?: Array<PageInfo>,
+): Section {
+  const parent: { id: string; itemCount: number } = { id: "", itemCount: 0 };
+  if (container instanceof Section) {
+    parent.id = container.id;
+    parent.itemCount = container.content.filter(
+      (item) => item instanceof Section,
+    ).length;
+  } else {
+    parent.itemCount = container.filter(
+      (item) => item instanceof Section,
+    ).length;
+  }
+  const sectionId = `${parent.id}/s${parent.itemCount}`;
+  const currentSection = new Section(
+    sectionId,
+    [],
+    new Header(`${sectionId}/h`, []),
+    new Footer(`${sectionId}/f`, []),
+  );
+
+  // parse section elements
+  for (const elementPath of section.elements ?? []) {
+    const { kind, index } = parseAzureElementPath(elementPath);
+    if (!kind) {
+      console.warn("Invalid element path: ", elementPath);
+      continue;
+    }
+    switch (kind) {
+      case "sections": {
+        // mark section as parsed to avoid double counting nested sections
+        const subSection = doc.sections?.[index];
+        if (subSection) {
+          currentSection.content.push(
+            parseAzureSection(
+              subSection,
+              currentSection,
+              doc,
+              parsedSections,
+              referenceMap,
+              pages,
+            ),
+          );
+        }
+        parsedSections.add(elementPath);
+        break;
+      }
+      case "paragraphs": {
+        const paragraph = doc.paragraphs?.[index];
+        if (paragraph) {
+          const { location, item } = parseAzureParagraph(
+            paragraph,
+            index,
+            currentSection,
+            referenceMap,
+            pages,
+          );
+          switch (location) {
+            case "header":
+              currentSection.header.content.push(item);
+              break;
+            case "footer":
+              currentSection.footer.content.push(item);
+              break;
+            case "content":
+              currentSection.content.push(item);
+              break;
+          }
+        }
+        break;
+      }
+      case "tables": {
+        const table = doc.tables?.[index];
+        if (table) {
+          currentSection.content.push(
+            parseAzureTable(table, index, currentSection, referenceMap, pages),
+          );
+        }
+        break;
+      }
+      case "figures": {
+        console.log("Found figure: ", elementPath);
+        // TODO: implement figure parsing
+        break;
+      }
+      default:
+        console.warn("Invalid element kind: ", kind);
+        break;
+    }
+  }
+
+  return currentSection;
+}
+
+/**
+ * Converts an Azure paragraph into a structured content element.
+ *
+ * Azure paragraphs have roles (pageHeader, pageFooter, or regular content) that
+ * determine the return type. This method creates the appropriate element type,
+ * registers it in the reference map with bounding regions, and assigns a unique ID.
+ *
+ * @param paragraph - Azure paragraph object
+ * @param index - Index of this paragraph in the document
+ * @param referenceMap - Map to register paragraph location
+ * @param pages - Page information for normalizing coordinates
+ * @returns Paragraph, PageHeader, or PageFooter based on the role
+ *
+ * @private
+ */
+function parseAzureParagraph(
+  paragraph: AzureParagraph,
+  paragraphIndex: number,
+  parentSection: Section,
+  referenceMap: ReferenceMap,
+  pages?: Array<PageInfo>,
+): { location: "header" | "footer" | "content"; item: Paragraph | Heading } {
+  const { role, content } = paragraph;
+  const depth = parentSection.id.split("/").length - 1;
+  switch (role) {
+    case "title":
+    case "sectionHeading": {
+      const headingId = `${parentSection.id}/h/h${paragraphIndex}`;
+      addItemToReferenceMap(paragraph, headingId, referenceMap, pages);
+      return {
+        location: "header",
+        item: new Heading(headingId, depth as 1 | 2 | 3 | 4 | 5, content),
+      };
+    }
+    case "pageHeader": {
+      const headingId = `${parentSection.id}/h/p${paragraphIndex}`;
+      addItemToReferenceMap(paragraph, headingId, referenceMap, pages);
+      return {
+        location: "header",
+        item: new Heading(headingId, 1, content),
+      };
+    }
+    case "pageFooter": {
+      const paragraphId = `${parentSection.id}/f/p${paragraphIndex}`;
+      addItemToReferenceMap(paragraph, paragraphId, referenceMap, pages);
+      return {
+        location: "footer",
+        item: new Paragraph(paragraphId, content),
+      };
+    }
+    default: {
+      const paragraphId = `${parentSection.id}/p${paragraphIndex}`;
+      addItemToReferenceMap(paragraph, paragraphId, referenceMap, pages);
+      return {
+        location: "content",
+        item: new Paragraph(paragraphId, content),
+      };
+    }
+  }
+}
+
+/**
+ * Converts an Azure table into the standardized Table structure.
+ *
+ * Azure tables provide cells with row/column indices that need to be grouped into rows.
+ * This method organizes cells by row, identifies header rows (containing columnHeader cells),
+ * and creates structured Row objects. All cells are registered in the reference map.
+ *
+ * @param table - Azure table object with cells
+ * @param index - Index of this table in the document
+ * @param referenceMap - Map to register cell locations
+ * @param pages - Page information for normalizing coordinates
+ * @returns Structured Table object with properly typed rows
+ *
+ * @private
+ */
+function parseAzureTable(
+  table: AzureTable,
+  tableIndex: number,
+  parentSection: Section,
+  referenceMap: ReferenceMap,
+  pages?: Array<PageInfo>,
+): Table {
+  const rowMap: Map<number, AzureTableCell[]> = new Map();
+  for (const cell of table.cells) {
+    if (!rowMap.has(cell.rowIndex)) {
+      rowMap.set(cell.rowIndex, []);
+    }
+    rowMap.get(cell.rowIndex)!.push(cell);
+  }
+
+  // convert to rows
+  const rows: Row[] = [];
+  const sortedRowIndices = Array.from(rowMap.keys()).sort();
+  const tableId = `${parentSection.id}/t${tableIndex}`;
+  for (const rowIndex of sortedRowIndices) {
+    const cells = rowMap.get(rowIndex) ?? [];
+    cells.sort((a, b) => a.columnIndex - b.columnIndex);
+
+    const rowType = cells.some((c) => c.kind === "columnHeader")
+      ? "head"
+      : "body";
+    const row = new Row(rowType, []);
+    for (const cell of cells) {
+      const cellId = `${tableId}/${cell.rowIndex}/${cell.columnIndex}`;
+      addItemToReferenceMap(cell, cellId, referenceMap, pages);
+      row.cells.push(
+        new Cell(cellId, cell.content ?? "", cell.columnSpan, cell.rowSpan),
+      );
+    }
+    rows.push(row);
+  }
+
+  return new Table(tableId, rows);
+}
+
+function addItemToReferenceMap(
+  item: AzureParagraph | AzureTableCell,
+  itemId: string,
+  referenceMap: ReferenceMap,
+  pages?: Array<PageInfo>,
+) {
+  if (!pages) {
+    console.warn(
+      "No pages provided for item: ",
+      itemId,
+      ", unable to add to reference map",
+    );
+    return;
+  }
+  if (!item.boundingRegions) {
+    console.warn("No bounding regions for item: ", itemId);
+    return;
+  }
+  referenceMap[itemId] = {
+    text: item.content ?? "",
+    boundingRegions: item.boundingRegions.map((r) => ({
+      pageNumber: r.pageNumber,
+      polygon: normalizePolygon(r.polygon as Polygon, r.pageNumber, pages),
+    })),
+  };
+}
+
+/**
+ * Normalizes Azure polygon coordinates to 0-1 range.
+ * Azure returns coordinates in inches (for PDFs) or pixels (for images).
+ * We need to normalize them by dividing by page width/height.
+ */
+function normalizePolygon(
+  polygon: number[],
+  pageNumber: number,
+  pages?: Array<PageInfo>,
+): Polygon {
+  if (!pages) {
+    throw new Error("Missing page metadata for polygon normalization");
+  }
+
+  const page = pages.find((p) => p.pageNumber === pageNumber);
+  if (!page?.width || !page?.height) {
+    return polygon as Polygon;
+  }
+
+  // Normalize coordinates to 0-1 range
+  const normalized: number[] = [];
+  for (const [i, coord] of polygon.entries()) {
+    const scaleBy = i % 2 === 0 ? page.width : page.height;
+    normalized.push(coord / scaleBy);
+  }
+
+  return normalized as Polygon;
+}
+
+function parseAzureElementPath(elementPath: string): {
+  kind: string | null;
+  index: number;
+} {
+  const [_, kind, indexStr] = elementPath.split("/");
+  if (!indexStr || !kind) {
+    console.warn(
+      "Invalid element path: ",
+      elementPath,
+      ", returning dummy path",
+    );
+    return { kind: null, index: 0 };
+  }
+  return { kind, index: parseInt(indexStr) };
+}
+
+/**
+ * Transforms Azure-specific extraction output into a standardized entity format.
+ */
+export function serializePretrainedLayoutAzure(content: AzureDocumentResult): {
+  entities: Record<string, PretrainedEntity>;
+} {
+  const entities: Record<string, PretrainedEntity> = {};
+  if (content.documents && Array.isArray(content.documents)) {
+    for (const document of content.documents) {
+      if (document.fields) {
+        for (const [fieldName, fieldValue] of Object.entries(document.fields)) {
+          const convertedEntity = convertAzureFieldValueToEntity(
+            fieldName,
+            fieldValue,
+            content.pages,
+          );
+
+          // If an entity with this name already exists, convert to array or append to existing array
+          if (entities[fieldName]) {
+            const existing = entities[fieldName];
+            if (existing.type === "array") {
+              existing.valueArray = existing.valueArray || [];
+              existing.valueArray.push(convertedEntity);
+            } else {
+              // Convert to array
+              entities[fieldName] = {
+                id: "dummy",
+                type: "array",
+                valueArray: [existing, convertedEntity],
+              };
+            }
+          } else {
+            entities[fieldName] = convertedEntity;
+          }
+        }
+      }
+    }
+  }
+
+  return { entities };
+}
+
+/**
+ * Converts an Azure field value into the standardized PretrainedEntity format.
+ *
+ * Azure's field structure already includes typed values, so this method primarily
+ * maps Azure types to our standardized types and handles nested structures (arrays
+ * and objects). It preserves confidence scores and bounding regions.
+ *
+ * @param fieldName - Name of the field being converted
+ * @param fieldValue - Azure field output from Document Intelligence
+ * @param pages - Page information for normalizing coordinates
+ * @returns Standardized PretrainedEntity with typed values
+ *
+ * @private
+ */
+function convertAzureFieldValueToEntity(
+  fieldName: string,
+  fieldValue: AzureFieldOutput,
+  pages?: Array<{ pageNumber?: number; width?: number; height?: number }>,
+): PretrainedEntity {
+  const entity: PretrainedEntity = {
+    id: "dummy",
+    type: fieldValue.type as EntityType,
+    content: fieldValue.content ?? "",
+    confidence: fieldValue.confidence ?? 0,
+    boundingRegions:
+      fieldValue.boundingRegions?.map((r) => ({
+        pageNumber: r.pageNumber || 0,
+        polygon: normalizePolygon(
+          r.polygon as number[],
+          r.pageNumber || 0,
+          pages,
+        ),
+      })) ?? [],
+  };
+
+  switch (fieldValue.type) {
+    case "array":
+      entity.valueArray = [];
+      if (fieldValue.valueArray) {
+        for (const arrayItem of fieldValue.valueArray) {
+          entity.valueArray.push(
+            convertAzureFieldValueToEntity(fieldName, arrayItem, pages),
+          );
+        }
+      }
+      break;
+    case "object":
+      if (!fieldValue.valueObject) {
+        console.warn("Object value is required");
+        return entity;
+      }
+      entity.valueObject = {};
+      for (const [key, value] of Object.entries(fieldValue.valueObject)) {
+        entity.valueObject[key] = convertAzureFieldValueToEntity(
+          key,
+          value,
+          pages,
+        );
+      }
+      break;
+    case "date":
+      if (fieldValue.valueDate) {
+        entity.valueDate = fieldValue.valueDate;
+      }
+      break;
+    case "number":
+      if (fieldValue.valueNumber !== undefined) {
+        entity.valueNumber = fieldValue.valueNumber;
+      }
+      break;
+    case "string":
+      if (fieldValue.valueString) {
+        entity.valueString = fieldValue.valueString;
+      }
+      break;
+    case "boolean":
+      if (fieldValue.valueBoolean !== undefined) {
+        entity.valueBoolean = fieldValue.valueBoolean;
+      }
+      break;
+  }
+  return entity;
+}