import { describe, expect, it, vi } from 'vitest'; import { structureDocument } from '../../src/document-workbench/structure-document'; import { Inf0GatewayClient, type Inf0GatewayReceipt } from '../../src/document-workbench/inf0-gateway-client'; type Req = Omit[0], 'messages' | 'inputTokens'>; const request = { context: { tenantId: 'a', projectId: 'p', actorId: 'u', capability: 'inf0.complete', correlationId: 'c' }, admission: { tenantId: 'a', projectId: 'p', jobId: 'j', modelId: 'local_only', dataClass: 'high', maxInputTokens: 200, maxOutputTokens: 100, maximumCredits: 5, }, policy: { tenantId: 'a', projectId: 'p', allowedModelIds: ['local_only'], allowedDataClasses: ['high'], destination: 'internal', externalTransferAllowed: false, maxInputTokens: 200, maxOutputTokens: 100, }, reservation: { tenantId: 'a', projectId: 'p', jobId: 'j', status: 'reserved', credits: 5 }, } satisfies Req; const response: Inf0GatewayReceipt = { status: 'model_stopped', content: JSON.stringify({ value: { name: 'Acme' }, evidence: [{ field: 'name', sourceId: 'page1', start: 6, end: 10, transform: 'exact' }], }), modelExecuted: 'm1', providerKind: 'local', promptTokens: 50, completionTokens: 20, totalTokens: 70, }; const base = { request, countTokens: vi.fn(async () => 50), instructions: 'Extract the customer name', fields: [{ name: 'name', type: 'string' as const, required: true }], sources: [{ id: 'page1', text: 'Name: Acme' }], }; describe('structure document workflow', () => { it('returns only source-supported values and carries model usage without raw prompt', async () => { const complete = vi.fn(async () => response); const out = await structureDocument({ ...base, gateway: { complete } }); expect(out).toMatchObject({ status: 'mechanically_validated', value: { name: 'Acme' }, model: { modelExecuted: 'm1' }, }); expect(out.model).not.toHaveProperty('content'); expect(complete).toHaveBeenCalledOnce(); }); it('sends hallucinated data, invalid JSON and truncated output to review', async () => { const fake = { ...response, content: JSON.stringify({ value: { name: 'Other' }, evidence: [{ field: 'name', sourceId: 'page1', start: 6, end: 10, transform: 'exact' }], }), }; expect( (await structureDocument({ ...base, gateway: { complete: vi.fn(async () => fake) } })).status, ).toBe('needs_review'); expect( ( await structureDocument({ ...base, gateway: { complete: vi.fn(async () => ({ ...response, content: 'not json' })) }, }) ).status, ).toBe('needs_review'); expect( ( await structureDocument({ ...base, gateway: { complete: vi.fn(async () => ({ ...response, status: 'needs_review' as const })) }, }) ).status, ).toBe('needs_review'); }); it('runs through the real INF0 adapter and rejects an ungrounded response', async () => { const fetchMock = vi.fn(async () => Response.json({ content: response.content, modelExecuted: 'm1', finishReason: 'stop', routeDecision: { route: 'local_only', providerKind: 'local' }, cost: { promptTokens: 50, completionTokens: 20, totalTokens: 70 }, }), ); const gateway = new Inf0GatewayClient({ baseUrl: 'https://inf0.example.test', serviceToken: 'test', fetchImpl: fetchMock, }); const out = await structureDocument({ ...base, gateway }); expect(out).toMatchObject({ status: 'mechanically_validated', value: { name: 'Acme' } }); const call = fetchMock.mock.calls[0] as unknown as [string, RequestInit]; const body = JSON.parse(String(call[1].body)); expect(body.messages[1].content).toContain('Name: Acme'); expect(body.data_classification).toBe('high'); }); it('computes evidence deterministically when a local model returns values only', async () => { const out = await structureDocument({ ...base, gateway: { complete: vi.fn(async () => ({ ...response, content: JSON.stringify({ value: { name: 'Acme' } }) })), }, }); expect(out).toMatchObject({ status: 'mechanically_validated', value: { name: 'Acme' } }); if (out.status !== 'mechanically_validated') throw new Error('expected validated'); expect(out.evidence).toEqual([ { field: 'name', sourceId: 'page1', start: 6, end: 10, transform: 'exact' }, ]); }); it('accepts a flat object or a fenced JSON object from a small local model', async () => { const flat = await structureDocument({ ...base, gateway: { complete: vi.fn(async () => ({ ...response, content: JSON.stringify({ name: 'Acme' }) })) }, }); expect(flat).toMatchObject({ status: 'mechanically_validated', value: { name: 'Acme' } }); const fenced = await structureDocument({ ...base, gateway: { complete: vi.fn(async () => ({ ...response, content: '```json\n{"value":{"name":"Acme"}}\n```' })), }, }); expect(fenced).toMatchObject({ status: 'mechanically_validated', value: { name: 'Acme' } }); }); it('refuses ambiguous or absent values when computing evidence itself', async () => { const twice = await structureDocument({ ...base, sources: [{ id: 'p1', text: 'Acme AB and Acme AB' }], gateway: { complete: vi.fn(async () => ({ ...response, content: JSON.stringify({ value: { name: 'Acme AB' } }), })), }, }); expect(twice.status).toBe('needs_review'); const absent = await structureDocument({ ...base, gateway: { complete: vi.fn(async () => ({ ...response, content: JSON.stringify({ value: { name: 'Globex' } }), })), }, }); expect(absent.status).toBe('needs_review'); }); it('derives evidence for decimals and dates from locale-formatted source text', async () => { const out = await structureDocument({ request, countTokens: vi.fn(async () => 50), instructions: 'Extract amount and due date', fields: [ { name: 'amount', type: 'decimal' as const, required: true }, { name: 'due', type: 'date' as const, required: true }, ], sources: [{ id: 'p1', text: 'Belopp: 12 500,00 SEK\nFörfallodatum: 14/10/2026' }], gateway: { complete: vi.fn(async () => ({ ...response, content: JSON.stringify({ value: { amount: '12500.00', due: '2026-10-14' } }), })), }, }); expect(out.status).toBe('mechanically_validated'); if (out.status !== 'mechanically_validated') throw new Error('expected validated'); expect(out.evidence.map((e) => e.transform)).toEqual(['decimal', 'date']); expect(out.evidence[0]).toMatchObject({ start: 8, end: 17 }); }); it('rewrites a JSON-number decimal to the exact source string (no rounding)', async () => { const out = await structureDocument({ request, countTokens: vi.fn(async () => 50), instructions: 'Extract the amount', fields: [{ name: 'amount', type: 'decimal' as const, required: true }], sources: [{ id: 'p1', text: 'Belopp: 12 500,00 SEK' }], gateway: { complete: vi.fn(async () => ({ ...response, content: JSON.stringify({ amount: 12500 }) })), }, }); expect(out.status).toBe('mechanically_validated'); if (out.status !== 'mechanically_validated') throw new Error('expected validated'); expect(out.value.amount).toBe('12500.00'); }); it('rejects malformed contracts before a model call', async () => { const complete = vi.fn(async () => response); await expect(structureDocument({ ...base, fields: [], gateway: { complete } })).rejects.toThrow( 'Invalid structure', ); expect(complete).not.toHaveBeenCalled(); }); it('validates against the exact source and field contract supplied before async tokenization', async () => { const sources = [{ id: 'page1', text: 'Name: Evil' }]; const fields = [{ name: 'name', type: 'string' as const, required: true }]; const countTokens = vi.fn(async () => { sources[0]!.text = 'Name: Acme'; fields[0]!.required = false; return 50; }); const complete = vi.fn(async () => response); const result = await structureDocument({ ...base, sources, fields, countTokens, gateway: { complete } }); expect(result).toMatchObject({ status: 'needs_review', issues: [{ field: 'name', code: 'evidence_mismatch' }], }); expect(complete.mock.calls[0]?.[0]?.messages[1]?.content).toContain('Name: Evil'); }); it('uses the original admission and policy even if the caller mutates its request during tokenization', async () => { const mutableRequest = structuredClone(request); const countTokens = vi.fn(async () => { mutableRequest.admission.modelId = 'changed'; mutableRequest.policy.allowedModelIds.push('changed'); return 50; }); const complete = vi.fn(async () => response); await structureDocument({ ...base, request: mutableRequest, countTokens, gateway: { complete } }); expect(complete.mock.calls[0]?.[0]?.admission.modelId).toBe('local_only'); expect(complete.mock.calls[0]?.[0]?.policy.allowedModelIds).toEqual(['local_only']); }); it('keeps tokenizer input identical to the message sent to the gateway', async () => { const countTokens = vi.fn(async (messages: Parameters[0]['messages']) => { expect(Object.isFrozen(messages)).toBe(true); expect(Object.isFrozen(messages[1])).toBe(true); return 50; }); const complete = vi.fn(async () => response); await structureDocument({ ...base, countTokens, gateway: { complete } }); expect(complete.mock.calls[0]?.[0]?.messages).toEqual(countTokens.mock.calls[0]?.[0]); }); });