ConusAI vendor v0.1.0-dev.20261027+9b1131fb

modules/documents/core/process.ts

@@ -0,0 +1,304 @@
+// 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,
+        )
+    },
+  }
+}