diff --git a/src/lib/documents/__tests__/extract-reads.test.ts b/src/lib/documents/__tests__/extract-reads.test.ts index 57fa54549..568f1eba1 100644 --- a/src/lib/documents/__tests__/extract-reads.test.ts +++ b/src/lib/documents/__tests__/extract-reads.test.ts @@ -201,16 +201,21 @@ describe("planInboundTextReads", () => { const line = (i: number) => `Analyt ${i} ${i},0 mg/dl 1 - 9`; const text = (lines: number) => Array.from({ length: lines }, (_, i) => line(i)).join("\n"); + const schedule = (calls: number) => ({ + maxModelCalls: () => calls, + arm: async () => undefined, + }); + const noDeadline = schedule(Number.POSITIVE_INFINITY); it("reads a short text as one read", () => { - const plan = planInboundTextReads(text(10), undefined); + const plan = planInboundTextReads(text(10), noDeadline); expect(plan.args.chunks).toHaveLength(1); expect(plan.partial).toBe(false); }); it("splits a long text at line boundaries into reads no longer than one read carries", () => { const whole = text(400); - const plan = planInboundTextReads(whole, undefined); + const plan = planInboundTextReads(whole, noDeadline); expect(plan.args.chunks.length).toBeGreaterThan(1); for (const chunk of plan.args.chunks) { expect(chunk.text!.length).toBeLessThanOrEqual(INBOUND_TEXT_READ_CHARS); @@ -221,10 +226,6 @@ describe("planInboundTextReads", () => { }); it("grows the parts when the deadline allows few reads, and reads from the start past twice the size", () => { - const schedule = (calls: number) => ({ - maxModelCalls: () => calls, - arm: async () => undefined, - }); // Two reads' worth of calls: the parts grow to hold the text. const grown = planInboundTextReads(text(200), schedule(4)); expect(grown.args.chunks).toHaveLength(2); @@ -239,7 +240,7 @@ describe("planInboundTextReads", () => { }); it("caps the reads of a run with no deadline", () => { - const plan = planInboundTextReads(text(4000), undefined); + const plan = planInboundTextReads(text(4000), noDeadline); expect(plan.args.chunks.length).toBe(OCR_MAX_PAGES); expect(plan.partial).toBe(true); }); diff --git a/src/lib/documents/extract.ts b/src/lib/documents/extract.ts index d48803484..361faa61a 100644 --- a/src/lib/documents/extract.ts +++ b/src/lib/documents/extract.ts @@ -462,7 +462,7 @@ export interface InboundReads { /** A document's text in parts at line boundaries, as many as the run's deadline allows. */ export function planInboundTextReads( text: string, - schedule: OcrSchedule | undefined, + schedule: OcrSchedule, ): InboundReads { const { chunks, unreadParts } = planTextReads( text, @@ -483,7 +483,7 @@ export function planInboundTextReads( */ export async function planInboundReads( source: InboundReadSource, - schedule: OcrSchedule | undefined, + schedule: OcrSchedule, ): Promise<({ ok: true } & InboundReads) | { ok: false }> { if (source.mode === "text") { return { ok: true, ...planInboundTextReads(source.text, schedule) }; diff --git a/src/lib/documents/ocr-plan.ts b/src/lib/documents/ocr-plan.ts index 3fec56366..f46242133 100644 --- a/src/lib/documents/ocr-plan.ts +++ b/src/lib/documents/ocr-plan.ts @@ -66,8 +66,7 @@ export const SYNCHRONOUS_READ_SCHEDULE: OcrSchedule = { }; /** How many reads the deadline leaves room for. */ -export function maxReadsFor(schedule: OcrSchedule | undefined): number { - if (!schedule) return Number.POSITIVE_INFINITY; +export function maxReadsFor(schedule: OcrSchedule): number { return Math.max( 1, Math.floor(schedule.maxModelCalls() / OCR_EXTRACT_MODEL_CALLS), @@ -130,7 +129,7 @@ export function splitOcrText( */ export function planTextReads( text: string, - schedule: OcrSchedule | undefined, + schedule: OcrSchedule, target: number, ): { chunks: OcrChunk[]; unreadParts: number } { const maxReads = Math.min(maxReadsFor(schedule), OCR_MAX_PAGES); @@ -186,7 +185,7 @@ export async function planVisionReads( bytes: Buffer, mime: OcrVisionMime, pdfSupported: boolean, - schedule: OcrSchedule | undefined, + schedule: OcrSchedule, ): Promise { if (mime !== "application/pdf") { return { diff --git a/src/lib/labs/__tests__/ocr-run.test.ts b/src/lib/labs/__tests__/ocr-run.test.ts index f26d226e7..2ca533791 100644 --- a/src/lib/labs/__tests__/ocr-run.test.ts +++ b/src/lib/labs/__tests__/ocr-run.test.ts @@ -94,6 +94,8 @@ const schedule = (maxModelCalls: number) => ({ maxModelCalls: () => maxModelCalls, arm: vi.fn().mockResolvedValue(undefined), }); +/** A run whose deadline leaves room for every read. */ +const noDeadline = () => schedule(Number.POSITIVE_INFINITY); beforeEach(() => { vi.clearAllMocks(); @@ -156,6 +158,18 @@ describe("planning the reads", () => { }); }); +describe("a scan is always planned against a schedule", () => { + it("cannot be started without one", () => { + // The scan only runs in the worker, which hands it the run's deadline. + // A caller with no schedule would read a twenty-page report as forty + // model calls, so the signature leaves no such call to write. + const start = () => + // @ts-expect-error -- `schedule` is required + executeOcrExtraction({ userId: "u1", input: text(), budget }); + expect(typeof start).toBe("function"); + }); +}); + describe("text mode", () => { it("charges what the provider reported to the provider that served", async () => { vi.mocked(runOcrExtraction).mockImplementation(extraction(2345)); @@ -163,6 +177,7 @@ describe("text mode", () => { userId: "u1", input: text(), budget, + schedule: noDeadline(), }); expect(out).toEqual({ ok: true, data: { rows: [] } }); expect(reconcileSpend).toHaveBeenCalledWith( @@ -177,7 +192,12 @@ describe("text mode", () => { it("charges the reservation when the provider reported nothing", async () => { vi.mocked(runOcrExtraction).mockImplementation(extraction(0)); - await executeOcrExtraction({ userId: "u1", input: text(), budget }); + await executeOcrExtraction({ + userId: "u1", + input: text(), + budget, + schedule: noDeadline(), + }); expect(reconcileSpend).toHaveBeenCalledWith( "u1", 1500, @@ -210,7 +230,12 @@ describe("text mode", () => { it("passes the cached share of what was spent to the ledger", async () => { vi.mocked(runOcrExtraction).mockImplementation(extraction(3000, 1200)); - await executeOcrExtraction({ userId: "u1", input: text(), budget }); + await executeOcrExtraction({ + userId: "u1", + input: text(), + budget, + schedule: noDeadline(), + }); expect(reconcileSpend).toHaveBeenCalledWith( "u1", 1500, @@ -231,7 +256,12 @@ describe("text mode", () => { expect(await args.reserveRead!()).toBe(true); return extraction(20_000)(args); }); - await executeOcrExtraction({ userId: "u1", input: text(), budget }); + await executeOcrExtraction({ + userId: "u1", + input: text(), + budget, + schedule: noDeadline(), + }); expect(reserveBudget).toHaveBeenCalledWith( "u1", 8000, @@ -265,6 +295,7 @@ describe("text mode", () => { userId: "u1", input: text(), budget: { ...budget, owner: "user" }, + schedule: noDeadline(), }); expect(vi.mocked(reserveBudget).mock.calls[0][3]).toBe(2_000_000); expect(vi.mocked(reserveBudget).mock.calls[0][4]).toBe("user"); @@ -293,7 +324,12 @@ describe("text mode", () => { answer = await args.reserveRead!(); return { data: { rows: [] } as never, spentTokens: 0, cachedTokens: 0 }; }); - await executeOcrExtraction({ userId: "u1", input: text(), budget }); + await executeOcrExtraction({ + userId: "u1", + input: text(), + budget, + schedule: noDeadline(), + }); // 18 500 caught up first (allowed, the day stood at 190 000); the next // read is then judged at 208 500 and refused. Booking only the 8 000 // ceiling, or both in one reservation, would have let it through. @@ -322,7 +358,12 @@ describe("text mode", () => { await args.reserveRead!(); return { data: { rows: [] } as never, spentTokens: 0, cachedTokens: 0 }; }); - await executeOcrExtraction({ userId: "u1", input: text(), budget }); + await executeOcrExtraction({ + userId: "u1", + input: text(), + budget, + schedule: noDeadline(), + }); // 20 000 spent, 4 000 of it cached: 16 000 net against 1 500 reserved. expect(vi.mocked(reserveBudget).mock.calls[0][1]).toBe(14_500); }); @@ -337,6 +378,7 @@ describe("text mode", () => { userId: "u1", input: text(), budget: { ...budget, owner: undefined as never }, + schedule: noDeadline(), }); expect(answer).toBe(false); expect(reserveBudget).not.toHaveBeenCalled(); @@ -347,7 +389,12 @@ describe("text mode", () => { new OcrExtractError("x", "budget"), ); expect( - await executeOcrExtraction({ userId: "u1", input: text(), budget }), + await executeOcrExtraction({ + userId: "u1", + input: text(), + budget, + schedule: noDeadline(), + }), ).toMatchObject({ status: 429, errorCode: "labs.ocr.budgetExceeded" }); vi.mocked(runOcrExtraction).mockRejectedValueOnce( @@ -358,14 +405,24 @@ describe("text mode", () => { ), ); expect( - await executeOcrExtraction({ userId: "u1", input: text(), budget }), + await executeOcrExtraction({ + userId: "u1", + input: text(), + budget, + schedule: noDeadline(), + }), ).toMatchObject({ ok: false, errorCode: "consent.ai.required" }); vi.mocked(runOcrExtraction).mockRejectedValueOnce( new OcrExtractError("x", "error"), ); expect( - await executeOcrExtraction({ userId: "u1", input: text(), budget }), + await executeOcrExtraction({ + userId: "u1", + input: text(), + budget, + schedule: noDeadline(), + }), ).toMatchObject({ status: 500, errorCode: "aiRuns.failed" }); }); @@ -379,7 +436,12 @@ describe("text mode", () => { answer = await args.reserveRead!(); return extraction(4000)(args); }); - await executeOcrExtraction({ userId: "u1", input: text(), budget }); + await executeOcrExtraction({ + userId: "u1", + input: text(), + budget, + schedule: noDeadline(), + }); expect(answer).toBe(false); // The refused reservation was handed back by reserveBudget itself; the // run settles against the first read's alone. @@ -399,6 +461,7 @@ describe("text mode", () => { userId: "u1", input: text("x".repeat(OCR_TEXT_READ_MAX_CHARS + 1)), budget, + schedule: noDeadline(), }); expect(out).toMatchObject({ ok: false, @@ -421,7 +484,12 @@ describe("text mode", () => { args.meter!.spent += 900; throw new OcrExtractError("x"); }); - await executeOcrExtraction({ userId: "u1", input: text(), budget }); + await executeOcrExtraction({ + userId: "u1", + input: text(), + budget, + schedule: noDeadline(), + }); expect(reconcileSpend).toHaveBeenCalledWith( "u1", 1500, @@ -438,6 +506,7 @@ describe("text mode", () => { userId: "u1", input: text(), budget, + schedule: noDeadline(), }); expect(out).toMatchObject({ ok: false, @@ -462,6 +531,7 @@ describe("text mode", () => { userId: "u1", input: text(), budget, + schedule: noDeadline(), }); expect(out).toMatchObject({ ok: false, @@ -482,6 +552,7 @@ describe("text mode", () => { userId: "u1", input: text(), budget, + schedule: noDeadline(), }); expect(out).toMatchObject({ ok: false, status: 502 }); expect(JSON.stringify(out)).toContain("configured AI provider"); @@ -563,6 +634,7 @@ describe("vision mode", () => { userId: "u1", input: pdf(anthropicPick), budget, + schedule: noDeadline(), }); const args = vi.mocked(runOcrExtraction).mock.calls[0][0]; expect(args.chunks.map((c) => c.pages)).toEqual([[1], [2], [3], [4]]); @@ -585,6 +657,7 @@ describe("vision mode", () => { userId: "u1", input: pdf(anthropicPick), budget, + schedule: noDeadline(), }); const args = vi.mocked(runOcrExtraction).mock.calls[i][0]; expect(args.chunks).toHaveLength(1); @@ -627,6 +700,7 @@ describe("vision mode", () => { userId: "u1", input: pdf(), budget, + schedule: noDeadline(), }); expect(out).toEqual({ ok: true, @@ -643,6 +717,7 @@ describe("vision mode", () => { userId: "u1", input: pdf(), budget, + schedule: noDeadline(), }); expect(out).toMatchObject({ ok: false, diff --git a/src/lib/labs/ocr-run.ts b/src/lib/labs/ocr-run.ts index 629151cc3..75b668be5 100644 --- a/src/lib/labs/ocr-run.ts +++ b/src/lib/labs/ocr-run.ts @@ -137,8 +137,11 @@ export async function executeOcrExtraction(args: { userId: string; input: OcrRunInput; budget: AiRunBudget | undefined; - /** The run's deadline; absent, the read is planned without one. */ - schedule?: OcrSchedule; + /** + * The run's deadline. Required: the scan only runs in the background + * worker, and a read planned without a bound could make forty model calls. + */ + schedule: OcrSchedule; }): Promise> { const { userId, input, budget, schedule } = args; const servedBy = input.pick.entry.providerType as ProviderChainType; @@ -205,7 +208,7 @@ export async function executeOcrExtraction(args: { errorCode: "labs.ocr.tooLarge", }); } - await schedule?.arm(parts.length * OCR_EXTRACT_MODEL_CALLS); + await schedule.arm(parts.length * OCR_EXTRACT_MODEL_CALLS); return read({ mode: "text", chunks: parts.map((text) => ({ text, pages: [] })), @@ -238,7 +241,7 @@ export async function executeOcrExtraction(args: { }); } - await schedule?.arm(chunks.length * OCR_EXTRACT_MODEL_CALLS); + await schedule.arm(chunks.length * OCR_EXTRACT_MODEL_CALLS); const outcome = await read({ mode: "vision", chunks, totalPages }); return outcome.ok && pageCoverage ? { ok: true, data: { ...outcome.data, pageCoverage } }