template-app/test/documents-extraction.test.ts

// SPDX-License-Identifier: LicenseRef-ConusAI-Template-1.0
// The extraction of the documents module through the ModelGateway (plan item AT2.3): what the model is
// sent, what is done with its answer, and what happens when it does not answer. The model is the fake
// platform's gateway route, scripted per test; the gateway client, the extractor, the verification, the
// cache, the proposals and the review all run for real.

import { type LineView, LOW_TRUST_NOTE } from "@app/modules/documents"
import { uploadViaProcedures } from "@conusai-kit/testing"
import type { FakeChatResponder } from "@conusai-kit/testing/fake-platform"
import { afterEach, describe, expect, it } from "vitest"
import {
  bytesOf,
  closeApps,
  type DocumentsApp,
  makeDocumentsApp,
  type SeededProject,
  uploadAndProcess,
} from "./support/documents-app.js"
import { deterministicModel } from "./support/fake-model.js"

afterEach(closeApps)

/** Every page image the platform drew, as the fake platform recorded it. */
async function pagesDrawn(app: DocumentsApp): Promise<Array<{ job: number; page: number; dpi: number }>> {
  const state = (await (await fetch(`${app.platform.url}/__e2e/state`)).json()) as {
    ocrPages: Array<{ job: number; page: number; dpi: number }>
  }
  return state.ocrPages
}

const PAGE = { width: 1000, height: 1400, unit: "px", dpi: 150, rotation: 0 }

/**
 * One page of an OCR document whose lines are `rows`, 40 px apart, one block per line (`p{n}-b{m}`), as
 * the platform returns it once the page is classified: a scan when it was recognised, text otherwise.
 * `repeated` lists the rows the reader flagged as repeating on most pages.
 */
function pageOf(
  index: number,
  rows: string[],
  source: "text_layer" | "ocr" = "text_layer",
  confidence = 0.97,
  options: { pageClass?: string; repeated?: readonly string[] } = {},
) {
  return {
    ...PAGE,
    index,
    class: options.pageClass ?? (source === "ocr" ? "scan" : "text"),
    source,
    text: rows.join("\n"),
    lines: rows.map((text, order) => ({
      text,
      bbox: [100, 100 + order * 40, 600, 130 + order * 40],
      polygon: null,
      confidence,
      order,
      words: [],
      repeated: options.repeated?.includes(text) ?? false,
    })),
    blocks: rows.map((_, order) => ({
      id: `p${index + 1}-b${order + 1}`,
      bbox: [90, 90 + order * 40, 700, 140 + order * 40],
      order,
      lines: [order],
    })),
  }
}

/** A row long enough that a page of it is not weak for want of text: its image is not asked for. */
const SHEET_ROW = "Reference: QR-2026-0417 (rev 2)"

const documentOf = (...pages: Array<ReturnType<typeof pageOf>>) => ({
  schema_version: "conusai.ocr.document/v1",
  engine: { name: "fixture-engine", version: "1", detection_model: null, recognition_model: null },
  page_count: pages.length,
  pages,
})

/** A one-page OCR document whose lines are `rows`. */
const onePage = (rows: string[], source: "text_layer" | "ocr" = "text_layer", confidence = 0.97) =>
  documentOf(pageOf(0, rows, source, confidence))

interface Detail {
  document: {
    status: string
    failureCode: string | null
    pending: number
    skippedPages: Array<{ page: number; pageClass: string }>
  }
  lines: LineView[]
}

async function run(
  responder: FakeChatResponder | undefined,
  document: Record<string, unknown>,
  upload: { filename: string; contentType: string } = { filename: "report.pdf", contentType: "application/pdf" },
): Promise<{ app: DocumentsApp; project: SeededProject; detail: Detail }> {
  const app = await makeDocumentsApp({ document, ...(responder ? { responder } : {}) })
  const project = await app.seedProject("acme", "owner@acme.test")
  const file = await uploadAndProcess(app, project, upload.filename, upload.contentType)
  const detail = (await project.call.callByName("documents.document.get", {
    ...project.slugs,
    fileId: file.fileId,
  })) as Detail
  return { app, project, detail }
}

const reply = (lines: Array<{ label: string; value: string; quote: string }>) => ({
  content: JSON.stringify({ lines }),
})

type Message = { role: string; content: string | Array<{ type: string; text?: string; image_url?: { url: string } }> }
const messagesOf = (request: Record<string, unknown>) => request.messages as Message[]
const userParts = (request: Record<string, unknown>) => {
  const content = messagesOf(request).find((message) => message.role === "user")?.content
  return typeof content === "string" ? [] : (content ?? [])
}

describe("what the model is sent", () => {
  it("asks once per page, with the page's text and the schema of the module's line", async () => {
    const { app } = await run(undefined, onePage(["Reference: QR-1", "Issued on: 2026-04-01"]))
    expect(app.model.requests).toHaveLength(1)
    const request = app.model.requests[0] as Record<string, unknown>
    expect(request.model).toBe("fake-model")
    expect(request.temperature).toBe(0)

    const [system] = messagesOf(request)
    expect(system?.role).toBe("system")
    const [text, ...rest] = userParts(request)
    expect(text?.text).toContain("Page 1 of 1.")
    expect(text?.text).toContain("Reference: QR-1\nIssued on: 2026-04-01")
    // a text-layer PDF page needs no image
    expect(rest).toEqual([])

    const format = request.response_format as {
      type: string
      json_schema: { schema: { properties: { lines: { items: { required: string[]; properties: object } } } } }
    }
    expect(format.type).toBe("json_schema")
    const item = format.json_schema.schema.properties.lines.items
    expect(item.required.sort()).toEqual(["label", "quote", "value"])
    // the page, the position and the confidence are the code's to set, so the model is not asked for them
    expect(Object.keys(item.properties).sort()).toEqual(["label", "quote", "value"])
  })

  it("adds the page image, as a content part, for a scanned page of an image upload", async () => {
    const { app } = await run(undefined, onePage(["Reference: QR-1"], "ocr", 0.7), {
      filename: "scan.png",
      contentType: "image/png",
    })
    const parts = userParts(app.model.requests[0] as Record<string, unknown>)
    expect(parts.map((part) => part.type)).toEqual(["text", "image_url"])
    expect(parts[1]?.image_url?.url).toMatch(/^data:image\/png;base64,/)
  })

  it("fetches the image of a scanned page of a PDF from the platform's page route and sends it with the text", async () => {
    const { app, detail } = await run(undefined, onePage(["Reference: QR-1"], "ocr", 0.7))
    const parts = userParts(app.model.requests[0] as Record<string, unknown>)
    expect(parts.map((part) => part.type)).toEqual(["text", "image_url"])
    expect(parts[0]?.text).toContain(LOW_TRUST_NOTE)
    expect(parts[1]?.image_url?.url).toMatch(/^data:image\/png;base64,iVBORw0KGgo/)
    // the page was drawn once, for the job that read the file, at the default dpi
    expect(await pagesDrawn(app)).toEqual([{ job: 1, page: 1, dpi: 150 }])
    expect(detail.lines.map((line) => line.value)).toEqual(["QR-1"])
  })

  it("sends a table page as image and text, and a text page as text alone without asking for its image", async () => {
    const { app } = await run(
      undefined,
      documentOf(pageOf(0, ["Item | Qty", "Bolts | 12"], "ocr", 0.97, { pageClass: "table" }), pageOf(1, [SHEET_ROW])),
    )
    const kinds = app.model.requests.map((request) =>
      userParts(request as Record<string, unknown>).map((part) => part.type),
    )
    expect(kinds).toContainEqual(["text", "image_url"])
    expect(kinds).toContainEqual(["text"])
    expect(await pagesDrawn(app)).toEqual([{ job: 1, page: 1, dpi: 150 }])
  })

  it("leaves the header that repeats on the pages out of the text the model is shown", async () => {
    const header = "ACME: Specification sheet"
    const { app, detail } = await run(
      undefined,
      documentOf(
        pageOf(0, [header, "Reference: QR-1"], "text_layer", 0.97, { repeated: [header] }),
        pageOf(1, [header, "Issued on: 2026-04-01"], "text_layer", 0.97, { repeated: [header] }),
      ),
    )
    for (const request of app.model.requests) {
      expect(userParts(request as Record<string, unknown>)[0]?.text).not.toContain("ACME")
    }
    expect(detail.lines.map((line) => line.label)).toEqual(["Reference", "Issued on"])
  })

  it("puts the reader's block id in the evidence of a line, and its block's box", async () => {
    const { project, detail } = await run(undefined, onePage(["Reference: QR-1", "Issued on: 2026-04-01"]))
    expect(detail.lines.map((line) => line.blockId)).toEqual(["p1-b1", "p1-b2"])
    expect(detail.lines[1]?.box).toEqual({ x: 0.09, y: 130 / 1400, width: 0.61, height: 50 / 1400 })
    const proposals = (await project.call.callByName("proposals.proposal.list", {
      ...project.slugs,
      status: "pending",
    })) as { proposals: Array<{ evidence: Array<{ blockId?: string }> }> }
    expect(proposals.proposals.flatMap((proposal) => proposal.evidence.map((item) => item.blockId)).sort()).toEqual([
      "p1-b1",
      "p1-b2",
    ])
  })
})

describe("the pages left out", () => {
  const labels = Array.from({ length: 300 }, (_, at) => `A-${at} 2400`)

  it("never asks the model about a drawing or an empty page, and reports them on the document", async () => {
    const { app, detail } = await run(
      undefined,
      documentOf(
        pageOf(0, [SHEET_ROW]),
        pageOf(1, labels, "text_layer", 0.97, { pageClass: "drawing" }),
        pageOf(2, [], "ocr", 0.97, { pageClass: "empty" }),
      ),
    )
    expect(app.model.requests).toHaveLength(1)
    expect(detail.document).toMatchObject({
      status: "extracted",
      skippedPages: [
        { page: 2, pageClass: "drawing" },
        { page: 3, pageClass: "empty" },
      ],
    })
    expect(detail.lines.map((line) => line.value)).toEqual(["QR-2026-0417 (rev 2)"])
    // no image was drawn for a page nobody reads
    expect(await pagesDrawn(app)).toEqual([])
  })

  it("reads a scan beside a drawing: the scan produces lines, the drawing none", async () => {
    const { detail } = await run(
      undefined,
      documentOf(
        pageOf(0, labels, "text_layer", 0.97, { pageClass: "drawing" }),
        pageOf(1, ["Delivery no: DN-77", "Received by: A. Sample"], "ocr", 0.8),
      ),
    )
    expect(detail.lines.map((line) => [line.page, line.value])).toEqual([
      [2, "DN-77"],
      [2, "A. Sample"],
    ])
    expect(detail.document.skippedPages).toEqual([{ page: 1, pageClass: "drawing" }])
  })

  it("reads a drawing when the deployment's list does not exclude it", async () => {
    const app = await makeDocumentsApp({
      document: documentOf(pageOf(0, ["Scale: 1:50"], "ocr", 0.9, { pageClass: "drawing" })),
      excludedPageClasses: [],
    })
    const project = await app.seedProject("acme", "owner@acme.test")
    const file = await uploadAndProcess(app, project, "plan.pdf", "application/pdf")
    const detail = (await project.call.callByName("documents.document.get", {
      ...project.slugs,
      fileId: file.fileId,
    })) as Detail
    expect(detail.document.skippedPages).toEqual([])
    expect(detail.lines.map((line) => line.value)).toEqual(["1:50"])
  })
})

describe("a quote taken from the image of a scanned page", () => {
  it("stays needs_review with the scanned-quote note when the page text does not hold it", async () => {
    const { detail } = await run(
      () => reply([{ label: "Qty", value: "12", quote: "Qty 12 pcs" }]),
      onePage(["Delivery note"], "ocr", 0.7),
    )
    expect(detail.lines).toHaveLength(1)
    expect(detail.lines[0]).toMatchObject({ verification: "needs_review", note: "scanned_quote", confidence: 0 })
  })

  it("carries no note for a quote the text holds", async () => {
    const { detail } = await run(undefined, onePage(["Reference: QR-1"], "ocr", 0.7))
    expect(detail.lines[0]).toMatchObject({ verification: "verified", note: null, blockId: "p1-b1" })
  })
})

describe("what is done with the answer", () => {
  const rows = ["Sąskaita: INV-2026-07", "Suma: 1 250,00 EUR", "Pastaba: apmokėti per 14 dienų"]

  it("verifies each quote on the page it was asked about, folding only whitespace, case and diacritics", async () => {
    const { detail } = await run(
      () =>
        reply([
          { label: "Sąskaita", value: "INV-2026-07", quote: "Sąskaita: INV-2026-07" },
          // case, diacritics and whitespace differ from the page: still the same words
          { label: "Suma", value: "1 250,00 EUR", quote: "suma:   1 250,00\n eur" },
          { label: "Pastaba", value: "apmokėti per 14 dienų", quote: "PASTABA: APMOKETI PER 14 DIENU" },
        ]),
      onePage(rows),
    )
    expect(detail.lines.map((line) => [line.label, line.verification])).toEqual([
      ["Sąskaita", "verified"],
      ["Suma", "verified"],
      ["Pastaba", "verified"],
    ])
    // a verified line is located from the reader's own lines, and carries the reader's confidence
    expect(detail.lines[0]?.box).not.toBeNull()
    expect(detail.lines[0]?.confidence).toBe(0.97)
    expect(detail.lines.every((line) => line.page === 1 && line.status === "pending")).toBe(true)
  })

  it("keeps a line whose quote is not on the page as needs_review, with no box and no confidence", async () => {
    const { detail } = await run(
      () =>
        reply([
          { label: "Sąskaita", value: "INV-2026-07", quote: "Sąskaita: INV-2026-07" },
          // a changed digit is a different word: folding does not forgive it
          { label: "Suma", value: "1 520,00 EUR", quote: "Suma: 1 520,00 EUR" },
          { label: "Made up", value: "x", quote: "Nothing like this is written here" },
        ]),
      onePage(rows),
    )
    expect(detail.lines.map((line) => [line.label, line.verification])).toEqual([
      ["Sąskaita", "verified"],
      ["Suma", "needs_review"],
      ["Made up", "needs_review"],
    ])
    const flagged = detail.lines.filter((line) => line.verification === "needs_review")
    expect(flagged.every((line) => line.box === null && line.confidence === 0 && line.status === "pending")).toBe(true)
    // kept, not dropped and not approved: the document is waiting for a person
    expect(detail.document).toMatchObject({ status: "extracted", pending: 3 })
  })

  it("verifies against the page the model was asked about, not against another page", async () => {
    const two = {
      ...onePage(["Reference: QR-1"]),
      page_count: 2,
      pages: [...onePage(["Reference: QR-1"]).pages, { ...onePage(["Total: 5 EUR"]).pages[0], index: 1 }],
    }
    // the model answers every page with page 1's line: on page 2 it is not there
    const { detail } = await run(() => reply([{ label: "Reference", value: "QR-1", quote: "Reference: QR-1" }]), two)
    expect(detail.lines.map((line) => [line.page, line.verification])).toEqual([
      [1, "verified"],
      [2, "needs_review"],
    ])
  })

  it("drops a line the model repeats", async () => {
    const line = { label: "Reference", value: "QR-1", quote: "Reference: QR-1" }
    const { detail } = await run(() => reply([line, line]), onePage(["Reference: QR-1"]))
    expect(detail.lines).toHaveLength(1)
  })
})

describe("when the model does not answer", () => {
  const document = onePage(["Reference: QR-1"])

  it("a refusal is asked again once, then ends the document with a typed code and no proposals", async () => {
    const { app, detail } = await run(() => ({ content: null, refusal: "I cannot help with that." }), document)
    expect(app.model.requests).toHaveLength(2)
    expect(detail.document).toMatchObject({ status: "failed", failureCode: "extraction_refused" })
    expect(detail.lines).toEqual([])
  })

  it("a reply cut off at the token ceiling is asked again with more room, and the second answer is used", async () => {
    let calls = 0
    const { app, detail } = await run((request) => {
      calls += 1
      return calls === 1 ? { content: '{"lines": [{"label"', finishReason: "length" } : deterministicModel(request)
    }, document)
    expect(app.model.requests.map((request) => request.max_tokens)).toEqual([4096, 8192])
    expect(detail.document.status).toBe("extracted")
    expect(detail.lines).toHaveLength(1)
  })

  it("a reply that is always cut off ends the document as truncated", async () => {
    const { app, detail } = await run(() => ({ content: '{"lines": [', finishReason: "length" }), document)
    expect(app.model.requests).toHaveLength(2)
    expect(detail.document).toMatchObject({ status: "failed", failureCode: "extraction_truncated" })
  })

  it("a reply outside the schema is asked again once, by the gateway, and then ends the document as invalid", async () => {
    const { app, detail } = await run(() => ({ content: JSON.stringify({ lines: "none" }) }), document)
    // two asks in all: the gateway's retry with the validation errors is the one retry
    expect(app.model.requests).toHaveLength(2)
    expect(detail.document).toMatchObject({ status: "failed", failureCode: "extraction_invalid" })
    const retry = messagesOf(app.model.requests[1] as Record<string, unknown>).at(-1)
    expect(retry?.content).toContain("did not match the required schema")
  })

  it("records no document text in the failure: the code is all the document keeps", async () => {
    const { detail } = await run(() => ({ content: null, refusal: "Reference: QR-1" }), document)
    expect(JSON.stringify(detail.document)).not.toContain("QR-1")
  })
})

describe("the cache", () => {
  it("a file with the same bytes is neither read nor asked about again", async () => {
    const { app, project } = await run(undefined, onePage(["Reference: QR-1", "Issued on: 2026-04-01"]))
    expect(app.ocr.submitted).toHaveLength(1)
    expect(app.model.requests).toHaveLength(1)
    const stored = app.store.keysOf({ level: "project", orgId: project.orgId } as never)
    // the OCR result, the job it came from (so its pages can still be drawn), and the page's model reply
    expect(stored.filter((key) => key.startsWith("documents/cache/"))).toHaveLength(3)

    // the same bytes under another name: a second file, the same content hash
    const second = await uploadViaProcedures(project.call, project.slugs, {
      filename: "copy.pdf",
      contentType: "application/pdf",
      bytes: bytesOf("%PDF-1.4 fixture report.pdf"),
    })
    await app.deliver()
    await app.runJobs()

    expect(app.ocr.submitted).toHaveLength(1)
    expect(app.model.requests).toHaveLength(1)
    const view = (await project.call.callByName("documents.document.get", {
      ...project.slugs,
      fileId: second.fileId,
    })) as Detail
    expect(view.lines.map((line) => line.value)).toEqual(["QR-1", "2026-04-01"])
  })

  it("a change of the page's text is a different request", async () => {
    const { app, project } = await run(undefined, onePage(["Reference: QR-1"]))
    app.platform.ocr.setDocument(onePage(["Reference: QR-2"]))
    await uploadViaProcedures(project.call, project.slugs, {
      filename: "other.pdf",
      contentType: "application/pdf",
      bytes: bytesOf("%PDF-1.4 something else"),
    })
    await app.deliver()
    await app.runJobs()
    expect(app.model.requests).toHaveLength(2)
  })
})