fix(compute): unify PDF layout result handling with object key support

Update compute backend interfaces and job logic to consistently handle parsed
PDF layout results as either direct data or S3 object key references. Refactor
result types, local and worker backends, and job processing to support this
pattern. Improves flexibility for large document parsing and object storage
integration.
This commit is contained in:
Richard R 2026-05-20 15:38:24 -06:00
parent 4ddedee8ec
commit 88809b826b
6 changed files with 63 additions and 14 deletions

View file

@ -36,10 +36,17 @@ export interface PdfLayoutJobRequest {
documentObjectKey: string;
}
export interface PdfLayoutJobResult {
parsed: ParsedPdfDocument;
timing?: WorkerJobTiming;
}
export type PdfLayoutJobResult =
| {
parsed: ParsedPdfDocument;
parsedObjectKey?: never;
timing?: WorkerJobTiming;
}
| {
parsed?: never;
parsedObjectKey: string;
timing?: WorkerJobTiming;
};
export type WorkerJobState = 'queued' | 'running' | 'succeeded' | 'failed';

View file

@ -33,7 +33,7 @@ import {
type WorkerJobStatusResponse,
type WorkerJobTiming,
} from '@openreader/compute-core/contracts';
import { GetObjectCommand, S3Client } from '@aws-sdk/client-s3';
import { GetObjectCommand, PutObjectCommand, S3Client } from '@aws-sdk/client-s3';
const JOBS_STREAM_NAME = 'compute_jobs';
const WHISPER_JOBS_SUBJECT = 'jobs.whisper';
@ -44,6 +44,8 @@ const JOB_STATES_BUCKET = 'job_states';
const JOB_STATES_TTL_MS = 24 * 60 * 60 * 1000;
const PULL_EXPIRES_MS = 1000;
const LOOP_ERROR_BACKOFF_MS = 500;
const DOCUMENT_ID_REGEX = /^[a-f0-9]{64}$/i;
const SAFE_NAMESPACE_REGEX = /^[a-zA-Z0-9._-]{1,128}$/;
interface QueuedJob<TPayload> {
jobId: string;
@ -104,6 +106,20 @@ function normalizeS3Prefix(prefix: string | undefined): string {
return value ? value.replace(/^\/+|\/+$/g, '') : 'openreader';
}
function sanitizeNamespace(namespace: string | null): string | null {
if (!namespace) return null;
return SAFE_NAMESPACE_REGEX.test(namespace) ? namespace : null;
}
function documentParsedKey(id: string, namespace: string | null, prefix: string): string {
if (!DOCUMENT_ID_REGEX.test(id)) {
throw new Error(`Invalid document id: ${id}`);
}
const ns = sanitizeNamespace(namespace);
const nsSegment = ns ? `ns/${ns}/` : '';
return `${prefix}/documents_v1/parsed_v1/${nsSegment}${id}.json`;
}
function buildS3Client(): S3Client {
const bucket = requireEnv('S3_BUCKET');
const region = requireEnv('S3_REGION');
@ -468,6 +484,19 @@ async function main(): Promise<void> {
return toArrayBuffer(new Uint8Array(bytes));
};
const putParsedObject = async (documentId: string, namespace: string | null, parsed: unknown): Promise<string> => {
const key = documentParsedKey(documentId, namespace, s3Prefix);
const body = Buffer.from(JSON.stringify(parsed));
await s3.send(new PutObjectCommand({
Bucket: s3Bucket,
Key: key,
Body: body,
ContentType: 'application/json',
ServerSideEncryption: 'AES256',
}));
return key;
};
if (prewarmModels) {
await ensureComputeModels();
}
@ -638,7 +667,6 @@ async function main(): Promise<void> {
return {
...result,
timing: {
...(result.timing ?? {}),
queueWaitMs,
s3FetchMs,
computeMs,
@ -667,10 +695,10 @@ async function main(): Promise<void> {
);
const computeMs = Date.now() - computeStartedAt;
const parsedObjectKey = await putParsedObject(parsed.documentId, parsed.namespace, result.parsed);
return {
...result,
parsedObjectKey,
timing: {
...(result.timing ?? {}),
queueWaitMs,
s3FetchMs,
computeMs,

View file

@ -35,6 +35,6 @@ export class LocalComputeBackend implements ComputeBackend {
if (!pdfBytes) {
throw new Error('Local compute PDF layout requires pdfBytes or (documentId + namespace)');
}
return (await runPdfLayoutFromPdfBuffer({ documentId: input.documentId, pdfBytes })).parsed;
return { parsed: (await runPdfLayoutFromPdfBuffer({ documentId: input.documentId, pdfBytes })).parsed };
}
}

View file

@ -21,8 +21,12 @@ export interface PdfLayoutInput {
pdfBytes?: ArrayBuffer;
}
export type PdfLayoutResult =
| { parsed: ParsedPdfDocument; parsedObjectKey?: never }
| { parsed?: never; parsedObjectKey: string };
export interface ComputeBackend {
mode: ComputeMode;
alignWords(input: WhisperAlignInput): Promise<WhisperAlignResult>;
parsePdfLayout(input: PdfLayoutInput): Promise<ParsedPdfDocument>;
parsePdfLayout(input: PdfLayoutInput): Promise<PdfLayoutResult>;
}

View file

@ -124,7 +124,13 @@ export class WorkerComputeBackend implements ComputeBackend {
if (status.status !== 'succeeded' || !status.result) {
throw new Error(status.error?.message || 'PDF layout worker job did not complete');
}
return status.result.parsed;
if (status.result.parsedObjectKey) {
return { parsedObjectKey: status.result.parsedObjectKey };
}
if (status.result.parsed) {
return { parsed: status.result.parsed };
}
throw new Error('PDF layout worker job completed without parsed output');
});
}

View file

@ -29,14 +29,18 @@ export async function parsePdfJob(input: ParsePdfJobInput): Promise<void> {
.where(and(eq(documents.id, input.documentId), eq(documents.userId, input.userId)));
const compute = await getCompute();
const parsed = await compute.parsePdfLayout({
const layout = await compute.parsePdfLayout({
documentId: input.documentId,
namespace: input.namespace,
documentObjectKey: documentKey(input.documentId, input.namespace),
});
const parsedJson = Buffer.from(JSON.stringify(parsed));
const parsedJsonKey = await putParsedDocumentBlob(input.documentId, parsedJson, input.namespace);
let parsedJsonKey = layout.parsedObjectKey ?? null;
if (!parsedJsonKey) {
if (!layout.parsed) throw new Error('Compute backend did not return parsed result');
const parsedJson = Buffer.from(JSON.stringify(layout.parsed));
parsedJsonKey = await putParsedDocumentBlob(input.documentId, parsedJson, input.namespace);
}
const cleared = await clearTtsSegmentCache({
userId: input.userId,