305 lines
12 KiB
TypeScript
305 lines
12 KiB
TypeScript
// What the OCR job does with one uploaded file: read it, extract candidate lines, and put each one
|
|
// in front of a person as a proposal. It runs in the worker, never in a request (AGENTS.md: a
|
|
// request enqueues and answers), and every step is safe to run again, because pg-boss retries a job
|
|
// that throws:
|
|
//
|
|
// - the document row is created once (unique on the file);
|
|
// - the OCR job id is stored the moment the platform accepts it, so a retry resumes that job and
|
|
// does not pay for the document twice (and the submit carries an idempotency key besides);
|
|
// - a candidate that already has a proposal, in any state, is not proposed again, so a line a
|
|
// person rejected does not come back and a crash half-way through adds only the rest;
|
|
// - what a stage produced (the OCR result, the job it came from, a model's reply for a page) is cached
|
|
// by the hash of its inputs, so the retry reads those instead of paying for them again, and finds the
|
|
// same lines;
|
|
// - a page image is asked of the platform only for a page the extractor decided to show a model, and
|
|
// only through the job that read the file.
|
|
//
|
|
// A failure no retry would change (the object is gone, the platform says the job failed, the model
|
|
// refused or never answered in the schema) is recorded on the document as a code and ends the job;
|
|
// anything else throws and is retried.
|
|
|
|
import { AppError } from "@conusai-kit/core"
|
|
import type { ObjectStore } from "@conusai-kit/data"
|
|
import { toInstant } from "@conusai-kit/data"
|
|
import { getFile, type ProcessorJobPayload, readAll, sha256Hex } from "@conusai-kit/files"
|
|
import { createProposal } from "@conusai-kit/proposals"
|
|
import { z } from "zod"
|
|
import { documents, type FailureCode, LINE_TARGET_KEY, type LineChange } from "../contract/index.js"
|
|
import { type DocumentScope, type DocumentsDeps, lineKeyOf, lineProposals } from "./lines.js"
|
|
import {
|
|
type DocumentOcr,
|
|
type Extraction,
|
|
ExtractionFailed,
|
|
type LineExtractor,
|
|
OcrJobFailed,
|
|
type SkippedPage,
|
|
} from "./ports.js"
|
|
import { projectDocument, readDocumentSchema } from "./read-document.js"
|
|
import { type StageCache, stageCache } from "./stage-cache.js"
|
|
|
|
const IMAGE_CONTENT_TYPES = ["image/png", "image/jpeg"]
|
|
|
|
/** The job that read a file, kept beside the OCR result so a file read once can still be shown page by page. */
|
|
const readJobSchema = z.object({ jobId: z.number().int().nonnegative() })
|
|
|
|
export interface ProcessDeps extends DocumentsDeps {
|
|
readonly store: ObjectStore
|
|
readonly ocr: DocumentOcr
|
|
readonly extractor: LineExtractor
|
|
/** Where stage artefacts and model replies are cached; by default the project's space in `store`. */
|
|
readonly cache?: StageCache
|
|
/** The most lines one document may propose: a review queue is for people, and a runaway reader must not fill it. */
|
|
readonly maxLines?: number
|
|
/** Maximum number of source bytes the job may buffer. */
|
|
readonly maxBytes?: number
|
|
readonly now?: () => number
|
|
/** Aborted when the queue takes the job back (a lapsed claim, a stopping worker): the wait ends with it. */
|
|
readonly signal?: AbortSignal
|
|
}
|
|
|
|
/** The database operations the extraction pipeline needs, each implemented as its own transaction. */
|
|
export interface DocumentProcessRecords {
|
|
getFile(scope: DocumentScope, fileId: string): ReturnType<typeof getFile>
|
|
ensureProcessing(
|
|
scope: DocumentScope,
|
|
file: Awaited<ReturnType<typeof getFile>>,
|
|
payload: ProcessorJobPayload,
|
|
): Promise<ProcessDocumentRow | undefined>
|
|
setOcrJob(scope: DocumentScope, document: ProcessDocumentRow, jobId: number): Promise<ProcessDocumentRow>
|
|
markFailed(scope: DocumentScope, document: ProcessDocumentRow, code: FailureCode, finishedAt: string): Promise<void>
|
|
proposals(scope: DocumentScope, documentId: string): ReturnType<typeof lineProposals>
|
|
createProposal(
|
|
scope: DocumentScope,
|
|
input: Parameters<typeof createProposal>[2],
|
|
idempotencyKey: string,
|
|
): Promise<void>
|
|
finish(
|
|
scope: DocumentScope,
|
|
document: ProcessDocumentRow,
|
|
result: { pageCount: number; proposedCount: number; skippedPages: SkippedPage[]; finishedAt: string },
|
|
): Promise<void>
|
|
}
|
|
|
|
export interface ProcessDocumentRow {
|
|
readonly id: string
|
|
readonly filename: string
|
|
readonly status: "processing" | "extracted" | "failed"
|
|
readonly ocrJobId: number | null
|
|
readonly version: number
|
|
}
|
|
|
|
export type ProcessOutcome = "extracted" | "failed" | "skipped"
|
|
|
|
const DEFAULT_MAX_LINES = 500
|
|
export async function processDocument(
|
|
deps: ProcessDeps,
|
|
scope: DocumentScope,
|
|
payload: ProcessorJobPayload,
|
|
): Promise<ProcessOutcome> {
|
|
return processDocumentWithRecords(
|
|
{ ...deps, maxBytes: deps.maxBytes ?? deps.files.policy.maxBytes },
|
|
scope,
|
|
payload,
|
|
createDataRecords(deps),
|
|
)
|
|
}
|
|
|
|
/**
|
|
* Run OCR and extraction over an external-effect port while record transitions remain explicit.
|
|
* The Convex profile implements each records method with one internal mutation/query; it never
|
|
* simulates a multi-request transaction around the whole pipeline.
|
|
*/
|
|
export async function processDocumentWithRecords(
|
|
deps: Omit<ProcessDeps, keyof DocumentsDeps>,
|
|
scope: DocumentScope,
|
|
payload: ProcessorJobPayload,
|
|
records: DocumentProcessRecords,
|
|
): Promise<ProcessOutcome> {
|
|
const clock = deps.now ?? Date.now
|
|
let file: Awaited<ReturnType<typeof getFile>>
|
|
try {
|
|
file = await records.getFile(scope, payload.fileId)
|
|
} catch (error) {
|
|
// Deleted between the upload and this job: nothing to read, and nothing a retry would find.
|
|
if (error instanceof AppError && error.status === 404) return "skipped"
|
|
throw error
|
|
}
|
|
const ensuredDocument = await records.ensureProcessing(scope, file, payload)
|
|
if (!ensuredDocument) return "skipped"
|
|
let activeDocument: ProcessDocumentRow = ensuredDocument
|
|
|
|
const fail = async (code: FailureCode): Promise<ProcessOutcome> => {
|
|
await records.markFailed(scope, activeDocument, code, toInstant(clock()))
|
|
return "failed"
|
|
}
|
|
|
|
const object = await deps.store.get(scope, payload.storageKey)
|
|
if (!object) return fail("file_missing")
|
|
// The most this job reads into memory is what an upload may be: the files policy's own ceiling.
|
|
const bytes = await readAll(object.body, deps.maxBytes ?? DEFAULT_FILE_BYTES)
|
|
// The bytes the job reads are the bytes the upload stored: a digest that moved is a different file.
|
|
if ((await sha256Hex(bytes)) !== payload.sha256) return fail("file_changed")
|
|
|
|
// Stage one, reading the file: a file already read (a retry, the same bytes uploaded again) is a
|
|
// cache hit and costs the platform nothing. The key is the file's own digest, which was just checked.
|
|
const cache = deps.cache ?? stageCache(deps.store, scope)
|
|
let read = await cache.get("ocr", payload.sha256, readDocumentSchema)
|
|
let readJobId = (await cache.get("ocr-job", payload.sha256, readJobSchema))?.jobId ?? activeDocument.ocrJobId
|
|
if (!read) {
|
|
let jobId = activeDocument.ocrJobId
|
|
if (jobId === null) {
|
|
const accepted = await deps.ocr.submit(bytes, {
|
|
filename: file.filename,
|
|
idempotencyKey: `documents.ocr.${payload.fileId}.${payload.sha256}`,
|
|
})
|
|
jobId = accepted.jobId
|
|
activeDocument = await records.setOcrJob(scope, activeDocument, jobId)
|
|
}
|
|
try {
|
|
read = projectDocument(await deps.ocr.wait(jobId, deps.signal ? { signal: deps.signal } : {}))
|
|
} catch (error) {
|
|
if (error instanceof OcrJobFailed) return fail("ocr_failed")
|
|
throw error
|
|
}
|
|
await cache.put("ocr", payload.sha256, read)
|
|
await cache.put("ocr-job", payload.sha256, { jobId })
|
|
readJobId = jobId
|
|
}
|
|
|
|
// Stage two, extraction: the extractor decides what is a line; a failure it calls permanent is the
|
|
// document's, recorded as its code, and any other is thrown for the queue to retry.
|
|
const isImage = IMAGE_CONTENT_TYPES.includes(file.contentType)
|
|
let extraction: Extraction
|
|
try {
|
|
extraction = await deps.extractor.extract({
|
|
document: read,
|
|
// An image upload is its own single page. A page of a PDF is drawn by the platform from the job
|
|
// that read it, on demand and only for a page the extractor asks for; a job the platform no
|
|
// longer holds has no image to give, and the page is then read from its text.
|
|
pageImage: async (index) => {
|
|
if (isImage) return index === 0 ? { bytes, mediaType: file.contentType } : undefined
|
|
if (readJobId === null) return undefined
|
|
return deps.ocr.page(readJobId, index + 1, deps.signal ? { signal: deps.signal } : {})
|
|
},
|
|
cache,
|
|
...(deps.signal ? { signal: deps.signal } : {}),
|
|
})
|
|
} catch (error) {
|
|
if (error instanceof ExtractionFailed) return fail(error.code)
|
|
throw error
|
|
}
|
|
const candidates = extraction.lines.slice(0, deps.maxLines ?? DEFAULT_MAX_LINES)
|
|
const proposed = new Set(
|
|
(await records.proposals(scope, activeDocument.id)).map(
|
|
(proposal) => (proposal.change as { lineKey?: string }).lineKey,
|
|
),
|
|
)
|
|
for (const [position, candidate] of candidates.entries()) {
|
|
const lineKey = await lineKeyOf(activeDocument.id, candidate.page, position)
|
|
if (proposed.has(lineKey)) continue
|
|
const change: LineChange = {
|
|
lineKey,
|
|
label: candidate.label,
|
|
value: candidate.value,
|
|
page: candidate.page,
|
|
position,
|
|
quote: candidate.quote,
|
|
confidence: candidate.confidence,
|
|
verification: candidate.verification,
|
|
...(candidate.note ? { note: candidate.note } : {}),
|
|
edited: false,
|
|
}
|
|
await records.createProposal(
|
|
scope,
|
|
{
|
|
targetKey: LINE_TARGET_KEY,
|
|
targetRecordId: activeDocument.id,
|
|
change,
|
|
evidence: [
|
|
{
|
|
fileId: payload.fileId,
|
|
page: candidate.page,
|
|
...(candidate.blockId ? { blockId: candidate.blockId } : {}),
|
|
...(candidate.box ? { box: candidate.box } : {}),
|
|
quote: candidate.quote,
|
|
},
|
|
],
|
|
model: `${read.engine.name}+${deps.extractor.name}`.slice(0, 128),
|
|
promptVersion: deps.extractor.version,
|
|
createdBy: scope.actor.userId,
|
|
},
|
|
`documents.line.${lineKey}`,
|
|
)
|
|
proposed.add(lineKey)
|
|
}
|
|
|
|
await records.finish(scope, activeDocument, {
|
|
pageCount: read.page_count,
|
|
proposedCount: proposed.size,
|
|
skippedPages: [...extraction.skipped],
|
|
finishedAt: toInstant(clock()),
|
|
})
|
|
return "extracted"
|
|
}
|
|
|
|
const DEFAULT_FILE_BYTES = 50 * 1024 * 1024
|
|
|
|
function createDataRecords(deps: ProcessDeps): DocumentProcessRecords {
|
|
return {
|
|
getFile: (scope, fileId) => getFile(deps.files, scope, fileId),
|
|
async ensureProcessing(scope, file, payload) {
|
|
const repo = deps.data.forScope(scope).of(documents)
|
|
const existing = (await repo.list({ where: [{ field: "fileId", op: "eq", value: payload.fileId }], limit: 1 }))
|
|
.items[0]
|
|
if (existing) return existing.status === "processing" ? existing : undefined
|
|
try {
|
|
return await repo.insert({
|
|
fileId: payload.fileId,
|
|
filename: file.filename,
|
|
status: "processing",
|
|
ocrJobId: null,
|
|
pageCount: null,
|
|
proposedCount: null,
|
|
failureCode: null,
|
|
skippedPages: null,
|
|
finishedAt: null,
|
|
})
|
|
} catch (error) {
|
|
if (!(error instanceof AppError && error.status === 409)) throw error
|
|
const raced = (await repo.list({ where: [{ field: "fileId", op: "eq", value: payload.fileId }], limit: 1 }))
|
|
.items[0]
|
|
if (!raced) throw error
|
|
return raced.status === "processing" ? raced : undefined
|
|
}
|
|
},
|
|
setOcrJob: (scope, document, jobId) =>
|
|
deps.data.forScope(scope).of(documents).update(document.id, { ocrJobId: jobId }, document.version),
|
|
async markFailed(scope, document, code, finishedAt) {
|
|
await deps.data
|
|
.forScope(scope)
|
|
.of(documents)
|
|
.update(document.id, { status: "failed", failureCode: code, finishedAt }, document.version)
|
|
},
|
|
proposals: (scope, documentId) => lineProposals(deps, scope, documentId),
|
|
async createProposal(scope, input, idempotencyKey) {
|
|
await createProposal(deps.proposals, scope, input, idempotencyKey)
|
|
},
|
|
async finish(scope, document, result) {
|
|
await deps.data
|
|
.forScope(scope)
|
|
.of(documents)
|
|
.update(
|
|
document.id,
|
|
{
|
|
status: "extracted",
|
|
pageCount: result.pageCount,
|
|
proposedCount: result.proposedCount,
|
|
skippedPages: result.skippedPages as never,
|
|
finishedAt: result.finishedAt,
|
|
},
|
|
document.version,
|
|
)
|
|
},
|
|
}
|
|
}
|