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:
parent
4ddedee8ec
commit
88809b826b
6 changed files with 63 additions and 14 deletions
|
|
@ -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';
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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 };
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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>;
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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');
|
||||
});
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Reference in a new issue