Sort and Extract a Mixed Document Pile

import fs from "node:fs";

const API_KEY = process.env.UNSILOED_API_KEY;
const BASE_URL = "https://prod.visionapi.unsiloed.ai";

// The document types we expect in the pile. The description guides the split.
const CATEGORIES = [
  { name: "Invoice", description: "A bill issued by a seller requesting payment for goods or services" },
  { name: "Receipt", description: "Proof of a completed purchase from a store or point of sale" },
  { name: "Purchase Order", description: "A buyer-issued document authorizing a purchase from a vendor" },
];

// One extraction schema per type. The keys match the category names above.
const SCHEMAS = {
  Invoice: {
    type: "object",
    properties: {
      vendor_name:    { type: "string", description: "Company issuing the invoice (the seller)" },
      invoice_number: { type: "string", description: "Invoice identifier" },
      total_due:      { type: "number", description: "Total amount due in US dollars" },
    },
    required: ["vendor_name", "invoice_number", "total_due"],
    additionalProperties: false,
  },
  Receipt: {
    type: "object",
    properties: {
      store_name:     { type: "string", description: "Name of the store" },
      transaction_id: { type: "string", description: "Transaction or receipt ID" },
      total:          { type: "number", description: "Total amount paid in US dollars" },
    },
    required: ["store_name", "transaction_id", "total"],
    additionalProperties: false,
  },
  "Purchase Order": {
    type: "object",
    properties: {
      po_number:   { type: "string", description: "Purchase order number" },
      vendor_name: { type: "string", description: "Vendor the order is placed with" },
      required_by: { type: "string", description: "Date the order is required by" },
    },
    required: ["po_number", "vendor_name", "required_by"],
    additionalProperties: false,
  },
};

async function waitFor(url) {
  // Poll a job URL until it finishes, then return the completed response.
  for (let i = 0; i < 90; i++) { // roughly 6 minutes at 4 seconds per poll
    const job = await (await fetch(url, { headers: { "api-key": API_KEY } })).json();
    if (["completed", "review"].includes(job.status)) return job;
    if (job.status === "failed") throw new Error(job.error || "job failed");
    await new Promise((r) => setTimeout(r, 4000));
  }
  throw new Error(`Job at ${url} did not finish in time`);
}

// Stage 1: split the pile into one file per document type.
const splitForm = new FormData();
splitForm.append("file", new Blob([fs.readFileSync("pile.pdf")]), "pile.pdf");
splitForm.append("categories", JSON.stringify(CATEGORIES));
const splitRes = await fetch(`${BASE_URL}/splitter`, {
  method: "POST",
  headers: { "api-key": API_KEY },
  body: splitForm,
});
if (!splitRes.ok) throw new Error(`split submit failed: HTTP ${splitRes.status} ${await splitRes.text()}`);
const splitJob = await splitRes.json();
const sections = (await waitFor(`${BASE_URL}/splitter/${splitJob.job_id}`)).result.files;
console.log(`Split into ${sections.length} documents: ${sections.map((s) => s.name).join(", ")}`);

// Stage 2: extract type-specific fields from each section, by URL.
const extractJobs = {};
for (const section of sections) {
  const category = section.name.replace(/\.pdf$/, "");
  const schema = SCHEMAS[category];
  if (!schema) {
    console.log(`No schema for '${category}', skipping.`);
    continue;
  }
  const extractForm = new FormData();
  extractForm.append("file_url", section.full_path);
  extractForm.append("schema_data", JSON.stringify(schema));
  extractForm.append("model", "gamma");
  const extractRes = await fetch(`${BASE_URL}/v2/extract`, {
    method: "POST",
    headers: { "api-key": API_KEY },
    body: extractForm,
  });
  if (!extractRes.ok) throw new Error(`extract submit failed: HTTP ${extractRes.status} ${await extractRes.text()}`);
  extractJobs[category] = (await extractRes.json()).job_id;
}

const results = {};
for (const [category, jobId] of Object.entries(extractJobs)) {
  const done = await waitFor(`${BASE_URL}/extract/${jobId}`);
  results[category] = Object.fromEntries(
    Object.entries(done.result).map(([field, cell]) => [field, cell.value]),
  );
}

console.log(JSON.stringify(results, null, 2));