115 lines
4.8 KiB
TypeScript
115 lines
4.8 KiB
TypeScript
// The extractor over a scripted ModelGateway: the mapping of the gateway's failures to the codes a
|
|
// document records, and the retry each one gets. The whole path through the real gateway client is
|
|
// proved by test/documents-extraction.test.ts.
|
|
|
|
import { AiError, type GenerateObjectOptions, type ModelGateway, ModelGatewayError } from "@conusai-kit/platform"
|
|
import { describe, expect, it } from "vitest"
|
|
import { modelLineExtractor } from "../adapters/model-extractor.js"
|
|
import { ExtractionFailed, type ExtractionInput, type ReadDocument } from "../core/ports.js"
|
|
import { noCache, type StageCache } from "../core/stage-cache.js"
|
|
|
|
const document: ReadDocument = {
|
|
engine: { name: "fixture", version: "1" },
|
|
page_count: 1,
|
|
pages: [
|
|
{
|
|
index: 0,
|
|
width: 100,
|
|
height: 100,
|
|
source: "text_layer",
|
|
text: "Reference: QR-1",
|
|
class: "text",
|
|
lines: [{ text: "Reference: QR-1", bbox: [10, 10, 60, 20], confidence: 0.9, order: 0, repeated: false }],
|
|
blocks: [{ id: "p1-b1", bbox: [10, 10, 60, 20], order: 0, lines: [0] }],
|
|
},
|
|
],
|
|
}
|
|
|
|
const input = (cache: StageCache = noCache): ExtractionInput => ({ document, pageImage: async () => undefined, cache })
|
|
|
|
/** A gateway that answers from a script, one entry per call: a value is the reply, an Error is thrown. */
|
|
function scripted(...script: Array<unknown>): { gateway: ModelGateway; options: GenerateObjectOptions[] } {
|
|
const options: GenerateObjectOptions[] = []
|
|
const gateway: ModelGateway = {
|
|
async generateObject(_schema, _messages, call) {
|
|
options.push(call ?? {})
|
|
const next = script.shift()
|
|
if (next instanceof Error) throw next
|
|
return next as never
|
|
},
|
|
}
|
|
return { gateway, options }
|
|
}
|
|
|
|
const good = { lines: [{ label: "Reference", value: "QR-1", quote: "Reference: QR-1" }] }
|
|
const extractorOf = (gateway: ModelGateway) => modelLineExtractor({ gateway, model: "m" })
|
|
|
|
describe("the model extractor", () => {
|
|
it("names the model and the prompt version, which every proposal records", () => {
|
|
const extractor = extractorOf(scripted().gateway)
|
|
expect(extractor.name).toBe("m")
|
|
expect(extractor.version).toMatch(/^[a-z-]+\.v\d+$/)
|
|
})
|
|
|
|
it.each([
|
|
["refusal", "extraction_refused", 2],
|
|
["truncated", "extraction_truncated", 2],
|
|
["empty", "extraction_invalid", 2],
|
|
// a schema failure was already asked again by the gateway: the extractor does not ask a third time
|
|
["schema", "extraction_invalid", 1],
|
|
] as const)("a %s is asked again as the rules say, then ends as %s", async (kind, code, calls) => {
|
|
const failure = () => new ModelGatewayError(kind, "x")
|
|
const { gateway, options } = scripted(failure(), failure())
|
|
await expect(extractorOf(gateway).extract(input())).rejects.toMatchObject({ code })
|
|
expect(options).toHaveLength(calls)
|
|
})
|
|
|
|
it("answers when the second attempt is good, giving a truncated reply twice the room", async () => {
|
|
const { gateway, options } = scripted(new ModelGatewayError("truncated", "x"), good)
|
|
const { lines } = await extractorOf(gateway).extract(input())
|
|
expect(lines.map((line) => line.value)).toEqual(["QR-1"])
|
|
expect(options.map((call) => call.maxTokens)).toEqual([4096, 8192])
|
|
})
|
|
|
|
it("treats a request the gateway will never accept as the document's failure", async () => {
|
|
const bad = () => new AiError("bad image", { status: 400 })
|
|
const { gateway } = scripted(bad(), bad())
|
|
await expect(extractorOf(gateway).extract(input())).rejects.toMatchObject({
|
|
code: "extraction_invalid",
|
|
name: ExtractionFailed.name,
|
|
})
|
|
})
|
|
|
|
it.each([
|
|
["a network error", new AiError("down", { code: "network_error" })],
|
|
["a 503", new AiError("busy", { status: 503 })],
|
|
["a missing token scope", new AiError("no", { status: 403 })],
|
|
["a missing model", new ModelGatewayError("no_model", "x")],
|
|
])("leaves %s to the queue, which retries the job", async (_name, error) => {
|
|
const { gateway, options } = scripted(error)
|
|
await expect(extractorOf(gateway).extract(input())).rejects.toBe(error)
|
|
expect(options).toHaveLength(1)
|
|
})
|
|
|
|
it("reads a cached reply without asking the model", async () => {
|
|
const stored = new Map<string, unknown>()
|
|
const cache: StageCache = {
|
|
async get(stage, key, schema) {
|
|
const value = stored.get(`${stage}/${key}`)
|
|
return value === undefined ? undefined : schema.parse(value)
|
|
},
|
|
async put(stage, key, value) {
|
|
stored.set(`${stage}/${key}`, value)
|
|
},
|
|
}
|
|
const { gateway, options } = scripted(good)
|
|
await extractorOf(gateway).extract(input(cache))
|
|
await extractorOf(gateway).extract(input(cache))
|
|
expect(options).toHaveLength(1)
|
|
// another model is another key
|
|
const other = scripted(good)
|
|
await modelLineExtractor({ gateway: other.gateway, model: "other" }).extract(input(cache))
|
|
expect(other.options).toHaveLength(1)
|
|
})
|
|
})
|