structured.ts
2,275 bytes
| 1 | /** |
|---|---|
| 2 | * Force a model to return schema-valid JSON, with one repair attempt. |
| 3 | * |
| 4 | * Pipeline: call → strip/extract/validate → on failure, re-ask with the raw |
| 5 | * output and the exact validation error → validate again. Tokens spent on the |
| 6 | * failed attempt are still billed (the user's key paid for them), so usage is |
| 7 | * accumulated across both calls and surfaced even when the repair ultimately |
| 8 | * fails. |
| 9 | */ |
| 10 | import type { z } from 'zod'; |
| 11 | import { parseStructured } from './json-repair'; |
| 12 | import type { LlmClient, LlmRequest } from './llm-client'; |
| 13 | import { buildRepairPrompt } from './prompts'; |
| 14 | import type { Usage } from './types'; |
| 15 | import { addUsage } from './usage'; |
| 16 | |
| 17 | export interface StructuredResult<T> { |
| 18 | value: T; |
| 19 | usage: Usage; |
| 20 | latencyMs: number; |
| 21 | raw: string; |
| 22 | repaired: boolean; |
| 23 | } |
| 24 | |
| 25 | /** Raised when JSON is still invalid after the single repair retry. */ |
| 26 | export class StructuredParseError extends Error { |
| 27 | readonly usage: Usage; |
| 28 | readonly latencyMs: number; |
| 29 | constructor(message: string, opts: { usage: Usage; latencyMs: number }) { |
| 30 | super(message); |
| 31 | this.name = 'StructuredParseError'; |
| 32 | this.usage = opts.usage; |
| 33 | this.latencyMs = opts.latencyMs; |
| 34 | } |
| 35 | } |
| 36 | |
| 37 | export async function requestStructured<S extends z.ZodTypeAny>( |
| 38 | llm: LlmClient, |
| 39 | req: LlmRequest, |
| 40 | schema: S, |
| 41 | ): Promise<StructuredResult<z.infer<S>>> { |
| 42 | const first = await llm.complete({ ...req, json: true }); |
| 43 | const parsed = parseStructured(first.text, schema); |
| 44 | if (parsed.ok) { |
| 45 | return { |
| 46 | value: parsed.value, |
| 47 | usage: first.usage, |
| 48 | latencyMs: first.latencyMs, |
| 49 | raw: first.text, |
| 50 | repaired: false, |
| 51 | }; |
| 52 | } |
| 53 | |
| 54 | // Single repair attempt: hand the model its own output + the exact error. |
| 55 | const second = await llm.complete({ |
| 56 | ...req, |
| 57 | json: true, |
| 58 | messages: buildRepairPrompt(first.text, parsed.error), |
| 59 | }); |
| 60 | const usage = addUsage(first.usage, second.usage); |
| 61 | const latencyMs = first.latencyMs + second.latencyMs; |
| 62 | const parsed2 = parseStructured(second.text, schema); |
| 63 | if (parsed2.ok) { |
| 64 | return { value: parsed2.value, usage, latencyMs, raw: second.text, repaired: true }; |
| 65 | } |
| 66 | |
| 67 | throw new StructuredParseError( |
| 68 | `Model output failed schema validation after one repair attempt: ${parsed2.error}`, |
| 69 | { usage, latencyMs }, |
| 70 | ); |
| 71 | } |
| 72 | |