structured.test.ts
2,991 bytes
| 1 | import { describe, expect, it } from 'vitest'; |
|---|---|
| 2 | import { z } from 'zod'; |
| 3 | import type { LlmClient, LlmRequest, LlmResult, LlmStreamHandle } from './llm-client'; |
| 4 | import { requestStructured, StructuredParseError } from './structured'; |
| 5 | |
| 6 | const schema = z.object({ answer: z.string(), score: z.number().min(0).max(10) }); |
| 7 | |
| 8 | /** A client that returns a scripted sequence of raw responses. */ |
| 9 | class ScriptedClient implements LlmClient { |
| 10 | calls: LlmRequest[] = []; |
| 11 | constructor(private readonly responses: string[]) {} |
| 12 | async complete(req: LlmRequest): Promise<LlmResult> { |
| 13 | this.calls.push(req); |
| 14 | const text = this.responses[this.calls.length - 1] ?? this.responses.at(-1) ?? ''; |
| 15 | return { |
| 16 | text, |
| 17 | usage: { promptTokens: 10, completionTokens: 5, totalTokens: 15, costUsd: 0.001 }, |
| 18 | model: req.model, |
| 19 | latencyMs: 1, |
| 20 | }; |
| 21 | } |
| 22 | streamComplete(): Promise<LlmStreamHandle> { |
| 23 | throw new Error('not used'); |
| 24 | } |
| 25 | } |
| 26 | |
| 27 | const req: LlmRequest = { model: 'x/y', messages: [{ role: 'user', content: 'hi' }] }; |
| 28 | |
| 29 | describe('requestStructured', () => { |
| 30 | it('returns the parsed value on a clean first attempt', async () => { |
| 31 | const client = new ScriptedClient(['{"answer":"ok","score":7}']); |
| 32 | const r = await requestStructured(client, req, schema); |
| 33 | expect(r.repaired).toBe(false); |
| 34 | expect(r.value).toEqual({ answer: 'ok', score: 7 }); |
| 35 | expect(client.calls).toHaveLength(1); |
| 36 | expect(r.usage.costUsd).toBeCloseTo(0.001); |
| 37 | }); |
| 38 | |
| 39 | it('repairs once when the first response is schema-invalid', async () => { |
| 40 | const client = new ScriptedClient([ |
| 41 | '{"answer":"ok","score":99}', // invalid: score > 10 |
| 42 | '{"answer":"ok","score":8}', // repaired |
| 43 | ]); |
| 44 | const r = await requestStructured(client, req, schema); |
| 45 | expect(r.repaired).toBe(true); |
| 46 | expect(r.value.score).toBe(8); |
| 47 | expect(client.calls).toHaveLength(2); |
| 48 | // usage accumulates across both calls |
| 49 | expect(r.usage.costUsd).toBeCloseTo(0.002); |
| 50 | // the repair call carries the repair system prompt |
| 51 | expect(client.calls[1]!.messages[0]!.content).toMatch(/JSON fixer/i); |
| 52 | }); |
| 53 | |
| 54 | it('repairs malformed (non-JSON-wrapped) output', async () => { |
| 55 | const client = new ScriptedClient([ |
| 56 | 'I think the answer is good.', // no JSON at all |
| 57 | '```json\n{"answer":"fixed","score":6}\n```', |
| 58 | ]); |
| 59 | const r = await requestStructured(client, req, schema); |
| 60 | expect(r.repaired).toBe(true); |
| 61 | expect(r.value.answer).toBe('fixed'); |
| 62 | }); |
| 63 | |
| 64 | it('throws StructuredParseError (with billed usage) after a failed repair', async () => { |
| 65 | const client = new ScriptedClient(['garbage', 'still garbage']); |
| 66 | await expect(requestStructured(client, req, schema)).rejects.toBeInstanceOf(StructuredParseError); |
| 67 | try { |
| 68 | await requestStructured(client, req, schema); |
| 69 | } catch (e) { |
| 70 | expect(e).toBeInstanceOf(StructuredParseError); |
| 71 | // tokens from both attempts are still billed |
| 72 | expect((e as StructuredParseError).usage.costUsd).toBeCloseTo(0.002); |
| 73 | } |
| 74 | }); |
| 75 | }); |
| 76 | |